add context for conn

internal/util/ws/ws.go

@@ -1,12 +1,13 @@
 package ws
 
 import (
+	"context"
 	"net"
 	"sync"
 	"time"
 
+	ctx_pkg "github.com/go-gost/x/ctx"
 	xio "github.com/go-gost/x/internal/io"
-	xnet "github.com/go-gost/x/internal/net"
 	"github.com/gorilla/websocket"
 )
 
@@ -14,28 +15,31 @@ type WebsocketConn interface {
 	net.Conn
 	WriteMessage(int, []byte) error
 	ReadMessage() (int, []byte, error)
-	xnet.SrcAddr
 	xio.CloseRead
 	xio.CloseWrite
 }
 
 type websocketConn struct {
 	*websocket.Conn
-	rb      []byte
-	srcAddr net.Addr
-	mux     sync.Mutex
+	rb  []byte
+	ctx context.Context
+	mux sync.Mutex
 }
 
 func Conn(conn *websocket.Conn) WebsocketConn {
-	return &websocketConn{
-		Conn: conn,
+	ctx := context.Background()
+	if cc, ok := conn.NetConn().(ctx_pkg.Context); ok {
+		if cv := cc.Context(); cv != nil {
+			ctx = cv
+		}
 	}
+	return ContextConn(ctx, conn)
 }
 
-func ConnWithSrcAddr(conn *websocket.Conn, srcAddr net.Addr) WebsocketConn {
+func ContextConn(ctx context.Context, conn *websocket.Conn) WebsocketConn {
 	return &websocketConn{
-		Conn:    conn,
-		srcAddr: srcAddr,
+		Conn: conn,
+		ctx:  ctx,
 	}
 }
 
@@ -80,13 +84,6 @@ func (c *websocketConn) SetWriteDeadline(t time.Time) error {
 	return c.Conn.SetWriteDeadline(t)
 }
 
-func (c *websocketConn) SrcAddr() net.Addr {
-	if sa, ok := c.Conn.NetConn().(xnet.SrcAddr); ok {
-		return sa.SrcAddr()
-	}
-	return c.srcAddr
-}
-
 func (c *websocketConn) CloseRead() error {
 	if sc, ok := c.Conn.NetConn().(xio.CloseRead); ok {
 		return sc.CloseRead()
@@ -100,3 +97,7 @@ func (c *websocketConn) CloseWrite() error {
 	}
 	return xio.ErrUnsupported
 }
+
+func (c *websocketConn) Context() context.Context {
+	return c.ctx
+}