udp.go

package netpipe

import (
	"fmt"
	"net"
	"sync"

	"github.com/google/uuid"
)

// UDP wire format (per packet - no length prefix needed):
//
//   [1 byte: flags][N bytes: body]
//
// Each UDP packet IS a message boundary, so no framing required.
// Flags byte is identical to TCP: bit 0 = encrypted.

const (
	// udpMaxPacket is the max safe UDP payload size.
	// 65507 is the theoretical max, but we cap at 64 KB for sanity.
	udpMaxPacket = 65507
)

// ---------------------------------------------------------------------------
// UDP Server
// ---------------------------------------------------------------------------

// udpClient tracks a UDP "client" by their remote address.
type udpClient struct {
	id   string
	addr *net.UDPAddr
}

// listenUDP starts the UDP server loop. Called by Server.Listen() when protocol is "udp".
func (s *Server) listenUDP() error {
	addr, err := net.ResolveUDPAddr("udp", fmt.Sprintf(":%d", s.config.Port))
	if err != nil {
		return fmt.Errorf("netpipe: resolve udp addr: %w", err)
	}

	conn, err := net.ListenUDP("udp", addr)
	if err != nil {
		return fmt.Errorf("netpipe: listen udp: %w", err)
	}
	s.udpConn = conn
	fmt.Printf("netpipe server listening on udp :%d\n", s.config.Port)

	buf := make([]byte, udpMaxPacket)
	for {
		n, remoteAddr, err := conn.ReadFromUDP(buf)
		if err != nil {
			if s.udpConn == nil {
				return nil // clean shutdown
			}
			continue
		}
		if n < 1 {
			continue // need at least the flags byte
		}

		// identify or register this client by address
		clientID := s.getOrRegisterUDPClient(remoteAddr)

		flags := buf[0]
		body := make([]byte, n-1)
		copy(body, buf[1:n])

		// Stage 7: pong from client
		if flags&flagPong != 0 {
			s.meta.touchPong(clientID)
			continue
		}

		// Ignore pings from clients
		if flags&flagPing != 0 {
			continue
		}

		// Stage 7: track data activity
		s.meta.touchData(clientID)

		// Stage 6: stream chunk - buffer and reassemble
		if flags&flagStream != 0 {
			streamID, index, total, chunk, err := parseStreamBody(body)
			if err != nil {
				continue
			}
			assembled, complete := s.streams.addChunkFrom(clientID, streamID, index, total, chunk)
			if complete && s.onStream != nil {
				s.onStream(clientID, assembled)
			}
			continue
		}

		data := body

		// auto-decrypt if flagged and key is set
		if flags&flagEncrypted != 0 && s.config.EncryptionKey != "" {
			plaintext, err := decrypt(body, s.config.EncryptionKey)
			if err != nil {
				continue // bad key or corrupt - skip
			}
			data = plaintext
		}

		if s.onData != nil {
			s.onData(clientID, data)
		}
	}
}

// getOrRegisterUDPClient finds or creates a client entry for a UDP remote address.
func (s *Server) getOrRegisterUDPClient(addr *net.UDPAddr) string {
	key := addr.String()

	s.udpMu.RLock()
	if client, ok := s.udpClients[key]; ok {
		s.udpMu.RUnlock()
		return client.id
	}
	s.udpMu.RUnlock()

	// new client
	id := uuid.New().String()
	s.udpMu.Lock()
	// double-check after acquiring write lock
	if client, ok := s.udpClients[key]; ok {
		s.udpMu.Unlock()
		return client.id
	}
	s.udpClients[key] = &udpClient{id: id, addr: addr}
	s.udpMu.Unlock()

	if s.onConnect != nil {
		s.onConnect(id)
	}
	return id
}

// sendUDP sends a UDP packet with flags to a specific client by UUID.
func (s *Server) sendUDP(clientID string, flags byte, body []byte) error {
	s.udpMu.RLock()
	var target *net.UDPAddr
	for _, c := range s.udpClients {
		if c.id == clientID {
			target = c.addr
			break
		}
	}
	s.udpMu.RUnlock()

	if target == nil {
		return fmt.Errorf("netpipe: udp client %s not found", clientID)
	}

	packet := make([]byte, 1+len(body))
	packet[0] = flags
	copy(packet[1:], body)

	_, err := s.udpConn.WriteToUDP(packet, target)
	return err
}

// kickUDP removes a UDP client from the registry and fires OnDisconnect.
func (s *Server) kickUDP(clientID string) error {
	s.udpMu.Lock()
	var found bool
	for key, c := range s.udpClients {
		if c.id == clientID {
			delete(s.udpClients, key)
			found = true
			break
		}
	}
	s.udpMu.Unlock()

	if !found {
		return fmt.Errorf("netpipe: udp client %s not found", clientID)
	}

	if s.onDisconnect != nil {
		s.onDisconnect(clientID)
	}
	return nil
}

// broadcastUDP sends a UDP packet to all known clients.
func (s *Server) broadcastUDP(flags byte, body []byte, excludeID string) {
	packet := make([]byte, 1+len(body))
	packet[0] = flags
	copy(packet[1:], body)

	s.udpMu.RLock()
	targets := make([]*net.UDPAddr, 0, len(s.udpClients))
	for _, c := range s.udpClients {
		if c.id != excludeID {
			targets = append(targets, c.addr)
		}
	}
	s.udpMu.RUnlock()

	for _, addr := range targets {
		s.udpConn.WriteToUDP(packet, addr)
	}
}

// ---------------------------------------------------------------------------
// UDP Client
// ---------------------------------------------------------------------------

// connectUDP sets up the UDP "connection" (it's really just a bound socket).
func (c *Client) connectUDP() error {
	addr, err := net.ResolveUDPAddr("udp", fmt.Sprintf("%s:%d", c.config.Host, c.config.Port))
	if err != nil {
		return fmt.Errorf("netpipe: resolve udp addr: %w", err)
	}

	conn, err := net.DialUDP("udp", nil, addr)
	if err != nil {
		return fmt.Errorf("netpipe: connect udp: %w", err)
	}
	c.conn = conn
	return nil
}

// sendUDP sends a UDP packet with flags.
func (c *Client) sendUDP(flags byte, body []byte) error {
	if c.conn == nil {
		return fmt.Errorf("netpipe: not connected")
	}
	packet := make([]byte, 1+len(body))
	packet[0] = flags
	copy(packet[1:], body)
	_, err := c.conn.Write(packet)
	return err
}

// listenUDP runs the UDP receive loop for the client.
func (c *Client) listenUDP() {
	if c.conn == nil {
		return
	}

	buf := make([]byte, udpMaxPacket)
	for {
		n, err := c.conn.Read(buf)
		if err != nil {
			break
		}
		if n < 1 {
			continue
		}

		flags := buf[0]
		body := make([]byte, n-1)
		copy(body, buf[1:n])

		// Stage 7: server ping - auto-respond with pong
		if flags&flagPing != 0 {
			c.sendPong()
			continue
		}

		// Ignore pong from server
		if flags&flagPong != 0 {
			continue
		}

		// Stage 6: stream chunk - buffer and reassemble
		if flags&flagStream != 0 {
			streamID, index, total, chunk, err := parseStreamBody(body)
			if err != nil {
				continue
			}
			assembled, complete := c.streams.addChunk(streamID, index, total, chunk)
			if complete && c.onStream != nil {
				c.onStream(assembled)
			}
			continue
		}

		data := body

		// auto-decrypt if flagged and key provided
		if flags&flagEncrypted != 0 && c.encryptionKey != "" {
			plaintext, err := decrypt(body, c.encryptionKey)
			if err != nil {
				continue
			}
			data = plaintext
		}

		if c.onData != nil {
			c.onData(data)
		}
	}

	c.conn = nil
	if c.onDisconnect != nil {
		c.onDisconnect()
	}
}

// ---------------------------------------------------------------------------
// UDP fields added to Server and Client - initialised in constructors
// ---------------------------------------------------------------------------

// udpFields holds UDP-specific server state. Embedded via composition
// in the Server struct (see server.go changes).
type udpFields struct {
	udpConn    *net.UDPConn
	udpClients map[string]*udpClient // key = "ip:port"
	udpMu      sync.RWMutex
}