rl544
108 lines
2.3 KiB
Go
108 lines
2.3 KiB
Go
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])
|
|
}
|
|
}
|