lpra/internal/server/tunnelserver.go

package server

import (
	"crypto/rand"
	"encoding/hex"
	"fmt"
	"io"
	"log"
	"net"
	"strings"
	"sync"
	"time"

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

// readLine reads a single '\n'-terminated line, up to max bytes, one byte at a
// time so it never buffers past the delimiter into following binary frames.
func readLine(r io.Reader, max int) (string, error) {
	var b []byte
	buf := make([]byte, 1)
	for len(b) < max {
		n, err := r.Read(buf)
		if n > 0 {
			if buf[0] == '\n' {
				return string(b), nil
			}
			b = append(b, buf[0])
		}
		if err != nil {
			return "", err
		}
	}
	return string(b), nil
}

// tunnelHub tracks the live direct-tunnel connection for each agent and
// multiplexes forwarded TCP streams over it.
type tunnelHub struct {
	db     *DB
	mu     sync.Mutex
	ln     net.Listener
	agents map[string]*agentTunnel
}

func newTunnelHub(db *DB) *tunnelHub {
	return &tunnelHub{db: db, agents: map[string]*agentTunnel{}}
}

// agentTunnel is one authenticated direct tunnel to an agent.
type agentTunnel struct {
	agentID string
	conn    net.Conn
	wmu     sync.Mutex // serialize frame writes

	mu      sync.Mutex
	nextID  uint32
	streams map[uint32]net.Conn
}

// listen runs the TCP tunnel acceptor until the listener closes. Listening
// again moves the tunnel: the previous listener is closed and its agents are
// dropped, so they redial the new address.
func (h *tunnelHub) listen(addr string) error {
	ln, err := net.Listen("tcp", addr)
	if err != nil {
		return err
	}
	h.mu.Lock()
	old := h.ln
	h.ln = ln
	var stale []*agentTunnel
	if old != nil {
		for _, at := range h.agents {
			stale = append(stale, at)
		}
	}
	h.mu.Unlock()
	if old != nil {
		old.Close()
	}
	for _, at := range stale {
		at.conn.Close()
	}
	log.Printf("tunnel: listening on %s (direct TCP forwarding)", addr)
	go func() {
		for {
			conn, err := ln.Accept()
			if err != nil {
				return
			}
			go h.handleConn(conn)
		}
	}()
	return nil
}

func (h *tunnelHub) handleConn(conn net.Conn) {
	if tc, ok := conn.(*net.TCPConn); ok {
		_ = tc.SetNoDelay(true)
	}
	agentID, err := h.authenticate(conn)
	if err != nil {
		conn.Close()
		return
	}
	at := &agentTunnel{agentID: agentID, conn: conn, streams: map[uint32]net.Conn{}}

	h.mu.Lock()
	if old := h.agents[agentID]; old != nil {
		old.conn.Close()
	}
	h.agents[agentID] = at
	h.mu.Unlock()
	log.Printf("tunnel: agent %s connected", agentID)

	go at.keepalive()
	at.readLoop()

	h.mu.Lock()
	if h.agents[agentID] == at {
		delete(h.agents, agentID)
	}
	h.mu.Unlock()
	at.shutdown()
	log.Printf("tunnel: agent %s disconnected", agentID)
}

// authenticate performs the nonce/HMAC handshake and returns the agent id.
func (h *tunnelHub) authenticate(conn net.Conn) (string, error) {
	_ = conn.SetDeadline(time.Now().Add(10 * time.Second))
	nonce := make([]byte, 32)
	if _, err := rand.Read(nonce); err != nil {
		return "", err
	}
	if _, err := fmt.Fprintf(conn, "LPRA-TUN1 %s\n", hex.EncodeToString(nonce)); err != nil {
		return "", err
	}
	line, err := readLine(conn, 512)
	if err != nil {
		return "", err
	}
	parts := strings.Fields(strings.TrimSpace(line))
	if len(parts) != 2 {
		return "", fmt.Errorf("bad handshake")
	}
	agentID, sig := parts[0], parts[1]
	agent, err := h.db.GetAgent(agentID)
	if err != nil || agent == nil {
		return "", fmt.Errorf("unknown agent")
	}
	key, err := tunnel.KeyFromHash(agent.TokenHash)
	if err != nil {
		return "", err
	}
	if !tunnel.EqualSig(sig, tunnel.Sign(key, nonce)) {
		return "", fmt.Errorf("bad signature")
	}
	_ = conn.SetDeadline(time.Time{}) // clear
	return agentID, nil
}

func (at *agentTunnel) writeFrame(f tunnel.Frame) error {
	at.wmu.Lock()
	defer at.wmu.Unlock()
	return tunnel.WriteFrame(at.conn, f)
}

func (at *agentTunnel) keepalive() {
	t := time.NewTicker(20 * time.Second)
	defer t.Stop()
	for range t.C {
		if at.writeFrame(tunnel.Frame{Type: tunnel.TypePing}) != nil {
			return
		}
	}
}

// readLoop demultiplexes frames from the agent to their streams.
func (at *agentTunnel) readLoop() {
	for {
		f, err := tunnel.ReadFrame(at.conn)
		if err != nil {
			return
		}
		switch f.Type {
		case tunnel.TypeData:
			at.mu.Lock()
			c := at.streams[f.Stream]
			at.mu.Unlock()
			if c != nil {
				_, _ = c.Write(f.Payload)
			}
		case tunnel.TypeClose:
			at.closeStream(f.Stream)
		}
	}
}

// openStream wires a freshly accepted user connection to a new agent stream
// targeting 127.0.0.1:port on the client.
func (at *agentTunnel) openStream(port int, userConn net.Conn) {
	at.mu.Lock()
	at.nextID++
	id := at.nextID
	at.streams[id] = userConn
	at.mu.Unlock()

	// Tell the agent to dial the client-side port.
	if at.writeFrame(tunnel.Frame{Type: tunnel.TypeOpen, Stream: id, Payload: []byte{byte(port >> 8), byte(port)}}) != nil {
		at.closeStream(id)
		return
	}

	// user -> agent
	go func() {
		buf := make([]byte, 32*1024)
		for {
			n, err := userConn.Read(buf)
			if n > 0 {
				if at.writeFrame(tunnel.Frame{Type: tunnel.TypeData, Stream: id, Payload: append([]byte(nil), buf[:n]...)}) != nil {
					break
				}
			}
			if err != nil {
				break
			}
		}
		_ = at.writeFrame(tunnel.Frame{Type: tunnel.TypeClose, Stream: id})
		at.closeStream(id)
	}()
}

func (at *agentTunnel) closeStream(id uint32) {
	at.mu.Lock()
	c := at.streams[id]
	delete(at.streams, id)
	at.mu.Unlock()
	if c != nil {
		c.Close()
	}
}

func (at *agentTunnel) shutdown() {
	at.mu.Lock()
	conns := at.streams
	at.streams = map[uint32]net.Conn{}
	at.mu.Unlock()
	for _, c := range conns {
		c.Close()
	}
	at.conn.Close()
}

// get returns the live tunnel for an agent, or nil.
func (h *tunnelHub) get(agentID string) *agentTunnel {
	h.mu.Lock()
	defer h.mu.Unlock()
	return h.agents[agentID]
}