redapricot/client/client.go

package client

import (
	"context"
	crand "crypto/rand"
	"encoding/binary"
	"fmt"
	"log"
	"net"
	"sync"
	"sync/atomic"
	"time"

	"github.com/iceBear67/redapricot/client/wire"
)

// Client is a redapricot client: it holds a control session with the hub and a
// pool of worker connections used to serve player streams.
type Client struct {
	cfg        *Config
	pskBytes   []byte
	pskAddr    string
	serverPort uint16

	mappings map[string]Mapping // normalized pattern -> mapping
	pool     *WorkerPool

	// ctx/cancel own every pool dial: Close cancels it so an in-flight dial
	// aborts instead of holding a goroutine for the whole handshake timeout.
	// The control path uses the caller's context from Start, which is the same
	// shutdown signal by convention.
	ctx    context.Context
	cancel context.CancelFunc

	// closing is set by Close; a conn-loss teardown checks it and closes
	// streams outright rather than parking them for a reattach that is never
	// coming. Allocate also consults it through the pool's own flag.
	closing atomic.Bool

	streamWnd int     // our advertised per-stream receive window (bytes)
	shaper    *Shaper // caps aggregate egress to the hub; nil when unlimited
	chunk     int     // DATA payload cap; shrinks below DataChunkSize at low rates

	mu   sync.Mutex
	ctrl *wire.FramedConn
}

// New builds a client from config.
func New(cfg *Config) *Client {
	ctx, cancel := context.WithCancel(context.Background())
	c := &Client{
		cfg:      cfg,
		pskBytes: []byte(cfg.PSK),
		pskAddr:  wire.PSKAddress([]byte(cfg.PSK)),
		mappings: make(map[string]Mapping),
		ctx:      ctx,
		cancel:   cancel,
	}
	if _, portStr, err := net.SplitHostPort(cfg.Server); err == nil {
		if p, err := net.LookupPort("tcp", portStr); err == nil {
			c.serverPort = uint16(p)
		}
	}
	for _, m := range cfg.Mappings {
		c.mappings[NormalizeAddress(m.Pattern)] = m
	}
	c.streamWnd = clampWindow(cfg.StreamWindowBytes)
	// Parsed here rather than in LoadConfig because a Config may also be built
	// directly (tests). LoadConfig has already rejected a malformed value on the
	// file path, so a failure here can only come from a hand-built Config.
	bps, err := parseBandwidth(cfg.MaxBandwidth)
	if err != nil {
		log.Printf("client: %v; continuing without a bandwidth limit", err)
	}
	c.shaper = NewShaper(bps)
	c.chunk = c.shaper.chunkSize()
	if c.shaper != nil {
		log.Printf("egress shaped to %d B/s (burst %d B, chunk %d B)",
			bps, int64(c.shaper.burst), c.shaper.chunk)
	}
	c.pool = newWorkerPool(c, cfg.MaxConn)
	return c
}

func clampWindow(w int) int {
	if w <= 0 {
		return DefaultStreamWindow
	}
	if w < MinStreamWindow {
		return MinStreamWindow
	}
	if w > MaxStreamWindow {
		return MaxStreamWindow
	}
	return w
}

// session is an established redapricot session: the frame transport plus what
// was negotiated during establishment.
type session struct {
	fc        *wire.FramedConn
	peerWnd   int  // hub's advertised per-stream receive window
	heartbeat bool // hub accepted mux-level PING/PONG on worker conns
	resume    bool // hub accepted stream resumption (§7.5)
	// hubGrace is how long the hub will hang a parked player, as advertised in
	// SessionReady. Zero when resumption was not negotiated.
	hubGrace time.Duration
}

// dialSession opens a TCP connection, performs the Intent-17 handshake, the
// Phase-A rekey, and reads SessionReady. Per-stream flow control is mandatory:
// a hub that does not echo the STREAM_FC flag is rejected.
//
// The whole exchange is bounded by HandshakeTimeout. A hub that accepts the
// socket but never answers (wedged event loop, a load balancer accepting on its
// behalf) must fail fast rather than park the caller forever.
//
// ctx bounds the dial: the control path passes the caller's context so a
// shutdown mid-handshake aborts the attempt, and the pool passes the client's
// own context so Close cancels in-flight dials. After the dial, a cancelled
// ctx keeps aborting by closing the conn underneath the deadline-bounded
// handshake I/O.
func (c *Client) dialSession(ctx context.Context, magic byte) (sess *session, err error) {
	d := &net.Dialer{Timeout: HandshakeTimeout}
	conn, err := d.DialContext(ctx, "tcp", c.cfg.Server)
	if err != nil {
		return nil, err
	}
	// When ctx ends (shutdown), close the conn so the handshake below fails
	// immediately instead of waiting out its deadline.
	stop := context.AfterFunc(ctx, func() { _ = conn.Close() })
	defer stop()
	if tcp, ok := conn.(*net.TCPConn); ok {
		_ = tcp.SetNoDelay(true)
		_ = tcp.SetKeepAlive(true)
		_ = tcp.SetKeepAlivePeriod(TCPKeepAlivePeriod)
	}
	ok := false
	defer func() {
		if !ok {
			_ = conn.Close()
		}
	}()
	if err := conn.SetDeadline(time.Now().Add(HandshakeTimeout)); err != nil {
		return nil, err
	}

	// 1. plaintext Minecraft Handshake, Intent 17, address = hex(SHA3-224(PSK)).
	hs := wire.BuildHandshake(ProtocolVersion, c.pskAddr, c.serverPort, IntentRedapricot)
	if _, err := conn.Write(hs); err != nil {
		return nil, err
	}

	// 2. Phase-A ciphers derived from the PSK.
	fc := wire.NewFramedConn(conn,
		wire.CipherFor(c.pskBytes, wire.DirS2C), // in: server -> client
		wire.CipherFor(c.pskBytes, wire.DirC2S), // out: client -> server
	)

	// 3. Rekey frame (Phase A), including the mandatory feature flags and our
	// per-stream receive window.
	rnd := make([]byte, 16)
	if _, err := crand.Read(rnd); err != nil {
		return nil, err
	}
	ts := time.Now().UnixMilli()
	offered := FlagStreamFC | FlagWorkerHeartbeat
	if c.cfg.resumeEnabled() {
		offered |= FlagStreamResume
	}
	rekeyMsg := wire.NewWriter().U8(magic).VarInt(len(rnd)).Bytes(rnd).I64(ts).
		VarInt(offered).VarInt(c.streamWnd).Out()
	if err := fc.WriteFrame(rekeyMsg); err != nil {
		return nil, err
	}

	// 4. Switch to Phase-B ciphers: REKEY = Rand || Timestamp(I64 BE).
	rekey := make([]byte, 0, len(rnd)+8)
	rekey = append(rekey, rnd...)
	var tsb [8]byte
	binary.BigEndian.PutUint64(tsb[:], uint64(ts))
	rekey = append(rekey, tsb[:]...)
	fc.SwitchCiphers(
		wire.CipherFor(rekey, wire.DirS2C),
		wire.CipherFor(rekey, wire.DirC2S),
	)

	// 5. SessionReady: the type byte followed by the hub's accepted flags and
	// its per-stream receive window. Both are required.
	payload, err := fc.ReadFrame()
	if err != nil {
		return nil, err
	}
	if len(payload) < 1 || payload[0] != CtlSessionReady {
		return nil, fmt.Errorf("expected SessionReady, got %v", payload)
	}
	r := wire.NewReader(payload[1:])
	flags, ferr := r.VarInt()
	hubWnd, werr := r.VarInt()
	if ferr != nil || werr != nil || flags&FlagStreamFC == 0 || hubWnd <= 0 {
		return nil, fmt.Errorf("hub did not accept per-stream flow control (unsupported hub version?)")
	}
	if hubWnd > MaxStreamWindow {
		hubWnd = MaxStreamWindow
	}
	// Resumption is negotiated per connection, and the hub's grace period rides
	// along when it accepts. Our own grace is clamped strictly under the hub's:
	// the client must always give up first, or the hub drops a hanging player
	// while we are still mid-reattach. A hub that accepts the flag but omits the
	// grace is treated as not supporting it at all rather than guessed at.
	resume := flags&FlagStreamResume != 0
	var hubGrace time.Duration
	if resume {
		graceMs, gerr := r.VarInt()
		if gerr != nil || graceMs <= 0 {
			log.Printf("hub accepted stream resume without advertising a grace period; disabling resume")
			resume = false
		} else {
			hubGrace = time.Duration(graceMs) * time.Millisecond
		}
	}

	// The session is live: drop the establishment deadline. From here on
	// liveness is the heartbeat's job (and WriteFrame bounds each write).
	if err := conn.SetDeadline(time.Time{}); err != nil {
		return nil, err
	}
	ok = true
	return &session{
		fc:        fc,
		peerWnd:   hubWnd,
		heartbeat: flags&FlagWorkerHeartbeat != 0,
		resume:    resume,
		hubGrace:  hubGrace,
	}, nil
}

// statsOn reports whether performance diagnostics are enabled. When off, no
// counter struct is ever allocated and the instrumentation is a single branch.
func (c *Client) statsOn() bool { return c.cfg.StatsIntervalMs > 0 }

// Start establishes the control session and registers all patterns. It returns
// once the initial connection succeeds; subsequent drops are handled in the
// background with reconnect.
func (c *Client) Start(ctx context.Context) error {
	if err := c.connectControl(ctx); err != nil {
		return err
	}
	if c.statsOn() {
		go c.statsLoop(ctx.Done())
	}
	return nil
}

func (c *Client) connectControl(ctx context.Context) error {
	sess, err := c.dialSession(ctx, MagicControl)
	if err != nil {
		return fmt.Errorf("control connect: %w", err)
	}
	ctrl := &ctrlSession{fc: sess.fc}
	ctrl.lastPong.Store(time.Now().UnixMilli())
	c.registerAll(sess.fc)
	c.mu.Lock()
	c.ctrl = sess.fc
	c.mu.Unlock()
	log.Printf("control session established with %s", c.cfg.Server)
	go c.serveControl(ctx, ctrl)
	go c.pingLoop(ctx, ctrl)
	return nil
}

// ctrlSession tracks liveness for one control connection. A control session
// whose path dies silently must be detected, otherwise the hub keeps routing
// players to a session the client will never read from and nobody can connect.
type ctrlSession struct {
	fc       *wire.FramedConn
	lastPong atomic.Int64 // unix ms of the most recent Pong
}

func (c *Client) registerAll(fc *wire.FramedConn) {
	for pattern := range c.mappings {
		msg := wire.NewWriter().U8(CtlRegister).String(pattern).Out()
		if err := fc.WriteFrame(msg); err != nil {
			log.Printf("register %q: %v", pattern, err)
			return
		}
		log.Printf("registered pattern %q", pattern)
	}
}

func (c *Client) serveControl(ctx context.Context, ctrl *ctrlSession) {
	for {
		payload, err := ctrl.fc.ReadFrame()
		if err != nil {
			break
		}
		c.dispatchControl(ctrl, payload)
	}
	_ = ctrl.fc.Close()
	if ctx.Err() != nil || c.closing.Load() {
		return
	}
	// Reconnect with backoff, but try immediately first. While the control
	// session is down the hub has no live route for this client, so every
	// millisecond of delay is a player arriving to be told there is no such
	// server — and a session usually dies to a transient blip that the very next
	// dial would have survived. Sleeping first spent that window unconditionally.
	//
	// The wait is on ctx rather than time.Sleep so shutdown is not held up by a
	// backoff that has grown to the cap. Close() is authoritative on its own:
	// it gates the loop directly because a caller may shut the client down by
	// calling Close() without cancelling the context it passed to Start.
	for backoff := time.Duration(0); ctx.Err() == nil && !c.closing.Load(); {
		if backoff > 0 {
			select {
			case <-ctx.Done():
				return
			case <-time.After(backoff):
			}
		}
		if c.closing.Load() {
			return
		}
		if err := c.connectControl(ctx); err == nil {
			return
		} else {
			log.Printf("control reconnect failed: %v", err)
		}
		switch {
		case backoff == 0:
			backoff = 500 * time.Millisecond
		case backoff < maxControlBackoff:
			backoff = min(backoff*2, maxControlBackoff)
		}
	}
}

func (c *Client) dispatchControl(ctrl *ctrlSession, payload []byte) {
	r := wire.NewReader(payload)
	t, err := r.U8()
	if err != nil {
		return
	}
	switch t {
	case CtlSessionReady:
		// ignore
	case CtlRegisterAck:
		pattern, _ := r.String()
		status, _ := r.U8()
		switch status {
		case RegisterOk:
			log.Printf("pattern %q registered", pattern)
		case RegisterErrPattern:
			log.Printf("pattern %q rejected: not a valid regular expression", pattern)
		default:
			log.Printf("pattern %q rejected: status=%d", pattern, status)
		}
	case CtlControlRequest:
		cid, err := r.Bytes(CIDLen)
		if err != nil {
			return
		}
		pattern, _ := r.String()
		ip, _ := r.String()
		port, _ := r.U16()
		go c.handleControlRequest(cid, pattern, ip, int(port))
	case CtlPong:
		ctrl.lastPong.Store(time.Now().UnixMilli())
	default:
		log.Printf("control: unknown message type %d", t)
	}
}

// pingLoop keeps the control session alive and, crucially, verifies that the
// hub is still answering. A path that dies silently (no FIN/RST) would
// otherwise leave the read loop parked forever: the client would believe it is
// still registered while the hub routes players into the void.
func (c *Client) pingLoop(ctx context.Context, ctrl *ctrlSession) {
	ticker := time.NewTicker(c.cfg.pingInterval())
	defer ticker.Stop()
	timeout := c.cfg.heartbeatTimeout()
	for {
		select {
		case <-ctx.Done():
			return
		case <-ticker.C:
			last := time.UnixMilli(ctrl.lastPong.Load())
			if time.Since(last) > timeout {
				log.Printf("control session silent for %s; dropping it to force a reconnect", time.Since(last).Round(time.Second))
				_ = ctrl.fc.Close() // unblocks serveControl, which reconnects
				return
			}
			msg := wire.NewWriter().U8(CtlPing).I64(time.Now().UnixMilli()).Out()
			if err := ctrl.fc.WriteFrame(msg); err != nil {
				return
			}
		}
	}
}

// handleControlRequest reacts to a matched player: allocate a worker stream,
// SYN it, and bridge it to the mapped destination.
func (c *Client) handleControlRequest(cid []byte, pattern, ip string, port int) {
	mapping, ok := c.mappings[NormalizeAddress(pattern)]
	if !ok {
		log.Printf("control-request for unmapped pattern %q; ignoring", pattern)
		return
	}
	log.Printf("player %s:%d joined via pattern %q -> %s", ip, port, pattern, mapping.Destination)

	// Allocate and publish must agree on a live conn: Allocate hands out a
	// (conn, sid) pair that can die before we register on it, which would strand
	// the stream in a map nothing iterates. registerStream reports that, and we
	// simply pick another conn.
	var st *Stream
	var lg *leg
	for attempt := 0; attempt < allocateAttempts; attempt++ {
		wc, sid, err := c.pool.Allocate()
		if err != nil {
			log.Printf("worker allocate failed: %v", err)
			return
		}
		st = newStream(c, wc, sid, cid, mapping, ip, port)
		// Register before SYN so inbound DATA can never race ahead of the table,
		// and start the pumps before the (bounded) SYN write so a failed or slow
		// SYN cannot strand a stream that nothing would ever tear down.
		if wc.registerStream(sid, st) {
			lg = st.conn()
			break
		}
		st = nil
	}
	if st == nil {
		log.Printf("worker allocate failed: no live conn after %d attempts", allocateAttempts)
		return
	}

	go st.writeLoop()
	go st.run()
	if err := lg.wc.sendSyn(lg.sid, cid); err != nil {
		log.Printf("stream %s: SYN failed: %v", lg, err)
		st.teardown(false)
	}
}

// WorkerConnCount reports the current number of open worker connections
// (exposed for tests/observability).
func (c *Client) WorkerConnCount() int { return c.pool.count() }

// Close tears down the control session and all worker connections. Idempotent:
// a second call (or a Close racing a reconnect) only re-closes what is still
// open.
func (c *Client) Close() {
	c.closing.Store(true)
	c.cancel() // aborts in-flight pool dials
	c.mu.Lock()
	fc := c.ctrl
	c.mu.Unlock()
	if fc != nil {
		_ = fc.Close()
	}
	c.pool.closeAll()
	c.shaper.Stop()
}