q3rcon-proxy/udpproxy.go

// Package udpproxy implements a simple UDP proxy server that forwards rcon and query packets between clients and a target server.
package udpproxy

import (
	"fmt"
	"net"
	"time"

	log "github.com/sirupsen/logrus"
)

// Client represents a UDP proxy server that forwards rcon and query packets between clients and a target server.
// It maintains a session cache to manage client sessions and handles packet forwarding between clients and the target server.
type Client struct {
	laddr *net.UDPAddr
	raddr *net.UDPAddr

	proxyConn *net.UDPConn

	sessionCache   sessionCache
	sessionTimeout time.Duration
}

// New creates a new Client with the specified proxy and target addresses, and applies any provided options.
func New(proxy, target string, options ...Option) (*Client, error) {
	laddr, err := net.ResolveUDPAddr("udp", proxy)
	if err != nil {
		return nil, fmt.Errorf("invalid proxy address: %w", err)
	}

	raddr, err := net.ResolveUDPAddr("udp", target)
	if err != nil {
		return nil, fmt.Errorf("invalid target address: %w", err)
	}

	c := &Client{
		laddr:          laddr,
		raddr:          raddr,
		sessionCache:   newSessionCache(),
		sessionTimeout: 20 * time.Minute,
	}

	for _, o := range options {
		o(c)
	}

	return c, nil
}

// ListenAndServe starts the UDP proxy server and listens for incoming packets from clients.
// It reads packets from the proxy connection, checks for existing sessions, and forwards packets to the target server.
func (c *Client) ListenAndServe() error {
	var err error
	c.proxyConn, err = net.ListenUDP("udp", c.laddr)
	if err != nil {
		return fmt.Errorf("failed to listen on proxy address: %w", err)
	}

	go c.pruneSessions()

	buf := make([]byte, 2048)
	for {
		n, caddr, err := c.proxyConn.ReadFromUDP(buf)
		if err != nil {
			log.Error(err)
		}

		session, ok := c.sessionCache.read(caddr.String())
		if !ok {
			session, err = newSession(caddr, c.raddr, c.proxyConn)
			if err != nil {
				log.Error(err)
				continue
			}

			c.sessionCache.insert(caddr.String(), session)
		}

		go session.proxyTo(buf[:n])
	}
}

func (c *Client) pruneSessions() {
	ticker := time.NewTicker(1 * time.Minute)

	for range ticker.C {
		for _, session := range c.sessionCache.data {
			if time.Since(session.updateTime) > c.sessionTimeout {
				c.sessionCache.delete(session.caddr.String())
				log.Tracef("session for %s deleted", session.caddr)
			}
		}
	}
}