redapricot/client/client.go

package client

import (
	"context"
	crand "crypto/rand"
	"encoding/binary"
	"fmt"
	"log"
	"net"
	"sync"
	"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

	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.pool = newWorkerPool(c, cfg.MaxConn)
	return c
}

// dialSession opens a TCP connection, performs the Intent-17 handshake, the
// Phase-A rekey, and reads SessionReady, returning an established frame conn.
func (c *Client) dialSession(magic byte) (*wire.FramedConn, error) {
	conn, err := net.DialTimeout("tcp", c.cfg.Server, 10*time.Second)
	if err != nil {
		return nil, err
	}
	if tcp, ok := conn.(*net.TCPConn); ok {
		_ = tcp.SetNoDelay(true)
	}
	ok := false
	defer func() {
		if !ok {
			_ = conn.Close()
		}
	}()

	// 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).
	rnd := make([]byte, 16)
	if _, err := crand.Read(rnd); err != nil {
		return nil, err
	}
	ts := time.Now().UnixMilli()
	rekeyMsg := wire.NewWriter().U8(magic).VarInt(len(rnd)).Bytes(rnd).I64(ts).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.
	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)
	}
	ok = true
	return fc, 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 {
	fc, err := c.dialSession(MagicControl)
	if err != nil {
		return fmt.Errorf("control connect: %w", err)
	}
	c.registerAll(fc)
	c.mu.Lock()
	c.ctrl = fc
	c.mu.Unlock()
	log.Printf("control session established with %s", c.cfg.Server)
	go c.serveControl(ctx, fc)
	go c.pingLoop(ctx, fc)
	return nil
}

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, fc *wire.FramedConn) {
	for {
		payload, err := fc.ReadFrame()
		if err != nil {
			break
		}
		c.dispatchControl(payload)
	}
	_ = 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(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()
		username, _ := r.String()
		go c.handleControlRequest(cid, pattern, ip, int(port), username)
	case CtlPong:
		// ignore
	default:
		log.Printf("control: unknown message type %d", t)
	}
}

func (c *Client) pingLoop(ctx context.Context, fc *wire.FramedConn) {
	ticker := time.NewTicker(time.Duration(c.cfg.PingIntervalMs) * time.Millisecond)
	defer ticker.Stop()
	for {
		select {
		case <-ctx.Done():
			return
		case <-ticker.C:
			msg := wire.NewWriter().U8(CtlPing).I64(time.Now().UnixMilli()).Out()
			if err := 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, username string) {
	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 as %q via pattern %q -> %s", ip, port, username, 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)
	wc.registerStream(sid, st)
	wc.sendSyn(sid, cid)
	go st.run()
}

// 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()
}