Added Probe/ProbeTLS functions and demo command

protocol/probe.go

@@ -0,0 +1,70 @@
+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()))
+}