sniffing protocol for port forwarding

handler/forward/internal/forward/forward.go

@@ -14,7 +14,7 @@ import (
 	xio "github.com/go-gost/x/internal/io"
 )
 
-func SniffHost(ctx context.Context, rdw io.ReadWriter) (rw io.ReadWriter, host string, err error) {
+func Sniffing(ctx context.Context, rdw io.ReadWriter) (rw io.ReadWriter, host string, protocol string, err error) {
 	rw = rdw
 
 	// try to sniff TLS traffic
@@ -24,7 +24,8 @@ func SniffHost(ctx context.Context, rdw io.ReadWriter) (rw io.ReadWriter, host s
 	if err == nil &&
 		hdr[0] == dissector.Handshake &&
 		binary.BigEndian.Uint16(hdr[1:3]) == tls.VersionTLS10 {
-		return sniffSNI(ctx, rw)
+		rw, host, err = sniffSNI(ctx, rw)
+		return
 	}
 
 	// try to sniff HTTP traffic
@@ -39,6 +40,8 @@ func SniffHost(ctx context.Context, rdw io.ReadWriter) (rw io.ReadWriter, host s
 		}
 	}
 
+	protocol = sniffProtocol(hdr[:])
+
 	return
 }
 
@@ -82,3 +85,14 @@ func isHTTP(s string) bool {
 		strings.HasPrefix(http.MethodConnect, s) ||
 		strings.HasPrefix(http.MethodTrace, s)
 }
+
+const (
+	SSHv2 = "SSH-2"
+)
+
+func sniffProtocol(hdr []byte) string {
+	if string(hdr) == SSHv2 {
+		return "ssh"
+	}
+	return ""
+}

handler/forward/local/handler.go

@@ -93,15 +93,20 @@ func (h *forwardHandler) Handle(ctx context.Context, conn net.Conn, opts ...hand
 
 	var rw io.ReadWriter = conn
 	var host string
+	var protocol string
 	if h.md.sniffing {
 		if network == "tcp" {
-			rw, host, _ = forward.SniffHost(ctx, conn)
+			rw, host, protocol, _ = forward.Sniffing(ctx, conn)
+			h.options.Logger.Debugf("sniffing: host=%s, protocol=%s", host, protocol)
 		}
 	}
 
 	var target *chain.Node
 	if h.hop != nil {
-		target = h.hop.Select(ctx, chain.HostSelectOption(host))
+		target = h.hop.Select(ctx,
+			chain.HostSelectOption(host),
+			chain.ProtocolSelectOption(protocol),
+		)
 	}
 	if target == nil {
 		err := errors.New("target not available")

handler/forward/remote/handler.go

@@ -84,14 +84,18 @@ func (h *forwardHandler) Handle(ctx context.Context, conn net.Conn, opts ...hand
 
 	var rw io.ReadWriter = conn
 	var host string
+	var protocol string
 	if h.md.sniffing {
 		if network == "tcp" {
-			rw, host, _ = forward.SniffHost(ctx, conn)
+			rw, host, protocol, _ = forward.Sniffing(ctx, conn)
 		}
 	}
 	var target *chain.Node
 	if h.hop != nil {
-		target = h.hop.Select(ctx, chain.HostSelectOption(host))
+		target = h.hop.Select(ctx,
+			chain.HostSelectOption(host),
+			chain.ProtocolSelectOption(protocol),
+		)
 	}
 	if target == nil {
 		err := errors.New("target not available")