lpra/internal/server/udp.go

package server

import (
	"encoding/binary"
	"log"
	"net"
	"strings"
	"sync"
	"time"

	"github.com/x-grk/longpolling-remote-access/internal/tunnel"
)

// udpRelay carries UDP forwards over the direct UDP channel. The agent
// registers its UDP return address; end-user datagrams arriving on a rule's
// host port are relayed to the agent and replies are sent back.
//
// Datagram framing on the tunnel control socket:
//
//	registration (agent->server): "R <agentID> <sig>"
//	data         (both):          'D' | sid(4) | port(2) | payload
type udpRelay struct {
	db     *DB
	mu     sync.Mutex
	ctrl   *net.UDPConn
	gcOnce sync.Once

	agentAddr map[string]*net.UDPAddr // agentID -> last registered addr
	sessions  map[uint32]*udpSession
	nextSID   uint32
	rulePorts map[int64]*net.UDPConn // ruleID -> per-rule listener
}

type udpSession struct {
	id         uint32
	agentID    string
	ruleConn   *net.UDPConn
	userAddr   *net.UDPAddr
	clientPort int
	last       time.Time
}

func newUDPRelay(db *DB) *udpRelay {
	return &udpRelay{
		db:        db,
		agentAddr: map[string]*net.UDPAddr{},
		sessions:  map[uint32]*udpSession{},
		rulePorts: map[int64]*net.UDPConn{},
	}
}

// listen opens the shared UDP control socket. Listening again moves it: the
// previous socket is closed and agents register anew on the new address.
func (u *udpRelay) listen(addr string) error {
	ua, err := net.ResolveUDPAddr("udp", addr)
	if err != nil {
		return err
	}
	c, err := net.ListenUDP("udp", ua)
	if err != nil {
		return err
	}
	u.mu.Lock()
	old := u.ctrl
	u.ctrl = c
	if old != nil {
		u.agentAddr = map[string]*net.UDPAddr{}
	}
	u.mu.Unlock()
	if old != nil {
		old.Close()
	}
	log.Printf("tunnel: listening on %s/udp (direct UDP forwarding)", addr)
	go u.ctrlLoop(c)
	u.gcOnce.Do(func() { go u.gc() })
	return nil
}

func (u *udpRelay) ctrlLoop(ctrl *net.UDPConn) {
	buf := make([]byte, 65535)
	for {
		n, src, err := ctrl.ReadFromUDP(buf)
		if err != nil {
			return
		}
		if n == 0 {
			continue
		}
		switch buf[0] {
		case 'R':
			u.handleReg(string(buf[1:n]), src)
		case 'D':
			u.handleAgentData(buf[1:n])
		}
	}
}

// handleReg validates "agentID sig" and records the agent's UDP addr.
// sig = HMAC(key, "agentID|minute") accepting the current or previous minute.
func (u *udpRelay) handleReg(line string, src *net.UDPAddr) {
	f := strings.Fields(strings.TrimSpace(line))
	if len(f) != 2 {
		return
	}
	agentID, sig := f[0], f[1]
	agent, err := u.db.GetAgent(agentID)
	if err != nil || agent == nil {
		return
	}
	key, err := tunnel.KeyFromHash(agent.TokenHash)
	if err != nil {
		return
	}
	now := time.Now().Unix() / 60
	ok := false
	for _, m := range []int64{now, now - 1} {
		msg := []byte(agentID + "|" + itoa64(m))
		if tunnel.EqualSig(sig, tunnel.Sign(key, msg)) {
			ok = true
			break
		}
	}
	if !ok {
		return
	}
	u.mu.Lock()
	u.agentAddr[agentID] = src
	u.mu.Unlock()
}

// handleAgentData routes a reply datagram from an agent back to the end user.
func (u *udpRelay) handleAgentData(b []byte) {
	if len(b) < 6 {
		return
	}
	sid := binary.BigEndian.Uint32(b[0:4])
	payload := b[6:]
	u.mu.Lock()
	s := u.sessions[sid]
	if s != nil {
		s.last = time.Now()
	}
	u.mu.Unlock()
	if s != nil {
		_, _ = s.ruleConn.WriteToUDP(payload, s.userAddr)
	}
}

// openRulePort starts the per-rule UDP listener on the host port.
func (u *udpRelay) openRulePort(rule ForwardRule) error {
	ua, err := net.ResolveUDPAddr("udp", rule.ListenAddr+":"+itoa(rule.HostPort))
	if err != nil {
		return err
	}
	c, err := net.ListenUDP("udp", ua)
	if err != nil {
		return err
	}
	u.mu.Lock()
	u.rulePorts[rule.ID] = c
	u.mu.Unlock()
	go u.rulePortLoop(rule, c)
	return nil
}

func (u *udpRelay) closeRulePort(ruleID int64) {
	u.mu.Lock()
	c := u.rulePorts[ruleID]
	delete(u.rulePorts, ruleID)
	u.mu.Unlock()
	if c != nil {
		c.Close()
	}
}

func (u *udpRelay) rulePortLoop(rule ForwardRule, c *net.UDPConn) {
	buf := make([]byte, 65535)
	for {
		n, userAddr, err := c.ReadFromUDP(buf)
		if err != nil {
			return
		}
		u.mu.Lock()
		aAddr := u.agentAddr[rule.AgentID]
		ctrl := u.ctrl
		// Find existing session for this (rule, user).
		var sess *udpSession
		for _, s := range u.sessions {
			if s.ruleConn == c && udpAddrEqual(s.userAddr, userAddr) {
				sess = s
				break
			}
		}
		if sess == nil {
			u.nextSID++
			sess = &udpSession{
				id: u.nextSID, agentID: rule.AgentID, ruleConn: c,
				userAddr: userAddr, clientPort: rule.ClientPort,
			}
			u.sessions[sess.id] = sess
		}
		sess.last = time.Now()
		u.mu.Unlock()

		if aAddr == nil || ctrl == nil {
			continue // agent has not registered its UDP endpoint yet
		}
		// 'D' | sid(4) | clientPort(2) | payload
		out := make([]byte, 7+n)
		out[0] = 'D'
		binary.BigEndian.PutUint32(out[1:5], sess.id)
		binary.BigEndian.PutUint16(out[5:7], uint16(rule.ClientPort))
		copy(out[7:], buf[:n])
		_, _ = ctrl.WriteToUDP(out, aAddr)
	}
}

func (u *udpRelay) gc() {
	t := time.NewTicker(30 * time.Second)
	defer t.Stop()
	for range t.C {
		cut := time.Now().Add(-90 * time.Second)
		u.mu.Lock()
		for id, s := range u.sessions {
			if s.last.Before(cut) {
				delete(u.sessions, id)
			}
		}
		u.mu.Unlock()
	}
}

func udpAddrEqual(a, b *net.UDPAddr) bool {
	return a != nil && b != nil && a.Port == b.Port && a.IP.Equal(b.IP)
}

func itoa(n int) string { return itoa64(int64(n)) }
func itoa64(n int64) string {
	if n == 0 {
		return "0"
	}
	neg := n < 0
	if neg {
		n = -n
	}
	var b [24]byte
	i := len(b)
	for n > 0 {
		i--
		b[i] = byte('0' + n%10)
		n /= 10
	}
	if neg {
		i--
		b[i] = '-'
	}
	return string(b[i:])
}