WireGuard proxy: Refactor (#6287)

And #6303 (comment)

Author: LjhAUMEM

app/proxyman/inbound/always.go

@@ -57,16 +57,23 @@ func NewAlwaysOnInboundHandler(ctx context.Context, tag string, receiverConfig *
 	if err != nil {
 		return nil, err
 	}
-
-	// Set tag and sniffing config in context before creating proxy
-	// This allows proxies like TUN to access these settings
-	ctx = session.ContextWithInbound(ctx, &session.Inbound{Tag: tag})
-	if receiverConfig.SniffingSettings != nil {
-		ctx = session.ContextWithContent(ctx, &session.Content{
-			SniffingRequest: sniffingRequest,
-		})
+	src := net.TCPDestination(net.AnyIP, 0)
+	if receiverConfig.Listen != nil {
+		src.Address = receiverConfig.Listen.AsAddress()
+	}
+	if receiverConfig.PortList != nil && len(receiverConfig.PortList.Range) > 0 {
+		src.Port = net.Port(receiverConfig.PortList.Range[0].From)
 	}
-	rawProxy, err := common.CreateObject(ctx, proxyConfig)
+	mss, err := internet.ToMemoryStreamConfig(receiverConfig.StreamSettings)
+	if err != nil {
+		return nil, errors.New("failed to parse stream config").Base(err).AtWarning()
+	}
+
+	newCtx := session.ContextWithInbound(ctx, &session.Inbound{Tag: tag, Source: src})
+	newCtx = session.ContextWithContent(newCtx, &session.Content{SniffingRequest: sniffingRequest})
+	newCtx = session.ContextWithStreamSettings(newCtx, mss)
+
+	rawProxy, err := common.CreateObject(newCtx, proxyConfig)
 	if err != nil {
 		return nil, err
 	}
@@ -92,11 +99,6 @@ func NewAlwaysOnInboundHandler(ctx context.Context, tag string, receiverConfig *
 		address = net.AnyIP
 	}
 
-	mss, err := internet.ToMemoryStreamConfig(receiverConfig.StreamSettings)
-	if err != nil {
-		return nil, errors.New("failed to parse stream config").Base(err).AtWarning()
-	}
-
 	if receiverConfig.ReceiveOriginalDestination {
 		if mss.SocketSettings == nil {
 			mss.SocketSettings = &internet.SocketConfig{}

app/proxyman/outbound/handler.go

@@ -108,7 +108,9 @@ func NewHandler(ctx context.Context, config *core.OutboundHandlerConfig) (outbou
 
 	ctx = session.ContextWithFullHandler(ctx, h)
 
-	rawProxyHandler, err := common.CreateObject(ctx, proxyConfig)
+	newCtx := session.ContextWithStreamSettings(ctx, h.streamSettings)
+
+	rawProxyHandler, err := common.CreateObject(newCtx, proxyConfig)
 	if err != nil {
 		return nil, err
 	}

common/session/context.go

@@ -26,6 +26,8 @@ const (
 	fullHandlerKey            ctx.SessionKey = 10 // outbound gets full handler
 	mitmAlpn11Key             ctx.SessionKey = 11 // used by TLS dialer
 	mitmServerNameKey         ctx.SessionKey = 12 // used by TLS dialer
+
+	streamSettingsKey ctx.SessionKey = 13
 )
 
 func ContextWithInbound(ctx context.Context, inbound *Inbound) context.Context {
@@ -192,3 +194,11 @@ func MitmServerNameFromContext(ctx context.Context) string {
 	}
 	return ""
 }
+
+func ContextWithStreamSettings(ctx context.Context, streamSettings any) context.Context {
+	return context.WithValue(ctx, streamSettingsKey, streamSettings)
+}
+
+func StreamSettingsFromContext(ctx context.Context) any {
+	return ctx.Value(streamSettingsKey)
+}

infra/conf/wireguard.go

@@ -3,6 +3,7 @@ package conf
 import (
 	"encoding/base64"
 	"encoding/hex"
+	"strconv"
 	"strings"
 
 	"github.com/xtls/xray-core/common/errors"
@@ -37,8 +38,9 @@ func (c *WireGuardPeerConfig) Build() (proto.Message, error) {
 	}
 
 	config.Endpoint = c.Endpoint
-	// default 0
-	config.KeepAlive = c.KeepAlive
+	if c.KeepAlive != 0 {
+		config.KeepAlive = strconv.FormatUint(uint64(c.KeepAlive), 10)
+	}
 	if c.AllowedIPs == nil {
 		config.AllowedIps = []string{"0.0.0.0/0", "::0/0"}
 	} else {
@@ -56,7 +58,6 @@ type WireGuardConfig struct {
 	Address        []string               `json:"address"`
 	Peers          []*WireGuardPeerConfig `json:"peers"`
 	MTU            int32                  `json:"mtu"`
-	NumWorkers     int32                  `json:"workers"`
 	Reserved       []byte                 `json:"reserved"`
 	DomainStrategy string                 `json:"domainStrategy"`
 }
@@ -93,9 +94,6 @@ func (c *WireGuardConfig) Build() (proto.Message, error) {
 	} else {
 		config.Mtu = c.MTU
 	}
-	// these a fallback code exists in wireguard-go code,
-	// we don't need to process fallback manually
-	config.NumWorkers = c.NumWorkers
 
 	if len(c.Reserved) != 0 && len(c.Reserved) != 3 {
 		return nil, errors.New(`"reserved" should be empty or 3 bytes`)

infra/conf/wireguard_test.go

@@ -38,12 +38,10 @@ func TestWireGuardConfig(t *testing.T) {
 						// also can read from hex form directly
 						PublicKey:  "6e65ce0be17517110c17d77288ad87e7fd5252dcc7d09b95a39d61db03df832a",
 						Endpoint:   "127.0.0.1:1234",
-						KeepAlive:  0,
 						AllowedIps: []string{"0.0.0.0/0", "::0/0"},
 					},
 				},
 				Mtu:            1300,
-				NumWorkers:     2,
 				DomainStrategy: wireguard.DeviceConfig_FORCE_IP64,
 				NoKernelTun:    false,
 			},

proxy/hysteria/client.go

@@ -29,6 +29,13 @@ type Client struct {
 }
 
 func NewClient(ctx context.Context, config *ClientConfig) (*Client, error) {
+	v := core.MustFromContext(ctx)
+	p := v.GetFeature(policy.ManagerType()).(policy.Manager)
+
+	streamSettings := session.StreamSettingsFromContext(ctx).(*internet.MemoryStreamConfig)
+	if _, ok := streamSettings.ProtocolSettings.(*hysteria.Config); !ok {
+		return nil, errors.New("not hysteria transport")
+	}
 	if config.Server == nil {
 		return nil, errors.New(`no target server found`)
 	}
@@ -37,12 +44,10 @@ func NewClient(ctx context.Context, config *ClientConfig) (*Client, error) {
 		return nil, errors.New("failed to get server spec").Base(err)
 	}
 
-	v := core.MustFromContext(ctx)
-	client := &Client{
+	return &Client{
 		server:        server,
-		policyManager: v.GetFeature(policy.ManagerType()).(policy.Manager),
-	}
-	return client, nil
+		policyManager: p,
+	}, nil
 }
 
 func (c *Client) Process(ctx context.Context, link *transport.Link, dialer internet.Dialer) error {

proxy/hysteria/server.go

@@ -16,6 +16,7 @@ import (
 	"github.com/xtls/xray-core/features/routing"
 	"github.com/xtls/xray-core/proxy/hysteria/account"
 	"github.com/xtls/xray-core/transport"
+	"github.com/xtls/xray-core/transport/internet"
 	"github.com/xtls/xray-core/transport/internet/hysteria"
 	"github.com/xtls/xray-core/transport/internet/stat"
 )
@@ -27,6 +28,14 @@ type Server struct {
 }
 
 func NewServer(ctx context.Context, config *ServerConfig) (*Server, error) {
+	v := core.MustFromContext(ctx)
+	p := v.GetFeature(policy.ManagerType()).(policy.Manager)
+
+	streamSettings := session.StreamSettingsFromContext(ctx).(*internet.MemoryStreamConfig)
+	if _, ok := streamSettings.ProtocolSettings.(*hysteria.Config); !ok {
+		return nil, errors.New("not hysteria transport")
+	}
+
 	validator := account.NewValidator()
 	for _, user := range config.Users {
 		u, err := user.ToMemoryUser()
@@ -39,14 +48,11 @@ func NewServer(ctx context.Context, config *ServerConfig) (*Server, error) {
 		}
 	}
 
-	v := core.MustFromContext(ctx)
-	s := &Server{
+	return &Server{
 		config:        config,
 		validator:     validator,
-		policyManager: v.GetFeature(policy.ManagerType()).(policy.Manager),
-	}
-
-	return s, nil
+		policyManager: p,
+	}, nil
 }
 
 func (s *Server) HysteriaInboundValidator() *account.Validator {