add context for conn

internal/net/proxyproto/conn.go

@@ -1,6 +1,7 @@
 package proxyproto
 
 import (
+	"context"
 	"net"
 	"strconv"
 
@@ -10,6 +11,11 @@ import (
 
 type serverConn struct {
 	net.Conn
+	ctx context.Context
+}
+
+func (c *serverConn) Context() context.Context {
+	return c.ctx
 }
 
 func (c *serverConn) RemoteAddr() net.Addr {
@@ -26,14 +32,6 @@ func (c *serverConn) LocalAddr() net.Addr {
 	return c.Conn.LocalAddr()
 }
 
-func (c *serverConn) SrcAddr() net.Addr {
-	return c.Conn.RemoteAddr()
-}
-
-func (c *serverConn) DstAddr() net.Addr {
-	return c.Conn.LocalAddr()
-}
-
 func (c *serverConn) CloseRead() error {
 	if sc, ok := c.Conn.(xio.CloseRead); ok {
 		return sc.CloseRead()

internal/net/proxyproto/listener.go

@@ -1,9 +1,11 @@
 package proxyproto
 
 import (
+	"context"
 	"net"
 	"time"
 
+	"github.com/go-gost/x/ctx"
 	proxyproto "github.com/pires/go-proxyproto"
 )
 
@@ -16,7 +18,18 @@ func (ln *listener) Accept() (net.Conn, error) {
 	if err != nil {
 		return nil, err
 	}
-	return &serverConn{Conn: conn}, nil
+
+	innerCtx := context.Background()
+	if c, ok := conn.(ctx.Context); ok {
+		if v := c.Context(); v != nil {
+			innerCtx = v
+		}
+	}
+
+	innerCtx = ctx.ContextWithSrcAddr(innerCtx, conn.RemoteAddr())
+	innerCtx = ctx.ContextWithDstAddr(innerCtx, conn.LocalAddr())
+
+	return &serverConn{Conn: conn, ctx: innerCtx}, nil
 }
 
 func WrapListener(ppv int, ln net.Listener, readHeaderTimeout time.Duration) net.Listener {