Refactor buffer usage

go.mod

@@ -3,7 +3,7 @@ module github.com/sagernet/sing-shadowsocks
 go 1.18
 
 require (
-	github.com/sagernet/sing v0.0.0-20220606055031-45523fbbe7ae
+	github.com/sagernet/sing v0.0.0-20220609091055-86d0144940e7
 	golang.org/x/crypto v0.0.0-20220525230936-793ad666bf5e
 	lukechampine.com/blake3 v1.1.7
 )

go.sum

@@ -1,8 +1,8 @@
 github.com/klauspost/cpuid/v2 v2.0.9/go.mod h1:FInQzS24/EEf25PyTYn52gqo7WaD8xa0213Md/qVLRg=
 github.com/klauspost/cpuid/v2 v2.0.12 h1:p9dKCg8i4gmOxtv35DvrYoWqYzQrvEVdjQ762Y0OqZE=
 github.com/klauspost/cpuid/v2 v2.0.12/go.mod h1:g2LTdtYhdyuGPqyWyv7qRAmj1WBqxuObKfj5c0PQa7c=
-github.com/sagernet/sing v0.0.0-20220606055031-45523fbbe7ae h1:SrqT/QhXJwAw7hfAxgDqEPF3J0yK2U6Y8lGZhAubs9E=
+github.com/sagernet/sing v0.0.0-20220609091055-86d0144940e7 h1:Q+uNKLNSKqpx+p96qcBTVFh8RUKiQFr4IrNVi5Q5yl0=
-github.com/sagernet/sing v0.0.0-20220606055031-45523fbbe7ae/go.mod h1:w2HnJzXKHpD6F5Z/9XlSD4qbcpHY2RSZuQnFzqgELMg=
+github.com/sagernet/sing v0.0.0-20220609091055-86d0144940e7/go.mod h1:w2HnJzXKHpD6F5Z/9XlSD4qbcpHY2RSZuQnFzqgELMg=
 golang.org/x/crypto v0.0.0-20220525230936-793ad666bf5e h1:T8NU3HyQ8ClP4SEE+KbFlg6n0NhuTsN4MyznaarGsZM=
 golang.org/x/crypto v0.0.0-20220525230936-793ad666bf5e/go.mod h1:IxCIyHEi3zRg3s0A5j5BB6A9Jmi73HwBIUl50j+osU4=
 golang.org/x/sys v0.0.0-20220520151302-bc2c85ada10a h1:dGzPydgVsqGcTRVwiLJ1jVbufYwmzD3LfVPLKsKg+0k=

shadowaead/protocol.go

@@ -250,6 +250,7 @@ func (c *clientPacketConn) WritePacket(buffer *buf.Buffer, destination M.Socksad
 		return err
 	}
 	writeCipher.Seal(buffer.Index(c.keySaltLength), rw.ZeroBytes[:writeCipher.NonceSize()], buffer.From(c.keySaltLength), nil)
+	buffer.Extend(Overhead)
 	return common.Error(c.Write(buffer.Bytes()))
 }
 
@@ -290,15 +291,17 @@ func (c *clientPacketConn) ReadFrom(p []byte) (n int, addr net.Addr, err error)
 		return
 	}
 	addr = destination.UDPAddr()
-	copy(p, buffer.Bytes())
+	n = copy(p, buffer.Bytes())
 	return
 }
 
 func (c *clientPacketConn) WriteTo(p []byte, addr net.Addr) (n int, err error) {
 	destination := M.SocksaddrFromNet(addr)
-	_buffer := buf.StackNewSize(c.keySaltLength + M.SocksaddrSerializer.AddrPortLen(destination) + len(p))
+	_buffer := buf.StackNewSize(c.keySaltLength + M.SocksaddrSerializer.AddrPortLen(destination) + len(p) + Overhead)
 	defer common.KeepAlive(_buffer)
 	buffer := common.Dup(_buffer)
+	buffer.Resize(c.keySaltLength+M.SocksaddrSerializer.AddrPortLen(destination), 0)
+	common.Must1(buffer.Write(p))
 	err = c.WritePacket(buffer, destination)
 	if err != nil {
 		return

shadowaead_2022/protocol.go

@@ -216,18 +216,20 @@ func (m *Method) writeExtendedIdentityHeaders(request *buf.Buffer, salt []byte)
 		keyMaterial := buf.Make(m.keySaltLength * 2)
 		copy(keyMaterial, psk)
 		copy(keyMaterial[m.keySaltLength:], salt)
-		_identitySubkey := buf.Make(m.keySaltLength)
+		_identitySubkey := buf.StackNewSize(m.keySaltLength)
 		identitySubkey := common.Dup(_identitySubkey)
-		blake3.DeriveKey(identitySubkey, "shadowsocks 2022 identity subkey", keyMaterial)
+		identitySubkey.Extend(identitySubkey.FreeLen())
+		blake3.DeriveKey(identitySubkey.Bytes(), "shadowsocks 2022 identity subkey", keyMaterial)
 
 		pskHash := m.pskHash[aes.BlockSize*i : aes.BlockSize*(i+1)]
 
 		header := request.Extend(16)
-		b, err := m.blockConstructor(identitySubkey)
+		b, err := m.blockConstructor(identitySubkey.Bytes())
 		if err != nil {
 			return err
 		}
 		b.Encrypt(header, pskHash)
+		identitySubkey.Release()
 		common.KeepAlive(_identitySubkey)
 		if i == pskLen-2 {
 			break
@@ -237,7 +239,7 @@ func (m *Method) writeExtendedIdentityHeaders(request *buf.Buffer, salt []byte)
 }
 
 func (c *clientConn) writeRequest(payload []byte) error {
-	salt := buf.Make(c.keySaltLength)
+	salt := make([]byte, c.keySaltLength)
 	common.Must1(io.ReadFull(rand.Reader, salt))
 
 	key := SessionKey(c.pskList[len(c.pskList)-1], salt, c.keySaltLength)
@@ -301,19 +303,28 @@ func (c *clientConn) readResponse() error {
 		return nil
 	}
 
-	_salt := buf.Make(c.keySaltLength)
+	_salt := buf.StackNewSize(c.keySaltLength)
 	salt := common.Dup(_salt)
-	_, err := io.ReadFull(c.Conn, salt)
+
+	_, err := salt.ReadFullFrom(c.Conn, salt.FreeLen())
 	if err != nil {
+		salt.Release()
+		common.KeepAlive(_salt)
+
 		return err
 	}
+	if !c.replayFilter.Check(salt.Bytes()) {
+		salt.Release()
+		common.KeepAlive(_salt)
 
-	if !c.replayFilter.Check(salt) {
 		return ErrSaltNotUnique
 	}
 
-	key := SessionKey(c.pskList[len(c.pskList)-1], salt, c.keySaltLength)
+	key := SessionKey(c.pskList[len(c.pskList)-1], salt.Bytes(), c.keySaltLength)
+
+	salt.Release()
 	common.KeepAlive(_salt)
+
 	readCipher, err := c.constructor(common.Dup(key))
 	if err != nil {
 		return err
@@ -349,16 +360,17 @@ func (c *clientConn) readResponse() error {
 		return E.Extend(ErrBadTimestamp, "received ", epoch, ", diff ", diff, "s")
 	}
 
-	_requestSalt := buf.Make(c.keySaltLength)
+	_requestSalt := buf.StackNewSize(c.keySaltLength)
 	requestSalt := common.Dup(_requestSalt)
-	_, err = io.ReadFull(reader, requestSalt)
+	_, err = requestSalt.ReadFullFrom(reader, requestSalt.FreeLen())
 	if err != nil {
 		return err
 	}
 
-	if bytes.Compare(requestSalt, c.requestSalt) > 0 {
+	if bytes.Compare(requestSalt.Bytes(), c.requestSalt) > 0 {
 		return ErrBadRequestSalt
 	}
+	requestSalt.Release()
 	common.KeepAlive(_requestSalt)
 
 	var length uint16

shadowaead_2022/relay.go

@@ -158,15 +158,17 @@ func (s *Relay[U]) newConnection(ctx context.Context, conn net.Conn, metadata M.
 	keyMaterial := buf.Make(s.keySaltLength * 2)
 	copy(keyMaterial, s.iPSK)
 	copy(keyMaterial[s.keySaltLength:], requestSalt)
-	_identitySubkey := buf.Make(s.keySaltLength)
+	_identitySubkey := buf.StackNewSize(s.keySaltLength)
 	identitySubkey := common.Dup(_identitySubkey)
-	blake3.DeriveKey(identitySubkey, "shadowsocks 2022 identity subkey", keyMaterial)
+	identitySubkey.Extend(identitySubkey.FreeLen())
-	b, err := s.blockConstructor(identitySubkey)
+	blake3.DeriveKey(identitySubkey.Bytes(), "shadowsocks 2022 identity subkey", keyMaterial)
+	b, err := s.blockConstructor(identitySubkey.Bytes())
+	identitySubkey.Release()
+	common.KeepAlive(_identitySubkey)
 	if err != nil {
 		return err
 	}
 	b.Decrypt(eiHeader, eiHeader)
-	common.KeepAlive(_identitySubkey)
 
 	var user U
 	if u, loaded := s.uPSKHash[_eiHeader]; loaded {
@@ -185,10 +187,7 @@ func (s *Relay[U]) newConnection(ctx context.Context, conn net.Conn, metadata M.
 	}
 	metadata.Protocol = "shadowsocks-relay"
 	metadata.Destination = s.uDestination[user]
-	conn = &bufio.BufferedConn{
+	conn = bufio.NewCachedConn(conn, requestHeader)
-		Conn:   conn,
-		Buffer: requestHeader,
-	}
 	return s.handler.NewConnection(ctx, conn, metadata)
 }
 

shadowaead_2022/service.go

@@ -125,7 +125,7 @@ func (s *Service) NewConnection(ctx context.Context, conn net.Conn, metadata M.M
 }
 
 func (s *Service) newConnection(ctx context.Context, conn net.Conn, metadata M.Metadata) error {
-	header := buf.Make(s.keySaltLength + shadowaead.Overhead + RequestHeaderFixedChunkLength)
+	header := make([]byte, s.keySaltLength+shadowaead.Overhead+RequestHeaderFixedChunkLength)
 
 	n, err := conn.Read(header)
 	if err != nil {
@@ -234,13 +234,16 @@ type serverConn struct {
 }
 
 func (c *serverConn) writeResponse(payload []byte) (n int, err error) {
-	_salt := buf.Make(c.keySaltLength)
+	_salt := buf.StackNewSize(c.keySaltLength)
-	salt := common.Dup(_salt[:])
+	salt := common.Dup(_salt)
-	common.Must1(io.ReadFull(rand.Reader, salt))
+	salt.WriteRandom(salt.FreeLen())
-	key := SessionKey(c.uPSK, salt, c.keySaltLength)
+
+	key := SessionKey(c.uPSK, salt.Bytes(), c.keySaltLength)
 	common.KeepAlive(_salt)
 	writeCipher, err := c.constructor(common.Dup(key))
 	if err != nil {
+		salt.Release()
+		common.KeepAlive(_salt)
 		return
 	}
 	writer := shadowaead.NewWriter(
@@ -250,16 +253,20 @@ func (c *serverConn) writeResponse(payload []byte) (n int, err error) {
 	)
 	common.KeepAlive(key)
 	header := writer.Buffer()
-	header.Write(salt)
+	header.Write(salt.Bytes())
 
-	_headerFixedChunk := buf.Make(1 + 8 + c.keySaltLength + 2)
+	salt.Release()
-	headerFixedChunk := buf.With(common.Dup(_headerFixedChunk))
+	common.KeepAlive(_salt)
+
+	_headerFixedChunk := buf.StackNewSize(1 + 8 + c.keySaltLength + 2)
+	headerFixedChunk := common.Dup(_headerFixedChunk)
 	common.Must(headerFixedChunk.WriteByte(HeaderTypeServer))
 	common.Must(binary.Write(headerFixedChunk, binary.BigEndian, uint64(time.Now().Unix())))
 	common.Must1(headerFixedChunk.Write(c.requestSalt))
 	common.Must(binary.Write(headerFixedChunk, binary.BigEndian, uint16(len(payload))))
 
 	writer.WriteChunk(header, headerFixedChunk.Slice())
+	headerFixedChunk.Release()
 	common.KeepAlive(_headerFixedChunk)
 	c.requestSalt = nil
 

shadowaead_2022/service_multi.go

@@ -129,15 +129,17 @@ func (s *MultiService[U]) newConnection(ctx context.Context, conn net.Conn, meta
 	keyMaterial := buf.Make(s.keySaltLength * 2)
 	copy(keyMaterial, s.psk)
 	copy(keyMaterial[s.keySaltLength:], requestSalt)
-	_identitySubkey := buf.Make(s.keySaltLength)
+	_identitySubkey := buf.StackNewSize(s.keySaltLength)
 	identitySubkey := common.Dup(_identitySubkey)
-	blake3.DeriveKey(identitySubkey, "shadowsocks 2022 identity subkey", keyMaterial)
+	identitySubkey.Extend(identitySubkey.FreeLen())
-	b, err := s.blockConstructor(identitySubkey)
+	blake3.DeriveKey(identitySubkey.Bytes(), "shadowsocks 2022 identity subkey", keyMaterial)
+	b, err := s.blockConstructor(identitySubkey.Bytes())
+	identitySubkey.Release()
+	common.KeepAlive(_identitySubkey)
 	if err != nil {
 		return err
 	}
 	b.Decrypt(eiHeader, eiHeader)
-	common.KeepAlive(_identitySubkey)
 
 	var user U
 	var uPSK []byte

shadowstream/protocol.go

@@ -299,9 +299,10 @@ func (c *clientPacketConn) ReadFrom(p []byte) (n int, addr net.Addr, err error)
 
 func (c *clientPacketConn) WriteTo(p []byte, addr net.Addr) (n int, err error) {
 	destination := M.SocksaddrFromNet(addr)
-	_buffer := buf.Make(c.saltLength + M.SocksaddrSerializer.AddrPortLen(destination) + len(p))
+	_buffer := buf.StackNewSize(c.saltLength + M.SocksaddrSerializer.AddrPortLen(destination) + len(p))
 	defer common.KeepAlive(_buffer)
-	buffer := buf.With(common.Dup(_buffer))
+	buffer := common.Dup(_buffer)
+	defer buffer.Release()
 	common.Must1(buffer.ReadFullFrom(rand.Reader, c.saltLength))
 	err = M.SocksaddrSerializer.WriteAddrPort(buffer, M.SocksaddrFromNet(addr))
 	if err != nil {