71 lines
2.3 KiB
Go
71 lines
2.3 KiB
Go
package protocol
|
|
|
|
import (
|
|
"context"
|
|
"crypto/tls"
|
|
"net"
|
|
"time"
|
|
)
|
|
|
|
// Dialer used by probes to establish a connection.
|
|
var Dialer net.Dialer
|
|
|
|
// Probe a network service by reading its banner and running protocol detection.
|
|
func Probe(network, address string) (proto *Protocol, confidence float64, err error) {
|
|
return ProbeContext(context.Background(), network, address)
|
|
}
|
|
|
|
// ProbeContext is like [Probe] with a [context.Context].
|
|
func ProbeContext(ctx context.Context, network, address string) (proto *Protocol, confidence float64, err error) {
|
|
var conn net.Conn
|
|
if conn, err = Dialer.DialContext(ctx, network, address); err != nil {
|
|
return
|
|
}
|
|
defer func() { _ = conn.Close() }()
|
|
return probeConn(conn)
|
|
}
|
|
|
|
// ProbeTimeout is like [Probe] but with a fixed timeout.
|
|
func ProbeTimeout(network, address string, timeout time.Duration) (proto *Protocol, confidence float64, err error) {
|
|
ctx, cancel := context.WithTimeout(context.Background(), timeout)
|
|
defer cancel()
|
|
return ProbeContext(ctx, network, address)
|
|
}
|
|
|
|
// ProbeTLS is like [Probe] but first establishes a TLS connection.
|
|
func ProbeTLS(network, address string, tlsConfig *tls.Config) (proto *Protocol, confidence float64, err error) {
|
|
return ProbeTLSContext(context.Background(), network, address, tlsConfig)
|
|
}
|
|
|
|
// ProbeTLSContext is like [ProbeTLS] with a [context.Context].
|
|
func ProbeTLSContext(ctx context.Context, network, address string, tlsConfig *tls.Config) (proto *Protocol, confidence float64, err error) {
|
|
var conn net.Conn
|
|
if conn, err = Dialer.DialContext(ctx, network, address); err != nil {
|
|
return
|
|
}
|
|
defer func() { _ = conn.Close() }()
|
|
secure := tls.Client(conn, tlsConfig)
|
|
if err = secure.Handshake(); err != nil {
|
|
return
|
|
}
|
|
return probeConn(secure)
|
|
}
|
|
|
|
// ProbeTLSTimeout is like [ProbeTLS] but with a fixed timeout.
|
|
func ProbeTLSTimeout(network, address string, tlsConfig *tls.Config, timeout time.Duration) (proto *Protocol, confidence float64, err error) {
|
|
ctx, cancel := context.WithTimeout(context.Background(), timeout)
|
|
defer cancel()
|
|
return ProbeTLSContext(ctx, network, address, tlsConfig)
|
|
}
|
|
|
|
func probeConn(conn net.Conn) (proto *Protocol, confidence float64, err error) {
|
|
var (
|
|
data = make([]byte, 2048)
|
|
n int
|
|
)
|
|
if n, err = conn.Read(data); err != nil {
|
|
return
|
|
}
|
|
return Detect(Client, data[:n], getPortFromAddr(conn.LocalAddr()), getPortFromAddr(conn.RemoteAddr()))
|
|
}
|