chore: better restls

adapter/outbound/shadowsocks.go

@@ -86,22 +86,14 @@ type restlsOption struct {
 
 // StreamConn implements C.ProxyAdapter
 func (ss *ShadowSocks) StreamConn(c net.Conn, metadata *C.Metadata) (net.Conn, error) {
-	switch ss.obfsMode {
+	// fix tls handshake not timeout
-	case shadowtls.Mode:
+	ctx, cancel := context.WithTimeout(context.Background(), C.DefaultTLSTimeout)
-		// fix tls handshake not timeout
+	defer cancel()
-		ctx, cancel := context.WithTimeout(context.Background(), C.DefaultTLSTimeout)
+	return ss.StreamConnContext(ctx, c, metadata)
-		defer cancel()
-		var err error
-		c, err = shadowtls.NewShadowTLS(ctx, c, ss.shadowTLSOption)
-		if err != nil {
-			return nil, err
-		}
-
-	}
-	return ss.streamConn(c, metadata)
 }
 
-func (ss *ShadowSocks) streamConn(c net.Conn, metadata *C.Metadata) (net.Conn, error) {
+func (ss *ShadowSocks) StreamConnContext(ctx context.Context, c net.Conn, metadata *C.Metadata) (net.Conn, error) {
+	useEarly := false
 	switch ss.obfsMode {
 	case "tls":
 		c = obfs.NewTLSObfs(c, ss.obfsOption.Host)
@@ -114,21 +106,30 @@ func (ss *ShadowSocks) streamConn(c net.Conn, metadata *C.Metadata) (net.Conn, e
 		if err != nil {
 			return nil, fmt.Errorf("%s connect error: %w", ss.addr, err)
 		}
+	case shadowtls.Mode:
+		var err error
+		c, err = shadowtls.NewShadowTLS(ctx, c, ss.shadowTLSOption)
+		if err != nil {
+			return nil, err
+		}
+		useEarly = true
 	case restls.Mode:
 		var err error
-		c, err = restls.NewRestls(c, ss.restlsConfig)
+		c, err = restls.NewRestls(ctx, c, ss.restlsConfig)
 		if err != nil {
 			return nil, fmt.Errorf("%s (restls) connect error: %w", ss.addr, err)
 		}
+		useEarly = true
 	}
+	useEarly = useEarly || N.NeedHandshake(c)
 	if metadata.NetWork == C.UDP && ss.option.UDPOverTCP {
-		if N.NeedHandshake(c) {
+		if useEarly {
 			return ss.method.DialEarlyConn(c, M.ParseSocksaddr(uot.UOTMagicAddress+":443")), nil
 		} else {
 			return ss.method.DialConn(c, M.ParseSocksaddr(uot.UOTMagicAddress+":443"))
 		}
 	}
-	if N.NeedHandshake(c) {
+	if useEarly {
 		return ss.method.DialEarlyConn(c, M.ParseSocksaddr(metadata.RemoteAddress())), nil
 	} else {
 		return ss.method.DialConn(c, M.ParseSocksaddr(metadata.RemoteAddress()))
@@ -152,15 +153,7 @@ func (ss *ShadowSocks) DialContextWithDialer(ctx context.Context, dialer C.Diale
 		safeConnClose(c, err)
 	}(c)
 
-	switch ss.obfsMode {
+	c, err = ss.StreamConnContext(ctx, c, metadata)
-	case shadowtls.Mode:
-		c, err = shadowtls.NewShadowTLS(ctx, c, ss.shadowTLSOption)
-		if err != nil {
-			return nil, err
-		}
-	}
-
-	c, err = ss.streamConn(c, metadata)
 	return NewConn(c, ss), err
 }
 

transport/restls/restls.go

@@ -1,6 +1,7 @@
 package restls
 
 import (
+	"context"
 	"net"
 
 	tls "github.com/3andne/restls-client-go"
@@ -10,48 +11,29 @@ const (
 	Mode string = "restls"
 )
 
-// Restls
 type Restls struct {
-	net.Conn
+	*tls.UConn
-	firstPacketCache []byte
-	firstPacket      bool
 }
 
-func (r *Restls) Read(b []byte) (int, error) {
+func (r *Restls) Upstream() any {
-	if err := r.Conn.(*tls.UConn).Handshake(); err != nil {
+	return r.UConn.NetConn()
-		return 0, err
-	}
-	n, err := r.Conn.(*tls.UConn).Read(b)
-	return n, err
-}
-
-func (r *Restls) Write(b []byte) (int, error) {
-	if r.firstPacket {
-		r.firstPacketCache = append([]byte(nil), b...)
-		r.firstPacket = false
-		return len(b), nil
-	}
-	if len(r.firstPacketCache) != 0 {
-		b = append(r.firstPacketCache, b...)
-		r.firstPacketCache = nil
-	}
-	n, err := r.Conn.(*tls.UConn).Write(b)
-	return n, err
 }
 
 // NewRestls return a Restls Connection
-func NewRestls(conn net.Conn, config *tls.Config) (net.Conn, error) {
+func NewRestls(ctx context.Context, conn net.Conn, config *tls.Config) (net.Conn, error) {
+	clientHellowID := tls.HelloChrome_Auto
 	if config != nil {
 		clientIDPtr := config.ClientID.Load()
 		if clientIDPtr != nil {
-			return &Restls{
+			clientHellowID = *clientIDPtr
-				Conn:        tls.UClient(conn, config, *clientIDPtr),
-				firstPacket: true,
-			}, nil
 		}
 	}
-	return &Restls{
+	restls := &Restls{
-		Conn:        tls.UClient(conn, config, tls.HelloChrome_Auto),
+		UConn: tls.UClient(conn, config, clientHellowID),
-		firstPacket: true,
+	}
-	}, nil
+	if err := restls.HandshakeContext(ctx); err != nil {
+		return nil, err
+	}
+
+	return restls, nil
 }