chore: make all net.Conn wrapper can pass through N.ExtendedConn

common/callback/callback.go

@@ -22,53 +22,23 @@ func (c *firstWriteCallBackConn) Write(b []byte) (n int, err error) {
 	return c.Conn.Write(b)
 }
 
+func (c *firstWriteCallBackConn) WriteBuffer(buffer *buf.Buffer) (err error) {
+	defer func() {
+		if !c.written {
+			c.written = true
+			c.callback(err)
+		}
+	}()
+	return c.Conn.WriteBuffer(buffer)
+}
+
 func (c *firstWriteCallBackConn) Upstream() any {
 	return c.Conn
 }
 
-type extendedConn interface {
-	C.Conn
-	N.ExtendedConn
-}
-
-type firstWriteCallBackExtendedConn struct {
-	extendedConn
-	callback func(error)
-	written  bool
-}
-
-func (c *firstWriteCallBackExtendedConn) Write(b []byte) (n int, err error) {
-	defer func() {
-		if !c.written {
-			c.written = true
-			c.callback(err)
-		}
-	}()
-	return c.extendedConn.Write(b)
-}
-
-func (c *firstWriteCallBackExtendedConn) WriteBuffer(buffer *buf.Buffer) (err error) {
-	defer func() {
-		if !c.written {
-			c.written = true
-			c.callback(err)
-		}
-	}()
-	return c.extendedConn.WriteBuffer(buffer)
-}
-
-func (c *firstWriteCallBackExtendedConn) Upstream() any {
-	return c.extendedConn
-}
+var _ N.ExtendedConn = (*firstWriteCallBackConn)(nil)
 
 func NewFirstWriteCallBackConn(c C.Conn, callback func(error)) C.Conn {
-	if c, ok := c.(extendedConn); ok {
-		return &firstWriteCallBackExtendedConn{
-			extendedConn: c,
-			callback:     callback,
-			written:      false,
-		}
-	}
 	return &firstWriteCallBackConn{
 		Conn:     c,
 		callback: callback,

common/net/refconn.go

@@ -4,10 +4,12 @@ import (
 	"net"
 	"runtime"
 	"time"
+
+	"github.com/Dreamacro/clash/common/buf"
 )
 
 type refConn struct {
-	conn net.Conn
+	conn ExtendedConn
 	ref  any
 }
 
@@ -55,8 +57,20 @@ func (c *refConn) Upstream() any {
 	return c.conn
 }
 
+func (c *refConn) ReadBuffer(buffer *buf.Buffer) error {
+	defer runtime.KeepAlive(c.ref)
+	return c.conn.ReadBuffer(buffer)
+}
+
+func (c *refConn) WriteBuffer(buffer *buf.Buffer) error {
+	defer runtime.KeepAlive(c.ref)
+	return c.conn.WriteBuffer(buffer)
+}
+
+var _ ExtendedConn = (*refConn)(nil)
+
 func NewRefConn(conn net.Conn, ref any) net.Conn {
-	return &refConn{conn: conn, ref: ref}
+	return &refConn{conn: NewExtendedConn(conn), ref: ref}
 }
 
 type refPacketConn struct {