Add UDP proxy support

README.md

@@ -4,7 +4,7 @@
 
 - a web UI for viewing configured TCP proxies
 - a Bitcask-backed data store for persisted proxy definitions
-- a TCP proxy runner powered by `inet.af/tcpproxy`
+- a TCP proxy runner built with Go's standard `net` package
 
 This README is written for working in a GitHub Codespace.
 
@@ -29,7 +29,7 @@ The project is functional enough to:
 - start the web app
 - load stored proxy definitions from `./infinite.db`
 - create new noodles from the UI
-- run active TCP proxies
+- run active TCP and UDP proxies
 - restrict a proxy to a specific source IP
 - update expiration values in the database every second while a proxy is active
 - close and delete expired noodles automatically
@@ -39,7 +39,7 @@ The project is functional enough to:
 Current limitations:
 
 - there is no REST API; management is currently through the server-rendered UI
-- proxy routing is TCP only in the current code path
+- protocols are limited to TCP and UDP
 
 ## Running In A Codespace
 
@@ -78,6 +78,7 @@ The main table includes an add row for creating a proxy with:
 
 - `Name`
 - `Allow From`: accepts `All` or a specific IP address, with the current client IP suggested
+- `Proto`: choose `TCP` or `UDP`
 - `Listen Port`
 - `Destination Port`
 - `Destination Host`

go.mod

@@ -3,14 +3,13 @@ module infinite-noodle
 go 1.19
 
 require (
+	github.com/google/uuid v1.6.0
 	go.mills.io/bitcask/v2 v2.1.3
-	inet.af/tcpproxy v0.0.0-20231102063150-2862066fc2a9
 )
 
 require (
 	github.com/abcum/lcp v0.0.0-20201209214815-7a3f3840be81 // indirect
 	github.com/gofrs/flock v0.8.1 // indirect
-	github.com/google/uuid v1.6.0 // indirect
 	github.com/hashicorp/go-immutable-radix/v2 v2.0.0 // indirect
 	github.com/hashicorp/golang-lru/v2 v2.0.7 // indirect
 	github.com/mattetti/filebuffer v1.0.1 // indirect

go.sum

@@ -1,6 +1,5 @@
 github.com/abcum/lcp v0.0.0-20201209214815-7a3f3840be81 h1:uHogIJ9bXH75ZYrXnVShHIyywFiUZ7OOabwd9Sfd8rw=
 github.com/abcum/lcp v0.0.0-20201209214815-7a3f3840be81/go.mod h1:6ZvnjTZX1LNo1oLpfaJK8h+MXqHxcBFBIwkgsv+xlv0=
-github.com/armon/go-proxyproto v0.0.0-20210323213023-7e956b284f0a/go.mod h1:QmP9hvJ91BbJmGVGSbutW19IC0Q9phDCLGaomwTJbgU=
 github.com/davecgh/go-spew v1.1.0/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
 github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
 github.com/davecgh/go-spew v1.1.2-0.20180830191138-d8f796af33cc h1:U9qPSI2PIWSS1VwoXQT9A3Wy9MM3WgvqSxFWenqJduM=
@@ -50,5 +49,3 @@ gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405 h1:yhCVgyC4o1eVCa2tZl7eS0r+
 gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0=
 gopkg.in/yaml.v3 v3.0.0-20200313102051-9f266ea9e77c/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM=
 gopkg.in/yaml.v3 v3.0.1 h1:fxVm/GzAzEWqLHuvctI91KS9hhNmmWOoWu0XTYJS7CA=
-inet.af/tcpproxy v0.0.0-20231102063150-2862066fc2a9 h1:zomTWJvjwLbKRgGameQtpK6DNFUbZ2oNJuWhgUkGp3M=
-inet.af/tcpproxy v0.0.0-20231102063150-2862066fc2a9/go.mod h1:Tojt5kmHpDIR2jMojxzZK2w2ZR7OILODmUo2gaSwjrk=

internal/app/app.go

@@ -2,15 +2,16 @@ package app
 
 import (
 	"fmt"
+	"io"
 	"log"
 	"net"
 	"net/http"
+	"strings"
+	"sync"
 	"time"
 
 	"infinite-noodle/internal/noodle"
 	"infinite-noodle/internal/web"
-
-	"inet.af/tcpproxy"
 )
 
 type Config struct {
@@ -33,7 +34,7 @@ func Run(cfg Config) error {
 
 	go systemCheck(db, noodleChannel)
 	go expirationCheck(db, noodleChannel)
-	go tcpProxify(noodleChannel)
+	go proxify(noodleChannel)
 
 	http.Handle("/static/", http.StripPrefix("/static/", http.FileServer(http.FS(web.StaticFiles()))))
 	http.HandleFunc("/", web.HandleMain(db, &noodleChannel))
@@ -55,26 +56,19 @@ func systemCheck(db *noodle.Database, noodleChannel chan noodle.Noodle) {
 	}
 }
 
-func startProxy(proxy *tcpproxy.Proxy) {
+func proxify(noodleChannel chan noodle.Noodle) {
-	log.Print(proxy.Run())
+	noodleMap := make(map[string]managedProxy)
-}
-
-func tcpProxify(noodleChannel chan noodle.Noodle) {
-	noodleMap := make(map[string]*tcpproxy.Proxy)
 	for {
 		item := <-noodleChannel
 		_, running := noodleMap[item.Id]
 		if item.IsUp && !running {
-			var proxy tcpproxy.Proxy
+			proxy, err := newProxy(item)
-			src := fmt.Sprintf("0.0.0.0:%d", item.ListenPort)
+			if err != nil {
-			dst := fmt.Sprintf("%s:%d", item.DestHost, item.DestPort)
+				log.Print(err)
-			log.Printf("Starting a noodle from %s to %s with source=%s", src, dst, item.Src)
+				continue
-			proxy.AddRoute(src, sourceRestrictedTarget{
+			}
-				allowedIP: item.Src,
+			noodleMap[item.Id] = proxy
-				target:    tcpproxy.To(dst),
+			go proxy.Run()
-			})
-			noodleMap[item.Id] = &proxy
-			go startProxy(&proxy)
 			continue
 		}
 
@@ -88,30 +82,357 @@ func tcpProxify(noodleChannel chan noodle.Noodle) {
 	}
 }
 
-type sourceRestrictedTarget struct {
+type managedProxy interface {
-	allowedIP string
+	Run()
-	target    tcpproxy.Target
+	Close() error
 }
 
-func (t sourceRestrictedTarget) HandleConn(conn net.Conn) {
+func newProxy(item noodle.Noodle) (managedProxy, error) {
-	if t.allowedIP == "" || t.allowedIP == "All" {
+	listenAddr := fmt.Sprintf("0.0.0.0:%d", item.ListenPort)
-		t.target.HandleConn(conn)
+	targetAddr := fmt.Sprintf("%s:%d", item.DestHost, item.DestPort)
-		return
+	proto := normalizeProto(item.Proto)
-	}
+	log.Printf("Starting a %s noodle from %s to %s with source=%s", proto, listenAddr, targetAddr, item.Src)
 
-	host, _, err := net.SplitHostPort(conn.RemoteAddr().String())
+	switch proto {
+	case "UDP":
+		return newUDPNoodleProxy(listenAddr, targetAddr, item.Src)
+	default:
+		return newTCPNoodleProxy(listenAddr, targetAddr, item.Src)
+	}
+}
+
+func normalizeProto(proto string) string {
+	switch strings.ToUpper(strings.TrimSpace(proto)) {
+	case "UDP":
+		return "UDP"
+	default:
+		return "TCP"
+	}
+}
+
+type tcpNoodleProxy struct {
+	listenAddr string
+	targetAddr string
+	allowedIP  string
+	listener   net.Listener
+	mu         sync.Mutex
+	conns      map[net.Conn]struct{}
+}
+
+func newTCPNoodleProxy(listenAddr, targetAddr, allowedIP string) (*tcpNoodleProxy, error) {
+	ln, err := net.Listen("tcp", listenAddr)
 	if err != nil {
-		log.Printf("Rejected noodle connection with invalid remote address %q", conn.RemoteAddr().String())
+		return nil, err
-		conn.Close()
-		return
 	}
-	if host != t.allowedIP {
+	return &tcpNoodleProxy{
-		log.Printf("Rejected noodle connection from %s; allowed source is %s", host, t.allowedIP)
+		listenAddr: listenAddr,
-		conn.Close()
+		targetAddr: targetAddr,
+		allowedIP:  allowedIP,
+		listener:   ln,
+		conns:      make(map[net.Conn]struct{}),
+	}, nil
+}
+
+func (p *tcpNoodleProxy) Run() {
+	for {
+		conn, err := p.listener.Accept()
+		if err != nil {
+			if isClosedNetworkError(err) {
+				return
+			}
+			log.Printf("Accept failed for %s: %v", p.listenAddr, err)
+			return
+		}
+		go p.handleConn(conn)
+	}
+}
+
+func (p *tcpNoodleProxy) handleConn(src net.Conn) {
+	if !allowSource(p.allowedIP, src.RemoteAddr()) {
+		log.Printf("Rejected noodle connection from %s; allowed source is %s", src.RemoteAddr().String(), p.allowedIP)
+		src.Close()
 		return
 	}
 
-	t.target.HandleConn(conn)
+	dst, err := net.Dial("tcp", p.targetAddr)
+	if err != nil {
+		log.Printf("Dial failed from %s to %s: %v", p.listenAddr, p.targetAddr, err)
+		src.Close()
+		return
+	}
+
+	p.trackConn(src)
+	p.trackConn(dst)
+	defer p.untrackAndClose(src)
+	defer p.untrackAndClose(dst)
+
+	var wg sync.WaitGroup
+	wg.Add(2)
+
+	go func() {
+		defer wg.Done()
+		if _, err := io.Copy(dst, src); err != nil && !isClosedNetworkError(err) {
+			log.Printf("Proxy copy src->dst failed for %s to %s: %v", p.listenAddr, p.targetAddr, err)
+		}
+		closeWrite(dst)
+	}()
+
+	go func() {
+		defer wg.Done()
+		if _, err := io.Copy(src, dst); err != nil && !isClosedNetworkError(err) {
+			log.Printf("Proxy copy dst->src failed for %s to %s: %v", p.targetAddr, p.listenAddr, err)
+		}
+		closeWrite(src)
+	}()
+
+	wg.Wait()
+}
+
+func (p *tcpNoodleProxy) Close() error {
+	err := p.listener.Close()
+	p.mu.Lock()
+	defer p.mu.Unlock()
+	for conn := range p.conns {
+		_ = conn.Close()
+	}
+	return err
+}
+
+func (p *tcpNoodleProxy) trackConn(conn net.Conn) {
+	p.mu.Lock()
+	defer p.mu.Unlock()
+	p.conns[conn] = struct{}{}
+}
+
+func (p *tcpNoodleProxy) untrackAndClose(conn net.Conn) {
+	p.mu.Lock()
+	delete(p.conns, conn)
+	p.mu.Unlock()
+	_ = conn.Close()
+}
+
+type udpNoodleProxy struct {
+	listenAddr string
+	targetAddr string
+	allowedIP  string
+	conn       *net.UDPConn
+	mu         sync.Mutex
+	sessions   map[string]*udpSession
+	closed     bool
+}
+
+type udpSession struct {
+	clientAddr *net.UDPAddr
+	backend    *net.UDPConn
+	lastSeen   time.Time
+}
+
+const udpSessionIdleTimeout = 2 * time.Minute
+
+func newUDPNoodleProxy(listenAddr, targetAddr, allowedIP string) (*udpNoodleProxy, error) {
+	addr, err := net.ResolveUDPAddr("udp", listenAddr)
+	if err != nil {
+		return nil, err
+	}
+	conn, err := net.ListenUDP("udp", addr)
+	if err != nil {
+		return nil, err
+	}
+	return &udpNoodleProxy{
+		listenAddr: listenAddr,
+		targetAddr: targetAddr,
+		allowedIP:  allowedIP,
+		conn:       conn,
+		sessions:   make(map[string]*udpSession),
+	}, nil
+}
+
+func (p *udpNoodleProxy) Run() {
+	go p.cleanupIdleSessions()
+
+	buf := make([]byte, 65535)
+	for {
+		n, clientAddr, err := p.conn.ReadFromUDP(buf)
+		if err != nil {
+			if isClosedNetworkError(err) {
+				return
+			}
+			log.Printf("UDP read failed for %s: %v", p.listenAddr, err)
+			return
+		}
+		if !allowSource(p.allowedIP, clientAddr) {
+			log.Printf("Rejected noodle datagram from %s; allowed source is %s", clientAddr.String(), p.allowedIP)
+			continue
+		}
+
+		session, err := p.sessionFor(clientAddr)
+		if err != nil {
+			log.Printf("UDP session setup failed for %s to %s: %v", clientAddr.String(), p.targetAddr, err)
+			continue
+		}
+		payload := append([]byte(nil), buf[:n]...)
+		if _, err := session.backend.Write(payload); err != nil && !isClosedNetworkError(err) {
+			log.Printf("UDP write failed from %s to %s: %v", clientAddr.String(), p.targetAddr, err)
+			p.removeSession(clientAddr.String(), true)
+		}
+	}
+}
+
+func (p *udpNoodleProxy) sessionFor(clientAddr *net.UDPAddr) (*udpSession, error) {
+	key := clientAddr.String()
+
+	p.mu.Lock()
+	if session, ok := p.sessions[key]; ok {
+		session.lastSeen = time.Now()
+		p.mu.Unlock()
+		return session, nil
+	}
+	p.mu.Unlock()
+
+	targetAddr, err := net.ResolveUDPAddr("udp", p.targetAddr)
+	if err != nil {
+		return nil, err
+	}
+	backend, err := net.DialUDP("udp", nil, targetAddr)
+	if err != nil {
+		return nil, err
+	}
+
+	session := &udpSession{
+		clientAddr: clientAddr,
+		backend:    backend,
+		lastSeen:   time.Now(),
+	}
+
+	p.mu.Lock()
+	if existing, ok := p.sessions[key]; ok {
+		p.mu.Unlock()
+		_ = backend.Close()
+		return existing, nil
+	}
+	p.sessions[key] = session
+	p.mu.Unlock()
+
+	go p.relayBackend(session)
+	return session, nil
+}
+
+func (p *udpNoodleProxy) relayBackend(session *udpSession) {
+	buf := make([]byte, 65535)
+	for {
+		n, err := session.backend.Read(buf)
+		if err != nil {
+			if !isClosedNetworkError(err) {
+				log.Printf("UDP backend read failed for %s: %v", p.targetAddr, err)
+			}
+			p.removeSession(session.clientAddr.String(), false)
+			return
+		}
+		if _, err := p.conn.WriteToUDP(buf[:n], session.clientAddr); err != nil {
+			if !isClosedNetworkError(err) {
+				log.Printf("UDP client write failed for %s: %v", session.clientAddr.String(), err)
+			}
+			p.removeSession(session.clientAddr.String(), true)
+			return
+		}
+
+		p.mu.Lock()
+		if current, ok := p.sessions[session.clientAddr.String()]; ok {
+			current.lastSeen = time.Now()
+		}
+		p.mu.Unlock()
+	}
+}
+
+func (p *udpNoodleProxy) cleanupIdleSessions() {
+	ticker := time.NewTicker(30 * time.Second)
+	defer ticker.Stop()
+
+	for range ticker.C {
+		if p.isClosed() {
+			return
+		}
+		cutoff := time.Now().Add(-udpSessionIdleTimeout)
+		var expired []string
+
+		p.mu.Lock()
+		for key, session := range p.sessions {
+			if session.lastSeen.Before(cutoff) {
+				expired = append(expired, key)
+			}
+		}
+		p.mu.Unlock()
+
+		for _, key := range expired {
+			p.removeSession(key, true)
+		}
+	}
+}
+
+func (p *udpNoodleProxy) removeSession(key string, closeBackend bool) {
+	p.mu.Lock()
+	session, ok := p.sessions[key]
+	if ok {
+		delete(p.sessions, key)
+	}
+	p.mu.Unlock()
+	if ok && closeBackend {
+		_ = session.backend.Close()
+	}
+}
+
+func (p *udpNoodleProxy) isClosed() bool {
+	p.mu.Lock()
+	defer p.mu.Unlock()
+	return p.closed
+}
+
+func (p *udpNoodleProxy) Close() error {
+	p.mu.Lock()
+	if p.closed {
+		p.mu.Unlock()
+		return nil
+	}
+	p.closed = true
+	sessions := p.sessions
+	p.sessions = make(map[string]*udpSession)
+	conn := p.conn
+	p.mu.Unlock()
+
+	for _, session := range sessions {
+		_ = session.backend.Close()
+	}
+	if conn == nil {
+		return nil
+	}
+	return conn.Close()
+}
+
+func allowSource(allowedIP string, addr net.Addr) bool {
+	if allowedIP == "" || allowedIP == "All" {
+		return true
+	}
+	host, _, err := net.SplitHostPort(addr.String())
+	if err != nil {
+		return false
+	}
+	return host == allowedIP
+}
+
+func closeWrite(conn net.Conn) {
+	if tcpConn, ok := conn.(*net.TCPConn); ok {
+		_ = tcpConn.CloseWrite()
+	}
+}
+
+func isClosedNetworkError(err error) bool {
+	if err == nil {
+		return false
+	}
+	if net.ErrClosed != nil && err == net.ErrClosed {
+		return true
+	}
+	return strings.Contains(err.Error(), "use of closed network connection")
 }
 
 func expirationCheck(db *noodle.Database, noodleChannel chan noodle.Noodle) {
@@ -164,7 +485,7 @@ func runTestSequence(db *noodle.Database) {
 		item := noodle.Noodle{
 			Id:         db.MakeID(),
 			Name:       "Name_Test",
-			Proto:      "Proto_Test",
+			Proto:      "TCP",
 			Src:        "All",
 			ListenPort: 1080 + i,
 			DestPort:   22,

internal/assets/templates/index.html

@@ -55,7 +55,14 @@
             <tbody>
                 <tr>
                     <td><input class="input is-link is-small" type="text" form="add-noodle" name="name" placeholder="Name" autocomplete="off" required/></td>
-                    <td>TCP</td>
+                    <td>
+                        <div class="select is-small is-link">
+                            <select form="add-noodle" name="proto" autocomplete="off" required>
+                                <option value="TCP">TCP</option>
+                                <option value="UDP">UDP</option>
+                            </select>
+                        </div>
+                    </td>
                     <td>
                         <input class="input is-link is-small" type="text" form="add-noodle" name="src" list="allow-from-options" placeholder="All or source IP" value="All" autocomplete="off" required/>
                         <datalist id="allow-from-options">

internal/web/handlers.go

@@ -86,6 +86,12 @@ func HandleAdd(db *noodle.Database, pc *chan noodle.Noodle) func(w http.Response
 			return
 		}
 
+		proto := strings.ToUpper(strings.TrimSpace(req.FormValue("proto")))
+		if proto != "TCP" && proto != "UDP" {
+			http.Error(w, "invalid protocol", http.StatusBadRequest)
+			return
+		}
+
 		clientIP := clientIPFromRequest(req)
 		src := strings.TrimSpace(req.FormValue("src"))
 		if src == clientIP {
@@ -99,7 +105,7 @@ func HandleAdd(db *noodle.Database, pc *chan noodle.Noodle) func(w http.Response
 		item := noodle.Noodle{
 			Id:         db.MakeID(),
 			Name:       strings.TrimSpace(req.FormValue("name")),
-			Proto:      "TCP",
+			Proto:      proto,
 			Src:        src,
 			ListenPort: listenPort,
 			DestPort:   destPort,