add credit based mux window

This commit is contained in:
iceBear67
2026-07-15 14:59:32 +00:00
parent ada07e0e36
commit bbe4efbe16
15 changed files with 651 additions and 139 deletions
+45 -14
View File
@@ -24,6 +24,8 @@ type Client struct {
mappings map[string]Mapping // normalized pattern -> mapping
pool *WorkerPool
streamWnd int // our advertised per-stream receive window (bytes)
mu sync.Mutex
ctrl *wire.FramedConn
}
@@ -44,16 +46,32 @@ func New(cfg *Config) *Client {
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
}
// 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) {
// Phase-A rekey, and reads SessionReady, returning an established frame conn
// and the hub's advertised per-stream receive window. Per-stream flow control
// is mandatory: a hub that does not echo the STREAM_FC flag is rejected.
func (c *Client) dialSession(magic byte) (fc *wire.FramedConn, peerWnd int, err error) {
conn, err := net.DialTimeout("tcp", c.cfg.Server, 10*time.Second)
if err != nil {
return nil, err
return nil, 0, err
}
if tcp, ok := conn.(*net.TCPConn); ok {
_ = tcp.SetNoDelay(true)
@@ -68,24 +86,26 @@ func (c *Client) dialSession(magic byte) (*wire.FramedConn, error) {
// 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
return nil, 0, err
}
// 2. Phase-A ciphers derived from the PSK.
fc := wire.NewFramedConn(conn,
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).
// 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
return nil, 0, err
}
ts := time.Now().UnixMilli()
rekeyMsg := wire.NewWriter().U8(magic).VarInt(len(rnd)).Bytes(rnd).I64(ts).Out()
rekeyMsg := wire.NewWriter().U8(magic).VarInt(len(rnd)).Bytes(rnd).I64(ts).
VarInt(FlagStreamFC).VarInt(c.streamWnd).Out()
if err := fc.WriteFrame(rekeyMsg); err != nil {
return nil, err
return nil, 0, err
}
// 4. Switch to Phase-B ciphers: REKEY = Rand || Timestamp(I64 BE).
@@ -99,16 +119,26 @@ func (c *Client) dialSession(magic byte) (*wire.FramedConn, error) {
wire.CipherFor(rekey, wire.DirC2S),
)
// 5. SessionReady.
// 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
return nil, 0, err
}
if len(payload) < 1 || payload[0] != CtlSessionReady {
return nil, fmt.Errorf("expected SessionReady, got %v", payload)
return nil, 0, 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, 0, fmt.Errorf("hub did not accept per-stream flow control (unsupported hub version?)")
}
if hubWnd > MaxStreamWindow {
hubWnd = MaxStreamWindow
}
ok = true
return fc, nil
return fc, hubWnd, nil
}
// Start establishes the control session and registers all patterns. It returns
@@ -119,7 +149,7 @@ func (c *Client) Start(ctx context.Context) error {
}
func (c *Client) connectControl(ctx context.Context) error {
fc, err := c.dialSession(MagicControl)
fc, _, err := c.dialSession(MagicControl)
if err != nil {
return fmt.Errorf("control connect: %w", err)
}
@@ -233,6 +263,7 @@ func (c *Client) handleControlRequest(cid []byte, pattern, ip string, port int,
st := newStream(wc, sid, cid, mapping, ip, port)
wc.registerStream(sid, st)
wc.sendSyn(sid, cid)
go st.writeLoop()
go st.run()
}
+1
View File
@@ -3,6 +3,7 @@
"psk": "change-me-to-a-long-random-passphrase",
"maxConn": 4,
"pingIntervalMs": 20000,
"streamWindowBytes": 262144,
"mappings": [
{
"pattern": "mc\\.example\\.com",
+20 -5
View File
@@ -29,10 +29,24 @@ const (
MuxData = 0x01
MuxFin = 0x02
MuxRst = 0x03
MuxWnd = 0x04
FrameError = 0x7F
SaturationThreshold = 8
// Session-establishment feature flags (trailing VarInt on the Rekey message).
FlagStreamFC = 0x01
// Per-stream flow-control window bounds (bytes). The advertised window is the
// receiver's promise of how much un-credited DATA it will buffer per stream.
DefaultStreamWindow = 256 * 1024
MinStreamWindow = 32 * 1024
MaxStreamWindow = 8 << 20
// DataChunkSize caps a single DATA frame's payload so no stream monopolizes
// the shared worker connection for long.
DataChunkSize = 32 * 1024
)
// Mapping routes a registered pattern to a real destination.
@@ -44,11 +58,12 @@ type Mapping struct {
// Config is the client configuration (PROTOCOL.md §9.2).
type Config struct {
Server string `json:"server"`
PSK string `json:"psk"`
MaxConn int `json:"maxConn"`
PingIntervalMs int `json:"pingIntervalMs"`
Mappings []Mapping `json:"mappings"`
Server string `json:"server"`
PSK string `json:"psk"`
MaxConn int `json:"maxConn"`
PingIntervalMs int `json:"pingIntervalMs"`
StreamWindowBytes int `json:"streamWindowBytes"` // per-stream receive window; 0 = default
Mappings []Mapping `json:"mappings"`
}
// LoadConfig reads and validates a JSON config file.
+168 -38
View File
@@ -55,18 +55,21 @@ func (p *WorkerPool) Allocate() (*WorkerConn, int, error) {
}
func (p *WorkerPool) dialWorker() (*WorkerConn, error) {
fc, err := p.client.dialSession(MagicWorker)
fc, peerWnd, err := p.client.dialSession(MagicWorker)
if err != nil {
return nil, err
}
wc := &WorkerConn{
pool: p,
fc: fc,
streams: make(map[int]*Stream),
nextSid: 1,
pool: p,
fc: fc,
sendWndInit: peerWnd,
recvWndInit: p.client.streamWnd,
streams: make(map[int]*Stream),
nextSid: 1,
}
go wc.readLoop()
log.Printf("opened worker conn (#%d in pool)", len(p.conns)+1)
log.Printf("opened worker conn (#%d in pool, send window %d, recv window %d)",
len(p.conns)+1, wc.sendWndInit, wc.recvWndInit)
return wc, nil
}
@@ -101,6 +104,9 @@ type WorkerConn struct {
pool *WorkerPool
fc *wire.FramedConn
sendWndInit int // hub's advertised per-stream receive window (our send budget)
recvWndInit int // our advertised per-stream receive window (bounds each recv queue)
mu sync.Mutex
streams map[int]*Stream
nextSid int
@@ -140,6 +146,9 @@ func (wc *WorkerConn) removeStream(sid int) *Stream {
return st
}
// readLoop dispatches inbound mux frames. It must never block on a stream's
// destination: DATA is only enqueued (the per-stream writeLoop does the actual
// destination writes), so one slow destination cannot stall other streams.
func (wc *WorkerConn) readLoop() {
for {
payload, err := wc.fc.ReadFrame()
@@ -160,9 +169,20 @@ func (wc *WorkerConn) readLoop() {
if st := wc.getStream(sid); st != nil {
st.deliverFromHub(r.Remaining())
}
case MuxFin, MuxRst:
case MuxWnd:
if delta, err := r.VarInt(); err == nil && delta > 0 {
if st := wc.getStream(sid); st != nil {
st.grantSendWnd(delta)
}
}
case MuxFin:
// Graceful: drain what is already queued to the destination first.
if st := wc.removeStream(sid); st != nil {
st.shutdown(false)
st.gracefulFin()
}
case MuxRst:
if st := wc.removeStream(sid); st != nil {
st.teardown(false)
}
default:
log.Printf("worker: unknown mux type %d", ftype)
@@ -178,7 +198,7 @@ func (wc *WorkerConn) readLoop() {
wc.streams = make(map[int]*Stream)
wc.mu.Unlock()
for _, st := range streams {
st.shutdown(false)
st.teardown(false)
}
}
@@ -198,7 +218,16 @@ func (wc *WorkerConn) sendRst(sid int) {
_ = wc.fc.WriteFrame(wire.NewWriter().U8(MuxRst).VarInt(sid).Out())
}
func (wc *WorkerConn) sendWndUpdate(sid, delta int) {
_ = wc.fc.WriteFrame(wire.NewWriter().U8(MuxWnd).VarInt(sid).VarInt(delta).Out())
}
// Stream bridges one player (via the hub) to one destination connection.
//
// Data from the hub is queued and written to the destination by a dedicated
// writeLoop goroutine. The queue is bounded by the advertised receive window —
// the hub never sends more un-credited bytes, so overflow is a protocol
// violation and resets the stream.
type Stream struct {
wc *WorkerConn
sid int
@@ -207,25 +236,33 @@ type Stream struct {
srcIP string
srcPort int
mu sync.Mutex
dest net.Conn
connected bool
preBuf []byte
closed bool
mu sync.Mutex
cond *sync.Cond
dest net.Conn
connected bool
closed bool
finPending bool // hub sent FIN; close the destination once the queue drains
q [][]byte // hub -> destination, waiting for writeLoop
qBytes int
sendWnd int // flow control: budget for destination -> hub DATA
consumed int // flow control: drained bytes not yet credited back to the hub
}
func newStream(wc *WorkerConn, sid int, cid []byte, m Mapping, ip string, port int) *Stream {
return &Stream{wc: wc, sid: sid, cid: cid, mapping: m, srcIP: ip, srcPort: port}
s := &Stream{wc: wc, sid: sid, cid: cid, mapping: m, srcIP: ip, srcPort: port, sendWnd: wc.sendWndInit}
s.cond = sync.NewCond(&s.mu)
return s
}
// run dials the destination, optionally writes the PROXY v2 header, flushes any
// buffered hub bytes, then pumps destination -> hub.
// run dials the destination, optionally writes the PROXY v2 header, then pumps
// destination -> hub (respecting the stream send window when negotiated).
func (s *Stream) run() {
dest, err := net.DialTimeout("tcp", s.mapping.Destination, 10*time.Second)
if err != nil {
log.Printf("stream %d: dial %s failed: %v", s.sid, s.mapping.Destination, err)
s.wc.removeStream(s.sid)
s.wc.sendRst(s.sid)
s.teardown(false)
return
}
if tcp, ok := dest.(*net.TCPConn); ok {
@@ -240,7 +277,6 @@ func (s *Stream) run() {
}
}
// Atomically flush pre-connect buffer and enable direct writes.
s.mu.Lock()
if s.closed {
s.mu.Unlock()
@@ -248,18 +284,18 @@ func (s *Stream) run() {
return
}
s.dest = dest
if len(s.preBuf) > 0 {
_, _ = dest.Write(s.preBuf)
s.preBuf = nil
}
s.connected = true
s.cond.Broadcast() // wake writeLoop: queued hub bytes can flow now
s.mu.Unlock()
// destination -> hub
buf := make([]byte, 32*1024)
buf := make([]byte, DataChunkSize)
for {
n, err := dest.Read(buf)
if n > 0 {
if !s.acquireSendWnd(n) {
break
}
if werr := s.wc.sendData(s.sid, buf[:n]); werr != nil {
break
}
@@ -268,7 +304,99 @@ func (s *Stream) run() {
break
}
}
s.shutdown(true)
s.teardown(true)
}
// writeLoop is the only writer to the destination. It drains the receive queue,
// credits the hub as bytes land on the destination socket, and performs the
// deferred graceful close when a FIN arrived with data still queued.
func (s *Stream) writeLoop() {
for {
s.mu.Lock()
for !s.closed && (!s.connected || (len(s.q) == 0 && !s.finPending)) {
s.cond.Wait()
}
if s.closed {
s.mu.Unlock()
return
}
if len(s.q) == 0 { // finPending and fully drained
s.mu.Unlock()
s.teardown(false)
return
}
data := s.q[0]
s.q = s.q[1:]
s.qBytes -= len(data)
dest := s.dest
s.mu.Unlock()
if _, err := dest.Write(data); err != nil {
s.teardown(true)
return
}
s.credit(len(data))
}
}
// deliverFromHub enqueues hub bytes for the destination. Called from the worker
// readLoop; it never blocks — a peer that exceeds the advertised window is a
// protocol violator and gets the stream reset.
func (s *Stream) deliverFromHub(data []byte) {
s.mu.Lock()
if s.closed || s.finPending {
s.mu.Unlock()
return
}
if s.qBytes+len(data) > s.wc.recvWndInit {
s.mu.Unlock()
log.Printf("stream %d: peer exceeded flow-control window; resetting", s.sid)
s.wc.removeStream(s.sid)
s.wc.sendRst(s.sid)
s.teardown(false)
return
}
s.q = append(s.q, data)
s.qBytes += len(data)
s.cond.Broadcast()
s.mu.Unlock()
}
// acquireSendWnd blocks until the stream may send n more bytes to the hub.
// Returns false if the stream closed while waiting.
func (s *Stream) acquireSendWnd(n int) bool {
s.mu.Lock()
defer s.mu.Unlock()
for !s.closed && s.sendWnd < n {
s.cond.Wait()
}
if s.closed {
return false
}
s.sendWnd -= n
return true
}
func (s *Stream) grantSendWnd(delta int) {
s.mu.Lock()
s.sendWnd += delta
s.cond.Broadcast()
s.mu.Unlock()
}
// credit accounts bytes drained to the destination and grants the hub more
// window once half of our receive window has been consumed.
func (s *Stream) credit(n int) {
s.mu.Lock()
s.consumed += n
if s.closed || s.consumed*2 < s.wc.recvWndInit {
s.mu.Unlock()
return
}
delta := s.consumed
s.consumed = 0
s.mu.Unlock()
s.wc.sendWndUpdate(s.sid, delta)
}
func (s *Stream) buildProxyHeader(dest net.Conn) []byte {
@@ -283,28 +411,29 @@ func (s *Stream) buildProxyHeader(dest net.Conn) []byte {
return BuildProxyV2(srcIP, s.srcPort, dstTCP.IP, dstTCP.Port)
}
// deliverFromHub writes bytes coming from the hub to the destination, buffering
// until the destination connection is established.
func (s *Stream) deliverFromHub(data []byte) {
// finDrainTimeout bounds how long a FIN'd stream may keep draining its queue
// into the destination, so a dead destination cannot hold the stream forever.
const finDrainTimeout = 10 * time.Second
// gracefulFin marks the hub-side close; writeLoop finishes the queue and then
// tears the stream down.
func (s *Stream) gracefulFin() {
s.mu.Lock()
if s.closed {
if s.closed || s.finPending {
s.mu.Unlock()
return
}
if !s.connected {
s.preBuf = append(s.preBuf, data...)
s.mu.Unlock()
return
s.finPending = true
if s.dest != nil {
_ = s.dest.SetWriteDeadline(time.Now().Add(finDrainTimeout))
}
dest := s.dest
s.cond.Broadcast()
s.mu.Unlock()
if _, err := dest.Write(data); err != nil {
s.shutdown(true)
}
}
// shutdown closes the stream; notifyHub sends a FIN to the hub when true.
func (s *Stream) shutdown(notifyHub bool) {
// teardown closes the stream immediately; notifyHub sends a FIN when true.
// Idempotent; wakes every goroutine parked on the stream.
func (s *Stream) teardown(notifyHub bool) {
s.mu.Lock()
if s.closed {
s.mu.Unlock()
@@ -312,6 +441,7 @@ func (s *Stream) shutdown(notifyHub bool) {
}
s.closed = true
dest := s.dest
s.cond.Broadcast()
s.mu.Unlock()
if dest != nil {