pub

internal/proxy/udp.go

@@ -0,0 +1,107 @@
+package proxy
+
+import (
+	"fmt"
+	"log"
+	"net"
+	"sync"
+	"time"
+)
+
+// udpSession tracks the state of a single client's UDP "connection"
+type udpSession struct {
+	clientAddr  *net.UDPAddr
+	backendConn *net.UDPConn
+	lastActive  time.Time
+}
+
+// RunUDPProxy starts a UDP reverse proxy, establishing a NAT-like session mapping for clients
+func RunUDPProxy(port int, targetBackend string) error {
+	addr := fmt.Sprintf(":%d", port)
+	proxyAddr, err := net.ResolveUDPAddr("udp", addr)
+	if err != nil {
+		return err
+	}
+
+	backendAddr, err := net.ResolveUDPAddr("udp", targetBackend)
+	if err != nil {
+		return err
+	}
+
+	proxyConn, err := net.ListenUDP("udp", proxyAddr)
+	if err != nil {
+		return err
+	}
+	defer proxyConn.Close()
+
+	log.Printf("Starting UDP proxy on %s, forwarding to %s\n", addr, targetBackend)
+
+	var mu sync.RWMutex
+	sessions := make(map[string]*udpSession)
+
+	// Background routine to clean up inactive sessions
+	go func() {
+		for {
+			time.Sleep(1 * time.Minute)
+			mu.Lock()
+			for k, s := range sessions {
+				// If a client hasn't sent or received packets in 3 minutes, close it
+				if time.Since(s.lastActive) > 3*time.Minute {
+					s.backendConn.Close()
+					delete(sessions, k)
+				}
+			}
+			mu.Unlock()
+		}
+	}()
+
+	buffer := make([]byte, 65535)
+	for {
+		n, clientAddr, err := proxyConn.ReadFromUDP(buffer)
+		if err != nil {
+			log.Printf("UDP Read error: %v", err)
+			continue
+		}
+
+		clientKey := clientAddr.String()
+		mu.RLock()
+		session, exists := sessions[clientKey]
+		mu.RUnlock()
+
+		if !exists {
+			// New client session detected, dial the backend
+			bConn, err := net.DialUDP("udp", nil, backendAddr)
+			if err != nil {
+				log.Printf("Failed to connect to UDP backend: %v", err)
+				continue
+			}
+
+			session = &udpSession{
+				clientAddr:  clientAddr,
+				backendConn: bConn,
+				lastActive:  time.Now(),
+			}
+
+			mu.Lock()
+			sessions[clientKey] = session
+			mu.Unlock()
+
+			// Start a goroutine to read from the backend and send back to this specific client
+			go func(s *udpSession) {
+				bBuf := make([]byte, 65535)
+				for {
+					bn, err := s.backendConn.Read(bBuf)
+					if err != nil {
+						break // Connection closed or error
+					}
+					s.lastActive = time.Now()
+					proxyConn.WriteToUDP(bBuf[:bn], s.clientAddr)
+				}
+			}(session)
+		}
+
+		// Forward the client's packet to the backend
+		session.lastActive = time.Now()
+		session.backendConn.Write(buffer[:n])
+	}
+}