initial commit

pkg/udpproxy/session.go

@@ -0,0 +1,66 @@
+package udpproxy
+
+import (
+	"log"
+	"net"
+	"time"
+)
+
+type Session struct {
+	serverConn *net.UDPConn
+	proxyConn  *net.UDPConn
+	caddr      *net.UDPAddr
+	updateTime time.Time
+}
+
+func createSession(caddr *net.UDPAddr, raddr *net.UDPAddr, proxyConn *net.UDPConn) (*Session, error) {
+	serverConn, err := net.DialUDP("udp", nil, raddr)
+	if err != nil {
+		return nil, err
+	}
+
+	session := &Session{
+		serverConn: serverConn,
+		proxyConn:  proxyConn,
+		caddr:      caddr,
+		updateTime: time.Now(),
+	}
+
+	go session.listen()
+
+	return session, nil
+}
+
+func (s *Session) listen() error {
+	for {
+		buf := make([]byte, 1024)
+		_, err := s.serverConn.Read(buf)
+		if err != nil {
+			log.Println(err)
+			continue
+		}
+
+		go s.proxyFrom(buf)
+	}
+}
+
+func (s *Session) proxyFrom(buf []byte) error {
+	s.updateTime = time.Now()
+	_, err := s.proxyConn.WriteToUDP(buf, s.caddr)
+	if err != nil {
+		return err
+	}
+
+	return nil
+}
+
+func (s *Session) proxyTo(buf []byte) error {
+	s.updateTime = time.Now()
+	_, err := s.serverConn.Write(buf)
+	if err != nil {
+		log.Println(err)
+		return err
+	}
+
+	return nil
+}

pkg/udpproxy/udpproxy.go

@@ -0,0 +1,84 @@
+package udpproxy
+
+import (
+	"log"
+	"net"
+	"sync"
+	"time"
+)
+
+type Client struct {
+	laddr *net.UDPAddr
+	raddr *net.UDPAddr
+
+	proxyConn *net.UDPConn
+
+	mutex    sync.RWMutex
+	sessions map[string]*Session
+}
+
+func New(port, target string) (*Client, error) {
+	laddr, err := net.ResolveUDPAddr("udp", port)
+	if err != nil {
+		return nil, err
+	}
+
+	raddr, err := net.ResolveUDPAddr("udp", target)
+	if err != nil {
+		return nil, err
+	}
+
+	return &Client{
+		laddr:    laddr,
+		raddr:    raddr,
+		mutex:    sync.RWMutex{},
+		sessions: map[string]*Session{},
+	}, nil
+}
+
+func (c *Client) ListenAndServe() error {
+	var err error
+	c.proxyConn, err = net.ListenUDP("udp", c.laddr)
+	if err != nil {
+		return err
+	}
+
+	go c.pruneSessions()
+
+	for {
+		buf := make([]byte, 2048)
+		_, caddr, err := c.proxyConn.ReadFromUDP(buf)
+		if err != nil {
+			log.Println(err)
+		}
+
+		session, found := c.sessions[caddr.String()]
+		if !found {
+			session, err = createSession(caddr, c.raddr, c.proxyConn)
+			if err != nil {
+				log.Println(err)
+				continue
+			}
+
+			c.sessions[caddr.String()] = session
+		}
+
+		go session.proxyTo(buf)
+	}
+}
+
+func (c *Client) pruneSessions() {
+	ticker := time.NewTicker(1 * time.Minute)
+
+	// the locks here could be abusive and i dont even know if this is a real
+	// problem but we definitely need to clean up stale sessions
+	for range ticker.C {
+		for _, session := range c.sessions {
+			c.mutex.RLock()
+			if time.Since(session.updateTime) > time.Minute*5 {
+				delete(c.sessions, session.caddr.String())
+			}
+			c.mutex.RUnlock()
+		}
+	}
+}

pkg/udpproxy/udpproxy_test.go

@@ -0,0 +1,83 @@
+package udpproxy
+
+import (
+	"log"
+	"net"
+	"testing"
+	"time"
+)
+
+func TestSendAndReceive(t *testing.T) {
+	go runLilProxy()
+	go runUDPServer()
+
+	paddr, err := net.ResolveUDPAddr("udp", "localhost:9000")
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	conn, err := net.DialUDP("udp", nil, paddr)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	go func() {
+		for {
+			buf := make([]byte, 2048)
+			_, _, err = conn.ReadFromUDP(buf)
+			if err != nil {
+				log.Fatal(err)
+			}
+
+			log.Printf("response received: %s", string(buf))
+		}
+	}()
+
+	for {
+		time.Sleep(1 * time.Second)
+		_, err = conn.Write([]byte("hi\n"))
+		if err != nil {
+			log.Fatal(err)
+		}
+	}
+}
+
+func runLilProxy() {
+	port := ":9000"
+	target := "localhost:9001"
+
+	c, err := New(port, target)
+	if err != nil {
+		log.Fatal(err)
+	}
+
+	log.Fatal(c.ListenAndServe())
+}
+
+func runUDPServer() {
+	taddr, err := net.ResolveUDPAddr("udp", ":9001")
+	if err != nil {
+		log.Fatal(err)
+	}
+
+	conn, err := net.ListenUDP("udp", taddr)
+	if err != nil {
+		log.Fatal(err)
+	}
+
+	for {
+		buf := make([]byte, 2048)
+		_, caddr, err := conn.ReadFromUDP(buf)
+		if err != nil {
+			log.Fatal(err)
+		}
+
+		log.Printf("request received: %s", string(buf))
+
+		_, err = conn.WriteToUDP([]byte("bye\n"), caddr)
+		if err != nil {
+			log.Fatal(err)
+		}
+
+	}
+}