From b60cf0260354dc96340fe9575c010ad08b71e141 Mon Sep 17 00:00:00 2001
From: RPRX <63339210+rprx@users.noreply.github.com>
Date: Wed, 20 Jan 2021 23:58:59 +0000
Subject: [PATCH] Optimize TPROXY Inbound UDP write back

Enhanced stability.
---
 proxy/dokodemo/dokodemo.go      |  8 ++--
 proxy/dokodemo/fakeudp_linux.go | 65 +++++++++++++--------------------
 2 files changed, 30 insertions(+), 43 deletions(-)

diff --git a/proxy/dokodemo/dokodemo.go b/proxy/dokodemo/dokodemo.go
index 91b0dc2f..5b6b0948 100644
--- a/proxy/dokodemo/dokodemo.go
+++ b/proxy/dokodemo/dokodemo.go
@@ -277,17 +277,17 @@ func (w *PacketWriter) WriteMultiBuffer(mb buf.MultiBuffer) error {
 					w.mark,
 				)
 				if err != nil {
+					newError(err).WriteToLog()
 					b.Release()
-					buf.ReleaseMulti(mb)
-					return err
+					continue
 				}
 				w.conns[*b.UDP] = conn
 			}
 			_, err = conn.WriteTo(b.Bytes(), w.back)
 			if err != nil {
-				conn.Close()
-				w.conns[*b.UDP] = nil
 				newError(err).WriteToLog()
+				w.conns[*b.UDP] = nil
+				conn.Close()
 			}
 			b.Release()
 		} else {
diff --git a/proxy/dokodemo/fakeudp_linux.go b/proxy/dokodemo/fakeudp_linux.go
index 7de71dab..455efb48 100644
--- a/proxy/dokodemo/fakeudp_linux.go
+++ b/proxy/dokodemo/fakeudp_linux.go
@@ -12,68 +12,55 @@ import (
 )
 
 func FakeUDP(addr *net.UDPAddr, mark int) (net.PacketConn, error) {
+	var af int
+	var sockaddr syscall.Sockaddr
 
-	localSocketAddress, af, err := udpAddrToSocketAddr(addr)
-	if err != nil {
-		return nil, &net.OpError{Op: "fake", Err: fmt.Errorf("build local socket address: %s", err)}
+	if len(addr.IP) == 4 {
+		af = syscall.AF_INET
+		sockaddr = &syscall.SockaddrInet4{Port: addr.Port}
+		copy(sockaddr.(*syscall.SockaddrInet4).Addr[:], addr.IP)
+	} else {
+		af = syscall.AF_INET6
+		sockaddr = &syscall.SockaddrInet6{Port: addr.Port}
+		copy(sockaddr.(*syscall.SockaddrInet6).Addr[:], addr.IP)
 	}
 
-	fileDescriptor, err := syscall.Socket(af, syscall.SOCK_DGRAM, 0)
-	if err != nil {
+	var fd int
+	var err error
+
+	if fd, err = syscall.Socket(af, syscall.SOCK_DGRAM, 0); err != nil {
 		return nil, &net.OpError{Op: "fake", Err: fmt.Errorf("socket open: %s", err)}
 	}
 
 	if mark != 0 {
-		if err = syscall.SetsockoptInt(fileDescriptor, syscall.SOL_SOCKET, syscall.SO_MARK, mark); err != nil {
-			syscall.Close(fileDescriptor)
+		if err = syscall.SetsockoptInt(fd, syscall.SOL_SOCKET, syscall.SO_MARK, mark); err != nil {
+			syscall.Close(fd)
 			return nil, &net.OpError{Op: "fake", Err: fmt.Errorf("set socket option: SO_MARK: %s", err)}
 		}
 	}
 
-	if err = syscall.SetsockoptInt(fileDescriptor, syscall.SOL_SOCKET, syscall.SO_REUSEADDR, 1); err != nil {
-		syscall.Close(fileDescriptor)
-		return nil, &net.OpError{Op: "fake", Err: fmt.Errorf("set socket option: SO_REUSEADDR: %s", err)}
-	}
-
-	if err = syscall.SetsockoptInt(fileDescriptor, syscall.SOL_SOCKET, unix.SO_REUSEPORT, 1); err != nil {
-		syscall.Close(fileDescriptor)
-		return nil, &net.OpError{Op: "fake", Err: fmt.Errorf("set socket option: SO_REUSEPORT: %s", err)}
-	}
-
-	if err = syscall.SetsockoptInt(fileDescriptor, syscall.SOL_IP, syscall.IP_TRANSPARENT, 1); err != nil {
-		syscall.Close(fileDescriptor)
+	if err = syscall.SetsockoptInt(fd, syscall.SOL_IP, syscall.IP_TRANSPARENT, 1); err != nil {
+		syscall.Close(fd)
 		return nil, &net.OpError{Op: "fake", Err: fmt.Errorf("set socket option: IP_TRANSPARENT: %s", err)}
 	}
 
-	if err = syscall.Bind(fileDescriptor, localSocketAddress); err != nil {
-		syscall.Close(fileDescriptor)
+	syscall.SetsockoptInt(fd, syscall.SOL_SOCKET, syscall.SO_REUSEADDR, 1)
+
+	syscall.SetsockoptInt(fd, syscall.SOL_SOCKET, unix.SO_REUSEPORT, 1)
+
+	if err = syscall.Bind(fd, sockaddr); err != nil {
+		syscall.Close(fd)
 		return nil, &net.OpError{Op: "fake", Err: fmt.Errorf("socket bind: %s", err)}
 	}
 
-	fdFile := os.NewFile(uintptr(fileDescriptor), fmt.Sprintf("net-udp-fake-%s", addr.String()))
+	fdFile := os.NewFile(uintptr(fd), fmt.Sprintf("net-udp-fake-%s", addr.String()))
 	defer fdFile.Close()
 
 	packetConn, err := net.FilePacketConn(fdFile)
 	if err != nil {
-		syscall.Close(fileDescriptor)
+		syscall.Close(fd)
 		return nil, &net.OpError{Op: "fake", Err: fmt.Errorf("convert file descriptor to connection: %s", err)}
 	}
 
 	return packetConn, nil
 }
-
-func udpAddrToSocketAddr(addr *net.UDPAddr) (syscall.Sockaddr, int, error) {
-	switch {
-	case addr.IP.To4() != nil:
-		ip := [4]byte{}
-		copy(ip[:], addr.IP.To4())
-
-		return &syscall.SockaddrInet4{Addr: ip, Port: addr.Port}, syscall.AF_INET, nil
-
-	default:
-		ip := [16]byte{}
-		copy(ip[:], addr.IP.To16())
-
-		return &syscall.SockaddrInet6{Addr: ip, Port: addr.Port, ZoneId: 0}, syscall.AF_INET6, nil
-	}
-}