Frame the VNC session metrics explicitly instead of inferring updates from payload bytes

client/vnc/server/metrics_conn.go

@@ -3,7 +3,6 @@
 package server
 
 import (
-	"encoding/binary"
 	"net"
 	"sync"
 	"sync/atomic"
@@ -42,11 +41,14 @@ type metricsConn struct {
 
 	recorder func(SessionTick)
 
-	bytesOut    atomic.Uint64
-	writes      atomic.Uint64
-	writeNanos  atomic.Uint64
-	largestPkt  atomic.Uint64
-	fbus        atomic.Uint64
+	bytesOut   atomic.Uint64
+	writes     atomic.Uint64
+	writeNanos atomic.Uint64
+	largestPkt atomic.Uint64
+	fbus       atomic.Uint64
+	// inFBU is true between beginFBU and endFBU, so only the writes that make
+	// up a FramebufferUpdate are counted against it.
+	inFBU       atomic.Bool
 	fbuBytes    atomic.Uint64
 	fbuRects    atomic.Uint64
 	maxFBUBytes atomic.Uint64
@@ -170,36 +172,44 @@ func (m *metricsConn) BusyFraction() float64 {
 	return m.busyFraction
 }
 
-// startsFBU reports whether the Write payload begins a FramebufferUpdate
-// message (message type byte 0). This holds both for the standalone 4-byte
-// header that sendDirtyAndMoves writes before its rect bodies and for the
-// single framed Write that sendFullUpdate / sendEmptyUpdate use to emit a
-// whole FBU (header plus body) at once. Either way the FBU boundary lines
-// up with this Write boundary.
-func startsFBU(p []byte) bool {
-	return len(p) >= 1 && p[0] == serverFramebufferUpdate
+// beginFBU records that the writes which follow belong to a new
+// FramebufferUpdate carrying rects rectangles, closing off the accounting for
+// the previous one.
+//
+// The encoder says this rather than the wrapper inferring it from the payload.
+// A FramebufferUpdate's message type is 0, and so is the leading byte of plenty
+// of other traffic: the four zero bytes of a successful security result, and
+// any rect body whose x coordinate is below 256. Sniffing counted each of those
+// as a new update, inflating the FBU count and cutting the byte and rect totals
+// of the update actually in flight into pieces.
+func (m *metricsConn) beginFBU(rects int) {
+	m.flushFBUMax()
+	m.inFBU.Store(true)
+	m.fbus.Add(1)
+	if rects > 0 {
+		m.fbuRects.Add(uint64(rects))
+	}
+}
+
+// endFBU closes the update beginFBU opened and folds its totals into the
+// per-tick maxima straight away, so a frame is accounted in the tick it was
+// sent in rather than whenever the next one happens to start. Writes outside a
+// begin/end pair — the handshake, clipboard traffic, bell — are not part of any
+// update and are left out of the FBU byte total.
+func (m *metricsConn) endFBU() {
+	m.flushFBUMax()
+	m.inFBU.Store(false)
 }
 
 func (m *metricsConn) Write(p []byte) (int, error) {
-	fbuStart := startsFBU(p)
-	if fbuStart {
-		m.flushFBUMax()
-		m.fbus.Add(1)
-	}
-
 	t0 := time.Now()
 	n, err := m.Conn.Write(p)
 	m.writeNanos.Add(uint64(time.Since(t0).Nanoseconds()))
 	m.bytesOut.Add(uint64(n))
 	m.writes.Add(1)
 
-	m.fbuBytes.Add(uint64(n))
-	if fbuStart {
-		// Rect count is carried in bytes 2:3 of the FBU header. A standalone
-		// header records it here; the rect bodies that follow only add bytes.
-		if len(p) >= 4 {
-			m.fbuRects.Add(uint64(binary.BigEndian.Uint16(p[2:4])))
-		}
+	if m.inFBU.Load() {
+		m.fbuBytes.Add(uint64(n))
 	}
 
 	if uint64(n) > m.largestPkt.Load() {

client/vnc/server/session_encode.go

@@ -296,6 +296,7 @@ func (s *session) sendDesktopSize(w, h int) error {
 	body := encodeDesktopSizeBody(w, h)
 	s.writeMu.Lock()
 	defer s.writeMu.Unlock()
+	defer s.markFBU(1)()
 	if _, err := s.conn.Write(header); err != nil {
 		return err
 	}
@@ -318,6 +319,7 @@ func (s *session) sendExtMouseAck() error {
 
 	s.writeMu.Lock()
 	defer s.writeMu.Unlock()
+	defer s.markFBU(1)()
 	if _, err := s.conn.Write(header); err != nil {
 		return err
 	}
@@ -448,13 +450,13 @@ func (s *session) sendEmptyUpdate() error {
 	if cursorRect == nil {
 		var buf [4]byte
 		buf[0] = serverFramebufferUpdate
-		return s.writeFramed(buf[:])
+		return s.writeFramed(buf[:], 0)
 	}
 	buf := make([]byte, 4+len(cursorRect))
 	buf[0] = serverFramebufferUpdate
 	binary.BigEndian.PutUint16(buf[2:4], 1)
 	copy(buf[4:], cursorRect)
-	return s.writeFramed(buf)
+	return s.writeFramed(buf, 1)
 }
 
 func (s *session) sendFullUpdate(img *image.RGBA) error {
@@ -486,7 +488,7 @@ func (s *session) sendFullUpdate(img *image.RGBA) error {
 		rectBuf = body
 	default:
 		if cursorRect == nil {
-			return s.writeFramed(encodeRawRect(img, pf, 0, 0, w, h))
+			return s.writeFramed(encodeRawRect(img, pf, 0, 0, w, h), 1)
 		}
 		rectBuf = encodeRawRect(img, pf, 0, 0, w, h)[4:]
 	}
@@ -497,7 +499,7 @@ func (s *session) sendFullUpdate(img *image.RGBA) error {
 	off := 4
 	off += copy(buf[off:], cursorRect)
 	copy(buf[off:], rectBuf)
-	return s.writeFramed(buf)
+	return s.writeFramed(buf, int(rectCount))
 }
 
 // encodeZlibSingle encodes one full-frame rect with Zlib. When cursorRect is
@@ -508,7 +510,7 @@ func (s *session) sendFullUpdate(img *image.RGBA) error {
 func (s *session) encodeZlibSingle(img *image.RGBA, pf clientPixelFormat, w, h int, zlib *zlibState, cursorRect []byte) (body []byte, done bool, err error) {
 	if zb, ok := encodeZlibRect(img, pf, 0, 0, w, h, zlib); ok {
 		if cursorRect == nil {
-			if werr := s.writeFramed(zb); werr != nil {
+			if werr := s.writeFramed(zb, 1); werr != nil {
 				return nil, true, werr
 			}
 			return nil, true, nil
@@ -516,7 +518,7 @@ func (s *session) encodeZlibSingle(img *image.RGBA, pf clientPixelFormat, w, h i
 		return zb[4:], false, nil
 	}
 	if cursorRect == nil {
-		if werr := s.writeFramed(encodeRawRect(img, pf, 0, 0, w, h)); werr != nil {
+		if werr := s.writeFramed(encodeRawRect(img, pf, 0, 0, w, h), 1); werr != nil {
 			return nil, true, werr
 		}
 		return nil, true, nil
@@ -524,15 +526,38 @@ func (s *session) encodeZlibSingle(img *image.RGBA, pf clientPixelFormat, w, h i
 	return encodeRawRect(img, pf, 0, 0, w, h)[4:], false, nil
 }
 
-func (s *session) writeFramed(buf []byte) error {
+// writeFramed writes one complete FramebufferUpdate, header and body together.
+// rects is the rectangle count in that header, reported to the metrics wrapper
+// so it knows where this update begins.
+func (s *session) writeFramed(buf []byte, rects int) error {
 	s.writeMu.Lock()
 	defer s.writeMu.Unlock()
+	defer s.markFBU(rects)()
 	if _, err := s.conn.Write(buf); err != nil {
 		return err
 	}
 	return nil
 }
 
+// fbuMarker is implemented by the metrics wrapper around the connection.
+type fbuMarker interface {
+	beginFBU(rects int)
+	endFBU()
+}
+
+// markFBU tells the connection wrapper that a FramebufferUpdate with rects
+// rectangles starts here, and returns the function that closes it. A plain
+// connection does not implement the interface, and the result is then a no-op.
+// Caller must hold writeMu, so the update and its writes stay one step.
+func (s *session) markFBU(rects int) func() {
+	m, ok := s.conn.(fbuMarker)
+	if !ok {
+		return func() {}
+	}
+	m.beginFBU(rects)
+	return m.endFBU
+}
+
 // sendDirtyAndMoves writes one FramebufferUpdate combining CopyRect moves
 // (cheap, 16 bytes each) and pixel-encoded dirty rects. Moves come first so
 // their source tiles are read from the client's pre-update framebuffer state,
@@ -561,6 +586,7 @@ func (s *session) sendDirtyAndMoves(img *image.RGBA, moves []copyRectMove, rects
 	s.writeMu.Lock()
 	defer s.writeMu.Unlock()
 
+	defer s.markFBU(total)()
 	if _, err := s.conn.Write(header); err != nil {
 		return err
 	}