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

	streamWnd int // our advertised per-stream receive window (bytes)

	mu   sync.Mutex
	ctrl *wire.FramedConn
}

// New builds a client from config.
func New(cfg *Config) *Client {
	c := &Client{
		cfg:      cfg,
		pskBytes: []byte(cfg.PSK),
		pskAddr:  wire.PSKAddress([]byte(cfg.PSK)),
		mappings: make(map[string]Mapping),
	}
	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)
	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
}

// 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.
func (c *Client) dialSession(magic byte) (sess *session, err error) {
	conn, err := net.DialTimeout("tcp", c.cfg.Server, HandshakeTimeout)
	if err != nil {
		return nil, err
	}
	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
	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
	}

	// 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}, nil
}

// 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 {
	return c.connectControl(ctx)
}

func (c *Client) connectControl(ctx context.Context) error {
	sess, err := c.dialSession(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 {
		return
	}
	// Reconnect with backoff.
	for backoff := 500 * time.Millisecond; ctx.Err() == nil; backoff *= 2 {
		if backoff > 10*time.Second {
			backoff = 10 * time.Second
		}
		time.Sleep(backoff)
		if err := c.connectControl(ctx); err == nil {
			return
		} else {
			log.Printf("control reconnect failed: %v", err)
		}
	}
}

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()
		log.Printf("register ack %q 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)
	wc, sid, err := c.pool.Allocate()
	if err != nil {
		log.Printf("worker allocate failed: %v", err)
		return
	}
	st := newStream(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.
	wc.registerStream(sid, st)
	go st.writeLoop()
	go st.run()
	if err := wc.sendSyn(sid, cid); err != nil {
		log.Printf("stream %d: SYN failed: %v", sid, 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.
func (c *Client) Close() {
	c.mu.Lock()
	fc := c.ctrl
	c.mu.Unlock()
	if fc != nil {
		_ = fc.Close()
	}
	c.pool.closeAll()
}