feat: enable fwmark (SO_MARK) for outgoing sockets

Jigsaw-Code · Aug 25, 2024 · 60aad45 · 60aad45
1 parent 0f6ad5b
commit 60aad45
Show file tree

Hide file tree

Showing 6 changed files with 118 additions and 45 deletions.
diff --git a/cmd/outline-ss-server/main.go b/cmd/outline-ss-server/main.go
@@ -27,17 +27,20 @@ import (
 	"time"
 
 	"github.com/Jigsaw-Code/outline-sdk/transport/shadowsocks"
-	"github.com/Jigsaw-Code/outline-ss-server/ipinfo"
-	"github.com/Jigsaw-Code/outline-ss-server/service"
 	"github.com/lmittmann/tint"
 	"github.com/prometheus/client_golang/prometheus"
 	"github.com/prometheus/client_golang/prometheus/promhttp"
 	"golang.org/x/term"
 	"gopkg.in/yaml.v2"
+
+	"github.com/Jigsaw-Code/outline-ss-server/ipinfo"
+	"github.com/Jigsaw-Code/outline-ss-server/service"
 )
 
-var logLevel = new(slog.LevelVar) // Info by default
-var logHandler slog.Handler
+var (
+	logLevel   = new(slog.LevelVar) // Info by default
+	logHandler slog.Handler
+)
 
 // Set by goreleaser default ldflags. See https://goreleaser.com/customization/build/
 var version = "dev"
@@ -68,7 +71,7 @@ type SSServer struct {
 	ports       map[int]*ssPort
 }
 
-func (s *SSServer) startPort(portNum int) error {
+func (s *SSServer) startPort(portNum int, fwmark uint) error {
 	listener, err := net.ListenTCP("tcp", &net.TCPAddr{Port: portNum})
 	if err != nil {
 		//lint:ignore ST1005 Shadowsocks is capitalized.
@@ -84,8 +87,8 @@ func (s *SSServer) startPort(portNum int) error {
 	port := &ssPort{tcpListener: listener, packetConn: packetConn, cipherList: service.NewCipherList()}
 	authFunc := service.NewShadowsocksStreamAuthenticator(port.cipherList, &s.replayCache, s.m)
 	// TODO: Register initial data metrics at zero.
-	tcpHandler := service.NewTCPHandler(authFunc, s.m, tcpReadTimeout)
-	packetHandler := service.NewPacketHandler(s.natTimeout, port.cipherList, s.m)
+	tcpHandler := service.NewTCPHandler(authFunc, s.m, tcpReadTimeout, fwmark)
+	packetHandler := service.NewPacketHandler(s.natTimeout, port.cipherList, s.m, fwmark)
 	s.ports[portNum] = port
 	go service.StreamServe(service.WrapStreamListener(listener.AcceptTCP), tcpHandler.Handle)
 	go packetHandler.Handle(port.packetConn)
@@ -144,15 +147,20 @@ func (s *SSServer) loadConfig(filename string) error {
 				return fmt.Errorf("failed to remove port %v: %w", portNum, err)
 			}
 		} else if count == +1 {
-			if err := s.startPort(portNum); err != nil {
+			if err := s.startPort(portNum, config.Fwmark); err != nil {
 				return err
 			}
 		}
 	}
 	for portNum, cipherList := range portCiphers {
 		s.ports[portNum].cipherList.Update(cipherList)
 	}
-	slog.Info("Loaded config.", "access keys", len(config.Keys), "ports", len(s.ports))
+	slog.Info("Loaded config.", "access keys", len(config.Keys), "ports", len(s.ports), "fwmark", func() any {
+		if config.Fwmark == 0 {
+			return "disabled"
+		}
+		return config.Fwmark
+	}())
 	s.m.SetNumAccessKeys(len(config.Keys), len(portCiphers))
 	return nil
 }
@@ -199,6 +207,7 @@ type Config struct {
 		Cipher string
 		Secret string
 	}
+	Fwmark uint
 }
 
 func readConfig(filename string) (*Config, error) {

diff --git a/internal/integration_test/integration_test.go b/internal/integration_test/integration_test.go
@@ -27,12 +27,13 @@ import (
 
 	"github.com/Jigsaw-Code/outline-sdk/transport"
 	"github.com/Jigsaw-Code/outline-sdk/transport/shadowsocks"
-	"github.com/Jigsaw-Code/outline-ss-server/ipinfo"
-	"github.com/Jigsaw-Code/outline-ss-server/service"
-	"github.com/Jigsaw-Code/outline-ss-server/service/metrics"
 	logging "github.com/op/go-logging"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
+
+	"github.com/Jigsaw-Code/outline-ss-server/ipinfo"
+	"github.com/Jigsaw-Code/outline-ss-server/service"
+	"github.com/Jigsaw-Code/outline-ss-server/service/metrics"
 )
 
 const maxUDPPacketSize = 64 * 1024
@@ -133,7 +134,7 @@ func TestTCPEcho(t *testing.T) {
 	const testTimeout = 200 * time.Millisecond
 	testMetrics := &service.NoOpTCPMetrics{}
 	authFunc := service.NewShadowsocksStreamAuthenticator(cipherList, &replayCache, testMetrics)
-	handler := service.NewTCPHandler(authFunc, testMetrics, testTimeout)
+	handler := service.NewTCPHandler(authFunc, testMetrics, testTimeout, 0)
 	handler.SetTargetDialer(&transport.TCPDialer{})
 	done := make(chan struct{})
 	go func() {
@@ -202,7 +203,7 @@ func TestRestrictedAddresses(t *testing.T) {
 	const testTimeout = 200 * time.Millisecond
 	testMetrics := &statusMetrics{}
 	authFunc := service.NewShadowsocksStreamAuthenticator(cipherList, nil, testMetrics)
-	handler := service.NewTCPHandler(authFunc, testMetrics, testTimeout)
+	handler := service.NewTCPHandler(authFunc, testMetrics, testTimeout, 0)
 	done := make(chan struct{})
 	go func() {
 		service.StreamServe(service.WrapStreamListener(proxyListener.AcceptTCP), handler.Handle)
@@ -262,18 +263,23 @@ var _ service.UDPMetrics = (*fakeUDPMetrics)(nil)
 func (m *fakeUDPMetrics) GetIPInfo(ip net.IP) (ipinfo.IPInfo, error) {
 	return ipinfo.IPInfo{CountryCode: "QQ"}, nil
 }
+
 func (m *fakeUDPMetrics) AddUDPPacketFromClient(clientInfo ipinfo.IPInfo, accessKey, status string, clientProxyBytes, proxyTargetBytes int) {
 	m.up = append(m.up, udpRecord{clientInfo, accessKey, status, clientProxyBytes, proxyTargetBytes})
 }
+
 func (m *fakeUDPMetrics) AddUDPPacketFromTarget(clientInfo ipinfo.IPInfo, accessKey, status string, targetProxyBytes, proxyClientBytes int) {
 	m.down = append(m.down, udpRecord{clientInfo, accessKey, status, targetProxyBytes, proxyClientBytes})
 }
+
 func (m *fakeUDPMetrics) AddUDPNatEntry(clientAddr net.Addr, accessKey string) {
 	m.natAdded++
 }
+
 func (m *fakeUDPMetrics) RemoveUDPNatEntry(clientAddr net.Addr, accessKey string) {
 	// Not tested because it requires waiting for a long timeout.
 }
+
 func (m *fakeUDPMetrics) AddUDPCipherSearch(accessKeyFound bool, timeToCipher time.Duration) {}
 
 func TestUDPEcho(t *testing.T) {
@@ -289,7 +295,7 @@ func TestUDPEcho(t *testing.T) {
 		t.Fatal(err)
 	}
 	testMetrics := &fakeUDPMetrics{}
-	proxy := service.NewPacketHandler(time.Hour, cipherList, testMetrics)
+	proxy := service.NewPacketHandler(time.Hour, cipherList, testMetrics, 0)
 	proxy.SetTargetIPValidator(allowAll)
 	done := make(chan struct{})
 	go func() {
@@ -384,7 +390,7 @@ func BenchmarkTCPThroughput(b *testing.B) {
 	const testTimeout = 200 * time.Millisecond
 	testMetrics := &service.NoOpTCPMetrics{}
 	authFunc := service.NewShadowsocksStreamAuthenticator(cipherList, nil, testMetrics)
-	handler := service.NewTCPHandler(authFunc, testMetrics, testTimeout)
+	handler := service.NewTCPHandler(authFunc, testMetrics, testTimeout, 0)
 	handler.SetTargetDialer(&transport.TCPDialer{})
 	done := make(chan struct{})
 	go func() {
@@ -448,7 +454,7 @@ func BenchmarkTCPMultiplexing(b *testing.B) {
 	const testTimeout = 200 * time.Millisecond
 	testMetrics := &service.NoOpTCPMetrics{}
 	authFunc := service.NewShadowsocksStreamAuthenticator(cipherList, &replayCache, testMetrics)
-	handler := service.NewTCPHandler(authFunc, testMetrics, testTimeout)
+	handler := service.NewTCPHandler(authFunc, testMetrics, testTimeout, 0)
 	handler.SetTargetDialer(&transport.TCPDialer{})
 	done := make(chan struct{})
 	go func() {
@@ -521,7 +527,7 @@ func BenchmarkUDPEcho(b *testing.B) {
 	if err != nil {
 		b.Fatal(err)
 	}
-	proxy := service.NewPacketHandler(time.Hour, cipherList, &service.NoOpUDPMetrics{})
+	proxy := service.NewPacketHandler(time.Hour, cipherList, &service.NoOpUDPMetrics{}, 0)
 	proxy.SetTargetIPValidator(allowAll)
 	done := make(chan struct{})
 	go func() {
@@ -565,7 +571,7 @@ func BenchmarkUDPManyKeys(b *testing.B) {
 	if err != nil {
 		b.Fatal(err)
 	}
-	proxy := service.NewPacketHandler(time.Hour, cipherList, &service.NoOpUDPMetrics{})
+	proxy := service.NewPacketHandler(time.Hour, cipherList, &service.NoOpUDPMetrics{}, 0)
 	proxy.SetTargetIPValidator(allowAll)
 	done := make(chan struct{})
 	go func() {

diff --git a/service/tcp.go b/service/tcp.go
@@ -24,16 +24,18 @@ import (
 	"log/slog"
 	"net"
 	"net/netip"
+	"os"
 	"sync"
 	"syscall"
 	"time"
 
 	"github.com/Jigsaw-Code/outline-sdk/transport"
 	"github.com/Jigsaw-Code/outline-sdk/transport/shadowsocks"
+	"github.com/shadowsocks/go-shadowsocks2/socks"
+
 	"github.com/Jigsaw-Code/outline-ss-server/ipinfo"
 	onet "github.com/Jigsaw-Code/outline-ss-server/net"
 	"github.com/Jigsaw-Code/outline-ss-server/service/metrics"
-	"github.com/shadowsocks/go-shadowsocks2/socks"
 )
 
 // TCPMetrics is used to report metrics on TCP connections.
@@ -170,7 +172,9 @@ type tcpHandler struct {
 }
 
 // NewTCPService creates a TCPService
-func NewTCPHandler(authenticate StreamAuthenticateFunc, m TCPMetrics, timeout time.Duration) TCPHandler {
+func NewTCPHandler(authenticate StreamAuthenticateFunc, m TCPMetrics, timeout time.Duration, fwmark uint) TCPHandler {
+	defaultDialer := makeValidatingTCPStreamDialer(onet.RequirePublicIP, fwmark)
+
 	return &tcpHandler{
 		m:            m,
 		readTimeout:  timeout,
@@ -179,10 +183,19 @@ func NewTCPHandler(authenticate StreamAuthenticateFunc, m TCPMetrics, timeout ti
 	}
 }
 
-var defaultDialer = makeValidatingTCPStreamDialer(onet.RequirePublicIP)
-
-func makeValidatingTCPStreamDialer(targetIPValidator onet.TargetIPValidator) transport.StreamDialer {
+func makeValidatingTCPStreamDialer(targetIPValidator onet.TargetIPValidator, fwmark uint) transport.StreamDialer {
 	return &transport.TCPDialer{Dialer: net.Dialer{Control: func(network, address string, c syscall.RawConn) error {
+		if fwmark > 0 {
+			err := c.Control(func(fd uintptr) {
+				err := syscall.SetsockoptInt(int(fd), syscall.SOL_SOCKET, syscall.SO_MARK, int(fwmark))
+				if err != nil {
+					slog.Error("Set fwmark failed.", "err", os.NewSyscallError("failed to set mark for TCP socket", err))
+				}
+			})
+			if err != nil {
+				slog.Error("Set TCPDialer Control func failed.", "err", err)
+			}
+		}
 		ip, _, _ := net.SplitHostPort(address)
 		return targetIPValidator(net.ParseIP(ip))
 	}}}
@@ -397,12 +410,17 @@ var _ TCPMetrics = (*NoOpTCPMetrics)(nil)
 
 func (m *NoOpTCPMetrics) AddClosedTCPConnection(clientInfo ipinfo.IPInfo, clientAddr net.Addr, accessKey string, status string, data metrics.ProxyMetrics, duration time.Duration) {
 }
+
 func (m *NoOpTCPMetrics) GetIPInfo(net.IP) (ipinfo.IPInfo, error) {
 	return ipinfo.IPInfo{}, nil
 }
+
 func (m *NoOpTCPMetrics) AddOpenTCPConnection(clientInfo ipinfo.IPInfo) {}
+
 func (m *NoOpTCPMetrics) AddAuthenticatedTCPConnection(clientAddr net.Addr, accessKey string) {
 }
+
 func (m *NoOpTCPMetrics) AddTCPProbe(status, drainResult string, listenerId string, clientProxyBytes int64) {
 }
+
 func (m *NoOpTCPMetrics) AddTCPCipherSearch(accessKeyFound bool, timeToCipher time.Duration) {}
diff --git a/service/tcp_test.go b/service/tcp_test.go
@@ -28,11 +28,12 @@ import (
 
 	"github.com/Jigsaw-Code/outline-sdk/transport"
 	"github.com/Jigsaw-Code/outline-sdk/transport/shadowsocks"
-	"github.com/Jigsaw-Code/outline-ss-server/ipinfo"
-	"github.com/Jigsaw-Code/outline-ss-server/service/metrics"
 	logging "github.com/op/go-logging"
 	"github.com/shadowsocks/go-shadowsocks2/socks"
 	"github.com/stretchr/testify/require"
+
+	"github.com/Jigsaw-Code/outline-ss-server/ipinfo"
+	"github.com/Jigsaw-Code/outline-ss-server/service/metrics"
 )
 
 func init() {
@@ -233,6 +234,7 @@ func (m *probeTestMetrics) AddClosedTCPConnection(clientInfo ipinfo.IPInfo, clie
 func (m *probeTestMetrics) GetIPInfo(net.IP) (ipinfo.IPInfo, error) {
 	return ipinfo.IPInfo{}, nil
 }
+
 func (m *probeTestMetrics) AddOpenTCPConnection(clientInfo ipinfo.IPInfo) {
 }
 
@@ -281,7 +283,7 @@ func TestProbeRandom(t *testing.T) {
 	require.NoError(t, err, "MakeTestCiphers failed: %v", err)
 	testMetrics := &probeTestMetrics{}
 	authFunc := NewShadowsocksStreamAuthenticator(cipherList, nil, testMetrics)
-	handler := NewTCPHandler(authFunc, testMetrics, 200*time.Millisecond)
+	handler := NewTCPHandler(authFunc, testMetrics, 200*time.Millisecond, 0)
 	done := make(chan struct{})
 	go func() {
 		StreamServe(WrapStreamListener(listener.AcceptTCP), handler.Handle)
@@ -358,8 +360,8 @@ func TestProbeClientBytesBasicTruncated(t *testing.T) {
 	cipher := firstCipher(cipherList)
 	testMetrics := &probeTestMetrics{}
 	authFunc := NewShadowsocksStreamAuthenticator(cipherList, nil, testMetrics)
-	handler := NewTCPHandler(authFunc, testMetrics, 200*time.Millisecond)
-	handler.SetTargetDialer(makeValidatingTCPStreamDialer(allowAll))
+	handler := NewTCPHandler(authFunc, testMetrics, 200*time.Millisecond, 0)
+	handler.SetTargetDialer(makeValidatingTCPStreamDialer(allowAll, 0))
 	done := make(chan struct{})
 	go func() {
 		StreamServe(WrapStreamListener(listener.AcceptTCP), handler.Handle)
@@ -393,8 +395,8 @@ func TestProbeClientBytesBasicModified(t *testing.T) {
 	cipher := firstCipher(cipherList)
 	testMetrics := &probeTestMetrics{}
 	authFunc := NewShadowsocksStreamAuthenticator(cipherList, nil, testMetrics)
-	handler := NewTCPHandler(authFunc, testMetrics, 200*time.Millisecond)
-	handler.SetTargetDialer(makeValidatingTCPStreamDialer(allowAll))
+	handler := NewTCPHandler(authFunc, testMetrics, 200*time.Millisecond, 0)
+	handler.SetTargetDialer(makeValidatingTCPStreamDialer(allowAll, 0))
 	done := make(chan struct{})
 	go func() {
 		StreamServe(WrapStreamListener(listener.AcceptTCP), handler.Handle)
@@ -429,8 +431,8 @@ func TestProbeClientBytesCoalescedModified(t *testing.T) {
 	cipher := firstCipher(cipherList)
 	testMetrics := &probeTestMetrics{}
 	authFunc := NewShadowsocksStreamAuthenticator(cipherList, nil, testMetrics)
-	handler := NewTCPHandler(authFunc, testMetrics, 200*time.Millisecond)
-	handler.SetTargetDialer(makeValidatingTCPStreamDialer(allowAll))
+	handler := NewTCPHandler(authFunc, testMetrics, 200*time.Millisecond, 0)
+	handler.SetTargetDialer(makeValidatingTCPStreamDialer(allowAll, 0))
 	done := make(chan struct{})
 	go func() {
 		StreamServe(WrapStreamListener(listener.AcceptTCP), handler.Handle)
@@ -472,7 +474,7 @@ func TestProbeServerBytesModified(t *testing.T) {
 	cipher := firstCipher(cipherList)
 	testMetrics := &probeTestMetrics{}
 	authFunc := NewShadowsocksStreamAuthenticator(cipherList, nil, testMetrics)
-	handler := NewTCPHandler(authFunc, testMetrics, 200*time.Millisecond)
+	handler := NewTCPHandler(authFunc, testMetrics, 200*time.Millisecond, 0)
 	done := make(chan struct{})
 	go func() {
 		StreamServe(WrapStreamListener(listener.AcceptTCP), handler.Handle)
@@ -503,7 +505,7 @@ func TestReplayDefense(t *testing.T) {
 	testMetrics := &probeTestMetrics{}
 	const testTimeout = 200 * time.Millisecond
 	authFunc := NewShadowsocksStreamAuthenticator(cipherList, &replayCache, testMetrics)
-	handler := NewTCPHandler(authFunc, testMetrics, testTimeout)
+	handler := NewTCPHandler(authFunc, testMetrics, testTimeout, 0)
 	snapshot := cipherList.SnapshotForClientIP(netip.Addr{})
 	cipherEntry := snapshot[0].Value.(*CipherEntry)
 	cipher := cipherEntry.CryptoKey
@@ -582,7 +584,7 @@ func TestReverseReplayDefense(t *testing.T) {
 	testMetrics := &probeTestMetrics{}
 	const testTimeout = 200 * time.Millisecond
 	authFunc := NewShadowsocksStreamAuthenticator(cipherList, &replayCache, testMetrics)
-	handler := NewTCPHandler(authFunc, testMetrics, testTimeout)
+	handler := NewTCPHandler(authFunc, testMetrics, testTimeout, 0)
 	snapshot := cipherList.SnapshotForClientIP(netip.Addr{})
 	cipherEntry := snapshot[0].Value.(*CipherEntry)
 	cipher := cipherEntry.CryptoKey
@@ -653,7 +655,7 @@ func probeExpectTimeout(t *testing.T, payloadSize int) {
 	require.NoError(t, err, "MakeTestCiphers failed: %v", err)
 	testMetrics := &probeTestMetrics{}
 	authFunc := NewShadowsocksStreamAuthenticator(cipherList, nil, testMetrics)
-	handler := NewTCPHandler(authFunc, testMetrics, testTimeout)
+	handler := NewTCPHandler(authFunc, testMetrics, testTimeout, 0)
 
 	done := make(chan struct{})
 	go func() {