Harden Wayland video relay

Ultraworked with [omo](https://github.com/code-yeongyu/oh-my-openagent)

Co-authored-by: sisyphus-dev-ai <sisyphus-dev-ai@users.noreply.github.com>

Author: minpeter

cmd/doubletake/main.go

@@ -371,6 +371,7 @@ func main() {
 	}
 	streamCfg.AutomaticHEVCAvailable = capturePreparation.AutomaticHEVCAvailable()
 	streamCfg.MeasuredVideoLatency = capturePreparation.MeasuredVideoLatency()
+	streamCfg.MinimumVideoLead = capturePreparation.MinimumVideoLead()
 
 	var capture *airplay.ScreenCapture
 	var broadcast *airplay.BroadcastCapture

internal/airplay/capture.go

@@ -97,6 +97,14 @@ func captureMinimumVideoLead(kind capturePreparationKind, measured time.Duration
 	return measured
 }
 
+func waylandRawVideoSize(receiverWidth, receiverHeight int, streamSize [2]int) (int, int) {
+	if receiverWidth <= 0 || receiverHeight <= 0 {
+		receiverWidth, receiverHeight = streamSize[0], streamSize[1]
+	}
+	receiverWidth, receiverHeight = fitVideoSize(receiverWidth, receiverHeight, 3840, 2160)
+	return receiverWidth &^ 1, receiverHeight &^ 1
+}
+
 // CapturePreparation performs the potentially interactive part of screen
 // capture before the receiver session starts. In particular, a Wayland
 // preparation completes the screencast portal request and retains its PipeWire
@@ -115,8 +123,9 @@ type CapturePreparation struct {
 	timestampedOutput  bool
 	automaticHEVCAvail bool
 	// measuredVideoLatency is the minimum screen lead required by local capture:
-	// either the 4K HEVC preflight or the isolated Wayland raw relay.
+	// the 4K HEVC preflight. minimumVideoLead also includes transport overhead.
 	measuredVideoLatency time.Duration
+	minimumVideoLead     time.Duration
 
 	pwNodeID   uint32
 	pwFd       *os.File
@@ -217,7 +226,7 @@ func PrepareCapture(ctx context.Context, cfg CaptureConfig) (*CapturePreparation
 	preparation.pwNodeID = nodeID
 	preparation.pwFd = pwFd
 	preparation.dbusConn = dbusConn
-	preparation.measuredVideoLatency = captureMinimumVideoLead(kind, preparation.measuredVideoLatency)
+	preparation.minimumVideoLead = captureMinimumVideoLead(kind, preparation.measuredVideoLatency)
 	return preparation, nil
 }
 
@@ -405,6 +414,14 @@ func (p *CapturePreparation) MeasuredVideoLatency() time.Duration {
 	return p.measuredVideoLatency
 }
 
+// MinimumVideoLead returns the full local capture/transport scheduling floor.
+func (p *CapturePreparation) MinimumVideoLead() time.Duration {
+	if p == nil {
+		return 0
+	}
+	return p.minimumVideoLead
+}
+
 // Close releases an unconsumed portal preparation. Once Start has taken
 // ownership, ScreenCapture.Stop owns the corresponding resources.
 func (p *CapturePreparation) Close() {
@@ -944,10 +961,7 @@ func startPreparedWaylandCapture(ctx context.Context, cfg CaptureConfig, encoder
 			lowLatencyVideoQueueStage(),
 		)
 	}
-	rawWidth, rawHeight := cfg.MaxWidth&^1, cfg.MaxHeight&^1
-	if rawWidth <= 0 || rawHeight <= 0 {
-		rawWidth, rawHeight = streamSize[0]&^1, streamSize[1]&^1
-	}
+	rawWidth, rawHeight := waylandRawVideoSize(cfg.MaxWidth, cfg.MaxHeight, streamSize)
 	if rawWidth <= 0 || rawHeight <= 0 {
 		cancel()
 		_ = pwFd.Close()

internal/airplay/capture_broadcast.go

@@ -48,6 +48,7 @@ type BroadcastCapture struct {
 	// adjacent PTS values. A leaky upstream queue can legitimately create large
 	// PTS gaps while only one encoded picture is pending.
 	frameDuration time.Duration
+	now           func() time.Time
 	mu            sync.Mutex
 	done          chan struct{}
 	err           error // set before done is closed
@@ -57,9 +58,9 @@ type BroadcastCapture struct {
 	// following sequence, which gives attachment an exact cutover even when a
 	// source read has completed but has not yet been fanned out.
 	sequence uint64
-	// primer is the latest complete parameter-set plus random-access AU. A
-	// receiver attached after capture starts needs it before live P-frames are
-	// decodable.
+	// Timestamped access-unit fan-out caches the latest complete parameter-set
+	// plus random-access AU. Legacy byte-stream fan-out retains its exact
+	// next-read cutover and therefore waits for the encoder's next keyframe.
 	primer VideoAccessUnit
 
 	drainTimeout time.Duration
@@ -143,6 +144,7 @@ func NewBroadcastCaptureWithFrameRate(src *ScreenCapture, fps int) *BroadcastCap
 		src:           src,
 		frames:        src != nil && src.frames != nil,
 		frameDuration: time.Second / time.Duration(fps),
+		now:           time.Now,
 		done:          make(chan struct{}),
 		drainTimeout:  broadcastSinkDrainTimeout,
 	}
@@ -162,6 +164,9 @@ func (bc *BroadcastCapture) AddSink() *BroadcastSink {
 	}
 	s.startSequence = bc.sequence + 1
 	s.primer = bc.primer
+	if len(s.primer.AnnexB) > 0 {
+		s.primer.PTS = bc.now()
+	}
 	bc.sinks = append(bc.sinks, s)
 	bc.mu.Unlock()
 	return s
@@ -188,6 +193,9 @@ func (bc *BroadcastCapture) AddBackpressuredSink() (*BroadcastSink, error) {
 	bc.exclusive = true
 	s.startSequence = bc.sequence + 1
 	s.primer = bc.primer
+	if len(s.primer.AnnexB) > 0 {
+		s.primer.PTS = bc.now()
+	}
 	bc.sinks = append(bc.sinks, s)
 	bc.mu.Unlock()
 	return s, nil
@@ -277,7 +285,7 @@ func (bc *BroadcastCapture) runFrames() error {
 			bc.mu.Lock()
 			if len(frame.AnnexB) <= broadcastSinkQueueBytes && isDecoderPrimer(frame.AnnexB) {
 				bc.primer = VideoAccessUnit{
-					AnnexB: append(bc.primer.AnnexB[:0], frame.AnnexB...),
+					AnnexB: append([]byte(nil), frame.AnnexB...),
 					PTS:    frame.PTS,
 				}
 			}
@@ -562,7 +570,7 @@ func (s *BroadcastSink) Read(p []byte) (int, error) {
 func (s *BroadcastSink) ReadVideoAccessUnit() (VideoAccessUnit, error) {
 	s.mu.Lock()
 	defer s.mu.Unlock()
-	for len(s.frameQueue) == 0 && !s.inputClosed && !s.closed {
+	for len(s.primer.AnnexB) == 0 && len(s.frameQueue) == 0 && !s.inputClosed && !s.closed {
 		s.cond.Wait()
 	}
 	if s.closed {
@@ -571,9 +579,6 @@ func (s *BroadcastSink) ReadVideoAccessUnit() (VideoAccessUnit, error) {
 	if len(s.primer.AnnexB) > 0 {
 		primer := s.primer
 		s.primer = VideoAccessUnit{}
-		if len(s.frameQueue) > 0 && !s.frameQueue[0].PTS.IsZero() {
-			primer.PTS = s.frameQueue[0].PTS.Add(-s.frameDuration)
-		}
 		if s.inputClosed && len(s.frameQueue) == 0 && len(s.queue) == 0 {
 			s.closeDoneLocked()
 		}

internal/airplay/capture_broadcast_test.go

@@ -123,6 +123,8 @@ func TestBroadcastCaptureReplaysDecoderPrimerToLateSink(t *testing.T) {
 		waitCh: make(chan struct{}),
 	}
 	broadcast := NewBroadcastCaptureWithFrameRate(capture, 30)
+	replayPTS := time.Unix(110, 0)
+	broadcast.now = func() time.Time { return replayPTS }
 	runDone := make(chan error, 1)
 	go func() { runDone <- broadcast.Run() }()
 
@@ -140,6 +142,34 @@ func TestBroadcastCaptureReplaysDecoderPrimerToLateSink(t *testing.T) {
 
 	sink := broadcast.AddSink()
 	defer sink.Close()
+	replayDone := make(chan struct {
+		frame VideoAccessUnit
+		err   error
+	}, 1)
+	go func() {
+		frame, err := sink.ReadVideoAccessUnit()
+		replayDone <- struct {
+			frame VideoAccessUnit
+			err   error
+		}{frame: frame, err: err}
+	}()
+	var replayed VideoAccessUnit
+	select {
+	case result := <-replayDone:
+		if result.err != nil {
+			t.Fatalf("read replayed decoder primer: %v", result.err)
+		}
+		replayed = result.frame
+	case <-time.After(time.Second):
+		t.Fatal("cached decoder primer was not available immediately")
+	}
+	if !bytes.Equal(replayed.AnnexB, primer.AnnexB) {
+		t.Fatalf("first late-sink frame = %x, want cached decoder primer %x", replayed.AnnexB, primer.AnnexB)
+	}
+	if !replayed.PTS.Equal(replayPTS) {
+		t.Fatalf("replayed primer PTS = %v, want attachment PTS %v", replayed.PTS, replayPTS)
+	}
+
 	boundary := VideoAccessUnit{
 		AnnexB: []byte{0, 0, 0, 1, 0x61, 0x40},
 		PTS:    primer.PTS.Add(10 * time.Second),
@@ -153,16 +183,6 @@ func TestBroadcastCaptureReplaysDecoderPrimerToLateSink(t *testing.T) {
 	frames <- live
 	close(frames)
 
-	replayed, err := sink.ReadVideoAccessUnit()
-	if err != nil {
-		t.Fatalf("read replayed decoder primer: %v", err)
-	}
-	if !bytes.Equal(replayed.AnnexB, primer.AnnexB) {
-		t.Fatalf("first late-sink frame = %x, want cached decoder primer %x", replayed.AnnexB, primer.AnnexB)
-	}
-	if wantPTS := live.PTS.Add(-time.Second / 30); !replayed.PTS.Equal(wantPTS) {
-		t.Fatalf("replayed primer PTS = %v, want one frame before live PTS %v", replayed.PTS, wantPTS)
-	}
 	next, err := sink.ReadVideoAccessUnit()
 	if err != nil {
 		t.Fatalf("read live frame after decoder primer: %v", err)
@@ -175,6 +195,49 @@ func TestBroadcastCaptureReplaysDecoderPrimerToLateSink(t *testing.T) {
 	}
 }
 
+func TestBroadcastCapturePrimerSnapshotSurvivesRefresh(t *testing.T) {
+	frames := make(chan VideoAccessUnit)
+	reads := make(chan struct{}, 4)
+	capture := &ScreenCapture{
+		frames: &signaledVideoAccessUnitReader{frames: frames, reads: reads},
+		waitCh: make(chan struct{}),
+	}
+	broadcast := NewBroadcastCapture(capture)
+	broadcast.now = func() time.Time { return time.Unix(120, 0) }
+	runDone := make(chan error, 1)
+	go func() { runDone <- broadcast.Run() }()
+
+	first := VideoAccessUnit{AnnexB: []byte{
+		0, 0, 0, 1, 0x67, 0x42,
+		0, 0, 0, 1, 0x68, 0xce,
+		0, 0, 0, 1, 0x65, 0xaa,
+	}}
+	second := VideoAccessUnit{AnnexB: []byte{
+		0, 0, 0, 1, 0x67, 0x64,
+		0, 0, 0, 1, 0x68, 0xee,
+		0, 0, 0, 1, 0x65, 0xbb,
+	}}
+	<-reads
+	frames <- first
+	<-reads
+	sink := broadcast.AddSink()
+	defer sink.Close()
+	frames <- second
+	<-reads
+
+	got, err := sink.ReadVideoAccessUnit()
+	if err != nil {
+		t.Fatalf("read primer snapshot: %v", err)
+	}
+	if !bytes.Equal(got.AnnexB, first.AnnexB) {
+		t.Fatalf("primer snapshot = %x, want %x", got.AnnexB, first.AnnexB)
+	}
+	close(frames)
+	if err := <-runDone; !errors.Is(err, io.EOF) {
+		t.Fatalf("broadcast run = %v, want EOF", err)
+	}
+}
+
 func TestBroadcastSinkBackpressuresWithOnePendingAccessUnit(t *testing.T) {
 	sink := newBroadcastSinkWithPolicy(nil, true)
 	base := time.Now()

internal/airplay/capture_test.go

@@ -151,6 +151,17 @@ func TestCaptureMinimumVideoLeadIncludesWaylandRawRelay(t *testing.T) {
 	}
 }
 
+func TestWaylandRawVideoSizeBoundsReceiverCanvas(t *testing.T) {
+	width, height := waylandRawVideoSize(1<<30, 1<<30, [2]int{2880, 1800})
+	if width != 2160 || height != 2160 {
+		t.Fatalf("hostile square receiver canvas = %dx%d, want bounded 2160x2160", width, height)
+	}
+	width, height = waylandRawVideoSize(0, 0, [2]int{2881, 1801})
+	if width != 2880 || height != 1800 {
+		t.Fatalf("portal fallback canvas = %dx%d, want even 2880x1800", width, height)
+	}
+}
+
 func TestLiveVideoProbeTimeoutTracksConfiguredFrameRate(t *testing.T) {
 	if got := liveVideoProbeTimeout(30); got != minimumLiveVideoProbeTimeout {
 		t.Fatalf("30fps live probe timeout = %v, want %v", got, minimumLiveVideoProbeTimeout)

internal/airplay/client.go

@@ -1203,6 +1203,7 @@ type StreamConfig struct {
 	VideoCodec             VideoCodec    // empty/h264, auto, or capability-gated hevc
 	AutomaticHEVCAvailable bool          // capture preflight found the hardware HEVC-4K path
 	MeasuredVideoLatency   time.Duration // measured minimum lead for the local HEVC capture path
+	MinimumVideoLead       time.Duration // known capture/transport lead before codec selection
 	NoEncrypt              bool          // Disable encryption for debugging
 	DirectKey              bool          // Use shk/shiv directly without SHA-512 derivation
 	NoAudio                bool          // Disable audio streaming

internal/airplay/mirror.go

@@ -504,6 +504,9 @@ func (c *AirPlayClient) setupMirrorSession(ctx context.Context, cfg StreamConfig
 	var receiverEventPort int
 	var attemptedEventPort int
 	audioControlLPort := audioCtrlConn.LocalAddr().(*net.UDPAddr).Port
+	if !targetLatencyIsExplicit() {
+		latencies = latencies.withMinimumVideoLead(cfg.MinimumVideoLead)
+	}
 	audioLatencySamples := samplesFor44k1(latencies.audio)
 	measuredLatencyApplied := false
 	audioSetupCommitted := false

internal/airplay/receiver_server_test.go

@@ -579,6 +579,37 @@ func TestMediaFirstAutoFallsBackBeforeCreatingLatencyMismatch(t *testing.T) {
 	}
 }
 
+func TestMediaFirstRelayLeadPreservesAutomaticHEVC(t *testing.T) {
+	SetTargetLatency(0)
+	t.Cleanup(func() { SetTargetLatency(0) })
+	_, client, ctx := newReceiverServerTestPair(t, ReceiverConfig{
+		Profile: ReceiverProfileRoku, DisplayWidth: 3840, DisplayHeight: 2160,
+	})
+	if err := client.Pair(ctx, ""); err != nil {
+		t.Fatalf("pair: %v", err)
+	}
+	session, err := client.SetupMirrorWithVideoCodecPreparation(ctx, StreamConfig{
+		VideoCodec:             VideoCodecAuto,
+		AutomaticHEVCAvailable: true,
+		MinimumVideoLead:       250 * time.Millisecond,
+	}, func(_, _ int, codec VideoCodec) error {
+		if codec != VideoCodecHEVC {
+			return fmt.Errorf("media-first relay codec = %s, want HEVC", codec)
+		}
+		return nil
+	})
+	if err != nil {
+		t.Fatalf("setup mirror: %v", err)
+	}
+	defer session.Close()
+	if session.videoCodec != VideoCodecHEVC || session.timestampBias != 250*time.Millisecond {
+		t.Fatalf("media-first relay session = codec %s lead %v, want HEVC/250ms", session.videoCodec, session.timestampBias)
+	}
+	if session.audioStream == nil || session.audioStream.latencySamples != samplesFor44k1(260*time.Millisecond) {
+		t.Fatalf("media-first relay audio lead = %#v, want %d samples", session.audioStream, samplesFor44k1(260*time.Millisecond))
+	}
+}
+
 func TestMediaFirstCalibratedAutoFallsBackBeforeLiveMeasurement(t *testing.T) {
 	SetTargetLatency(0)
 	t.Cleanup(func() { SetTargetLatency(0) })

internal/daemon/daemon.go

@@ -1135,6 +1135,7 @@ func (d *Daemon) connectAndStream(ctx context.Context, entry *activeStream, targ
 	streamCfg := d.mirrorStreamConfig()
 	streamCfg.AutomaticHEVCAvailable = capturePreparation.AutomaticHEVCAvailable()
 	streamCfg.MeasuredVideoLatency = capturePreparation.MeasuredVideoLatency()
+	streamCfg.MinimumVideoLead = capturePreparation.MinimumVideoLead()
 	var broadcast *airplay.BroadcastCapture
 	selectedCaptureKey := videoCaptureKey{maxWidth: -1, maxHeight: -1}
 	prepareVideo := func(width, height int, codec airplay.VideoCodec) (airplay.VideoPreparationResult, error) {