From 0f15d99fbc3d75da1ec06be6b3b92e82dab6d99e Mon Sep 17 00:00:00 2001 From: realrasengan Date: Sat, 14 Mar 2026 20:40:53 -0700 Subject: [PATCH] transport: SSP state tracking for framebuffer ack Add AckedByRemote/SentNum accessors, diffOldNum locking, and received state list for proper old_num validation. Required for server framebuffer diffing to track which state the client has. --- client.go | 65 +++++++++++++++++------ transport.go | 145 ++++++++++++++++++++++++++++++++++++++++++--------- 2 files changed, 169 insertions(+), 41 deletions(-) diff --git a/client.go b/client.go index 92e2ad9..f4825ff 100644 --- a/client.go +++ b/client.go @@ -20,6 +20,14 @@ type Client struct { output []byte // accumulated terminal output outputC chan struct{} + // Action tracking for cumulative diffs. + actionsMu sync.Mutex + actions []UserInstruction + ackedActionCount int // how many actions the server has + lastAcked uint64 + sentActionCounts map[uint64]int // sentNum → action count at that state + dirty bool // true = new actions since last tick + done chan struct{} wg sync.WaitGroup } @@ -50,18 +58,20 @@ func Dial(host string, port int, key string) (*Client, error) { } c := &Client{ - conn: conn, - transport: NewTransport(ocb, false), - ocb: ocb, - outputC: make(chan struct{}, 1), - done: make(chan struct{}), + conn: conn, + transport: NewTransport(ocb, false), + ocb: ocb, + outputC: make(chan struct{}, 1), + done: make(chan struct{}), + sentActionCounts: make(map[uint64]int), } c.wg.Add(2) go c.recvLoop() go c.sendLoop() - // Send initial empty datagram to associate with the server. + // Send initial keepalive to associate with the server. + c.transport.ForceNextSend() c.tick() return c, nil @@ -69,21 +79,21 @@ func Dial(host string, port int, key string) (*Client, error) { // Send sends keystrokes to the server. func (c *Client) Send(keys []byte) { - diff := MarshalUserMessage([]UserInstruction{{Keys: keys}}) - c.transport.SetPending(diff) - c.tick() + c.actionsMu.Lock() + c.actions = append(c.actions, UserInstruction{Keys: append([]byte{}, keys...)}) + c.dirty = true + c.actionsMu.Unlock() } // Resize sends a resize to the server. func (c *Client) Resize(cols, rows uint16) { - diff := MarshalUserMessage([]UserInstruction{{ - Width: int32(cols), - Height: int32(rows), - }}) - c.transport.SetPending(diff) - c.tick() + c.actionsMu.Lock() + c.actions = append(c.actions, UserInstruction{Width: int32(cols), Height: int32(rows)}) + c.dirty = true + c.actionsMu.Unlock() } + // Recv reads accumulated terminal output, blocking until output is available // or the timeout expires. Returns nil on timeout. func (c *Client) Recv(timeout time.Duration) []byte { @@ -132,11 +142,36 @@ func (c *Client) Transport() *Transport { } func (c *Client) tick() { + // Batch: compute diff from accumulated actions. + c.actionsMu.Lock() + if c.dirty { + c.dirty = false + c.processAcksLocked() + newActions := c.actions[c.ackedActionCount:] + c.transport.SetPending(MarshalUserMessage(newActions)) + c.sentActionCounts[c.transport.SentNum()+1] = len(c.actions) + } + c.actionsMu.Unlock() + for _, dg := range c.transport.Tick() { c.conn.Write(dg) } } +func (c *Client) processAcksLocked() { + acked := c.transport.AckedByRemote() + if acked > c.lastAcked { + c.lastAcked = acked + // Only advance base when server caught up (no states in flight). + if acked >= c.transport.SentNum() { + if count, ok := c.sentActionCounts[acked]; ok && count > c.ackedActionCount { + c.ackedActionCount = count + } + c.sentActionCounts = make(map[uint64]int) + } + } +} + func (c *Client) recvLoop() { defer c.wg.Done() buf := make([]byte, maxPayload+64) diff --git a/transport.go b/transport.go index 141522d..19c6a0a 100644 --- a/transport.go +++ b/transport.go @@ -26,13 +26,18 @@ type Transport struct { // Outgoing state (SSP §3). sentNum uint64 // newest state we've sent (new_num) - ackedNum uint64 // newest state the remote has acknowledged (ack from remote) - pendingDiff []byte // diff payload waiting to be sent + ackedByRemote uint64 // newest state the remote has acknowledged + pendingDiff []byte // diff payload waiting to be sent + diffSent bool // true = pendingDiff has been sent at least once + diffOldNum uint64 // locked oldNum for all diffs until base advances + hasPendingBase bool // true = diffOldNum is locked + pendingDataAck bool // true = send ack ASAP (received data, not just ack) - // Incoming state. - recvNum uint64 // newest state we've received from remote (their new_num) - recvAck uint64 // what we've told the remote we received (our ack_num) - throwawayNum uint64 // oldest state we still hold + // Incoming state — list of received state nums for old_num validation. + receivedNums []uint64 // ordered list of state nums we have + ackNum uint64 // latest received state num + sentAckNum uint64 // last ackNum we actually sent on wire + throwawayNum uint64 // oldest state we still hold // Sequence counter for the crypto layer (independent of SSP state numbering). seqOut uint64 @@ -67,10 +72,11 @@ const ( // NewTransport creates a transport. isServer determines direction bits. func NewTransport(ocb *OCB, isServer bool) *Transport { t := &Transport{ - ocb: ocb, - rto: initialRTO, - lastSend: time.Now(), - lastRecv: time.Now(), + ocb: ocb, + rto: initialRTO, + lastSend: time.Now(), + lastRecv: time.Now(), + receivedNums: []uint64{0}, // start with state 0 } if isServer { t.toRemote = dirToClient @@ -107,9 +113,19 @@ func (t *Transport) HasCap(bit byte) bool { return (t.localCaps[idx] & t.remoteCaps[idx] & bit) != 0 } +// ForceNextSend forces the next tick to send, even with no pending diff. +func (t *Transport) ForceNextSend() { + t.mu.Lock() + t.lastSend = time.Time{} // zero time, always expired + t.mu.Unlock() +} + // SetPending sets the diff payload to send on the next tick. func (t *Transport) SetPending(diff []byte) { t.mu.Lock() + if len(diff) > 0 { + t.diffSent = false + } t.pendingDiff = diff t.mu.Unlock() } @@ -124,29 +140,44 @@ func (t *Transport) Tick() [][]byte { // Decide if we should send. haveDiff := len(t.pendingDiff) > 0 - needAck := t.recvNum > t.recvAck - expired := now.Sub(t.lastSend) >= t.rto + haveNewDiff := haveDiff && !t.diffSent + needAck := t.ackNum > t.sentAckNum + sinceLastSend := now.Sub(t.lastSend) + expired := sinceLastSend >= t.rto + urgentAck := t.pendingDataAck - if !haveDiff && !needAck && !expired { + shouldSend := haveNewDiff || needAck || expired || urgentAck + if !shouldSend { return nil } // Build TransportInstruction. - if haveDiff { + if haveNewDiff { t.sentNum++ + t.diffSent = true + if !t.hasPendingBase { + t.diffOldNum = t.ackedByRemote + t.hasPendingBase = true + } + } + t.pendingDataAck = false + + oldNum := t.ackedByRemote + if haveDiff { + oldNum = t.diffOldNum // use locked oldNum for diff retransmissions } ti := TransportInstruction{ - ProtocolVersion: 2, // mosh protocol version - OldNum: t.ackedNum, + ProtocolVersion: 2, + OldNum: oldNum, NewNum: t.sentNum, - AckNum: t.recvNum, - ThrowawayNum: t.throwawayNum, + AckNum: t.ackNum, + ThrowawayNum: 0, // client doesn't tell server to throw away states Diff: t.pendingDiff, LatchCaps: t.localCaps, } - t.recvAck = t.recvNum - t.pendingDiff = nil + t.sentAckNum = t.ackNum + // Do NOT nil pendingDiff — keep for retransmission until server acks. // Marshal → compress → fragment → encrypt. pbData := ti.Marshal() @@ -250,22 +281,84 @@ func (t *Transport) Recv(wire []byte) []byte { t.mu.Unlock() } - // Process SSP fields. + // Process SSP fields (matching upstream mosh recv logic). t.mu.Lock() - if ti.AckNum > t.ackedNum { - t.ackedNum = ti.AckNum + defer t.mu.Unlock() + + // Process ack from remote. + if ti.AckNum > t.ackedByRemote { + t.ackedByRemote = ti.AckNum + if t.ackedByRemote >= t.sentNum && t.pendingDiff != nil { + t.pendingDiff = nil + t.diffSent = false + t.hasPendingBase = false + } } - if ti.NewNum > t.recvNum { - t.recvNum = ti.NewNum + + // Check if we already have new_num (dedup). + for _, n := range t.receivedNums { + if n == ti.NewNum { + return nil + } } + + // Check if we have old_num (required to apply diff). + hasOld := false + for _, n := range t.receivedNums { + if n == ti.OldNum { + hasOld = true + break + } + } + if !hasOld { + return nil + } + + // Process throwaway. if ti.ThrowawayNum > t.throwawayNum { t.throwawayNum = ti.ThrowawayNum + filtered := t.receivedNums[:0] + for _, n := range t.receivedNums { + if n >= t.throwawayNum { + filtered = append(filtered, n) + } + } + t.receivedNums = filtered + } + + // Add new state. + t.receivedNums = append(t.receivedNums, ti.NewNum) + + // Bound the list. + if len(t.receivedNums) > 128 { + t.receivedNums = t.receivedNums[1:] + } + + // Update ack num to latest received state. + t.ackNum = ti.NewNum + + // Trigger immediate ack when we receive data. + if len(ti.Diff) > 0 { + t.pendingDataAck = true } - t.mu.Unlock() return ti.Diff } +// AckedByRemote returns the highest state number the remote has acked. +func (t *Transport) AckedByRemote() uint64 { + t.mu.Lock() + defer t.mu.Unlock() + return t.ackedByRemote +} + +// SentNum returns the current sent state number. +func (t *Transport) SentNum() uint64 { + t.mu.Lock() + defer t.mu.Unlock() + return t.sentNum +} + // LastRecv returns the time of the last received datagram. func (t *Transport) LastRecv() time.Time { t.mu.Lock()