add context for conn

admission/wrapper/conn.go

@@ -8,7 +8,7 @@ import (
 	"syscall"
 
 	"github.com/go-gost/core/admission"
-	"github.com/go-gost/core/metadata"
+	"github.com/go-gost/x/ctx"
 	xio "github.com/go-gost/x/internal/io"
 	xnet "github.com/go-gost/x/internal/net"
 	"github.com/go-gost/x/internal/net/udp"
@@ -51,27 +51,6 @@ func (c *serverConn) SyscallConn() (rc syscall.RawConn, err error) {
 	return
 }
 
-func (c *serverConn) Metadata() metadata.Metadata {
-	if md, ok := c.Conn.(metadata.Metadatable); ok {
-		return md.Metadata()
-	}
-	return nil
-}
-
-func (c *serverConn) SrcAddr() net.Addr {
-	if sc, ok := c.Conn.(xnet.SrcAddr); ok {
-		return sc.SrcAddr()
-	}
-	return nil
-}
-
-func (c *serverConn) DstAddr() net.Addr {
-	if sc, ok := c.Conn.(xnet.DstAddr); ok {
-		return sc.DstAddr()
-	}
-	return nil
-}
-
 func (c *serverConn) CloseRead() error {
 	if sc, ok := c.Conn.(xio.CloseRead); ok {
 		return sc.CloseRead()
@@ -86,6 +65,13 @@ func (c *serverConn) CloseWrite() error {
 	return xio.ErrUnsupported
 }
 
+func (c *serverConn) Context() context.Context {
+	if innerCtx, ok := c.Conn.(ctx.Context); ok {
+		return innerCtx.Context()
+	}
+	return nil
+}
+
 type packetConn struct {
 	net.PacketConn
 	admission admission.Admission
@@ -117,9 +103,9 @@ func (c *packetConn) ReadFrom(p []byte) (n int, addr net.Addr, err error) {
 	}
 }
 
-func (c *packetConn) Metadata() metadata.Metadata {
-	if md, ok := c.PacketConn.(metadata.Metadatable); ok {
-		return md.Metadata()
+func (c *packetConn) Context() context.Context {
+	if innerCtx, ok := c.PacketConn.(ctx.Context); ok {
+		return innerCtx.Context()
 	}
 	return nil
 }
@@ -263,9 +249,9 @@ func (c *udpConn) SetDSCP(n int) error {
 	return nil
 }
 
-func (c *udpConn) Metadata() metadata.Metadata {
-	if md, ok := c.PacketConn.(metadata.Metadatable); ok {
-		return md.Metadata()
+func (c *udpConn) Context() context.Context {
+	if innerCtx, ok := c.PacketConn.(ctx.Context); ok {
+		return innerCtx.Context()
 	}
 	return nil
 }

admission/wrapper/listener.go

@@ -6,6 +6,7 @@ import (
 
 	"github.com/go-gost/core/admission"
 	"github.com/go-gost/core/logger"
+	xctx "github.com/go-gost/x/ctx"
 )
 
 type listener struct {
@@ -30,8 +31,21 @@ func (ln *listener) Accept() (net.Conn, error) {
 		if err != nil {
 			return nil, err
 		}
+
+		ctx := context.Background()
+		if cc, ok := c.(xctx.Context); ok {
+			if cv := cc.Context(); cv != nil {
+				ctx = cv
+			}
+		}
+
+		clientAddr := c.RemoteAddr()
+		if addr := xctx.SrcAddrFromContext(ctx); addr != nil {
+			clientAddr = addr
+		}
+
 		if ln.admission != nil &&
-			!ln.admission.Admit(context.Background(), c.RemoteAddr().String()) {
+			!ln.admission.Admit(ctx, clientAddr.String()) {
 			c.Close()
 			continue
 		}