From be8e0555e5033b5862e1d7ce7ef8b3fd154a0008 Mon Sep 17 00:00:00 2001
From: ginuerzh <ginuerzh@gmail.com>
Date: Mon, 23 Dec 2024 21:05:55 +0800
Subject: [PATCH] handle tls for tunnel entrypoint

---
 handler/tunnel/entrypoint.go | 72 ++++++++++++++++++++++++++++++++++++
 handler/tunnel/handler.go    |  1 +
 2 files changed, 73 insertions(+)

diff --git a/handler/tunnel/entrypoint.go b/handler/tunnel/entrypoint.go
index ddd5397e..d7f02faa 100644
--- a/handler/tunnel/entrypoint.go
+++ b/handler/tunnel/entrypoint.go
@@ -4,6 +4,7 @@ import (
 	"bufio"
 	"bytes"
 	"context"
+	"encoding/hex"
 	"errors"
 	"fmt"
 	"io"
@@ -24,6 +25,7 @@ import (
 	"github.com/go-gost/core/recorder"
 	"github.com/go-gost/core/sd"
 	"github.com/go-gost/relay"
+	dissector "github.com/go-gost/tls-dissector"
 	admission "github.com/go-gost/x/admission/wrapper"
 	ctxvalue "github.com/go-gost/x/ctx"
 	ctx_internal "github.com/go-gost/x/internal/ctx"
@@ -32,6 +34,7 @@ import (
 	xhttp "github.com/go-gost/x/internal/net/http"
 	"github.com/go-gost/x/internal/net/proxyproto"
 	"github.com/go-gost/x/internal/util/sniffing"
+	tls_util "github.com/go-gost/x/internal/util/tls"
 	ws_util "github.com/go-gost/x/internal/util/ws"
 	climiter "github.com/go-gost/x/limiter/conn/wrapper"
 	metrics "github.com/go-gost/x/metrics/wrapper"
@@ -58,6 +61,8 @@ type entrypoint struct {
 
 	sniffingWebsocket   bool
 	websocketSampleRate float64
+
+	readTimeout time.Duration
 }
 
 func (ep *entrypoint) Handle(ctx context.Context, conn net.Conn) (err error) {
@@ -112,6 +117,9 @@ func (ep *entrypoint) Handle(ctx context.Context, conn net.Conn) (err error) {
 	if v[0] == relay.Version1 {
 		return ep.handleConnect(ctx, conn, ro, log)
 	}
+	if v[0] == dissector.Handshake {
+		return ep.HandleTLS(ctx, conn, ro, log)
+	}
 	return ep.handleHTTP(ctx, conn, ro, log)
 }
 
@@ -547,6 +555,70 @@ func (ep *entrypoint) copyWebsocketFrame(w io.Writer, r io.Reader, buf *bytes.Bu
 	return nil
 }
 
+func (ep *entrypoint) HandleTLS(ctx context.Context, conn net.Conn, ro *xrecorder.HandlerRecorderObject, log logger.Logger) error {
+	buf := new(bytes.Buffer)
+	clientHello, err := dissector.ParseClientHello(io.TeeReader(conn, buf))
+	if err != nil {
+		return err
+	}
+
+	ro.TLS = &xrecorder.TLSRecorderObject{
+		ServerName:  clientHello.ServerName,
+		ClientHello: hex.EncodeToString(buf.Bytes()),
+	}
+	if len(clientHello.SupportedProtos) > 0 {
+		ro.TLS.Proto = clientHello.SupportedProtos[0]
+	}
+
+	host := clientHello.ServerName
+	if host != "" {
+		if _, _, err := net.SplitHostPort(host); err != nil {
+			host = net.JoinHostPort(strings.Trim(host, "[]"), "443")
+		}
+		ro.Host = host
+	}
+
+	ctx = ctxvalue.ContextWithClientAddr(ctx, ctxvalue.ClientAddr(ro.RemoteAddr))
+	ctx = ctx_internal.ContextWithRecorderObject(ctx, ro)
+	ctx = ctxvalue.ContextWithLogger(ctx, log)
+
+	cc, err := ep.dial(ctx, "tcp", host)
+	if err != nil {
+		return err
+	}
+	defer cc.Close()
+
+	if _, err := buf.WriteTo(cc); err != nil {
+		return err
+	}
+
+	xio.SetReadDeadline(cc, time.Now().Add(ep.readTimeout))
+	serverHello, err := dissector.ParseServerHello(io.TeeReader(cc, buf))
+	xio.SetReadDeadline(cc, time.Time{})
+
+	if serverHello != nil {
+		ro.TLS.CipherSuite = tls_util.CipherSuite(serverHello.CipherSuite).String()
+		ro.TLS.CompressionMethod = serverHello.CompressionMethod
+		if serverHello.Proto != "" {
+			ro.TLS.Proto = serverHello.Proto
+		}
+		if serverHello.Version > 0 {
+			ro.TLS.Version = tls_util.Version(serverHello.Version).String()
+		}
+	}
+
+	if buf.Len() > 0 {
+		ro.TLS.ServerHello = hex.EncodeToString(buf.Bytes())
+	}
+
+	if _, err := buf.WriteTo(conn); err != nil {
+		return err
+	}
+
+	xnet.Transport(conn, cc)
+	return err
+}
+
 func (ep *entrypoint) handleConnect(ctx context.Context, conn net.Conn, ro *xrecorder.HandlerRecorderObject, log logger.Logger) (err error) {
 	req := relay.Request{}
 	if _, err := req.ReadFrom(conn); err != nil {
diff --git a/handler/tunnel/handler.go b/handler/tunnel/handler.go
index 4a23ebab..7ea36f2e 100644
--- a/handler/tunnel/handler.go
+++ b/handler/tunnel/handler.go
@@ -92,6 +92,7 @@ func (h *tunnelHandler) Init(md md.Metadata) (err error) {
 		}),
 		sniffingWebsocket:   h.md.sniffingWebsocket,
 		websocketSampleRate: h.md.sniffingWebsocketSampleRate,
+		readTimeout:         h.md.entryPointReadTimeout,
 	}
 	h.ep.transport = &http.Transport{
 		DialContext:           h.ep.dial,