From 05f8de1b8f9568c9b6007827baf6ca9a3fbcb49f Mon Sep 17 00:00:00 2001
From: Darien Raymond <admin@v2ray.com>
Date: Mon, 3 Dec 2018 17:41:24 +0100
Subject: [PATCH] allow custom controller function

---
 transport/internet/system_listener.go      | 55 +++++++++++++++-------
 transport/internet/system_listener_test.go | 29 ++++++++++++
 2 files changed, 67 insertions(+), 17 deletions(-)
 create mode 100644 transport/internet/system_listener_test.go

diff --git a/transport/internet/system_listener.go b/transport/internet/system_listener.go
index 364d2d401..26e42dae4 100644
--- a/transport/internet/system_listener.go
+++ b/transport/internet/system_listener.go
@@ -12,36 +12,57 @@ var (
 	effectiveListener = DefaultListener{}
 )
 
-type DefaultListener struct{}
+type controller func(network, address string, fd uintptr) error
 
-func (*DefaultListener) Listen(ctx context.Context, addr net.Addr, sockopt *SocketConfig) (net.Listener, error) {
-	var lc net.ListenConfig
+type DefaultListener struct {
+	contollers []controller
+}
 
-	if sockopt != nil {
-		lc.Control = func(network, address string, c syscall.RawConn) error {
-			return c.Control(func(fd uintptr) {
+func getControlFunc(ctx context.Context, sockopt *SocketConfig, contollers []controller) func(network, address string, c syscall.RawConn) error {
+	return func(network, address string, c syscall.RawConn) error {
+		return c.Control(func(fd uintptr) {
+			if sockopt != nil {
 				if err := applyInboundSocketOptions(network, fd, sockopt); err != nil {
 					newError("failed to apply socket options to incoming connection").Base(err).WriteToLog(session.ExportIDToError(ctx))
 				}
-			})
-		}
+			}
+
+			for _, controller := range contollers {
+				if err := controller(network, address, fd); err != nil {
+					newError("failed to apply external controller").Base(err).WriteToLog(session.ExportIDToError(ctx))
+				}
+			}
+		})
+	}
+}
+
+func (dl *DefaultListener) Listen(ctx context.Context, addr net.Addr, sockopt *SocketConfig) (net.Listener, error) {
+	var lc net.ListenConfig
+
+	if sockopt != nil || len(dl.contollers) > 0 {
+		lc.Control = getControlFunc(ctx, sockopt, dl.contollers)
 	}
 
 	return lc.Listen(ctx, addr.Network(), addr.String())
 }
 
-func (*DefaultListener) ListenPacket(ctx context.Context, addr net.Addr, sockopt *SocketConfig) (net.PacketConn, error) {
+func (dl *DefaultListener) ListenPacket(ctx context.Context, addr net.Addr, sockopt *SocketConfig) (net.PacketConn, error) {
 	var lc net.ListenConfig
 
-	if sockopt != nil {
-		lc.Control = func(network, address string, c syscall.RawConn) error {
-			return c.Control(func(fd uintptr) {
-				if err := applyInboundSocketOptions(network, fd, sockopt); err != nil {
-					newError("failed to apply socket options to incoming connection").Base(err).WriteToLog(session.ExportIDToError(ctx))
-				}
-			})
-		}
+	if sockopt != nil || len(dl.contollers) > 0 {
+		lc.Control = getControlFunc(ctx, sockopt, dl.contollers)
 	}
 
 	return lc.ListenPacket(ctx, addr.Network(), addr.String())
 }
+
+// RegisterListenerController adds a controller to the effective system listener.
+// The controller can be used to operate on file descriptors before they are put into use.
+func RegisterListenerController(controller func(network, address string, fd uintptr) error) error {
+	if controller == nil {
+		return newError("nil listener controller")
+	}
+
+	effectiveListener.contollers = append(effectiveListener.contollers, controller)
+	return nil
+}
diff --git a/transport/internet/system_listener_test.go b/transport/internet/system_listener_test.go
new file mode 100644
index 000000000..ae4136155
--- /dev/null
+++ b/transport/internet/system_listener_test.go
@@ -0,0 +1,29 @@
+package internet_test
+
+import (
+	"context"
+	"net"
+	"testing"
+
+	"v2ray.com/core/common"
+	"v2ray.com/core/transport/internet"
+)
+
+func TestRegisterListenerController(t *testing.T) {
+	var gotFd uintptr
+
+	common.Must(internet.RegisterListenerController(func(network string, addr string, fd uintptr) error {
+		gotFd = fd
+		return nil
+	}))
+
+	conn, err := internet.ListenSystemPacket(context.Background(), &net.UDPAddr{
+		IP: net.IPv4zero,
+	}, nil)
+	common.Must(err)
+	common.Must(conn.Close())
+
+	if gotFd == 0 {
+		t.Error("expected none-zero fd, but actually 0")
+	}
+}