add -dynamic-destination flag

path-network/go-mmproxy path-network#15

Author: yuki0308

README.md

@@ -57,6 +57,8 @@ Usage of ./go-mmproxy:
     	Path to a file that contains allowed subnets of the proxy servers
   -close-after int
     	Number of seconds after which UDP socket will be cleaned up (default 60)
+  -dynamic-destination
+		Traffic wil be forwarded to the destination specified in the PROXY protocol header
   -l string
     	Address the proxy listens on (default "0.0.0.0:8443")
   -listeners int

main.go

@@ -31,6 +31,7 @@ type options struct {
 	Logger             *slog.Logger
 	udpCloseAfter      int
 	UDPCloseAfter      time.Duration
+	DynamicDestination bool
 }
 
 var Opts options
@@ -49,6 +50,7 @@ func init() {
 	flag.IntVar(&Opts.Listeners, "listeners", 1,
 		"Number of listener sockets that will be opened for the listen address (Linux 3.9+)")
 	flag.IntVar(&Opts.udpCloseAfter, "close-after", 60, "Number of seconds after which UDP socket will be cleaned up")
+	flag.BoolVar(&Opts.DynamicDestination, "dynamic-destination", false, "Traffic will be forwarded to the destination specified in the PROXY protocol header")
 }
 
 func listen(listenerNum int, errors chan<- error) {

tcp.go

@@ -44,19 +44,23 @@ func tcpHandleConnection(conn net.Conn, logger *slog.Logger) {
 		return
 	}
 
-	saddr, _, restBytes, err := PROXYReadRemoteAddr(buffer[:n], TCP)
+	saddr, daddr, restBytes, err := PROXYReadRemoteAddr(buffer[:n], TCP)
 	if err != nil {
 		logger.Debug("failed to parse PROXY header", "error", err, slog.Bool("dropConnection", true))
 		return
 	}
 
 	targetAddr := Opts.TargetAddr6
-	if saddr == nil {
-		if netip.MustParseAddrPort(conn.RemoteAddr().String()).Addr().Is4() {
+	if !Opts.DynamicDestination {
+		if saddr == nil {
+			if netip.MustParseAddrPort(conn.RemoteAddr().String()).Addr().Is4() {
+				targetAddr = Opts.TargetAddr4
+			}
+		} else if netip.MustParseAddrPort(saddr.String()).Addr().Is4() {
 			targetAddr = Opts.TargetAddr4
 		}
-	} else if netip.MustParseAddrPort(saddr.String()).Addr().Is4() {
-		targetAddr = Opts.TargetAddr4
+	} else {
+		targetAddr = netip.MustParseAddrPort(daddr.String())
 	}
 
 	clientAddr := "UNKNOWN"

udp.go

@@ -82,7 +82,7 @@ func udpCopyFromUpstream(downstream net.PacketConn, conn *udpConnection) {
 	}
 }
 
-func udpGetSocketFromMap(downstream net.PacketConn, downstreamAddr, saddr net.Addr, logger *slog.Logger,
+func udpGetSocketFromMap(downstream net.PacketConn, downstreamAddr, saddr net.Addr, daddr net.Addr, logger *slog.Logger,
 	connMap map[string]*udpConnection, socketClosures chan<- string) (*udpConnection, error) {
 	connKey := ""
 	if saddr != nil {
@@ -94,8 +94,12 @@ func udpGetSocketFromMap(downstream net.PacketConn, downstreamAddr, saddr net.Ad
 	}
 
 	targetAddr := Opts.TargetAddr6
-	if netip.MustParseAddr(downstreamAddr.String()).Is4() {
-		targetAddr = Opts.TargetAddr4
+	if !Opts.DynamicDestination {
+		if netip.MustParseAddr(downstreamAddr.String()).Is4() {
+			targetAddr = Opts.TargetAddr4
+		}
+	} else {
+		targetAddr = netip.MustParseAddrPort(daddr.String())
 	}
 
 	logger = logger.With(slog.String("downstreamAddr", downstreamAddr.String()), slog.String("targetAddr", targetAddr.String()))
@@ -159,7 +163,7 @@ func UDPListen(listenConfig *net.ListenConfig, logger *slog.Logger, errors chan<
 			continue
 		}
 
-		saddr, _, restBytes, err := PROXYReadRemoteAddr(buffer[:n], UDP)
+		saddr, daddr, restBytes, err := PROXYReadRemoteAddr(buffer[:n], UDP)
 		if err != nil {
 			logger.Debug("failed to parse PROXY header", "error", err, slog.String("remoteAddr", remoteAddr.String()))
 			continue
@@ -178,7 +182,7 @@ func UDPListen(listenConfig *net.ListenConfig, logger *slog.Logger, errors chan<
 			}
 		}
 
-		conn, err := udpGetSocketFromMap(ln, remoteAddr, saddr, logger, connectionMap, socketClosures)
+		conn, err := udpGetSocketFromMap(ln, remoteAddr, saddr, daddr, logger, connectionMap, socketClosures)
 		if err != nil {
 			continue
 		}