Add shadowsocks 2022 relay service

common/udpnat/service.go

@@ -40,22 +40,24 @@ func New[K comparable](maxAge int64, handler Handler) *Service[K] {
 }
 
 func (s *Service[T]) NewPacket(key T, writer func() N.PacketWriter, buffer *buf.Buffer, metadata M.Metadata) {
-	s.NewContextPacket(context.Background(), key, writer, buffer, metadata)
+	s.NewContextPacket(context.Background(), key, func() (context.Context, N.PacketWriter) {
+		return context.Background(), writer()
+	}, buffer, metadata)
 }
 
-func (s *Service[T]) NewContextPacket(ctx context.Context, key T, writer func() N.PacketWriter, buffer *buf.Buffer, metadata M.Metadata) {
+func (s *Service[T]) NewContextPacket(ctx context.Context, key T, init func() (context.Context, N.PacketWriter), buffer *buf.Buffer, metadata M.Metadata) {
 	c, loaded := s.nat.LoadOrStore(key, func() *conn {
 		c := &conn{
 			data:       make(chan packet),
 			localAddr:  metadata.Source,
 			remoteAddr: metadata.Destination,
-			source:     writer(),
 			fastClose:  metadata.Destination.Port == 53,
 		}
 		c.ctx, c.cancel = context.WithCancel(ctx)
 		return c
 	})
 	if !loaded {
+		ctx, c.source = init()
 		go func() {
 			err := s.handler.NewPacketConnection(ctx, c, metadata)
 			if err != nil {
@@ -69,7 +71,7 @@ func (s *Service[T]) NewContextPacket(ctx context.Context, key T, writer func()
 	if common.Done(c.ctx) {
 		s.nat.Delete(key)
 		c.access.Unlock()
-		s.NewContextPacket(ctx, key, writer, buffer, metadata)
+		s.NewContextPacket(ctx, key, init, buffer, metadata)
 		return
 	}
 	packetCtx, done := context.WithCancel(c.ctx)

protocol/shadowsocks/shadowaead_2022/protocol.go

@@ -106,9 +106,7 @@ func New(method string, pskList [][]byte, password string, secureRNG io.Reader)
 	}
 
 	switch method {
-	case "2022-blake3-aes-128-gcm":
-		m.udpBlockCipher = newAES(pskList[0])
-	case "2022-blake3-aes-256-gcm":
+	case "2022-blake3-aes-128-gcm", "2022-blake3-aes-256-gcm":
 		m.udpBlockCipher = newAES(pskList[0])
 	case "2022-blake3-chacha20-poly1305":
 		m.udpCipher = newXChacha20Poly1305(pskList[0])
@@ -297,7 +295,7 @@ func (c *clientConn) readResponse() error {
 	}
 
 	if !c.replayFilter.Check(salt) {
-		return E.New("salt not unique")
+		return ErrSaltNotUnique
 	}
 
 	key := SessionKey(c.pskList[len(c.pskList)-1], salt, c.keySaltLength)

protocol/shadowsocks/shadowaead_2022/relay.go

@@ -0,0 +1,233 @@
+package shadowaead_2022
+
+import (
+	"context"
+	"crypto/aes"
+	"crypto/cipher"
+	"encoding/binary"
+	"io"
+	"net"
+	"os"
+	"runtime"
+
+	"github.com/sagernet/sing/common"
+	"github.com/sagernet/sing/common/buf"
+	"github.com/sagernet/sing/common/cache"
+	E "github.com/sagernet/sing/common/exceptions"
+	M "github.com/sagernet/sing/common/metadata"
+	N "github.com/sagernet/sing/common/network"
+	"github.com/sagernet/sing/common/rw"
+	"github.com/sagernet/sing/common/udpnat"
+	"github.com/sagernet/sing/protocol/shadowsocks"
+	"github.com/sagernet/sing/protocol/shadowsocks/shadowaead"
+	"lukechampine.com/blake3"
+)
+
+type Relay[U comparable] struct {
+	name          string
+	secureRNG     io.Reader
+	keySaltLength int
+	handler       shadowsocks.Handler
+
+	constructor      func(key []byte) cipher.AEAD
+	blockConstructor func(key []byte) cipher.Block
+	udpBlockCipher   cipher.Block
+
+	iPSK         []byte
+	uPSKHash     map[U][aes.BlockSize]byte
+	uPSKHashR    map[[aes.BlockSize]byte]U
+	uDestination map[U]M.Socksaddr
+	uCipher      map[U]cipher.Block
+	udpNat       *udpnat.Service[uint64]
+	udpSessions  *cache.LruCache[uint64, *relayUDPSession]
+}
+
+func (s *Relay[U]) AddUser(user U, key []byte, destination M.Socksaddr) error {
+	if len(key) < s.keySaltLength {
+		return shadowsocks.ErrBadKey
+	} else if len(key) > s.keySaltLength {
+		key = Key(key, s.keySaltLength)
+	}
+
+	var uPSKHash [aes.BlockSize]byte
+	hash512 := blake3.Sum512(key)
+	copy(uPSKHash[:], hash512[:])
+
+	if oldHash, loaded := s.uPSKHash[user]; loaded {
+		delete(s.uPSKHashR, oldHash)
+	}
+
+	s.uPSKHash[user] = uPSKHash
+	s.uPSKHashR[uPSKHash] = user
+	s.uDestination[user] = destination
+	s.uCipher[user] = s.blockConstructor(key)
+
+	return nil
+}
+
+func (s *Relay[U]) RemoveUser(user U) {
+	if hash, loaded := s.uPSKHash[user]; loaded {
+		delete(s.uPSKHashR, hash)
+	}
+	delete(s.uPSKHash, user)
+	delete(s.uCipher, user)
+}
+
+func NewRelay[U comparable](method string, psk []byte, secureRNG io.Reader, udpTimeout int64, handler shadowsocks.Handler) (*Relay[U], error) {
+	s := &Relay[U]{
+		name:      method,
+		secureRNG: secureRNG,
+		handler:   handler,
+
+		uPSKHash:     make(map[U][aes.BlockSize]byte),
+		uPSKHashR:    make(map[[aes.BlockSize]byte]U),
+		uDestination: make(map[U]M.Socksaddr),
+		uCipher:      make(map[U]cipher.Block),
+
+		udpNat: udpnat.New[uint64](udpTimeout, handler),
+		udpSessions: cache.New(
+			cache.WithAge[uint64, *relayUDPSession](udpTimeout),
+			cache.WithUpdateAgeOnGet[uint64, *relayUDPSession](),
+		),
+	}
+
+	switch method {
+	case "2022-blake3-aes-128-gcm":
+		s.keySaltLength = 16
+		s.constructor = newAESGCM
+		s.blockConstructor = newAES
+	case "2022-blake3-aes-256-gcm":
+		s.keySaltLength = 32
+		s.constructor = newAESGCM
+		s.blockConstructor = newAES
+	default:
+		return nil, os.ErrInvalid
+	}
+	if len(psk) != s.keySaltLength {
+		if len(psk) < s.keySaltLength {
+			return nil, shadowsocks.ErrBadKey
+		} else {
+			psk = Key(psk, s.keySaltLength)
+		}
+	}
+	s.udpBlockCipher = s.blockConstructor(psk)
+	return s, nil
+}
+
+func (s *Relay[U]) NewConnection(ctx context.Context, conn net.Conn, metadata M.Metadata) error {
+	err := s.newConnection(ctx, conn, metadata)
+	if err != nil {
+		err = &shadowsocks.ServerConnError{Conn: conn, Source: metadata.Source, Cause: err}
+	}
+	return err
+}
+
+func (s *Relay[U]) newConnection(ctx context.Context, conn net.Conn, metadata M.Metadata) error {
+	_requestHeader := buf.StackNew()
+	defer runtime.KeepAlive(_requestHeader)
+	requestHeader := common.Dup(_requestHeader)
+	n, err := requestHeader.ReadFrom(conn)
+	if err != nil {
+		return err
+	} else if int(n) < s.keySaltLength+aes.BlockSize {
+		return shadowaead.ErrBadHeader
+	}
+	requestSalt := requestHeader.To(s.keySaltLength)
+	var _eiHeader [aes.BlockSize]byte
+	eiHeader := common.Dup(_eiHeader[:])
+	copy(eiHeader, requestHeader.Range(s.keySaltLength, s.keySaltLength+aes.BlockSize))
+
+	keyMaterial := buf.Make(s.keySaltLength * 2)
+	copy(keyMaterial, s.iPSK)
+	copy(keyMaterial[s.keySaltLength:], requestSalt)
+	_identitySubkey := buf.Make(s.keySaltLength)
+	identitySubkey := common.Dup(_identitySubkey)
+	blake3.DeriveKey(identitySubkey, "shadowsocks 2022 identity subkey", keyMaterial)
+	s.blockConstructor(identitySubkey).Decrypt(eiHeader, eiHeader)
+	runtime.KeepAlive(_identitySubkey)
+
+	var user U
+	if u, loaded := s.uPSKHashR[_eiHeader]; loaded {
+		user = u
+	} else {
+		return E.New("invalid request")
+	}
+	runtime.KeepAlive(_eiHeader)
+
+	copy(requestHeader.Range(aes.BlockSize, aes.BlockSize+s.keySaltLength), requestHeader.To(s.keySaltLength))
+	requestHeader.Advance(aes.BlockSize)
+
+	ctx = shadowsocks.UserContext[U]{
+		ctx,
+		user,
+	}
+	metadata.Protocol = "shadowsocks-relay"
+	metadata.Destination = s.uDestination[user]
+	conn = &rw.BufferedConn{
+		Conn:   conn,
+		Buffer: requestHeader,
+	}
+	return s.handler.NewConnection(ctx, conn, metadata)
+}
+
+func (s *Relay[U]) NewPacket(conn N.PacketConn, buffer *buf.Buffer, metadata M.Metadata) error {
+	err := s.newPacket(conn, buffer, metadata)
+	if err != nil {
+		err = &shadowsocks.ServerPacketError{Source: metadata.Source, Cause: err}
+	}
+	return err
+}
+
+func (s *Relay[U]) newPacket(conn N.PacketConn, buffer *buf.Buffer, metadata M.Metadata) error {
+	packetHeader := buffer.To(aes.BlockSize)
+	s.udpBlockCipher.Decrypt(packetHeader, packetHeader)
+
+	sessionId := binary.BigEndian.Uint64(packetHeader)
+
+	var _eiHeader [aes.BlockSize]byte
+	eiHeader := common.Dup(_eiHeader[:])
+	s.udpBlockCipher.Decrypt(eiHeader, buffer.Range(aes.BlockSize, 2*aes.BlockSize))
+
+	for i := range eiHeader {
+		eiHeader[i] = eiHeader[i] ^ packetHeader[i]
+	}
+
+	var user U
+	if u, loaded := s.uPSKHashR[_eiHeader]; loaded {
+		user = u
+	} else {
+		return E.New("invalid request")
+	}
+
+	session, _ := s.udpSessions.LoadOrStore(sessionId, func() *relayUDPSession {
+		return new(relayUDPSession)
+	})
+	session.sourceAddr = metadata.Source
+
+	s.uCipher[user].Encrypt(packetHeader, packetHeader)
+	copy(buffer.Range(aes.BlockSize, 2*aes.BlockSize), packetHeader)
+	buffer.Advance(aes.BlockSize)
+
+	metadata.Protocol = "shadowsocks-relay"
+	metadata.Destination = s.uDestination[user]
+	s.udpNat.NewContextPacket(context.Background(), sessionId, func() (context.Context, N.PacketWriter) {
+		return &shadowsocks.UserContext[U]{
+			context.Background(),
+			user,
+		}, &relayPacketWriter[U]{conn, session}
+	}, buffer, metadata)
+	return nil
+}
+
+type relayUDPSession struct {
+	sourceAddr M.Socksaddr
+}
+
+type relayPacketWriter[U comparable] struct {
+	N.PacketConn
+	session *relayUDPSession
+}
+
+func (w *relayPacketWriter[U]) WritePacket(buffer *buf.Buffer, _ M.Socksaddr) error {
+	return w.PacketConn.WritePacket(buffer, w.session.sourceAddr)
+}

protocol/shadowsocks/shadowaead_2022/service.go

@@ -29,33 +29,37 @@ import (
 )
 
 var (
-	ErrNoPadding  = E.New("bad request: missing payload or padding")
-	ErrBadPadding = E.New("bad request: damaged padding")
+	ErrSaltNotUnique = E.New("bad request: salt not unique")
+	ErrNoPadding     = E.New("bad request: missing payload or padding")
+	ErrBadPadding    = E.New("bad request: damaged padding")
 )
 
 type Service struct {
-	name             string
-	secureRNG        io.Reader
-	keySaltLength    int
+	name          string
+	secureRNG     io.Reader
+	keySaltLength int
+	handler       shadowsocks.Handler
+
 	constructor      func(key []byte) cipher.AEAD
 	blockConstructor func(key []byte) cipher.Block
 	udpCipher        cipher.AEAD
 	udpBlockCipher   cipher.Block
 	psk              []byte
-	replayFilter     replay.Filter
-	handler          shadowsocks.Handler
-	udpNat           *udpnat.Service[uint64]
-	sessions         *cache.LruCache[uint64, *serverUDPSession]
+
+	replayFilter replay.Filter
+	udpNat       *udpnat.Service[uint64]
+	udpSessions  *cache.LruCache[uint64, *serverUDPSession]
 }
 
 func NewService(method string, psk []byte, password string, secureRNG io.Reader, udpTimeout int64, handler shadowsocks.Handler) (shadowsocks.Service, error) {
 	s := &Service{
-		name:         method,
-		secureRNG:    secureRNG,
+		name:      method,
+		secureRNG: secureRNG,
+		handler:   handler,
+
 		replayFilter: replay.NewCuckoo(60),
-		handler:      handler,
 		udpNat:       udpnat.New[uint64](udpTimeout, handler),
-		sessions: cache.New[uint64, *serverUDPSession](
+		udpSessions: cache.New[uint64, *serverUDPSession](
 			cache.WithAge[uint64, *serverUDPSession](udpTimeout),
 			cache.WithUpdateAgeOnGet[uint64, *serverUDPSession](),
 		),
@@ -90,9 +94,7 @@ func NewService(method string, psk []byte, password string, secureRNG io.Reader,
 	}
 
 	switch method {
-	case "2022-blake3-aes-128-gcm":
-		s.udpBlockCipher = newAES(psk)
-	case "2022-blake3-aes-256-gcm":
+	case "2022-blake3-aes-128-gcm", "2022-blake3-aes-256-gcm":
 		s.udpBlockCipher = newAES(psk)
 	case "2022-blake3-chacha20-poly1305":
 		s.udpCipher = newXChacha20Poly1305(psk)
@@ -123,7 +125,7 @@ func (s *Service) newConnection(ctx context.Context, conn net.Conn, metadata M.M
 	requestSalt := header[:s.keySaltLength]
 
 	if !s.replayFilter.Check(requestSalt) {
-		return E.New("salt not unique")
+		return ErrSaltNotUnique
 	}
 
 	requestKey := SessionKey(s.psk, requestSalt, s.keySaltLength)
@@ -315,7 +317,7 @@ func (s *Service) newPacket(conn N.PacketConn, buffer *buf.Buffer, metadata M.Me
 		return err
 	}
 
-	session, loaded := s.sessions.LoadOrStore(sessionId, s.newUDPSession)
+	session, loaded := s.udpSessions.LoadOrStore(sessionId, s.newUDPSession)
 	if !loaded {
 		session.remoteSessionId = sessionId
 		if packetHeader != nil {
@@ -328,7 +330,7 @@ func (s *Service) newPacket(conn N.PacketConn, buffer *buf.Buffer, metadata M.Me
 
 returnErr:
 	if !loaded {
-		s.sessions.Delete(sessionId)
+		s.udpSessions.Delete(sessionId)
 	}
 	return err
 

protocol/shadowsocks/shadowaead_2022/service_multi.go

@@ -7,6 +7,7 @@ import (
 	"io"
 	"math"
 	"net"
+	"os"
 	"runtime"
 	"time"
 
@@ -64,7 +65,7 @@ func NewMultiService[U comparable](method string, iPSK []byte, secureRNG io.Read
 	case "2022-blake3-aes-128-gcm":
 	case "2022-blake3-aes-256-gcm":
 	default:
-		return nil, E.New("unsupported method ", method)
+		return nil, os.ErrInvalid
 	}
 
 	ss, err := NewService(method, iPSK, "", secureRNG, udpTimeout, handler)
@@ -100,7 +101,7 @@ func (s *MultiService[U]) newConnection(ctx context.Context, conn net.Conn, meta
 	}
 	requestSalt := requestHeader[:s.keySaltLength]
 	if !s.replayFilter.Check(requestSalt) {
-		return E.New("salt not unique")
+		return ErrSaltNotUnique
 	}
 
 	var _eiHeader [aes.BlockSize]byte
@@ -243,7 +244,7 @@ func (s *MultiService[U]) newPacket(conn N.PacketConn, buffer *buf.Buffer, metad
 		return err
 	}
 
-	session, loaded := s.sessions.LoadOrStore(sessionId, func() *serverUDPSession {
+	session, loaded := s.udpSessions.LoadOrStore(sessionId, func() *serverUDPSession {
 		return s.newUDPSession(uPSK)
 	})
 	if !loaded {
@@ -257,7 +258,7 @@ func (s *MultiService[U]) newPacket(conn N.PacketConn, buffer *buf.Buffer, metad
 
 returnErr:
 	if !loaded {
-		s.sessions.Delete(sessionId)
+		s.udpSessions.Delete(sessionId)
 	}
 	return err
 
@@ -314,12 +315,11 @@ process:
 	metadata.Destination = destination
 	session.remoteAddr = metadata.Source
 
-	var userCtx shadowsocks.UserContext[U]
-	userCtx.Context = context.Background()
-	userCtx.User = user
-
-	s.udpNat.NewContextPacket(&userCtx, sessionId, func() N.PacketWriter {
-		return &serverPacketWriter{s.Service, conn, session}
+	s.udpNat.NewContextPacket(context.Background(), sessionId, func() (context.Context, N.PacketWriter) {
+		return &shadowsocks.UserContext[U]{
+			context.Background(),
+			user,
+		}, &serverPacketWriter{s.Service, conn, session}
 	}, buffer, metadata)
 	return nil
 }