drover-go/internal/redirect/tcp.go

package redirect

import (
	"context"
	"errors"
	"fmt"
	"io"
	"net"
	"sync"
	"time"

	"git.okcu.io/root/drover-go/internal/socks5"
)

// Config configures the TCP redirector.
type Config struct {
	SOCKS5 socks5.Config
	Bind   string // "127.0.0.1:0" — listener bind addr
}

type mapping struct {
	dstIP   net.IP
	dstPort uint16
	added   time.Time
}

// Redirector is the loopback listener that catches NAT-rewritten SYNs
// from divert and tunnels them through SOCKS5.
type Redirector struct {
	cfg   Config
	ln    net.Listener
	mu    sync.RWMutex
	flows map[uint16]mapping // src_port → mapping
	wg    sync.WaitGroup
	ctx   context.Context
	cnl   context.CancelFunc
}

// New starts a Redirector. It binds the listener but does not yet
// have any mappings; SetMapping is called by the divert layer when
// it sees an outbound SYN from a target PID.
func New(cfg Config) (*Redirector, error) {
	bind := cfg.Bind
	if bind == "" {
		bind = "127.0.0.1:0"
	}
	ln, err := net.Listen("tcp", bind)
	if err != nil {
		return nil, fmt.Errorf("listen %s: %w", bind, err)
	}
	ctx, cnl := context.WithCancel(context.Background())
	r := &Redirector{
		cfg:   cfg,
		ln:    ln,
		flows: map[uint16]mapping{},
		ctx:   ctx,
		cnl:   cnl,
	}
	r.wg.Add(1)
	go r.acceptLoop()
	r.wg.Add(1)
	go r.sweepLoop()
	return r, nil
}

func (r *Redirector) LocalAddr() string { return r.ln.Addr().String() }

func (r *Redirector) LocalPort() uint16 {
	return uint16(r.ln.Addr().(*net.TCPAddr).Port)
}

// SetMapping records that future TCP connections originating from
// src_port should be tunneled to dstIP:dstPort. Called by the divert
// layer at SYN time.
func (r *Redirector) SetMapping(srcPort uint16, dstIP net.IP, dstPort uint16) {
	r.mu.Lock()
	r.flows[srcPort] = mapping{dstIP: dstIP, dstPort: dstPort, added: time.Now()}
	r.mu.Unlock()
}

// GetMapping returns the original (dstIP, dstPort) for a recorded flow
// keyed by src port, or ok=false if no mapping exists. Used by the
// engine's diverterLoop on the return path to look up the original
// target port when rewriting packets going from the listener back to
// the client.
func (r *Redirector) GetMapping(srcPort uint16) (net.IP, uint16, bool) {
	r.mu.RLock()
	defer r.mu.RUnlock()
	m, ok := r.flows[srcPort]
	if !ok {
		return nil, 0, false
	}
	return m.dstIP, m.dstPort, true
}

// Close stops accepting and tears down active flows.
func (r *Redirector) Close() error {
	r.cnl()
	err := r.ln.Close()
	r.wg.Wait()
	return err
}

func (r *Redirector) acceptLoop() {
	defer r.wg.Done()
	for {
		c, err := r.ln.Accept()
		if err != nil {
			return
		}
		r.wg.Add(1)
		go r.handle(c)
	}
}

func (r *Redirector) handle(c net.Conn) {
	defer r.wg.Done()
	defer c.Close()

	srcPort := uint16(c.RemoteAddr().(*net.TCPAddr).Port)
	r.mu.RLock()
	m, ok := r.flows[srcPort]
	r.mu.RUnlock()
	if !ok {
		return // unknown flow; drop quietly
	}

	ctx, cancel := context.WithTimeout(r.ctx, 10*time.Second)
	defer cancel()

	host := m.dstIP.String()
	upstream, err := socks5.Dial(ctx, r.cfg.SOCKS5, host, m.dstPort)
	if err != nil {
		return
	}
	defer upstream.Close()

	pump(c, upstream)
}

func pump(a, b net.Conn) {
	var wg sync.WaitGroup
	wg.Add(2)
	go func() {
		defer wg.Done()
		// G1: upstream → client (read from b, write to a).
		_, _ = io.Copy(a, b)
		if cw, ok := a.(closeWriter); ok {
			cw.CloseWrite()
		}
		// We are exiting — force G2's read of a to unblock so the
		// pump tears down even if the peer half never closes.
		_ = a.SetReadDeadline(time.Now())
	}()
	go func() {
		defer wg.Done()
		// G2: client → upstream (read from a, write to b).
		_, _ = io.Copy(b, a)
		if cw, ok := b.(closeWriter); ok {
			cw.CloseWrite()
		}
		_ = b.SetReadDeadline(time.Now())
	}()
	wg.Wait()
}

type closeWriter interface{ CloseWrite() error }

// sweepLoop removes mappings older than 30 minutes (T-6 in spec).
func (r *Redirector) sweepLoop() {
	defer r.wg.Done()
	tk := time.NewTicker(time.Minute)
	defer tk.Stop()
	for {
		select {
		case <-r.ctx.Done():
			return
		case <-tk.C:
			cutoff := time.Now().Add(-30 * time.Minute)
			r.mu.Lock()
			for k, m := range r.flows {
				if m.added.Before(cutoff) {
					delete(r.flows, k)
				}
			}
			r.mu.Unlock()
		}
	}
}

// Sentinel for callers.
var ErrNotMapped = errors.New("redirector: source port has no mapping")