adapt newest go-shadowsocks2 lib

internal/util/ss/conn.go

@@ -1,121 +1,38 @@
 package ss
 
 import (
-	"errors"
 	"net"
-	"net/netip"
-	"sync"
 
-	"github.com/go-gost/core/common/bufpool"
 	"github.com/go-gost/go-shadowsocks2/core"
 	"github.com/go-gost/go-shadowsocks2/socks"
 )
 
-const (
-	defaultBufferSize = 4096
-)
-
-var (
-	_ net.PacketConn = (*UDPConn)(nil)
-	_ net.Conn       = (*UDPConn)(nil)
-)
-
-// This wrapped connection has different behavior than ordinary connection:
-// 1. ReadFrom will return target addr of shadowsocks instead of remote addr
-type UDPConn struct {
-	client     *core.UDPClient
-	server     *core.UDPServer
-	sessionMap *sync.Map
+type udpConn struct {
 	net.PacketConn
-	raddr      net.Addr
-	taddr      net.Addr
-	bufferSize int
+	targetAddr net.Addr
 }
 
-func UDPClientConn(c net.PacketConn, remoteAddr, targetAddr net.Addr, bufferSize int, client *core.UDPClient, sessionMap *sync.Map) *UDPConn {
-	if bufferSize <= 0 {
-		bufferSize = defaultBufferSize
-	}
-
-	return &UDPConn{
-		PacketConn: c,
-		raddr:      remoteAddr,
-		taddr:      targetAddr,
-		bufferSize: bufferSize,
-		client:     client,
-		sessionMap: sessionMap,
+func UDPClientConn(c net.PacketConn, targetAddr net.Addr, client *core.UDPClient) net.Conn {
+	return &udpConn{
+		PacketConn: client.WrapConn(c),
+		targetAddr: targetAddr,
 	}
 }
 
-func (c *UDPConn) ReadFrom(b []byte) (n int, addr net.Addr, err error) {
-	buf := bufpool.Get(c.bufferSize)
-	defer bufpool.Put(buf)
-
-	clientAddr, err := netip.ParseAddrPort(c.LocalAddr().String())
-	if err != nil {
-		return
-	}
-
-	n, _, err = c.PacketConn.ReadFrom(buf)
-	if err != nil {
-		return
-	}
-
-	var payload []byte
-	var session core.UDPSession
-	if c.client != nil {
-		s, ok := c.sessionMap.Load(core.SessionHashFromAddrPort(clientAddr))
-		if !ok {
-			return 0, nil, errors.New("udp session cannot find")
-		}
-		session = s.(core.UDPSession)
-		payload, err = c.client.Outbound(buf[:n], session)
-		if err != nil {
-			return
-		}
-	} else {
-		return 0, nil, errors.New("UDPConn must be client ")
-	}
-
-	n = copy(b, payload)
-	addr, err = net.ResolveUDPAddr("udp", session.Target().String())
-
-	return
+func (c *udpConn) Read(b []byte) (int, error) {
+	n, _, err := c.PacketConn.ReadFrom(b)
+	return n, err
 }
 
-func (c *UDPConn) Read(b []byte) (n int, err error) {
-	n, _, err = c.ReadFrom(b)
-	return
+func (c *udpConn) Write(b []byte) (int, error) {
+	return c.WriteTo(b, c.targetAddr)
 }
 
-func (c *UDPConn) WriteTo(b []byte, addr net.Addr) (n int, err error) {
-	target := socks.ParseAddr(c.taddr.String())
-	clientAddr, err := netip.ParseAddrPort(c.LocalAddr().String())
-	if err != nil {
-		return
-	}
-
-	var session core.UDPSession
-	var encrypted []byte
-	if c.client != nil {
-		session, encrypted, err = c.client.Inbound(b, clientAddr, target)
-		if err != nil {
-			return
-		}
-		c.sessionMap.Store(session.Hash(), session)
-	} else {
-		return 0, errors.New("UDPConn must be client")
-	}
-
-	_, err = c.PacketConn.WriteTo(encrypted, c.raddr)
-	n = len(b)
-	return
+func (c *udpConn) WriteTo(b []byte, addr net.Addr) (int, error) {
+	target := socks.ParseAddr(addr.String())
+	return c.PacketConn.WriteTo(b, core.NewUDPClientPacketAddr(target, c.LocalAddr()))
 }
 
-func (c *UDPConn) Write(b []byte) (n int, err error) {
-	return c.WriteTo(b, c.taddr)
-}
-
-func (c *UDPConn) RemoteAddr() net.Addr {
-	return c.raddr
+func (c *udpConn) RemoteAddr() net.Addr {
+	return c.targetAddr
 }

internal/util/ss/ss.go

@@ -34,9 +34,16 @@ func (c *shadowConn) Write(b []byte) (n int, err error) {
 	n = len(b) // force byte length consistent
 	if c.wbuf.Len() > 0 {
 		c.wbuf.Write(b) // append the data to the cached header
-		_, err = c.Conn.Write(c.wbuf.Bytes())
+		written, err := c.Conn.Write(c.wbuf.Bytes())
+		if err != nil {
+			// If short write, keep unwritten bytes
+			if written > 0 && written < c.wbuf.Len() {
+				c.wbuf.Next(written)
+			}
+			return written, err
+		}
 		c.wbuf.Reset()
-		return
+		return n, nil
 	}
 	_, err = c.Conn.Write(b)
 	return