lpra/internal/agent/tunnel.go

package agent

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

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

// --- direct TCP mux tunnel ---

// runTunnel keeps a persistent connection to the host's tunnel at addr, dialing
// client-side ports on demand for "tcp" forward rules, until ctx is cancelled.
func (a *Agent) runTunnel(ctx context.Context, addr string) {
	backoff := time.Second
	for ctx.Err() == nil {
		if err := a.tunnelOnce(ctx, addr); err != nil {
			select {
			case <-time.After(backoff):
			case <-ctx.Done():
				return
			}
			if backoff < 15*time.Second {
				backoff *= 2
			}
			continue
		}
		backoff = time.Second
	}
}

func (a *Agent) tunnelOnce(ctx context.Context, addr string) error {
	d := net.Dialer{Timeout: 10 * time.Second}
	conn, err := d.DialContext(ctx, "tcp", addr)
	if err != nil {
		return err
	}
	defer conn.Close()
	// Unblock the read loop when the tunnel is stopped or moved.
	defer context.AfterFunc(ctx, func() { conn.Close() })()
	if tc, ok := conn.(*net.TCPConn); ok {
		_ = tc.SetNoDelay(true)
	}

	// Handshake: server sends "LPRA-TUN1 <nonceHex>", we reply "<agentID> <sig>".
	_ = conn.SetDeadline(time.Now().Add(10 * time.Second))
	line, err := readLineConn(conn, 512)
	if err != nil {
		return err
	}
	f := strings.Fields(line)
	if len(f) != 2 || f[0] != "LPRA-TUN1" {
		return fmt.Errorf("bad tunnel greeting")
	}
	nonce, err := hexDecode(f[1])
	if err != nil {
		return err
	}
	key := tunnel.AuthKey(a.state.Token)
	if _, err := fmt.Fprintf(conn, "%s %s\n", a.state.AgentID, tunnel.Sign(key, nonce)); err != nil {
		return err
	}
	_ = conn.SetDeadline(time.Time{})
	log.Printf("tunnel: connected to %s", addr)

	tc := &tunClient{conn: conn, streams: map[uint32]*agentStream{}}
	return tc.readLoop()
}

type tunClient struct {
	conn    net.Conn
	wmu     sync.Mutex
	mu      sync.Mutex
	streams map[uint32]*agentStream
}

// agentStream buffers downstream data in a channel so bytes that arrive before
// the client-side dial completes are not lost (the OPEN/DATA race).
type agentStream struct {
	in     chan []byte
	closed chan struct{}
	once   sync.Once
}

func (s *agentStream) close() { s.once.Do(func() { close(s.closed) }) }

func (t *tunClient) write(f tunnel.Frame) error {
	t.wmu.Lock()
	defer t.wmu.Unlock()
	return tunnel.WriteFrame(t.conn, f)
}

func (t *tunClient) readLoop() error {
	for {
		f, err := tunnel.ReadFrame(t.conn)
		if err != nil {
			t.shutdown()
			return err
		}
		switch f.Type {
		case tunnel.TypeOpen:
			if len(f.Payload) >= 2 {
				port := int(f.Payload[0])<<8 | int(f.Payload[1])
				t.mu.Lock()
				s := &agentStream{in: make(chan []byte, 256), closed: make(chan struct{})}
				t.streams[f.Stream] = s
				t.mu.Unlock()
				go t.openStream(f.Stream, port, s)
			}
		case tunnel.TypeData:
			t.mu.Lock()
			s := t.streams[f.Stream]
			t.mu.Unlock()
			if s != nil {
				select {
				case s.in <- f.Payload:
				case <-s.closed:
				}
			}
		case tunnel.TypeClose:
			t.closeStream(f.Stream)
		}
	}
}

func (t *tunClient) openStream(id uint32, port int, s *agentStream) {
	conn, err := net.DialTimeout("tcp", net.JoinHostPort("127.0.0.1", itoa(port)), 10*time.Second)
	if err != nil {
		_ = t.write(tunnel.Frame{Type: tunnel.TypeClose, Stream: id})
		t.closeStream(id)
		return
	}
	if tc, ok := conn.(*net.TCPConn); ok {
		_ = tc.SetNoDelay(true)
	}
	defer conn.Close()

	// downstream buffer -> local socket
	go func() {
		for {
			select {
			case b := <-s.in:
				if _, err := conn.Write(b); err != nil {
					s.close()
					return
				}
			case <-s.closed:
				return
			}
		}
	}()

	// local socket -> upstream frames
	buf := make([]byte, 32*1024)
	for {
		n, err := conn.Read(buf)
		if n > 0 {
			if t.write(tunnel.Frame{Type: tunnel.TypeData, Stream: id, Payload: append([]byte(nil), buf[:n]...)}) != nil {
				break
			}
		}
		if err != nil {
			break
		}
	}
	_ = t.write(tunnel.Frame{Type: tunnel.TypeClose, Stream: id})
	t.closeStream(id)
}

func (t *tunClient) closeStream(id uint32) {
	t.mu.Lock()
	s := t.streams[id]
	delete(t.streams, id)
	t.mu.Unlock()
	if s != nil {
		s.close()
	}
}

func (t *tunClient) shutdown() {
	t.mu.Lock()
	streams := t.streams
	t.streams = map[uint32]*agentStream{}
	t.mu.Unlock()
	for _, s := range streams {
		s.close()
	}
	t.conn.Close()
}

// --- direct UDP forwarding ---

// runUDP maintains the agent's UDP endpoint on the host's tunnel at addr and
// relays datagrams to/from client-side UDP ports for "udp" forward rules.
func (a *Agent) runUDP(ctx context.Context, addr string) {
	backoff := time.Second
	for ctx.Err() == nil {
		if err := a.udpOnce(ctx, addr); err != nil {
			select {
			case <-time.After(backoff):
			case <-ctx.Done():
				return
			}
			if backoff < 15*time.Second {
				backoff *= 2
			}
			continue
		}
		backoff = time.Second
	}
}

func (a *Agent) udpOnce(ctx context.Context, addr string) error {
	raddr, err := net.ResolveUDPAddr("udp", addr)
	if err != nil {
		return err
	}
	conn, err := net.DialUDP("udp", nil, raddr)
	if err != nil {
		return err
	}
	defer conn.Close()
	// Unblock the read loop when the tunnel is stopped or moved.
	defer context.AfterFunc(ctx, func() { conn.Close() })()

	key := tunnel.AuthKey(a.state.Token)
	register := func() {
		m := time.Now().Unix() / 60
		sig := tunnel.Sign(key, []byte(a.state.AgentID+"|"+itoa64(m)))
		_, _ = conn.Write([]byte("R" + a.state.AgentID + " " + sig))
	}
	register()

	// Keepalive registration so the host keeps our NAT mapping fresh.
	regCtx, cancel := context.WithCancel(ctx)
	defer cancel()
	go func() {
		t := time.NewTicker(20 * time.Second)
		defer t.Stop()
		for {
			select {
			case <-regCtx.Done():
				return
			case <-t.C:
				register()
			}
		}
	}()

	uc := &udpClient{server: conn, dials: map[uint32]*net.UDPConn{}}
	defer uc.shutdown()

	buf := make([]byte, 65535)
	for {
		n, err := conn.Read(buf)
		if err != nil {
			return err
		}
		if n < 7 || buf[0] != 'D' {
			continue
		}
		sid := binary.BigEndian.Uint32(buf[1:5])
		port := int(binary.BigEndian.Uint16(buf[5:7]))
		uc.send(sid, port, append([]byte(nil), buf[7:n]...))
	}
}

type udpClient struct {
	server *net.UDPConn // connection back to the host tunnel
	mu     sync.Mutex
	dials  map[uint32]*net.UDPConn // sid -> client-side socket
}

// send forwards a datagram to 127.0.0.1:port and ensures replies are relayed.
func (u *udpClient) send(sid uint32, port int, data []byte) {
	u.mu.Lock()
	c := u.dials[sid]
	u.mu.Unlock()
	if c == nil {
		raddr, err := net.ResolveUDPAddr("udp", net.JoinHostPort("127.0.0.1", itoa(port)))
		if err != nil {
			return
		}
		c, err = net.DialUDP("udp", nil, raddr)
		if err != nil {
			return
		}
		u.mu.Lock()
		u.dials[sid] = c
		u.mu.Unlock()
		go u.replyLoop(sid, c)
	}
	_, _ = c.Write(data)
}

func (u *udpClient) replyLoop(sid uint32, c *net.UDPConn) {
	buf := make([]byte, 65535)
	hdr := make([]byte, 7)
	hdr[0] = 'D'
	binary.BigEndian.PutUint32(hdr[1:5], sid)
	for {
		_ = c.SetReadDeadline(time.Now().Add(90 * time.Second))
		n, err := c.Read(buf)
		if err != nil {
			u.mu.Lock()
			delete(u.dials, sid)
			u.mu.Unlock()
			c.Close()
			return
		}
		out := append(append([]byte(nil), hdr...), buf[:n]...)
		_, _ = u.server.Write(out)
	}
}

func (u *udpClient) shutdown() {
	u.mu.Lock()
	dials := u.dials
	u.dials = map[uint32]*net.UDPConn{}
	u.mu.Unlock()
	for _, c := range dials {
		c.Close()
	}
}

// --- small helpers ---

func readLineConn(c net.Conn, max int) (string, error) {
	var b []byte
	one := make([]byte, 1)
	for len(b) < max {
		n, err := c.Read(one)
		if n > 0 {
			if one[0] == '\n' {
				return string(b), nil
			}
			b = append(b, one[0])
		}
		if err != nil {
			return "", err
		}
	}
	return string(b), nil
}

func hexDecode(s string) ([]byte, error) {
	out := make([]byte, len(s)/2)
	for i := 0; i < len(out); i++ {
		hi, err := hexVal(s[i*2])
		if err != nil {
			return nil, err
		}
		lo, err := hexVal(s[i*2+1])
		if err != nil {
			return nil, err
		}
		out[i] = hi<<4 | lo
	}
	return out, nil
}

func hexVal(c byte) (byte, error) {
	switch {
	case c >= '0' && c <= '9':
		return c - '0', nil
	case c >= 'a' && c <= 'f':
		return c - 'a' + 10, nil
	case c >= 'A' && c <= 'F':
		return c - 'A' + 10, nil
	}
	return 0, fmt.Errorf("bad hex")
}