feat: initial commit

internal/server/tunnelserver.go

@@ -0,0 +1,241 @@
+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
+	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 8446 TCP tunnel acceptor until the listener closes.
+func (h *tunnelHub) listen(addr string) error {
+	ln, err := net.Listen("tcp", addr)
+	if err != nil {
+		return err
+	}
+	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]
+}