Use the new WS Upgrade() function from the SDK.

Author: sbruens

caddy/outline_handler.go

@@ -18,12 +18,17 @@ import (
 	"errors"
 	"fmt"
 	"log/slog"
+	"net"
 
+	"github.com/Jigsaw-Code/outline-sdk/transport"
 	"github.com/caddyserver/caddy/v2"
 	"github.com/mholt/caddy-l4/layer4"
 )
 
-const outlineHandlerModuleName = "layer4.handlers.outline"
+const (
+	outlineHandlerModuleName    = "layer4.handlers.outline"
+	outlineConnectionTypeCtxKey = "layer4.handlers.outline.cxtype"
+)
 
 func init() {
 	caddy.RegisterModule(ModuleRegistration{
@@ -32,6 +37,13 @@ func init() {
 	})
 }
 
+type ConnectionType string
+
+const (
+	StreamConnectionType = ConnectionType("stream")
+	PacketConnectionType = ConnectionType("packet")
+)
+
 // OutlineHandler implements a Caddy layer4 plugin for Outline connections.
 type OutlineHandler struct {
 	ConnectionHandler string `json:"connection_handler,omitempty"`
@@ -85,5 +97,11 @@ func (h *OutlineHandler) Validate() error {
 
 // Handle implements layer4.NextHandler.
 func (h *OutlineHandler) Handle(cx *layer4.Connection, next layer4.Handler) error {
+	switch cx.Conn.(type) {
+	case transport.StreamConn:
+		cx.SetVar(outlineConnectionTypeCtxKey, StreamConnectionType)
+	case net.Conn:
+		cx.SetVar(outlineConnectionTypeCtxKey, PacketConnectionType)
+	}
 	return h.compiledHandler.Handle(cx, next)
 }

caddy/shadowsocks_handler.go

@@ -18,7 +18,6 @@ import (
 	"container/list"
 	"fmt"
 	"log/slog"
-	"net"
 	"time"
 
 	"github.com/Jigsaw-Code/outline-sdk/transport"
@@ -116,13 +115,18 @@ func (h *ShadowsocksHandler) Provision(ctx caddy.Context) error {
 
 // Handle implements layer4.NextHandler.
 func (h *ShadowsocksHandler) Handle(cx *layer4.Connection, _ layer4.Handler) error {
-	switch conn := cx.Conn.(type) {
-	case transport.StreamConn:
-		h.streamHandler.HandleStream(cx.Context, conn, h.metrics.AddOpenTCPConnection(conn))
-	case net.Conn:
-		h.associationHandler.HandleAssociation(cx.Context, conn, h.metrics.AddOpenUDPAssociation(conn))
-	default:
-		return fmt.Errorf("failed to handle unknown connection type: %T", conn)
+	connType, ok := cx.GetVar(outlineConnectionTypeCtxKey).(ConnectionType)
+	if !ok {
+		// Likely if the Shadowsocks handler was used directly instead of through
+		// the Outline connection handler.
+		return fmt.Errorf("unknown outline connection type")
+	}
+
+	switch connType {
+	case StreamConnectionType:
+		h.streamHandler.HandleStream(cx.Context, cx.Conn.(transport.StreamConn), h.metrics.AddOpenTCPConnection(cx))
+	case PacketConnectionType:
+		h.associationHandler.HandleAssociation(cx.Context, cx.Conn, h.metrics.AddOpenUDPAssociation(cx))
 	}
 	return nil
 }

caddy/ws2outline_handler.go

@@ -22,31 +22,24 @@ import (
 	"net/http"
 
 	"github.com/Jigsaw-Code/outline-sdk/transport"
+	"github.com/Jigsaw-Code/outline-sdk/x/websocket"
 	"github.com/caddyserver/caddy/v2"
 	"github.com/caddyserver/caddy/v2/modules/caddyhttp"
 	"github.com/mholt/caddy-l4/layer4"
 	"go.uber.org/zap"
-	"golang.org/x/net/websocket"
-)
 
-const (
-	wsModuleName = "http.handlers.ws2outline"
+	onet "github.com/Jigsaw-Code/outline-ss-server/net"
 )
 
+const wsModuleName = "http.handlers.ws2outline"
+
 func init() {
 	caddy.RegisterModule(ModuleRegistration{
 		ID:  wsModuleName,
 		New: func() caddy.Module { return new(WebSocketHandler) },
 	})
 }
 
-type ConnectionType string
-
-const (
-	connectionTypeStream ConnectionType = "stream"
-	connectionTypePacket ConnectionType = "packet"
-)
-
 // WebSocketHandler implements a middleware Caddy web handler that proxies
 // WebSockets Outline connections.
 type WebSocketHandler struct {
@@ -75,7 +68,7 @@ func (h *WebSocketHandler) Provision(ctx caddy.Context) error {
 	h.zlogger = ctx.Logger()
 	if h.Type == "" {
 		// The default connection type if not provided is a stream.
-		h.Type = connectionTypeStream
+		h.Type = StreamConnectionType
 	}
 
 	mod, err := ctx.AppIfConfigured(outlineModuleName)
@@ -101,7 +94,7 @@ func (h *WebSocketHandler) Provision(ctx caddy.Context) error {
 
 // Validate implements caddy.Validator.
 func (h *WebSocketHandler) Validate() error {
-	if h.Type != "" && h.Type != connectionTypeStream && h.Type != connectionTypePacket {
+	if h.Type != "" && h.Type != StreamConnectionType && h.Type != PacketConnectionType {
 		return fmt.Errorf("unsupported `type`: %v", h.Type)
 	}
 	if h.ConnectionHandler == "" {
@@ -117,55 +110,35 @@ func (h WebSocketHandler) ServeHTTP(w http.ResponseWriter, r *http.Request, _ ca
 		slog.String("path", r.URL.Path),
 		slog.Any("remote_addr", r.RemoteAddr))
 
+	conn, err := websocket.Upgrade(w, r, nil)
+	if err != nil {
+		h.logger.Error("failed to upgrade", "err", err)
+	}
+	defer conn.Close()
+	clientAddrPort, clientIpErr := onet.ParseAddrPortOrIP(r.RemoteAddr)
 	switch h.Type {
-	case connectionTypeStream:
-		handler := func(wsConn *websocket.Conn) {
-			raddr, err := net.ResolveTCPAddr("tcp", r.RemoteAddr)
-			if err != nil {
-				h.logger.Error("failed to upgrade", "err", err)
-				w.WriteHeader(http.StatusBadGateway)
-				return
-			}
-
-			conn := &wsToStreamConn{&wrappedConn{Conn: wsConn, raddr: raddr}}
-
-			if err = h.compiledHandler.Handle(layer4.WrapConnection(conn, []byte{}, h.zlogger), nil); err != nil {
-				h.logger.Error("failed to upgrade", "err", err)
-				w.WriteHeader(http.StatusBadGateway)
-				return
-			}
+	case StreamConnectionType:
+		if clientIpErr == nil {
+			conn = &replaceAddrConn{StreamConn: conn, raddr: net.TCPAddrFromAddrPort(clientAddrPort)}
 		}
-		websocket.Handler(handler).ServeHTTP(w, r)
-	case connectionTypePacket:
-		handler := func(wsConn *websocket.Conn) {
-			raddr, err := net.ResolveUDPAddr("udp", r.RemoteAddr)
-			if err != nil {
-				h.logger.Error("failed to upgrade", "err", err)
-				w.WriteHeader(http.StatusBadGateway)
-				return
-			}
-
-			conn := &wrappedConn{Conn: wsConn, raddr: raddr}
-
-			if err = h.compiledHandler.Handle(layer4.WrapConnection(conn, []byte{}, h.zlogger), nil); err != nil {
-				h.logger.Error("failed to upgrade", "err", err)
-				w.WriteHeader(http.StatusBadGateway)
-				return
-			}
+	case PacketConnectionType:
+		if clientIpErr == nil {
+			conn = &replaceAddrConn{StreamConn: conn, raddr: net.UDPAddrFromAddrPort(clientAddrPort)}
 		}
-		websocket.Handler(handler).ServeHTTP(w, r)
 	}
-
-	return nil
+	cx := layer4.WrapConnection(conn, []byte{}, h.zlogger)
+	cx.SetVar(outlineConnectionTypeCtxKey, h.Type)
+	return h.compiledHandler.Handle(cx, nil)
 }
 
-// wrappedConn overrides [websocket.Conn]'s remote address handling.
-type wrappedConn struct {
-	*websocket.Conn
+// TODO: Create a dedicated `ClientConn` struct with `ClientAddr` and `Conn`.
+// replaceAddrConn overrides a [transport.StreamConn]'s remote address handling.
+type replaceAddrConn struct {
+	transport.StreamConn
 	raddr net.Addr
 }
 
-func (c wrappedConn) RemoteAddr() net.Addr {
+func (c replaceAddrConn) RemoteAddr() net.Addr {
 	return c.raddr
 }
 
@@ -184,4 +157,3 @@ func (c *wsToStreamConn) CloseRead() error {
 func (c *wsToStreamConn) CloseWrite() error {
 	return c.Conn.Close()
 }
-