Merge remote-tracking branch 'origin/pr/96'

connector/ss/connector.go

@@ -2,6 +2,7 @@ package ss
 
 import (
 	"context"
+	"errors"
 	"fmt"
 	"net"
 	"time"
@@ -43,18 +44,20 @@ func (c *ssConnector) Init(md md.Metadata) (err error) {
 		return
 	}
 
-	if c.options.Auth != nil {
-		method := c.options.Auth.Username()
-		password, _ := c.options.Auth.Password()
-
-		clientConfig, err := utils.NewClientConfig(method, password)
-		if err != nil {
-			return err
-		}
-
-		c.client = core.NewTCPClient(clientConfig)
+	if c.options.Auth == nil {
+		return errors.New("ss: auth is required")
 	}
 
+	method := c.options.Auth.Username()
+	password, _ := c.options.Auth.Password()
+
+	clientConfig, err := utils.NewClientConfig(method, password)
+	if err != nil {
+		return err
+	}
+
+	c.client = core.NewTCPClient(clientConfig)
+
 	return
 }
 
@@ -106,14 +109,5 @@ func (c *ssConnector) Connect(ctx context.Context, conn net.Conn, network, addre
 		return nil, err
 	}
 
-	var sc net.Conn
-	if c.md.noDelay {
-		err := conn.(core.TCPConn).ClientFirstWrite()
-		if err != nil {
-			return nil, err
-		}
-	}
-	sc = conn
-
-	return sc, nil
+	return conn, nil
 }

connector/ss/udp/connector.go

@@ -2,9 +2,10 @@ package ss
 
 import (
 	"context"
+	"errors"
 	"fmt"
 	"net"
-	"sync"
+	"net/netip"
 	"time"
 
 	"github.com/go-gost/core/connector"
@@ -23,11 +24,10 @@ func init() {
 }
 
 type ssuConnector struct {
-	client     core.UDPClient
-	tcpClient  core.TCPClient
-	md         metadata
-	options    connector.Options
-	sessionMap sync.Map
+	clientCfg core.ClientConfig
+	tcpClient core.TCPClient
+	md        metadata
+	options   connector.Options
 }
 
 func NewConnector(opts ...connector.Option) connector.Connector {
@@ -46,17 +46,26 @@ func (c *ssuConnector) Init(md md.Metadata) (err error) {
 		return
 	}
 
-	if c.options.Auth != nil {
-		method := c.options.Auth.Username()
-		password, _ := c.options.Auth.Password()
-		clientConfig, err := utils.NewClientConfig(method, password)
-		if err != nil {
-			return err
-		}
-		c.client = core.NewUDPClient(clientConfig, 60)
-		c.tcpClient = core.NewTCPClient(clientConfig)
+	if c.options.Auth == nil {
+		return errors.New("ss: auth is required")
 	}
 
+	method := c.options.Auth.Username()
+	password, _ := c.options.Auth.Password()
+	clientConfig, err := utils.NewClientConfig(method, password)
+	if err != nil {
+		return err
+	}
+	clientConfig.UDPTimeout = time.Minute
+
+	tcpClientConfig, err := utils.NewClientConfig(method, password)
+	if err != nil {
+		return err
+	}
+
+	c.clientCfg = clientConfig
+	c.tcpClient = core.NewTCPClient(tcpClientConfig)
+
 	return
 }
 
@@ -87,15 +96,29 @@ func (c *ssuConnector) Connect(ctx context.Context, conn net.Conn, network, addr
 	if taddr == nil {
 		taddr = &net.UDPAddr{}
 	}
+	serverAddr, err := netip.ParseAddrPort(conn.RemoteAddr().String())
+	if err != nil {
+		return nil, err
+	}
+	clientCfg := c.clientCfg
+	clientCfg.ServerAddr = serverAddr
+	client := core.NewUDPClient(core.ClientConfig{
+		Cipher:     clientCfg.Cipher,
+		ServerAddr: serverAddr,
+		UDPTimeout: time.Minute,
+	})
+	if err := client.Init(); err != nil {
+		return nil, err
+	}
 
 	pc, ok := conn.(net.PacketConn)
 	if ok {
 		// standard UDP relay
-		return ss.UDPClientConn(pc, conn.RemoteAddr(), taddr, c.md.udpBufferSize, &c.client, &c.sessionMap), nil
+		return ss.UDPClientConn(pc, taddr, &client), nil
 	}
 
 	target := socks.ParseAddr(taddr.String())
-	conn, err := c.tcpClient.WrapConn(conn, target)
+	conn, err = c.tcpClient.WrapConn(conn, target)
 	if err != nil {
 		return nil, err
 	}

handler/ss/handler.go

@@ -5,6 +5,7 @@ import (
 	"bytes"
 	"context"
 	"crypto/tls"
+	"errors"
 	"net"
 	"time"
 
@@ -14,6 +15,7 @@ import (
 	"github.com/go-gost/core/observer/stats"
 	"github.com/go-gost/core/recorder"
 	"github.com/go-gost/go-shadowsocks2/core"
+	"github.com/go-gost/go-shadowsocks2/socks"
 	"github.com/go-gost/go-shadowsocks2/utils"
 	xctx "github.com/go-gost/x/ctx"
 	ictx "github.com/go-gost/x/internal/ctx"
@@ -54,19 +56,22 @@ func (h *ssHandler) Init(md md.Metadata) (err error) {
 	if err = h.parseMetadata(md); err != nil {
 		return
 	}
-	if h.options.Auth != nil {
-		method := h.options.Auth.Username()
-		password, _ := h.options.Auth.Password()
 
-		serverConfig, err := utils.NewServerConfig(method, password, h.md.users)
-		if err != nil {
-			return err
-		}
-		h.server = core.NewTCPServer(serverConfig)
-		err = h.server.Init()
-		if err != nil {
-			return err
-		}
+	if h.options.Auth == nil {
+		return errors.New("ss: auth is required")
+	}
+
+	method := h.options.Auth.Username()
+	password, _ := h.options.Auth.Password()
+
+	serverConfig, err := utils.NewServerConfig(method, password, h.md.users)
+	if err != nil {
+		return err
+	}
+	h.server = core.NewTCPServer(serverConfig)
+	err = h.server.Init()
+	if err != nil {
+		return err
 	}
 
 	for _, ro := range h.options.Recorders {
@@ -139,7 +144,13 @@ func (h *ssHandler) Handle(ctx context.Context, conn net.Conn, opts ...handler.H
 	if err != nil {
 		return
 	}
-	target := wrappedConn.Target()
+	var target socks.Addr
+	if tc, ok := wrappedConn.(interface{ Target() socks.Addr }); ok {
+		target = tc.Target()
+	}
+	if len(target) == 0 {
+		return
+	}
 	conn = wrappedConn
 
 	conn.SetReadDeadline(time.Now().Add(h.md.readTimeout))

handler/ss/udp/handler.go

@@ -5,12 +5,10 @@ import (
 	"context"
 	"errors"
 	"net"
-	"net/netip"
 	"sync"
 	"time"
 
 	"github.com/go-gost/core/bypass"
-	"github.com/go-gost/core/common/bufpool"
 	"github.com/go-gost/core/handler"
 	"github.com/go-gost/core/logger"
 	md "github.com/go-gost/core/metadata"
@@ -19,6 +17,7 @@ import (
 	"github.com/go-gost/go-shadowsocks2/utils"
 	xctx "github.com/go-gost/x/ctx"
 	ictx "github.com/go-gost/x/internal/ctx"
+	"github.com/go-gost/x/internal/net/udp"
 	"github.com/go-gost/x/internal/util/relay"
 	rate_limiter "github.com/go-gost/x/limiter/rate"
 	xrecorder "github.com/go-gost/x/recorder"
@@ -54,25 +53,28 @@ func (h *ssuHandler) Init(md md.Metadata) (err error) {
 		return
 	}
 
-	if h.options.Auth != nil {
-		method := h.options.Auth.Username()
-		password, _ := h.options.Auth.Password()
+	if h.options.Auth == nil {
+		return errors.New("ss: auth is required")
+	}
 
-		serverConfig, err := utils.NewServerConfig(method, password, h.md.users)
-		if err != nil {
-			return err
-		}
-		h.udpServer = core.NewUDPServer(serverConfig, time.Duration(60*time.Second))
-		err = h.udpServer.Init()
-		if err != nil {
-			return err
-		}
+	method := h.options.Auth.Username()
+	password, _ := h.options.Auth.Password()
 
-		h.tcpServer = core.NewTCPServer(serverConfig)
-		err = h.udpServer.Init()
-		if err != nil {
-			return err
-		}
+	serverConfig, err := utils.NewServerConfig(method, password, h.md.users)
+	if err != nil {
+		return err
+	}
+	serverConfig.UDPTimeout = time.Minute
+	h.udpServer = core.NewUDPServer(serverConfig)
+	err = h.udpServer.Init()
+	if err != nil {
+		return err
+	}
+
+	h.tcpServer = core.NewTCPServer(serverConfig)
+	err = h.tcpServer.Init()
+	if err != nil {
+		return err
 	}
 
 	for _, ro := range h.options.Recorders {
@@ -141,198 +143,148 @@ func (h *ssuHandler) Handle(ctx context.Context, conn net.Conn, opts ...handler.
 		// WriteTo will write data to original remote addr instead of the addr parameter passed to
 		pc = relay.UDPTunServerConn(conn)
 
-		return h.relayPacketTCP(ctx, pc, ro, log)
-	} else {
-		return h.relayPacketUDP(ctx, pc, ro, log)
-	}
-}
-
-// UDP over TCP
-func (h *ssuHandler) relayPacketTCP(ctx context.Context, src net.PacketConn, ro *xrecorder.HandlerRecorderObject, log logger.Logger) error {
-	bufferSize := h.md.udpBufferSize
-	if bufferSize <= 0 {
-		bufferSize = defaultBufferSize
-	}
-
-	b := bufpool.Get(bufferSize)
-	defer bufpool.Put(b)
-
-	var dstConn net.Conn
-	defer func() {
-		if dstConn != nil {
-			dstConn.Close()
-		}
-	}()
-
-	for {
-		n, targetAddr, err := src.ReadFrom(b)
+		// UDP over TCP still behaves like a regular packet tunnel.
+		var buf bytes.Buffer
+		c, err := h.options.Router.Dial(ictx.ContextWithBuffer(ctx, &buf), "udp", "")
+		ro.Route = buf.String()
 		if err != nil {
+			log.Error(err)
+			return err
+		}
+		defer c.Close()
+
+		cc, ok := c.(net.PacketConn)
+		if !ok {
+			err := errors.New("ss: wrong connection type")
+			log.Error(err)
 			return err
 		}
 
-		if ro.Host == "" {
-			ro.Host = targetAddr.String()
-		}
+		r := udp.NewRelay(pc, cc).
+			WithService(h.options.Service).
+			WithBypass(h.options.Bypass).
+			WithBufferSize(h.md.udpBufferSize).
+			WithLogger(log)
 
-		if h.options.Bypass != nil && h.options.Bypass.Contains(ctx, targetAddr.Network(), targetAddr.String(), bypass.WithService(h.options.Service)) {
-			log.Warn("bypass: ", targetAddr)
-			return nil
-		}
-
-		if dstConn == nil {
-			// obtain a udp connection
-			var buf bytes.Buffer
-			dstConn, err := h.options.Router.Dial(ictx.ContextWithBuffer(ctx, &buf), "udp", "") // UDP association
-			ro.Route = buf.String()
-			if err != nil {
-				log.Error(err)
-				return err
-			}
-
-			cc, ok := dstConn.(net.PacketConn)
-			if !ok {
-				err := errors.New("ss: wrong connection type")
-				log.Error(err)
-				return err
-			}
-
-			log.Infof("%s <-> %s", src.LocalAddr(), cc.LocalAddr())
-
-			go func() {
-				defer dstConn.Close()
-				for {
-					n, raddr, err := cc.ReadFrom(b)
-					if err != nil {
-						log.Warnf("failed to read response from %v: %v", raddr, err)
-						return
-					}
-
-					if h.options.Bypass != nil && h.options.Bypass.Contains(ctx, raddr.Network(), raddr.String(), bypass.WithService(h.options.Service)) {
-						log.Warn("bypass: ", raddr)
-						return
-					}
-
-					if _, err = src.WriteTo(b[:n], targetAddr); err != nil {
-						log.Warnf("failed to write response from %v: %v", targetAddr, err)
-						return
-					}
-
-					log.Tracef("%s <<< %s data: %d",
-						cc.LocalAddr(), raddr, n)
-				}
-			}()
-		}
-
-		if _, err = dstConn.(net.PacketConn).WriteTo(b[:n], targetAddr); err != nil {
-			return err
-		}
-
-		log.Tracef("%s >>> %s data: %d",
-			dstConn.(net.PacketConn).LocalAddr(), targetAddr, n)
+		return r.Run(ctx)
 	}
+
+	return h.relayPacketUDP(ctx, h.udpServer.WrapConn(pc), ro, log)
 }
 
-// standard udp relay
-// Dataflow: src (encrypted data) <-> dst (plaintext data)
 func (h *ssuHandler) relayPacketUDP(ctx context.Context, src net.PacketConn, ro *xrecorder.HandlerRecorderObject, log logger.Logger) error {
 	bufferSize := h.md.udpBufferSize
 	if bufferSize <= 0 {
 		bufferSize = defaultBufferSize
 	}
 
-	b := bufpool.Get(bufferSize)
-	defer bufpool.Put(b)
-
+	b := make([]byte, bufferSize)
 	for {
 		n, addr, err := src.ReadFrom(b)
 		if err != nil {
 			return err
 		}
 
-		clientAddr, err := netip.ParseAddrPort(addr.String())
-		if err != nil {
-			return err
-		}
-		session, payload, err := h.udpServer.Inbound(b[:n], clientAddr)
-		if err != nil {
-			return err
+		ctxAddr, ok := addr.(interface {
+			Session() core.UDPSession
+			ClientAddr() net.Addr
+			TargetAddr() net.Addr
+		})
+		if !ok {
+			return errors.New("ss: missing udp session context")
 		}
 
-		targetAddr, err := net.ResolveUDPAddr("udp", session.Target().String())
-		if ro.Host == "" {
+		targetAddr := ctxAddr.TargetAddr()
+		if ro.Host == "" && targetAddr != nil {
 			ro.Host = targetAddr.String()
 		}
 
-		if h.options.Bypass != nil && h.options.Bypass.Contains(ctx, targetAddr.Network(), targetAddr.String(), bypass.WithService(h.options.Service)) {
-			log.Warn("bypass: ", addr)
+		if h.options.Bypass != nil && targetAddr != nil && h.options.Bypass.Contains(ctx, "udp", targetAddr.String(), bypass.WithService(h.options.Service)) {
+			log.Warn("bypass: ", targetAddr)
 			return nil
 		}
+
+		session := ctxAddr.Session()
+		if session == nil {
+			return errors.New("ss: udp session not found")
+		}
+
+		dstConn, err := h.packetConnForSession(ctx, src, session, ro, log)
 		if err != nil {
 			return err
 		}
 
-		dstConn, ok := h.connMap.Load(session.Hash())
-		if !ok {
-			// obtain a udp connection
-			var buf bytes.Buffer
-			c, err := h.options.Router.Dial(ictx.ContextWithBuffer(ctx, &buf), "udp", "") // UDP association
-			ro.Route = buf.String()
-			if err != nil {
-				log.Error(err)
-				return err
-			}
-
-			cc, ok := c.(net.PacketConn)
-			if !ok {
-				err := errors.New("ss: wrong connection type")
-				log.Error(err)
-				return err
-			}
-			h.connMap.Store(session.Hash(), c)
-
-			dstConn = cc
-
-			log.Infof("%s <-> %s", src.LocalAddr(), cc.LocalAddr())
-
-			go func() {
-				defer func() {
-					c.Close()
-					h.connMap.Delete(session.Hash())
-				}()
-
-				for {
-					n, raddr, err := dstConn.(net.PacketConn).ReadFrom(b)
-					if err != nil {
-						log.Warnf("failed to read response: %v", err)
-						return
-					}
-
-					clientAddr := session.ClientAddr()
-					if h.options.Bypass != nil && h.options.Bypass.Contains(ctx, raddr.Network(), raddr.String(), bypass.WithService(h.options.Service)) {
-						log.Warn("bypass: ", raddr)
-						return
-					}
-					encrypted, err := h.udpServer.Outbound(b[:n], session)
-					if _, err = src.WriteTo(encrypted, net.UDPAddrFromAddrPort(clientAddr)); err != nil {
-						log.Warnf("failed to write response to %v: %v", clientAddr, err)
-						return
-					}
-
-					log.Tracef("%s <<< %s data: %d",
-						dstConn.(net.PacketConn).LocalAddr(), raddr, n)
-				}
-			}()
-		}
-
-		if n, err = dstConn.(net.PacketConn).WriteTo(payload, targetAddr); err != nil {
+		if _, err = dstConn.WriteTo(b[:n], targetAddr); err != nil {
 			return err
 		}
 
-		log.Tracef("%s >>> %s data: %d",
-			dstConn.(net.PacketConn).LocalAddr(), targetAddr, n)
+		log.Tracef("%s >>> %s data: %d", dstConn.LocalAddr(), targetAddr, n)
 	}
 }
 
+func (h *ssuHandler) packetConnForSession(ctx context.Context, src net.PacketConn, session core.UDPSession, ro *xrecorder.HandlerRecorderObject, log logger.Logger) (net.PacketConn, error) {
+	if cc, ok := h.connMap.Load(session.Hash()); ok {
+		return cc.(net.PacketConn), nil
+	}
+
+	var buf bytes.Buffer
+	c, err := h.options.Router.Dial(ictx.ContextWithBuffer(ctx, &buf), "udp", "")
+	ro.Route = buf.String()
+	if err != nil {
+		log.Error(err)
+		return nil, err
+	}
+
+	cc, ok := c.(net.PacketConn)
+	if !ok {
+		c.Close()
+		err := errors.New("ss: wrong connection type")
+		log.Error(err)
+		return nil, err
+	}
+
+	actual, loaded := h.connMap.LoadOrStore(session.Hash(), cc)
+	if loaded {
+		c.Close()
+		return actual.(net.PacketConn), nil
+	}
+
+	go func() {
+		defer func() {
+			cc.Close()
+			h.connMap.Delete(session.Hash())
+		}()
+
+		b := make([]byte, h.md.udpBufferSize)
+		if len(b) == 0 {
+			b = make([]byte, defaultBufferSize)
+		}
+
+		for {
+			n, raddr, err := cc.ReadFrom(b)
+			if err != nil {
+				log.Warnf("failed to read response: %v", err)
+				return
+			}
+
+			if h.options.Bypass != nil && h.options.Bypass.Contains(ctx, "udp", raddr.String(), bypass.WithService(h.options.Service)) {
+				log.Warn("bypass: ", raddr)
+				return
+			}
+
+			addr := core.NewUDPServerPacketAddr(raddr, net.UDPAddrFromAddrPort(session.ClientAddr()), session)
+			if _, err = src.WriteTo(b[:n], addr); err != nil {
+				log.Warnf("failed to write response to %v: %v", session.ClientAddr(), err)
+				return
+			}
+
+			log.Tracef("%s <<< %s data: %d", cc.LocalAddr(), raddr, n)
+		}
+	}()
+
+	return cc, nil
+}
+
 func (h *ssuHandler) checkRateLimit(addr net.Addr) bool {
 	if h.options.RateLimiter == nil {
 		return true

handler/ss/udp/metadata.go

@@ -24,6 +24,9 @@ func (h *ssuHandler) parseMetadata(md mdata.Metadata) (err error) {
 
 	h.md.key = mdutil.GetString(md, "key")
 	h.md.readTimeout = mdutil.GetDuration(md, "readTimeout")
+	if h.md.readTimeout == 0 {
+		h.md.readTimeout = time.Minute // Default read timeout
+	}
 	h.md.udpBufferSize = mdutil.GetInt(md, "udpBufferSize", "udp.bufferSize")
 
 	usersMap := mdutil.GetStringMapString(md, "users")

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