update Listener interface

go.mod

@@ -11,7 +11,7 @@ require (
 	github.com/go-gost/gosocks4 v0.0.1
 	github.com/go-gost/gosocks5 v0.3.1-0.20211109033403-d894d75b7f09
 	github.com/go-gost/metrics v0.0.0-20220314135054-2263ae431a5f
-	github.com/go-gost/x v0.0.0-20220314140325-7915c0401c2e
+	github.com/go-gost/x v0.0.0-20220314141508-422e07f4e7b5
 	github.com/gobwas/glob v0.2.3
 	github.com/miekg/dns v1.1.45
 	github.com/sirupsen/logrus v1.8.1

go.sum

@@ -136,6 +136,8 @@ github.com/go-gost/tls-dissector v0.0.2-0.20211125135007-2b5d5bd9c07e h1:73NGqAs
 github.com/go-gost/tls-dissector v0.0.2-0.20211125135007-2b5d5bd9c07e/go.mod h1:/9QfdewqmHdaE362Hv5nDaSWLx3pCmtD870d6GaquXs=
 github.com/go-gost/x v0.0.0-20220314140325-7915c0401c2e h1:sQjjbnEQBgVWYvJREYb7wdVqFabHzaevIUaOjZzSv8U=
 github.com/go-gost/x v0.0.0-20220314140325-7915c0401c2e/go.mod h1:kA9Ndel1O6xj/Iva5LkhaOYpHmjYpr0HBoZmsnDGKso=
+github.com/go-gost/x v0.0.0-20220314141508-422e07f4e7b5 h1:PaJHhGtn/tlcze5mJHHMQbK/1vhpB2B8ebUBImPQG5I=
+github.com/go-gost/x v0.0.0-20220314141508-422e07f4e7b5/go.mod h1:hJLG+4EBF2sXXr+nE6MoDhNESdostnkSQKgq3Yw6naY=
 github.com/go-kit/kit v0.8.0/go.mod h1:xBxKIO96dXMWWy0MnWVtmwkA9/13aqxPnvrjFYMA2as=
 github.com/go-kit/kit v0.9.0/go.mod h1:xBxKIO96dXMWWy0MnWVtmwkA9/13aqxPnvrjFYMA2as=
 github.com/go-kit/log v0.1.0/go.mod h1:zbhenjAZHb184qTLMA9ZjW7ThYL0H2mk7Q6pNt4vbaY=

pkg/common/admission/conn.go

@@ -1,223 +0,0 @@
-package admission
-
-import (
-	"errors"
-	"io"
-	"net"
-	"syscall"
-
-	"github.com/go-gost/gost/v3/pkg/admission"
-)
-
-var (
-	errUnsupport = errors.New("unsupported operation")
-)
-
-type packetConn struct {
-	net.PacketConn
-	admission admission.Admission
-}
-
-func WrapPacketConn(admission admission.Admission, pc net.PacketConn) net.PacketConn {
-	if admission == nil {
-		return pc
-	}
-	return &packetConn{
-		PacketConn: pc,
-		admission:  admission,
-	}
-}
-
-func (c *packetConn) ReadFrom(p []byte) (n int, addr net.Addr, err error) {
-	for {
-		n, addr, err = c.PacketConn.ReadFrom(p)
-		if err != nil {
-			return
-		}
-
-		if c.admission != nil &&
-			!c.admission.Admit(addr.String()) {
-			continue
-		}
-
-		return
-	}
-}
-
-type udpConn struct {
-	net.PacketConn
-	admission admission.Admission
-}
-
-func WrapUDPConn(admission admission.Admission, pc net.PacketConn) UDPConn {
-	return &udpConn{
-		PacketConn: pc,
-		admission:  admission,
-	}
-}
-
-func (c *udpConn) RemoteAddr() net.Addr {
-	if nc, ok := c.PacketConn.(remoteAddr); ok {
-		return nc.RemoteAddr()
-	}
-	return nil
-}
-
-func (c *udpConn) SetReadBuffer(n int) error {
-	if nc, ok := c.PacketConn.(setBuffer); ok {
-		return nc.SetReadBuffer(n)
-	}
-	return errUnsupport
-}
-
-func (c *udpConn) SetWriteBuffer(n int) error {
-	if nc, ok := c.PacketConn.(setBuffer); ok {
-		return nc.SetWriteBuffer(n)
-	}
-	return errUnsupport
-}
-
-func (c *udpConn) Read(b []byte) (n int, err error) {
-	if nc, ok := c.PacketConn.(io.Reader); ok {
-		n, err = nc.Read(b)
-		return
-	}
-	err = errUnsupport
-	return
-}
-
-func (c *udpConn) ReadFrom(p []byte) (n int, addr net.Addr, err error) {
-	for {
-		n, addr, err = c.PacketConn.ReadFrom(p)
-		if err != nil {
-			return
-		}
-		if c.admission != nil &&
-			!c.admission.Admit(addr.String()) {
-			continue
-		}
-		return
-	}
-}
-
-func (c *udpConn) ReadFromUDP(b []byte) (n int, addr *net.UDPAddr, err error) {
-	if nc, ok := c.PacketConn.(readUDP); ok {
-		for {
-			n, addr, err = nc.ReadFromUDP(b)
-			if err != nil {
-				return
-			}
-			if c.admission != nil &&
-				!c.admission.Admit(addr.String()) {
-				continue
-			}
-			return
-		}
-	}
-	err = errUnsupport
-	return
-}
-
-func (c *udpConn) ReadMsgUDP(b, oob []byte) (n, oobn, flags int, addr *net.UDPAddr, err error) {
-	if nc, ok := c.PacketConn.(readUDP); ok {
-		for {
-			n, oobn, flags, addr, err = nc.ReadMsgUDP(b, oob)
-			if err != nil {
-				return
-			}
-			if c.admission != nil &&
-				!c.admission.Admit(addr.String()) {
-				continue
-			}
-			return
-		}
-	}
-	err = errUnsupport
-	return
-}
-
-func (c *udpConn) Write(b []byte) (n int, err error) {
-	if nc, ok := c.PacketConn.(io.Writer); ok {
-		n, err = nc.Write(b)
-		return
-	}
-	err = errUnsupport
-	return
-}
-
-func (c *udpConn) WriteTo(p []byte, addr net.Addr) (n int, err error) {
-	n, err = c.PacketConn.WriteTo(p, addr)
-	return
-}
-
-func (c *udpConn) WriteToUDP(b []byte, addr *net.UDPAddr) (n int, err error) {
-	if nc, ok := c.PacketConn.(writeUDP); ok {
-		n, err = nc.WriteToUDP(b, addr)
-		return
-	}
-	err = errUnsupport
-	return
-}
-
-func (c *udpConn) WriteMsgUDP(b, oob []byte, addr *net.UDPAddr) (n, oobn int, err error) {
-	if nc, ok := c.PacketConn.(writeUDP); ok {
-		n, oobn, err = nc.WriteMsgUDP(b, oob, addr)
-		return
-	}
-	err = errUnsupport
-	return
-}
-
-func (c *udpConn) SyscallConn() (rc syscall.RawConn, err error) {
-	if nc, ok := c.PacketConn.(syscallConn); ok {
-		return nc.SyscallConn()
-	}
-	err = errUnsupport
-	return
-}
-
-func (c *udpConn) SetDSCP(n int) error {
-	if nc, ok := c.PacketConn.(setDSCP); ok {
-		return nc.SetDSCP(n)
-	}
-	return nil
-}
-
-type UDPConn interface {
-	net.PacketConn
-	io.Reader
-	io.Writer
-	readUDP
-	writeUDP
-	setBuffer
-	syscallConn
-	remoteAddr
-}
-
-type setBuffer interface {
-	SetReadBuffer(bytes int) error
-	SetWriteBuffer(bytes int) error
-}
-
-type readUDP interface {
-	ReadFromUDP(b []byte) (n int, addr *net.UDPAddr, err error)
-	ReadMsgUDP(b, oob []byte) (n, oobn, flags int, addr *net.UDPAddr, err error)
-}
-
-type writeUDP interface {
-	WriteToUDP(b []byte, addr *net.UDPAddr) (int, error)
-	WriteMsgUDP(b, oob []byte, addr *net.UDPAddr) (n, oobn int, err error)
-}
-
-type syscallConn interface {
-	SyscallConn() (syscall.RawConn, error)
-}
-
-type remoteAddr interface {
-	RemoteAddr() net.Addr
-}
-
-// tcpraw.TCPConn
-type setDSCP interface {
-	SetDSCP(int) error
-}

pkg/common/admission/listener.go

@@ -1,37 +0,0 @@
-package admission
-
-import (
-	"net"
-
-	"github.com/go-gost/gost/v3/pkg/admission"
-)
-
-type listener struct {
-	net.Listener
-	admission admission.Admission
-}
-
-func WrapListener(admission admission.Admission, ln net.Listener) net.Listener {
-	if admission == nil {
-		return ln
-	}
-	return &listener{
-		Listener:  ln,
-		admission: admission,
-	}
-}
-
-func (ln *listener) Accept() (net.Conn, error) {
-	for {
-		c, err := ln.Listener.Accept()
-		if err != nil {
-			return nil, err
-		}
-		if ln.admission != nil &&
-			!ln.admission.Admit(c.RemoteAddr().String()) {
-			c.Close()
-			continue
-		}
-		return c, err
-	}
-}

pkg/handler/option.go

@@ -8,6 +8,7 @@ import (
 	"github.com/go-gost/gost/v3/pkg/bypass"
 	"github.com/go-gost/gost/v3/pkg/chain"
 	"github.com/go-gost/gost/v3/pkg/logger"
+	"github.com/go-gost/gost/v3/pkg/metadata"
 )
 
 type Options struct {
@@ -58,6 +59,13 @@ func LoggerOption(logger logger.Logger) Option {
 }
 
 type HandleOptions struct {
+	Metadata metadata.Metadata
 }
 
 type HandleOption func(opts *HandleOptions)
+
+func MetadataHandleOption(md metadata.Metadata) HandleOption {
+	return func(opts *HandleOptions) {
+		opts.Metadata = md
+	}
+}

pkg/listener/listener.go

@@ -14,7 +14,9 @@ var (
 // Listener is a server listener, just like a net.Listener.
 type Listener interface {
 	Init(metadata.Metadata) error
-	net.Listener
+	Accept() (net.Conn, metadata.Metadata, error)
+	Addr() net.Addr
+	Close() error
 }
 
 type AcceptError struct {

pkg/listener/rtcp/listener.go

@@ -57,10 +57,10 @@ func (l *rtcpListener) Init(md md.Metadata) (err error) {
 	return
 }
 
-func (l *rtcpListener) Accept() (conn net.Conn, err error) {
+func (l *rtcpListener) Accept() (conn net.Conn, md md.Metadata, err error) {
 	select {
 	case <-l.closed:
-		return nil, net.ErrClosed
+		return nil, nil, net.ErrClosed
 	default:
 	}
 
@@ -70,7 +70,7 @@ func (l *rtcpListener) Accept() (conn net.Conn, err error) {
 			connector.MuxBindOption(true),
 		)
 		if err != nil {
-			return nil, listener.NewAcceptError(err)
+			return nil, nil, listener.NewAcceptError(err)
 		}
 		l.ln = metrics.WrapListener(l.options.Service, l.ln)
 	}
@@ -78,7 +78,7 @@ func (l *rtcpListener) Accept() (conn net.Conn, err error) {
 	if err != nil {
 		l.ln.Close()
 		l.ln = nil
-		return nil, listener.NewAcceptError(err)
+		return nil, nil, listener.NewAcceptError(err)
 	}
 	return
 }

pkg/listener/rudp/listener.go

@@ -57,10 +57,10 @@ func (l *rudpListener) Init(md md.Metadata) (err error) {
 	return
 }
 
-func (l *rudpListener) Accept() (conn net.Conn, err error) {
+func (l *rudpListener) Accept() (conn net.Conn, md md.Metadata, err error) {
 	select {
 	case <-l.closed:
-		return nil, net.ErrClosed
+		return nil, nil, net.ErrClosed
 	default:
 	}
 
@@ -73,14 +73,14 @@ func (l *rudpListener) Accept() (conn net.Conn, err error) {
 			connector.UDPDataQueueSizeBindOption(l.md.readQueueSize),
 		)
 		if err != nil {
-			return nil, listener.NewAcceptError(err)
+			return nil, nil, listener.NewAcceptError(err)
 		}
 	}
 	conn, err = l.ln.Accept()
 	if err != nil {
 		l.ln.Close()
 		l.ln = nil
-		return nil, listener.NewAcceptError(err)
+		return nil, nil, listener.NewAcceptError(err)
 	}
 
 	if pc, ok := conn.(net.PacketConn); ok {

pkg/listener/tcp/listener.go

@@ -15,7 +15,7 @@ func init() {
 }
 
 type tcpListener struct {
-	net.Listener
+	ln      net.Listener
 	logger  logger.Logger
 	md      metadata
 	options listener.Options
@@ -42,8 +42,20 @@ func (l *tcpListener) Init(md md.Metadata) (err error) {
 		return
 	}
 
-	ln = metrics.WrapListener(l.options.Service, ln)
+	l.ln = metrics.WrapListener(l.options.Service, ln)
-	l.Listener = ln
 
 	return
 }
+
+func (l *tcpListener) Accept() (conn net.Conn, md md.Metadata, err error) {
+	conn, err = l.ln.Accept()
+	return
+}
+
+func (l *tcpListener) Addr() net.Addr {
+	return l.ln.Addr()
+}
+
+func (l *tcpListener) Close() error {
+	return l.ln.Close()
+}

pkg/listener/tls/listener.go

@@ -17,7 +17,7 @@ func init() {
 }
 
 type tlsListener struct {
-	net.Listener
+	ln      net.Listener
 	logger  logger.Logger
 	md      metadata
 	options listener.Options
@@ -46,7 +46,20 @@ func (l *tlsListener) Init(md md.Metadata) (err error) {
 	ln = metrics.WrapListener(l.options.Service, ln)
 	ln = admission.WrapListener(l.options.Admission, ln)
 
-	l.Listener = tls.NewListener(ln, l.options.TLSConfig)
+	l.ln = tls.NewListener(ln, l.options.TLSConfig)
 
 	return
 }
+
+func (l *tlsListener) Accept() (conn net.Conn, md md.Metadata, err error) {
+	conn, err = l.ln.Accept()
+	return
+}
+
+func (l *tlsListener) Addr() net.Addr {
+	return l.ln.Addr()
+}
+
+func (l *tlsListener) Close() error {
+	return l.ln.Close()
+}

pkg/listener/udp/listener.go

@@ -16,7 +16,7 @@ func init() {
 }
 
 type udpListener struct {
-	net.Listener
+	ln      net.Listener
 	logger  logger.Logger
 	md      metadata
 	options listener.Options
@@ -50,7 +50,7 @@ func (l *udpListener) Init(md md.Metadata) (err error) {
 	}
 	conn = metrics.WrapPacketConn(l.options.Service, conn)
 
-	l.Listener = udp.NewListener(
+	l.ln = udp.NewListener(
 		conn,
 		laddr,
 		l.md.backlog,
@@ -59,3 +59,16 @@ func (l *udpListener) Init(md md.Metadata) (err error) {
 		l.logger)
 	return
 }
+
+func (l *udpListener) Accept() (conn net.Conn, md md.Metadata, err error) {
+	conn, err = l.ln.Accept()
+	return
+}
+
+func (l *udpListener) Addr() net.Addr {
+	return l.ln.Addr()
+}
+
+func (l *udpListener) Close() error {
+	return l.ln.Close()
+}

pkg/service/service.go

@@ -71,7 +71,7 @@ func (s *service) Serve() error {
 
 	var tempDelay time.Duration
 	for {
-		conn, e := s.listener.Accept()
+		conn, md, e := s.listener.Accept()
 		if e != nil {
 			if ne, ok := e.(net.Error); ok && ne.Temporary() {
 				if tempDelay == 0 {
@@ -108,7 +108,11 @@ func (s *service) Serve() error {
 				metrics.RequestSeconds(s.name).Observe(time.Since(start).Seconds())
 			}()
 
-			if err := s.handler.Handle(context.Background(), conn); err != nil {
+			if err := s.handler.Handle(
+				context.Background(),
+				conn,
+				handler.MetadataHandleOption(md),
+			); err != nil {
 				s.options.logger.Error(err)
 				metrics.HandlerErrors(s.name).Inc()
 			}