add traffic sniffing for handler

internal/io/io.go

@@ -1,6 +1,9 @@
 package io
 
-import "io"
+import (
+	"io"
+	"time"
+)
 
 type readWriter struct {
 	io.Reader
@@ -27,3 +30,14 @@ func NewReadWriteCloser(r io.Reader, w io.Writer, c io.Closer) io.ReadWriteClose
 		Closer: c,
 	}
 }
+
+type setReadDeadline interface {
+	SetReadDeadline(t time.Time) error
+}
+
+func SetReadDeadline(rw io.ReadWriter, t time.Time) error {
+	if v, _ := rw.(setReadDeadline); v != nil {
+		return v.SetReadDeadline(t)
+	}
+	return nil
+}

internal/util/forward/forward.go

@@ -1,90 +0,0 @@
-package forward
-
-import (
-	"bufio"
-	"bytes"
-	"context"
-	"crypto/tls"
-	"encoding/binary"
-	"io"
-	"net/http"
-	"strings"
-
-	dissector "github.com/go-gost/tls-dissector"
-	xio "github.com/go-gost/x/internal/io"
-)
-
-func Sniffing(ctx context.Context, rdw io.ReadWriter) (rw io.ReadWriter, host string, protocol string, err error) {
-	rw = rdw
-
-	// try to sniff TLS traffic
-	var hdr [dissector.RecordHeaderLen]byte
-	n, err := io.ReadFull(rw, hdr[:])
-	rw = xio.NewReadWriter(io.MultiReader(bytes.NewReader(hdr[:n]), rw), rw)
-	tlsVersion := binary.BigEndian.Uint16(hdr[1:3])
-	if err == nil &&
-		hdr[0] == dissector.Handshake &&
-		(tlsVersion >= tls.VersionTLS10 && tlsVersion <= tls.VersionTLS13) {
-		rw, host, err = sniffSNI(rw)
-		protocol = ProtoTLS
-		return
-	}
-
-	// try to sniff HTTP traffic
-	if isHTTP(string(hdr[:])) {
-		buf := new(bytes.Buffer)
-		var r *http.Request
-		r, err = http.ReadRequest(bufio.NewReader(io.TeeReader(rw, buf)))
-		rw = xio.NewReadWriter(io.MultiReader(buf, rw), rw)
-		if err == nil {
-			host = r.Host
-			protocol = ProtoHTTP
-			return
-		}
-	}
-
-	protocol = sniffProtocol(hdr[:])
-
-	return
-}
-
-func sniffSNI(rw io.ReadWriter) (io.ReadWriter, string, error) {
-	buf := new(bytes.Buffer)
-	host, err := getServerName(io.TeeReader(rw, buf))
-	rw = xio.NewReadWriter(io.MultiReader(buf, rw), rw)
-	return rw, host, err
-}
-
-func getServerName(r io.Reader) (host string, err error) {
-	clientHello, err := dissector.ParseClientHello(r)
-	if err != nil {
-		return
-	}
-	host = clientHello.ServerName
-	return
-}
-
-func isHTTP(s string) bool {
-	return strings.HasPrefix(http.MethodGet, s[:3]) ||
-		strings.HasPrefix(http.MethodPost, s[:4]) ||
-		strings.HasPrefix(http.MethodPut, s[:3]) ||
-		strings.HasPrefix(http.MethodDelete, s) ||
-		strings.HasPrefix(http.MethodOptions, s) ||
-		strings.HasPrefix(http.MethodPatch, s) ||
-		strings.HasPrefix(http.MethodHead, s[:4]) ||
-		strings.HasPrefix(http.MethodConnect, s) ||
-		strings.HasPrefix(http.MethodTrace, s)
-}
-
-const (
-	ProtoHTTP  = "http"
-	ProtoTLS   = "tls"
-	ProtoSSHv2 = "SSH-2"
-)
-
-func sniffProtocol(hdr []byte) string {
-	if string(hdr) == ProtoSSHv2 {
-		return "ssh"
-	}
-	return ""
-}

internal/util/sniffing/sniff.go

@@ -0,0 +1,55 @@
+package sniffing
+
+import (
+	"bufio"
+	"context"
+	"crypto/tls"
+	"encoding/binary"
+	"net/http"
+	"strings"
+
+	dissector "github.com/go-gost/tls-dissector"
+)
+
+const (
+	ProtoHTTP = "http"
+	ProtoTLS  = "tls"
+	ProtoSSH  = "ssh"
+)
+
+func Sniff(ctx context.Context, r *bufio.Reader) (proto string, err error) {
+	hdr, err := r.Peek(dissector.RecordHeaderLen)
+	if err != nil {
+		return
+	}
+
+	// try to sniff TLS traffic
+	tlsVersion := binary.BigEndian.Uint16(hdr[1:3])
+	if hdr[0] == dissector.Handshake &&
+		(tlsVersion >= tls.VersionTLS10 && tlsVersion <= tls.VersionTLS13) {
+		return ProtoTLS, nil
+	}
+
+	// try to sniff HTTP traffic
+	if isHTTP(string(hdr[:])) {
+		return ProtoHTTP, nil
+	}
+
+	if string(hdr) == "SSH-2" {
+		return ProtoSSH, nil
+	}
+
+	return
+}
+
+func isHTTP(s string) bool {
+	return strings.HasPrefix(http.MethodGet, s[:3]) ||
+		strings.HasPrefix(http.MethodPost, s[:4]) ||
+		strings.HasPrefix(http.MethodPut, s[:3]) ||
+		strings.HasPrefix(http.MethodDelete, s) ||
+		strings.HasPrefix(http.MethodOptions, s) ||
+		strings.HasPrefix(http.MethodPatch, s) ||
+		strings.HasPrefix(http.MethodHead, s[:4]) ||
+		strings.HasPrefix(http.MethodConnect, s) ||
+		strings.HasPrefix(http.MethodTrace, s)
+}