From a52838a1ce935392a8ab053ed3448fb5a1a20a19 Mon Sep 17 00:00:00 2001 From: gmegidish Date: Fri, 18 Sep 2026 17:49:15 +0200 Subject: [PATCH 1/4] feat(ios): share one avc stream between screencapture and screenrecord The broadcast extension sends video to its newest tcp client, so a second screencapture, or a screenrecord started while a capture was live, stole or restarted the stream. Real iOS devices now own a single avc source per device (devices/avc_hub.go) that fans the same bytes out to every subscriber: one goes to disk as mp4, one goes to the websocket/WebRTC consumer. - first subscriber starts the broadcast; later ones attach to it and their scale/fps/quality/bitrate are ignored - a late joiner gets a key frame request and starts at the head of the next key frame (SPS, PPS, timecode SEI), failing after 5s if none arrives - a slow subscriber is dropped instead of stalling the others - the source stops when the last subscriber leaves and restarts cleanly - Ctrl-C or a stop request detaches only that subscriber Android and simulators are unchanged. --- commands/screenrecord.go | 14 +- devices/avc_hub.go | 532 ++++++++++++++++++++++++++++++ devices/avc_hub_test.go | 687 +++++++++++++++++++++++++++++++++++++++ devices/common.go | 1 + devices/ios.go | 232 ++++++------- pkg/avc2mp4/nalparser.go | 8 +- 6 files changed, 1355 insertions(+), 119 deletions(-) create mode 100644 devices/avc_hub.go create mode 100644 devices/avc_hub_test.go diff --git a/commands/screenrecord.go b/commands/screenrecord.go index 2ee02483..6fea65ec 100644 --- a/commands/screenrecord.go +++ b/commands/screenrecord.go @@ -105,9 +105,9 @@ func ScreenRecordCommand(req ScreenRecordRequest) *CommandResponse { return dev.ScreenRecord(req.OutputPath, req.TimeLimit, req.StopChan) }, req, progress) case targetDevice.Platform() == "ios" && targetDevice.DeviceType() == "real": - // real iOS devices route through DeviceKit + ReplayKit; screenRecordIOSDevice + // real iOS devices route through DeviceKit + ReplayKit; screenRecordAvc // signals req.Ready itself once the broadcast picker is confirmed started. - return screenRecordIOSDevice(targetDevice, req, progress) + return screenRecordAvc(targetDevice, req, progress) default: err := fmt.Errorf("screen recording is not supported for this device type") req.signalReady(err) @@ -205,7 +205,11 @@ func (p *screenRecordProgress) downloaded(speedMBps float64) { fmt.Fprintf(p.out, "\nDownloading done, %.3f MB/sec\n", speedMBps) } -func screenRecordIOSDevice(targetDevice devices.ControllableDevice, req ScreenRecordRequest, progress *screenRecordProgress) *CommandResponse { +// screenRecordAvc records through the device's shared H.264 stream: subscribe, +// spool the elementary stream to a temp .avc, then mux it to mp4. Used by real +// iOS devices (DeviceKit + ReplayKit), so a concurrent screencapture and +// screenrecord share one broadcast instead of stealing it from each other. +func screenRecordAvc(targetDevice devices.ControllableDevice, req ScreenRecordRequest, progress *screenRecordProgress) *CommandResponse { tempFile, err := os.CreateTemp("", "screenrecord-*.avc") if err != nil { req.signalReady(err) @@ -241,6 +245,10 @@ func screenRecordIOSDevice(targetDevice devices.ControllableDevice, req ScreenRe _, writeErr := tempFile.Write(data) return writeErr == nil }, req.TimeLimit, req.StopChan), + // OnData only runs when a frame arrives, and a static screen emits very + // few; hand the stop channel down so leaving the shared stream doesn't + // wait for the next frame. + StopChan: req.StopChan, }) progress.recordingEnded() diff --git a/devices/avc_hub.go b/devices/avc_hub.go new file mode 100644 index 00000000..cfce03e8 --- /dev/null +++ b/devices/avc_hub.go @@ -0,0 +1,532 @@ +package devices + +import ( + "bytes" + "errors" + "os" + "os/signal" + "sync" + "sync/atomic" + "syscall" + "time" + + "github.com/mobile-next/mobilecli/pkg/avc2mp4" + "github.com/mobile-next/mobilecli/utils" +) + +// avcSource describes the single per-device H.264 source an avcHub fans out. +// start must not emit before it returns; it is expected to read in its own +// goroutine, call emit for every chunk and ended exactly once when the source +// is gone. stop makes the source end (close the conn, kill the process) and +// must not return before that teardown is complete; ended likewise must not run +// until it is. That is what lets the hub promise one live source per device. +type avcSource struct { + start func(emit func([]byte), ended func(error)) (stop func(), err error) + ready func() // optional: this subscriber is attached and bytes are coming + requestKeyFrame func() error // optional: ask the encoder for an immediate sync frame +} + +// avcSubscriber is one consumer of the shared stream. Each has its own buffered +// channel so a slow websocket cannot stall the disk writer. +type avcSubscriber struct { + ch chan []byte + joined chan struct{} // late joiners only: closed once their first chunk is queued + waiting bool // late joiner, still waiting for the next IDR + dropped atomic.Bool // set when the hub dropped this one for falling behind + err error // why this subscriber ended, set under the hub lock +} + +// avcHub shares one encoder/source between every concurrent avc capture of a +// device. The first subscriber starts the source and sees every byte from the +// start; later subscribers join the running one, so their scale/fps/quality/ +// bitrate are ignored. The last one leaving stops the source. +type avcHub struct { + mu sync.Mutex + subs map[*avcSubscriber]struct{} + stop func() + starting chan struct{} // closed when the in-flight src.start finished + stopping chan struct{} // closed when the in-flight stop() returned + endErr error // why the last source ended + epoch uint64 // bumped per source, so a stopped source's leftovers are ignored + sps []byte // latest parameter sets, with their start codes, for late joiners + pps []byte + tail []byte // incomplete NAL unit carried over from the previous chunk + + // set before the hub is used, so tests can run the key frame wait in + // milliseconds; zero means the avcKeyFrame* defaults + keyFrameTimeout time.Duration + keyFrameRetry time.Duration +} + +const ( + nalTypeIDR = 5 + nalTypeSPS = 7 + nalTypePPS = 8 +) + +// avcSubscriberQueue is how many chunks a subscriber may fall behind. +// ponytail: 256 chunks (~16MB at the 64KB source reads) is the ceiling; a +// consumer that far behind is broken rather than slow, so it is dropped instead +// of being buffered forever. +const avcSubscriberQueue = 256 + +// avcMaxPartialNAL bounds the incomplete NAL unit carried between chunks, so a +// source that never emits another start code cannot grow the hub without limit. +const avcMaxPartialNAL = 8 << 20 + +// How long a late joiner waits for its first key frame, and how often the +// request is repeated meanwhile. +const ( + avcKeyFrameTimeout = 5 * time.Second + avcKeyFrameRetry = 1 * time.Second +) + +var ( + errAvcSubscriberTooSlow = errors.New("avc stream consumer fell too far behind") + errAvcNoKeyFrame = errors.New("timed out waiting for a key frame on the shared avc stream") + errAvcSourceEnded = errors.New("avc source ended") + errAvcCancelled = errors.New("avc capture cancelled") +) + +// watchStop reports Ctrl+C or the caller's stop channel on the returned channel, +// so one capture can leave the shared stream — promptly, even while no frames are +// arriving — without stopping the other subscribers. The returned func stops +// watching. A nil stop channel means "signals only". +func watchStop(stop <-chan struct{}) (<-chan struct{}, func()) { + sigChan := make(chan os.Signal, 1) + signal.Notify(sigChan, syscall.SIGINT, syscall.SIGTERM) + + cancelled := make(chan struct{}) + watching := make(chan struct{}) + go func() { + defer signal.Stop(sigChan) + select { + case <-sigChan: + close(cancelled) + case <-stop: + close(cancelled) + case <-watching: + } + }() + + return cancelled, func() { close(watching) } +} + +// subscribe attaches onData to the device's shared avc source, starting it when +// nobody else is capturing. Blocks until onData returns false, cancel is closed, +// or the source ends. +func (h *avcHub) subscribe(src avcSource, onData func([]byte) bool, cancel <-chan struct{}) error { + sub, err := h.attach(src, cancel) + if err != nil || sub == nil { + return err + } + + // a late joiner is only live once its first key frame is queued, so ready() + // waits for that rather than firing on a stream it cannot decode yet + if sub.joined != nil { + if err := h.awaitFirstChunk(sub, src, cancel); err != nil { + if errors.Is(err, errAvcCancelled) { + return h.detach(sub, nil) + } + return h.detach(sub, err) + } + } + + if src.ready != nil { + src.ready() + } + + for { + select { + case chunk, ok := <-sub.ch: + if !ok || sub.dropped.Load() { + // a dropped subscriber leaves now instead of writing out a + // queue full of stale chunks first + return h.detach(sub, nil) + } + if !onData(chunk) { + return h.detach(sub, nil) + } + case <-cancel: + return h.detach(sub, nil) + } + } +} + +// attach registers a new subscriber, starting the source when nobody else is +// capturing. It returns (nil, nil) when cancel fired while this caller was +// waiting for another subscriber's start or stop to finish. +func (h *avcHub) attach(src avcSource, cancel <-chan struct{}) (*avcSubscriber, error) { + for { + h.mu.Lock() + if h.subs == nil { + h.subs = make(map[*avcSubscriber]struct{}) + } + + // somebody else is starting or stopping the one source; wait for them + // instead of racing a second one. on iOS a second dial would steal the + // broadcast extension's stream. + if busy := h.busyLocked(); busy != nil { + h.mu.Unlock() + if !waitOrCancel(busy, cancel) { + return nil, nil + } + continue + } + + if h.stop != nil { + // join the running source: hand this one the cached parameter sets + // and start its bytes at the next IDR + sub := &avcSubscriber{ch: make(chan []byte, avcSubscriberQueue), joined: make(chan struct{}), waiting: true} + h.subs[sub] = struct{}{} + h.mu.Unlock() + utils.Verbose("avc hub: joined a running stream, this request's capture settings are ignored") + return sub, nil + } + + started := make(chan struct{}) + h.starting = started + h.epoch++ + epoch := h.epoch + // registered before the start so nothing the source emits in the window + // between src.start returning and this subscriber being installed is lost + sub := &avcSubscriber{ch: make(chan []byte, avcSubscriberQueue)} + h.subs[sub] = struct{}{} + h.mu.Unlock() + + return h.startSource(src, sub, epoch, started) + } +} + +// busyLocked reports the channel to wait on while another subscriber holds the +// hub for a start or a stop, or nil when the hub is free. +func (h *avcHub) busyLocked() chan struct{} { + if h.stopping != nil { + return h.stopping + } + return h.starting +} + +// startSource runs src.start outside the hub lock — on iOS it takes ~10s for +// DeviceKit and the broadcast picker — then installs the stop func. Subscribers +// waiting on the start are released either way; when it failed, one of them +// becomes the next starter and tries again. +func (h *avcHub) startSource(src avcSource, sub *avcSubscriber, epoch uint64, started chan struct{}) (*avcSubscriber, error) { + stop, err := src.start( + func(chunk []byte) { h.publish(epoch, chunk) }, + func(err error) { h.sourceEnded(epoch, err) }, + ) + + h.mu.Lock() + h.starting = nil + defer close(started) + + if err != nil { + h.finishLocked(sub, err) + h.mu.Unlock() + return nil, err + } + if stop == nil { + stop = func() {} + } + if epoch != h.epoch { + // the source already ended while it was still starting + endErr := h.endErr + if endErr == nil { + endErr = errAvcSourceEnded + } + h.finishLocked(sub, endErr) + h.mu.Unlock() + stop() + return nil, endErr + } + + h.stop = stop + h.mu.Unlock() + return sub, nil +} + +// awaitFirstChunk blocks until a late joiner's first chunk (parameter sets plus +// key frame) is queued. It re-asks for a key frame every avcKeyFrameRetry and +// gives up after avcKeyFrameTimeout, so a joiner is never stranded on an +// encoder that ignored the request. +func (h *avcHub) awaitFirstChunk(sub *avcSubscriber, src avcSource, cancel <-chan struct{}) error { + giveUpAfter, retryEvery := h.keyFrameWaits() + timeout := time.NewTimer(giveUpAfter) + defer timeout.Stop() + retry := time.NewTicker(retryEvery) + defer retry.Stop() + + requestAvcKeyFrameFrom(src) + for { + select { + case <-sub.joined: + return nil + case <-retry.C: + requestAvcKeyFrameFrom(src) + case <-timeout.C: + return errAvcNoKeyFrame + case <-cancel: + return errAvcCancelled + } + } +} + +// keyFrameWaits returns how long a late joiner waits for its first key frame and +// how often it re-asks for one. +func (h *avcHub) keyFrameWaits() (giveUpAfter, retryEvery time.Duration) { + giveUpAfter, retryEvery = avcKeyFrameTimeout, avcKeyFrameRetry + if h.keyFrameTimeout > 0 { + giveUpAfter = h.keyFrameTimeout + } + if h.keyFrameRetry > 0 { + retryEvery = h.keyFrameRetry + } + return giveUpAfter, retryEvery +} + +// detach removes a subscriber, failing it with cause unless it already ended, +// and stops the source when it was the last one. +func (h *avcHub) detach(sub *avcSubscriber, cause error) error { + h.mu.Lock() + h.finishLocked(sub, cause) + err := sub.err + stop, stopping := h.takeStopIfIdleLocked() + h.mu.Unlock() + + if stop != nil { + utils.Verbose("avc hub: last subscriber left, stopping source") + h.runStop(stop, stopping) + } + return err +} + +// publish delivers one source chunk to every subscriber. +func (h *avcHub) publish(epoch uint64, chunk []byte) { + h.mu.Lock() + if epoch != h.epoch { + // bytes from a source we already stopped + h.mu.Unlock() + return + } + + buf, idrAt := h.scanLocked(chunk) + data := buf[len(buf)-len(chunk):] + sets := h.parameterSetsLocked() + for sub := range h.subs { + if !sub.waiting { + h.sendLocked(sub, data) + continue + } + // never hand a late joiner a stream it cannot decode: both parameter + // sets must be cached and a key frame must have arrived + if idrAt < 0 || len(h.sps) == 0 || len(h.pps) == 0 { + continue + } + // most encoders repeat SPS/PPS in front of every key frame; the cached + // copies are only for the ones that do not + if units := avc2mp4.ParseNALUnits(buf[idrAt:]); len(units) == 0 || units[0].Type != nalTypeSPS { + h.sendLocked(sub, sets) + } + h.sendLocked(sub, buf[idrAt:]) + h.markJoinedLocked(sub) + } + + stop, stopping := h.takeStopIfIdleLocked() + h.mu.Unlock() + + if stop != nil { + h.runStop(stop, stopping) + } +} + +// sourceEnded fails every subscriber with the source's error and clears the hub +// so a later subscribe starts a fresh source. +func (h *avcHub) sourceEnded(epoch uint64, err error) { + h.mu.Lock() + defer h.mu.Unlock() + + if epoch != h.epoch { + return + } + for sub := range h.subs { + h.finishLocked(sub, err) + } + h.endErr = err + h.clearSourceLocked() + // a late emit from the dead source must not repopulate sps/pps/tail + h.epoch++ +} + +// scanLocked splits the incoming bytes into NAL units, carrying an incomplete +// unit over to the next chunk, and caches the parameter sets. It returns that +// carried-over buffer (the previous partial NAL followed by chunk) and the +// offset of the first IDR start code in it, or -1. +func (h *avcHub) scanLocked(chunk []byte) ([]byte, int) { + buf := make([]byte, 0, len(h.tail)+len(chunk)) + buf = append(append(buf, h.tail...), chunk...) + + // headAt is where the current picture's non-slice units (SPS, PPS, SEI) + // began. a joiner starts there rather than at the IDR slice itself: the + // timecode SEI that stamps a picture precedes its slice, and without it the + // mp4 muxer has to drop the keyframe. + idrAt, headAt := -1, -1 + pos := avc2mp4.FindStartCode(buf, 0) + for pos >= 0 { + next := avc2mp4.FindStartCode(buf, pos+3) + if next < 0 { + // the unit is still incomplete, wait for the next chunk + break + } + if units := avc2mp4.ParseNALUnits(buf[pos:next]); len(units) > 0 { + nalType := units[0].Type + isSlice := nalType >= 1 && nalType <= nalTypeIDR + if !isSlice && headAt < 0 { + headAt = pos + } + switch nalType { + case nalTypeSPS: + h.sps = bytes.Clone(buf[pos:next]) + case nalTypePPS: + h.pps = bytes.Clone(buf[pos:next]) + case nalTypeIDR: + if idrAt < 0 { + idrAt = pos + if headAt >= 0 { + idrAt = headAt + } + } + } + if isSlice { + headAt = -1 + } + } + pos = next + } + + if pos < 0 { + // no start code yet: keep only what could be the head of one + pos = max(len(buf)-3, 0) + } + if headAt >= 0 { + // the buffer ends inside a picture's head: carry it so the slice in the + // next chunk still finds its SEI + pos = min(pos, headAt) + } + h.tail = buf[pos:len(buf):len(buf)] + if len(h.tail) > avcMaxPartialNAL { + utils.Verbose("avc hub: dropping %d bytes without a start code", len(h.tail)) + h.tail = nil + } + return buf, idrAt +} + +// parameterSetsLocked returns the cached SPS and PPS a late joiner needs before +// its first keyframe. +func (h *avcHub) parameterSetsLocked() []byte { + if len(h.sps) == 0 && len(h.pps) == 0 { + return nil + } + sets := make([]byte, 0, len(h.sps)+len(h.pps)) + return append(append(sets, h.sps...), h.pps...) +} + +// sendLocked queues data for one subscriber, dropping the subscriber when its +// queue is full rather than blocking the source. +func (h *avcHub) sendLocked(sub *avcSubscriber, data []byte) { + if len(data) == 0 { + return + } + if _, ok := h.subs[sub]; !ok { + return + } + + select { + case sub.ch <- data: + default: + utils.Verbose("avc hub: subscriber fell %d chunks behind, dropping it", cap(sub.ch)) + sub.dropped.Store(true) + h.finishLocked(sub, errAvcSubscriberTooSlow) + } +} + +// markJoinedLocked records that a late joiner's first chunk is on its way, which +// releases its awaitFirstChunk. +func (h *avcHub) markJoinedLocked(sub *avcSubscriber) { + if !sub.waiting { + return + } + sub.waiting = false + close(sub.joined) +} + +// finishLocked ends one subscriber with err; its subscribe call returns that error. +func (h *avcHub) finishLocked(sub *avcSubscriber, err error) { + if _, ok := h.subs[sub]; !ok { + return + } + delete(h.subs, sub) + sub.err = err + // a joiner that never got its key frame must stop waiting for one + h.markJoinedLocked(sub) + close(sub.ch) +} + +// takeStopIfIdleLocked hands back the source's stop func once no subscriber is +// left, together with the channel reserving the hub while it runs. The caller +// must pass both to runStop after releasing the lock. +func (h *avcHub) takeStopIfIdleLocked() (func(), chan struct{}) { + if h.stop == nil || len(h.subs) > 0 { + return nil, nil + } + stop := h.stop + h.clearSourceLocked() + // anything the stopped source still emits belongs to the old epoch + h.epoch++ + // hold the hub until stop() has returned, so a new subscriber cannot start a + // second source while this one is still tearing down + h.stopping = make(chan struct{}) + return stop, h.stopping +} + +// runStop tears the source down and then releases the hub for the next subscriber. +func (h *avcHub) runStop(stop func(), stopping chan struct{}) { + defer func() { + h.mu.Lock() + if h.stopping == stopping { + h.stopping = nil + } + h.mu.Unlock() + close(stopping) + }() + stop() +} + +func (h *avcHub) clearSourceLocked() { + h.stop = nil + h.sps = nil + h.pps = nil + h.tail = nil +} + +// requestAvcKeyFrameFrom asks the encoder for an immediate sync frame, if the +// source supports it. +func requestAvcKeyFrameFrom(src avcSource) { + if src.requestKeyFrame == nil { + return + } + if err := src.requestKeyFrame(); err != nil { + utils.Verbose("avc hub: keyframe request failed: %v", err) + } +} + +// waitOrCancel blocks until done is closed, reporting false when the caller's +// cancel channel fired first. +func waitOrCancel(done <-chan struct{}, cancel <-chan struct{}) bool { + select { + case <-done: + return true + case <-cancel: + return false + } +} diff --git a/devices/avc_hub_test.go b/devices/avc_hub_test.go new file mode 100644 index 00000000..75a67d13 --- /dev/null +++ b/devices/avc_hub_test.go @@ -0,0 +1,687 @@ +package devices + +import ( + "bytes" + "errors" + "sync" + "sync/atomic" + "testing" + "time" +) + +// --- annex-b builders ------------------------------------------------------- + +// nal builds one annex-b NAL unit: a 4-byte start code, the type byte, payload. +func nal(nalType byte, payload ...byte) []byte { + return append([]byte{0x00, 0x00, 0x00, 0x01, nalType}, payload...) +} + +func spsNAL() []byte { return nal(nalTypeSPS, 0xaa) } +func ppsNAL() []byte { return nal(nalTypePPS, 0xbb) } +func keyFrameNAL(marker byte) []byte { return nal(nalTypeIDR, marker) } +func deltaFrameNAL(marker byte) []byte { return nal(1, marker) } +func timecodeNAL(marker byte) []byte { return nal(6, marker) } +func concat(parts ...[]byte) []byte { return bytes.Join(parts, nil) } +func splitIntoSingleBytes(b []byte) [][]byte { + chunks := make([][]byte, 0, len(b)) + for i := range b { + chunks = append(chunks, b[i:i+1]) + } + return chunks +} + +// --- fake source ------------------------------------------------------------ + +// fakeAvcSource stands in for a device encoder: the test feeds it bytes and +// observes when the hub starts it, stops it, or asks for a key frame. +type fakeAvcSource struct { + mu sync.Mutex + starts int + keyFrameReqs int + live int // sources currently between start and a completed stop + maxLive int + emit func([]byte) + ended func(error) + keyFrameErr error + startGate chan struct{} // when set, start blocks on it + stopGate chan struct{} // when set, stop blocks on it + startBegan chan struct{} + stopBegan chan struct{} + started chan struct{} + stopped chan struct{} + keyFrames chan struct{} +} + +func newFakeAvcSource() *fakeAvcSource { + return &fakeAvcSource{ + startBegan: make(chan struct{}, 16), + stopBegan: make(chan struct{}, 16), + started: make(chan struct{}, 16), + stopped: make(chan struct{}, 16), + keyFrames: make(chan struct{}, 16), + } +} + +func (f *fakeAvcSource) source() avcSource { + return avcSource{ + start: f.start, + requestKeyFrame: f.requestKeyFrame, + } +} + +func (f *fakeAvcSource) start(emit func([]byte), ended func(error)) (func(), error) { + f.mu.Lock() + gate := f.startGate + f.mu.Unlock() + + signalNonBlocking(f.startBegan) + if gate != nil { + <-gate + } + + f.mu.Lock() + f.starts++ + f.live++ + if f.live > f.maxLive { + f.maxLive = f.live + } + f.emit, f.ended = emit, ended + stopGate := f.stopGate + f.mu.Unlock() + signalNonBlocking(f.started) + + return func() { + signalNonBlocking(f.stopBegan) + if stopGate != nil { + <-stopGate + } + f.mu.Lock() + f.live-- + f.mu.Unlock() + signalNonBlocking(f.stopped) + }, nil +} + +func (f *fakeAvcSource) requestKeyFrame() error { + f.mu.Lock() + f.keyFrameReqs++ + err := f.keyFrameErr + f.mu.Unlock() + signalNonBlocking(f.keyFrames) + return err +} + +// blockStart makes the next start hang until the returned func is called, the +// way DeviceKit plus the broadcast picker hangs for ~10s on a real iOS device. +func (f *fakeAvcSource) blockStart() func() { + gate := make(chan struct{}) + f.mu.Lock() + f.startGate = gate + f.mu.Unlock() + return func() { close(gate) } +} + +// blockStop makes stop hang until the returned func is called, standing in for a +// source that is still tearing down. +func (f *fakeAvcSource) blockStop() func() { + gate := make(chan struct{}) + f.mu.Lock() + f.stopGate = gate + f.mu.Unlock() + return func() { close(gate) } +} + +// neverProducesKeyFrames makes every key frame request fail, like an encoder +// that ignores the request. +func (f *fakeAvcSource) neverProducesKeyFrames(err error) { + f.mu.Lock() + f.keyFrameErr = err + f.mu.Unlock() +} + +func (f *fakeAvcSource) keyFrameRequestCount() int { + f.mu.Lock() + defer f.mu.Unlock() + return f.keyFrameReqs +} + +func (f *fakeAvcSource) mostSourcesAliveAtOnce() int { + f.mu.Lock() + defer f.mu.Unlock() + return f.maxLive +} + +func signalNonBlocking(ch chan struct{}) { + select { + case ch <- struct{}{}: + default: + } +} + +func (f *fakeAvcSource) feed(chunks ...[]byte) { + f.mu.Lock() + emit := f.emit + f.mu.Unlock() + for _, chunk := range chunks { + emit(chunk) + } +} + +func (f *fakeAvcSource) die(err error) { + f.mu.Lock() + ended := f.ended + f.mu.Unlock() + ended(err) +} + +func (f *fakeAvcSource) startCount() int { + f.mu.Lock() + defer f.mu.Unlock() + return f.starts +} + +func (f *fakeAvcSource) waitUntilStarted(t *testing.T) { waitForSignal(t, f.started, "source start") } +func (f *fakeAvcSource) waitUntilStopped(t *testing.T) { waitForSignal(t, f.stopped, "source stop") } +func (f *fakeAvcSource) waitUntilStartBegan(t *testing.T) { + waitForSignal(t, f.startBegan, "the source start to begin") +} +func (f *fakeAvcSource) waitUntilStopBegan(t *testing.T) { + waitForSignal(t, f.stopBegan, "the source stop to begin") +} + +func (f *fakeAvcSource) expectNotStartedAgain(t *testing.T) { + t.Helper() + select { + case <-f.startBegan: + t.Fatal("a second source was started while the first one was still stopping") + case <-time.After(50 * time.Millisecond): + } +} +func (f *fakeAvcSource) waitForKeyFrameRequest(t *testing.T) { + waitForSignal(t, f.keyFrames, "key frame request") +} + +func (f *fakeAvcSource) expectStillRunning(t *testing.T) { + t.Helper() + select { + case <-f.stopped: + t.Fatal("source was stopped while subscribers were still attached") + case <-time.After(50 * time.Millisecond): + } +} + +// --- fake subscriber -------------------------------------------------------- + +// capture is one subscriber: it runs subscribe() in the background and records +// every byte the hub hands it. +type capture struct { + mu sync.Mutex + received []byte + gate chan struct{} // when set, onData waits on it before recording + stopAfterFirstChunk bool + readyCalls atomic.Int32 + done chan error + cancel chan struct{} +} + +func startCapture(hub *avcHub, src avcSource) *capture { + return newCapture().run(hub, src) +} + +func startCaptureThatStopsAfterFirstChunk(hub *avcHub, src avcSource) *capture { + c := newCapture() + c.stopAfterFirstChunk = true + return c.run(hub, src) +} + +// startBlockedCapture starts a subscriber whose onData does not return until the +// returned release func is called. +func startBlockedCapture(hub *avcHub, src avcSource) (*capture, func()) { + c := newCapture() + c.gate = make(chan struct{}) + return c.run(hub, src), func() { close(c.gate) } +} + +func newCapture() *capture { + return &capture{done: make(chan error, 1), cancel: make(chan struct{})} +} + +func (c *capture) run(hub *avcHub, src avcSource) *capture { + inner := src.ready + src.ready = func() { + c.readyCalls.Add(1) + if inner != nil { + inner() + } + } + go func() { c.done <- hub.subscribe(src, c.onData, c.cancel) }() + return c +} + +// wasReportedReady tells whether the hub told this capture its stream is live. +func (c *capture) wasReportedReady() bool { return c.readyCalls.Load() > 0 } + +func (c *capture) onData(chunk []byte) bool { + if c.gate != nil { + <-c.gate + } + c.mu.Lock() + defer c.mu.Unlock() + c.received = append(c.received, chunk...) + return !c.stopAfterFirstChunk +} + +func (c *capture) interrupt() { close(c.cancel) } + +func (c *capture) bytes() []byte { + c.mu.Lock() + defer c.mu.Unlock() + return bytes.Clone(c.received) +} + +func (c *capture) waitForBytes(t *testing.T, n int) []byte { + t.Helper() + deadline := time.Now().Add(2 * time.Second) + for { + got := c.bytes() + if len(got) >= n { + return got + } + if time.Now().After(deadline) { + t.Fatalf("timed out waiting for %d bytes, got %d", n, len(got)) + } + time.Sleep(time.Millisecond) + } +} + +func (c *capture) waitUntilFinished(t *testing.T) error { + t.Helper() + select { + case err := <-c.done: + return err + case <-time.After(2 * time.Second): + t.Fatal("timed out waiting for the subscriber to finish") + return nil + } +} + +func waitForSignal(t *testing.T, ch <-chan struct{}, what string) { + t.Helper() + select { + case <-ch: + case <-time.After(2 * time.Second): + t.Fatalf("timed out waiting for %s", what) + } +} + +func expectBytes(t *testing.T, got, want []byte, what string) { + t.Helper() + if !bytes.Equal(got, want) { + t.Fatalf("%s:\n got %x\nwant %x", what, got, want) + } +} + +// --- tests ------------------------------------------------------------------ + +func TestFirstSubscriberReceivesEveryByteFromTheStart(t *testing.T) { + hub, source := &avcHub{}, newFakeAvcSource() + + first := startCapture(hub, source.source()) + source.waitUntilStarted(t) + + stream := concat(deltaFrameNAL(1), spsNAL(), ppsNAL(), keyFrameNAL(2), deltaFrameNAL(3)) + source.feed(stream) + + expectBytes(t, first.waitForBytes(t, len(stream)), stream, "first subscriber") +} + +func TestLateJoinerStartsAtTheParameterSetsAndNextKeyFrame(t *testing.T) { + hub, source := &avcHub{}, newFakeAvcSource() + + first := startCapture(hub, source.source()) + source.waitUntilStarted(t) + firstGOP := concat(spsNAL(), ppsNAL(), keyFrameNAL(1), deltaFrameNAL(2)) + source.feed(firstGOP) + first.waitForBytes(t, len(firstGOP)) + + late := startCapture(hub, source.source()) + source.waitForKeyFrameRequest(t) + + // mid-GOP bytes must never reach the late joiner + source.feed(deltaFrameNAL(3)) + source.feed(concat(keyFrameNAL(4), deltaFrameNAL(5))) + + wantLate := concat(spsNAL(), ppsNAL(), keyFrameNAL(4), deltaFrameNAL(5)) + expectBytes(t, late.waitForBytes(t, len(wantLate)), wantLate, "late joiner") +} + +func TestBothSubscribersReceiveIdenticalBytesAfterTheJoin(t *testing.T) { + hub, source := &avcHub{}, newFakeAvcSource() + + first := startCapture(hub, source.source()) + source.waitUntilStarted(t) + beforeJoin := concat(spsNAL(), ppsNAL(), keyFrameNAL(1), deltaFrameNAL(2)) + source.feed(beforeJoin) + + late := startCapture(hub, source.source()) + source.waitForKeyFrameRequest(t) + + afterJoin := concat(keyFrameNAL(3), deltaFrameNAL(4), deltaFrameNAL(5), deltaFrameNAL(6)) + source.feed(afterJoin) + + parameterSets := concat(spsNAL(), ppsNAL()) + firstBytes := first.waitForBytes(t, len(beforeJoin)+len(afterJoin)) + lateBytes := late.waitForBytes(t, len(parameterSets)+len(afterJoin)) + + expectBytes(t, firstBytes, concat(beforeJoin, afterJoin), "first subscriber") + expectBytes(t, lateBytes[len(parameterSets):], afterJoin, "late joiner from the key frame on") +} + +func TestLastSubscriberLeavingStopsTheSource(t *testing.T) { + hub, source := &avcHub{}, newFakeAvcSource() + + staying := startCapture(hub, source.source()) + source.waitUntilStarted(t) + + leaving := startCapture(hub, source.source()) + source.waitForKeyFrameRequest(t) + source.feed(concat(spsNAL(), ppsNAL(), keyFrameNAL(1), deltaFrameNAL(2))) + leaving.waitForBytes(t, 1) + + leaving.interrupt() + if err := leaving.waitUntilFinished(t); err != nil { + t.Fatalf("interrupted subscriber returned %v, want nil", err) + } + source.expectStillRunning(t) + + staying.interrupt() + if err := staying.waitUntilFinished(t); err != nil { + t.Fatalf("last subscriber returned %v, want nil", err) + } + source.waitUntilStopped(t) +} + +func TestASilentSourceStillLetsASubscriberStopPromptly(t *testing.T) { + hub, source := &avcHub{}, newFakeAvcSource() + + only := startCapture(hub, source.source()) + source.waitUntilStarted(t) + + // an idle android virtual display emits very few frames, so stopping must + // not wait for the next one to arrive + only.interrupt() + + if err := only.waitUntilFinished(t); err != nil { + t.Fatalf("stopped subscriber returned %v, want nil", err) + } + source.waitUntilStopped(t) +} + +func TestWatchStopCancelsWhenTheCallerStops(t *testing.T) { + stop := make(chan struct{}) + cancelled, stopWatching := watchStop(stop) + defer stopWatching() + + close(stop) + waitForSignal(t, cancelled, "cancellation") +} + +func TestWatchStopWithoutAStopChannelWaitsForASignalOnly(t *testing.T) { + cancelled, stopWatching := watchStop(nil) + defer stopWatching() + + select { + case <-cancelled: + t.Fatal("a capture with no stop channel was cancelled without a signal") + case <-time.After(50 * time.Millisecond): + } +} + +func TestResubscribingAfterTheSourceStoppedStartsItAgain(t *testing.T) { + hub, source := &avcHub{}, newFakeAvcSource() + + first := startCaptureThatStopsAfterFirstChunk(hub, source.source()) + source.waitUntilStarted(t) + source.feed(deltaFrameNAL(1)) + if err := first.waitUntilFinished(t); err != nil { + t.Fatalf("first subscriber returned %v, want nil", err) + } + source.waitUntilStopped(t) + + // the fresh source has no cached parameter sets, so this subscriber is a + // first subscriber again and gets everything from byte 0 + second := startCapture(hub, source.source()) + source.waitUntilStarted(t) + if got := source.startCount(); got != 2 { + t.Fatalf("source started %d times, want 2", got) + } + + source.feed(deltaFrameNAL(9)) + expectBytes(t, second.waitForBytes(t, len(deltaFrameNAL(9))), deltaFrameNAL(9), "restarted stream") +} + +func TestSlowSubscriberIsDroppedWithoutStallingTheOther(t *testing.T) { + hub, source := &avcHub{}, newFakeAvcSource() + + slow, releaseSlow := startBlockedCapture(hub, source.source()) + source.waitUntilStarted(t) + + fast := startCapture(hub, source.source()) + source.waitForKeyFrameRequest(t) + source.feed(concat(spsNAL(), ppsNAL(), keyFrameNAL(1))) + + // overflow the slow subscriber's queue while it sits inside onData + frames := make([][]byte, 0, avcSubscriberQueue+16) + for i := 0; i < avcSubscriberQueue+16; i++ { + frames = append(frames, deltaFrameNAL(byte(i))) + } + source.feed(frames...) + + // the fast subscriber kept flowing the whole time + fast.waitForBytes(t, len(keyFrameNAL(1))+len(frames)*len(frames[0])) + + releaseSlow() + if err := slow.waitUntilFinished(t); !errors.Is(err, errAvcSubscriberTooSlow) { + t.Fatalf("slow subscriber returned %v, want %v", err, errAvcSubscriberTooSlow) + } +} + +func TestStartCodeSplitAcrossChunksStillAlignsALateJoiner(t *testing.T) { + hub, source := &avcHub{}, newFakeAvcSource() + + first := startCapture(hub, source.source()) + source.waitUntilStarted(t) + + late := startCapture(hub, source.source()) + source.waitForKeyFrameRequest(t) + + gop := concat(spsNAL(), ppsNAL(), keyFrameNAL(1), deltaFrameNAL(2)) + source.feed(splitIntoSingleBytes(gop)...) + source.feed(deltaFrameNAL(3)) // completes the trailing NAL unit + + wantLate := concat(spsNAL(), ppsNAL(), keyFrameNAL(1), deltaFrameNAL(2), deltaFrameNAL(3)) + expectBytes(t, late.waitForBytes(t, len(wantLate)), wantLate, "late joiner over split chunks") + expectBytes(t, first.waitForBytes(t, len(gop)+len(deltaFrameNAL(3))), concat(gop, deltaFrameNAL(3)), "first subscriber over split chunks") +} + +// the timecode SEI that stamps a picture comes before its slice; a joiner that +// started at the slice would hand the mp4 muxer a key frame with no timestamp +func TestLateJoinerReceivesTheTimecodeThatPrecedesItsKeyFrame(t *testing.T) { + hub, source := &avcHub{}, newFakeAvcSource() + + first := startCapture(hub, source.source()) + source.waitUntilStarted(t) + source.feed(concat(spsNAL(), ppsNAL(), timecodeNAL(1), keyFrameNAL(1), timecodeNAL(2), deltaFrameNAL(2))) + first.waitForBytes(t, 1) + + late := startCapture(hub, source.source()) + source.waitForKeyFrameRequest(t) + + // the head of the key frame arrives in its own chunk, as it does over tcp + source.feed(concat(spsNAL(), ppsNAL(), timecodeNAL(3))) + source.feed(concat(keyFrameNAL(3), timecodeNAL(4))) + source.feed(deltaFrameNAL(4)) // completes the trailing NAL unit + + wantLate := concat(spsNAL(), ppsNAL(), timecodeNAL(3), keyFrameNAL(3), timecodeNAL(4), deltaFrameNAL(4)) + expectBytes(t, late.waitForBytes(t, len(wantLate)), wantLate, "late joiner starts at its key frame's head") +} + +func TestSourceErrorEndsEverySubscriber(t *testing.T) { + hub, source := &avcHub{}, newFakeAvcSource() + + first := startCapture(hub, source.source()) + source.waitUntilStarted(t) + + late := startCapture(hub, source.source()) + source.waitForKeyFrameRequest(t) + + sourceErr := errors.New("stream died") + source.die(sourceErr) + + if err := first.waitUntilFinished(t); !errors.Is(err, sourceErr) { + t.Fatalf("first subscriber returned %v, want %v", err, sourceErr) + } + if err := late.waitUntilFinished(t); !errors.Is(err, sourceErr) { + t.Fatalf("late joiner returned %v, want %v", err, sourceErr) + } +} + +// hubWithQuickKeyFrameWait is a hub whose late joiners give up on a missing key +// frame in milliseconds instead of seconds. +func hubWithQuickKeyFrameWait() *avcHub { + return &avcHub{keyFrameTimeout: 200 * time.Millisecond, keyFrameRetry: 10 * time.Millisecond} +} + +func TestASubscriberCanGiveUpWhileAnotherSubscribersStartIsStillRunning(t *testing.T) { + hub, source := &avcHub{}, newFakeAvcSource() + finishTheSlowStart := source.blockStart() + + // on a real iOS device this start sits in the broadcast picker for ~10s + first := startCapture(hub, source.source()) + source.waitUntilStartBegan(t) + + impatient := startCapture(hub, source.source()) + impatient.interrupt() + if err := impatient.waitUntilFinished(t); err != nil { + t.Fatalf("subscriber cancelled during someone else's start returned %v, want nil", err) + } + + finishTheSlowStart() + source.waitUntilStarted(t) + if got := source.startCount(); got != 1 { + t.Fatalf("source started %d times, want 1", got) + } + + source.feed(deltaFrameNAL(1)) + expectBytes(t, first.waitForBytes(t, len(deltaFrameNAL(1))), deltaFrameNAL(1), "the subscriber that waited out the start") +} + +func TestALateJoinerGivesUpWhenNoKeyFrameEverArrives(t *testing.T) { + hub, source := hubWithQuickKeyFrameWait(), newFakeAvcSource() + source.neverProducesKeyFrames(errors.New("encoder refused")) + + staying := startCapture(hub, source.source()) + source.waitUntilStarted(t) + source.feed(concat(spsNAL(), ppsNAL(), keyFrameNAL(1))) + staying.waitForBytes(t, 1) + + // no further key frame is ever produced, so the joiner must fail rather than + // block its caller forever + late := startCapture(hub, source.source()) + if err := late.waitUntilFinished(t); !errors.Is(err, errAvcNoKeyFrame) { + t.Fatalf("late joiner returned %v, want %v", err, errAvcNoKeyFrame) + } + if late.wasReportedReady() { + t.Fatal("a late joiner that never received a key frame was reported ready") + } + if got := source.keyFrameRequestCount(); got < 2 { + t.Fatalf("key frame was requested %d times, want it retried at least twice", got) + } + source.expectStillRunning(t) +} + +func TestALateJoinerIsOnlyReportedReadyOnceItsKeyFrameArrives(t *testing.T) { + hub, source := hubWithQuickKeyFrameWait(), newFakeAvcSource() + + startCapture(hub, source.source()) + source.waitUntilStarted(t) + source.feed(concat(spsNAL(), ppsNAL(), keyFrameNAL(1))) + + late := startCapture(hub, source.source()) + source.waitForKeyFrameRequest(t) + if late.wasReportedReady() { + t.Fatal("a late joiner was reported ready before its first key frame") + } + + source.feed(concat(keyFrameNAL(2), deltaFrameNAL(3))) + late.waitForBytes(t, 1) + if !late.wasReportedReady() { + t.Fatal("a late joiner receiving bytes was never reported ready") + } +} + +func TestANewSubscriberWaitsForTheOldSourceToFinishStopping(t *testing.T) { + hub, source := &avcHub{}, newFakeAvcSource() + finishTheSlowStop := source.blockStop() + + first := startCapture(hub, source.source()) + source.waitUntilStartBegan(t) + source.waitUntilStarted(t) + first.interrupt() + source.waitUntilStopBegan(t) + + // a second iOS dial here would steal the broadcast extension's stream from + // the source that is still closing its conn + second := startCapture(hub, source.source()) + source.expectNotStartedAgain(t) + + finishTheSlowStop() + source.waitUntilStarted(t) + if got := source.mostSourcesAliveAtOnce(); got != 1 { + t.Fatalf("%d sources were alive at once, want 1", got) + } + + source.feed(deltaFrameNAL(7)) + expectBytes(t, second.waitForBytes(t, len(deltaFrameNAL(7))), deltaFrameNAL(7), "the subscriber that waited out the stop") +} + +func TestADroppedSubscriberStopsWithoutDrainingItsStaleQueue(t *testing.T) { + hub, source := &avcHub{}, newFakeAvcSource() + + slow, releaseSlow := startBlockedCapture(hub, source.source()) + source.waitUntilStarted(t) + + firstChunk := concat(spsNAL(), ppsNAL(), keyFrameNAL(1)) + source.feed(firstChunk) + + backlog := make([][]byte, 0, avcSubscriberQueue+16) + for i := 0; i < avcSubscriberQueue+16; i++ { + backlog = append(backlog, deltaFrameNAL(byte(i))) + } + source.feed(backlog...) + + releaseSlow() + if err := slow.waitUntilFinished(t); !errors.Is(err, errAvcSubscriberTooSlow) { + t.Fatalf("slow subscriber returned %v, want %v", err, errAvcSubscriberTooSlow) + } + // it left right after the chunk it was already inside, instead of writing + // out a queue full of stale ones + expectBytes(t, slow.bytes(), firstChunk, "dropped subscriber") +} + +func TestConcurrentFirstSubscribersStartTheSourceOnce(t *testing.T) { + hub, source := &avcHub{}, newFakeAvcSource() + + const subscribers = 5 + for i := 0; i < subscribers; i++ { + startCapture(hub, source.source()) + } + + source.waitUntilStarted(t) + for i := 0; i < subscribers-1; i++ { + source.waitForKeyFrameRequest(t) + } + if got := source.startCount(); got != 1 { + t.Fatalf("source started %d times, want 1", got) + } +} diff --git a/devices/common.go b/devices/common.go index 5319525a..3598ba64 100644 --- a/devices/common.go +++ b/devices/common.go @@ -120,6 +120,7 @@ type ScreenCaptureConfig struct { OnProgress func(message string) // optional progress callback OnReady func() // optional: called once capture is confirmed live (e.g. after the ReplayKit broadcast picker is clicked), before streaming begins OnData func([]byte) bool // data callback - return false to stop + StopChan <-chan struct{} // optional (avc only): stops this capture when closed, even while no frames are arriving } // StartAgentConfig contains configuration for agent startup operations diff --git a/devices/ios.go b/devices/ios.go index 4d83462c..f1525a84 100644 --- a/devices/ios.go +++ b/devices/ios.go @@ -9,12 +9,10 @@ import ( "io" "net" "os" - "os/signal" "path/filepath" "strconv" "strings" "sync" - "syscall" "time" goios "github.com/danielpaulus/go-ios/ios" @@ -88,6 +86,7 @@ type IOSDevice struct { locationSimulationService *instruments.LocationSimulationService // open while an ios 17+ location override is held avcWriteMu sync.Mutex // serializes control writes on avcStreamConn + avcStream avcHub // fans the single h264 stream out to every concurrent capture } func (d *IOSDevice) ID() string { @@ -990,136 +989,143 @@ func (d *IOSDevice) sendAvcControl(method string, params map[string]any) error { return nil } -func (d *IOSDevice) StartScreenCapture(config ScreenCaptureConfig) error { - // handle avc format via DeviceKit - if config.Format == "avc" { +// ensureDeviceKitAvcRunning returns the ports of a DeviceKit session ready to +// stream H.264, starting it (and its broadcast) when it is not up yet. +func (d *IOSDevice) ensureDeviceKitAvcRunning(config ScreenCaptureConfig) (*DeviceKitInfo, error) { + if config.OnProgress != nil { + config.OnProgress("Checking DeviceKit status") + } + + // DeviceKit not running, start it normally + if !d.isDeviceKitRunning() { if config.OnProgress != nil { - config.OnProgress("Checking DeviceKit status") + config.OnProgress("Starting DeviceKit for H.264 streaming") } - var deviceKitInfo *DeviceKitInfo - var err error + // start DeviceKit + // Note: passing nil registry since this is internal call from StartScreenCapture + // ScreenCapture callers should have already registered the device via StartAgent + deviceKitInfo, err := d.StartDeviceKitAvc(nil) + if err != nil { + return nil, fmt.Errorf("failed to start DeviceKit: %w", err) + } + return deviceKitInfo, nil + } - // check if DeviceKit is already running - if d.isDeviceKitRunning() { - utils.Verbose("DeviceKit already running, reusing existing session") + utils.Verbose("DeviceKit already running, reusing existing session") - // check if we need to create port forwarders - d.mu.Lock() - hasHTTPForwarder := d.portForwarderDeviceKit != nil && d.portForwarderDeviceKit.IsRunning() - hasStreamForwarder := d.portForwarderAvc != nil && d.portForwarderAvc.IsRunning() - d.mu.Unlock() + // check if we need to create port forwarders + d.mu.Lock() + hasHTTPForwarder := d.portForwarderDeviceKit != nil && d.portForwarderDeviceKit.IsRunning() + hasStreamForwarder := d.portForwarderAvc != nil && d.portForwarderAvc.IsRunning() + d.mu.Unlock() - if hasHTTPForwarder && hasStreamForwarder { - // reuse existing forwarders - d.mu.Lock() - httpPort, _ := d.portForwarderDeviceKit.GetPorts() - streamPort, _ := d.portForwarderAvc.GetPorts() - d.mu.Unlock() + if !hasHTTPForwarder || !hasStreamForwarder { + // DeviceKit running but we need to create forwarders + deviceKitInfo, err := d.ensureDeviceKitPortForwarders() + if err != nil { + return nil, fmt.Errorf("failed to create port forwarders: %w", err) + } + if config.OnProgress != nil { + config.OnProgress("Using existing DeviceKit session") + } + return deviceKitInfo, nil + } - deviceKitInfo = &DeviceKitInfo{ - HTTPPort: httpPort, - StreamPort: streamPort, - } - } else { - // DeviceKit running but we need to create forwarders - deviceKitInfo, err = d.ensureDeviceKitPortForwarders() - if err != nil { - return fmt.Errorf("failed to create port forwarders: %w", err) - } - } + // reuse existing forwarders + d.mu.Lock() + httpPort, _ := d.portForwarderDeviceKit.GetPorts() + streamPort, _ := d.portForwarderAvc.GetPorts() + d.mu.Unlock() - if config.OnProgress != nil { - config.OnProgress("Using existing DeviceKit session") - } - } else { - // DeviceKit not running, start it normally - if config.OnProgress != nil { - config.OnProgress("Starting DeviceKit for H.264 streaming") - } + if config.OnProgress != nil { + config.OnProgress("Using existing DeviceKit session") + } + return &DeviceKitInfo{HTTPPort: httpPort, StreamPort: streamPort}, nil +} - // start DeviceKit - // Note: passing nil registry since this is internal call from StartScreenCapture - // ScreenCapture callers should have already registered the device via StartAgent - deviceKitInfo, err = d.StartDeviceKitAvc(nil) - if err != nil { - return fmt.Errorf("failed to start DeviceKit: %w", err) - } - } +// startAvcStream starts DeviceKit, dials its H.264 stream and reads it into +// emit until the conn dies. It is the avcHub source for real iOS devices. +func (d *IOSDevice) startAvcStream(config ScreenCaptureConfig, emit func([]byte), ended func(error)) (func(), error) { + deviceKitInfo, err := d.ensureDeviceKitAvcRunning(config) + if err != nil { + return nil, err + } - // DeviceKit is confirmed running (either reused or freshly started and the - // broadcast picker was clicked) — safe to tell the caller capture is live. - if config.OnReady != nil { - config.OnReady() - } + if config.OnProgress != nil { + config.OnProgress(fmt.Sprintf("Connecting to H.264 stream on localhost:%d", deviceKitInfo.StreamPort)) + } - if config.OnProgress != nil { - config.OnProgress(fmt.Sprintf("Connecting to H.264 stream on localhost:%d", deviceKitInfo.StreamPort)) - } + // connect to the TCP stream + conn, err := net.Dial("tcp", fmt.Sprintf("localhost:%d", deviceKitInfo.StreamPort)) + if err != nil { + return nil, fmt.Errorf("failed to connect to stream port: %w", err) + } - // connect to the TCP stream - conn, err := net.Dial("tcp", fmt.Sprintf("localhost:%d", deviceKitInfo.StreamPort)) - if err != nil { - return fmt.Errorf("failed to connect to stream port: %w", err) - } + // Expose the stream conn as the live encoder control channel. Control + // must ride this exact conn: the extension's TCPServer redirects video + // output to its newest client, so a separate control connection would + // steal the stream. + d.mu.Lock() + d.avcStreamConn = conn + d.mu.Unlock() - // Expose the stream conn as the live encoder control channel. Control - // must ride this exact conn: the extension's TCPServer redirects video - // output to its newest client, so a separate control connection would - // steal the stream. + // closing must be complete by the time the hub's stop() returns, and by the + // time ended() runs: the hub only ever allows one live stream per device, + // and a lingering conn would still be the extension's newest client. + closeStream := sync.OnceFunc(func() { + _ = conn.Close() d.mu.Lock() - d.avcStreamConn = conn + if d.avcStreamConn == conn { + d.avcStreamConn = nil + } d.mu.Unlock() - defer func() { - d.mu.Lock() - if d.avcStreamConn == conn { - d.avcStreamConn = nil - } - d.mu.Unlock() - }() - - // setup signal handling for Ctrl+C - sigChan := make(chan os.Signal, 1) - signal.Notify(sigChan, syscall.SIGINT, syscall.SIGTERM) - - // channel to signal when streaming is done - done := make(chan error, 1) - - // stream data in a goroutine - go func() { - defer func() { _ = conn.Close() }() - buffer := make([]byte, 65536) - for { - n, err := conn.Read(buffer) - if err != nil { - if err != io.EOF { - done <- fmt.Errorf("error reading from stream: %w", err) - } else { - done <- nil - } - return - } + }) - if n > 0 { - if !config.OnData(buffer[:n]) { - // client wants to stop the stream - done <- nil - return - } + go func() { + buffer := make([]byte, 65536) + var readErr error + for { + n, err := conn.Read(buffer) + if n > 0 { + emit(buffer[:n]) + } + if err != nil { + if err != io.EOF { + readErr = fmt.Errorf("error reading from stream: %w", err) } + break } - }() - - // wait for either signal or stream completion - select { - case <-sigChan: - _ = conn.Close() - utils.Verbose("stream closed by user") - return nil - case err := <-done: - utils.Verbose("stream ended") - return err } + + closeStream() + ended(readErr) + }() + + return closeStream, nil +} + +func (d *IOSDevice) StartScreenCapture(config ScreenCaptureConfig) error { + // handle avc format via DeviceKit. all concurrent captures share one + // broadcast: dialing a second conn would make the extension's TCPServer + // redirect the video to it and starve the first one. + if config.Format == "avc" { + cancel, stopWatching := watchStop(config.StopChan) + defer stopWatching() + + err := d.avcStream.subscribe(avcSource{ + start: func(emit func([]byte), ended func(error)) (func(), error) { + return d.startAvcStream(config, emit, ended) + }, + // DeviceKit is confirmed running (either reused or freshly started and + // the broadcast picker was clicked) — safe to tell the caller capture + // is live. a late joiner is live right away. + ready: config.OnReady, + requestKeyFrame: func() error { return RequestAvcKeyFrame(d) }, + }, config.OnData, cancel) + + utils.Verbose("stream ended") + return err } // mjpeg is served on the same port as the agent HTTP server at /mjpeg diff --git a/pkg/avc2mp4/nalparser.go b/pkg/avc2mp4/nalparser.go index b8d14d10..756f9bc5 100644 --- a/pkg/avc2mp4/nalparser.go +++ b/pkg/avc2mp4/nalparser.go @@ -14,7 +14,7 @@ func ParseNALUnits(data []byte) []NALUnit { i := 0 // find the first start code - i = findStartCode(data, i) + i = FindStartCode(data, i) if i < 0 { return nil } @@ -30,7 +30,7 @@ func ParseNALUnits(data []byte) []NALUnit { nalStart := i // find the next start code (or end of data) - next := findStartCode(data, i) + next := FindStartCode(data, i) if next < 0 { next = n } @@ -49,7 +49,9 @@ func ParseNALUnits(data []byte) []NALUnit { return units } -func findStartCode(data []byte, pos int) int { +// FindStartCode returns the index of the first Annex B start code at or after +// pos, or -1 when data holds none from there on. +func FindStartCode(data []byte, pos int) int { n := len(data) for i := pos; i+2 < n; i++ { if data[i] == 0x00 && data[i+1] == 0x00 { From f180d5ff4b8c949752b1be2a0289d096540b62dc Mon Sep 17 00:00:00 2001 From: gmegidish Date: Sat, 19 Sep 2026 12:48:27 +0200 Subject: [PATCH 2/4] test: make avc hub slow-subscriber tests independent of scheduling On a small CI runner the dropped-subscriber test flooded the queue before the subscriber had taken its first chunk, and on a single cpu the feeder outran the fast subscriber so the hub dropped it too. Wait until the subscriber is inside onData, and feed in batches the fast subscriber can drain. --- devices/avc_hub_test.go | 39 ++++++++++++++++++++++++++++++--------- 1 file changed, 30 insertions(+), 9 deletions(-) diff --git a/devices/avc_hub_test.go b/devices/avc_hub_test.go index 75a67d13..a59f59ce 100644 --- a/devices/avc_hub_test.go +++ b/devices/avc_hub_test.go @@ -218,6 +218,8 @@ type capture struct { mu sync.Mutex received []byte gate chan struct{} // when set, onData waits on it before recording + entered chan struct{} // closed the first time onData is called + enteredOnce sync.Once stopAfterFirstChunk bool readyCalls atomic.Int32 done chan error @@ -243,7 +245,7 @@ func startBlockedCapture(hub *avcHub, src avcSource) (*capture, func()) { } func newCapture() *capture { - return &capture{done: make(chan error, 1), cancel: make(chan struct{})} + return &capture{done: make(chan error, 1), cancel: make(chan struct{}), entered: make(chan struct{})} } func (c *capture) run(hub *avcHub, src avcSource) *capture { @@ -261,7 +263,19 @@ func (c *capture) run(hub *avcHub, src avcSource) *capture { // wasReportedReady tells whether the hub told this capture its stream is live. func (c *capture) wasReportedReady() bool { return c.readyCalls.Load() > 0 } +// waitUntilHandlingAChunk returns once the subscriber is inside onData, so a +// test knows that chunk left the queue before it floods the rest. +func (c *capture) waitUntilHandlingAChunk(t *testing.T) { + t.Helper() + select { + case <-c.entered: + case <-time.After(2 * time.Second): + t.Fatal("subscriber never started handling a chunk") + } +} + func (c *capture) onData(chunk []byte) bool { + c.enteredOnce.Do(func() { close(c.entered) }) if c.gate != nil { <-c.gate } @@ -470,15 +484,21 @@ func TestSlowSubscriberIsDroppedWithoutStallingTheOther(t *testing.T) { source.waitForKeyFrameRequest(t) source.feed(concat(spsNAL(), ppsNAL(), keyFrameNAL(1))) - // overflow the slow subscriber's queue while it sits inside onData - frames := make([][]byte, 0, avcSubscriberQueue+16) - for i := 0; i < avcSubscriberQueue+16; i++ { - frames = append(frames, deltaFrameNAL(byte(i))) + // overflow the slow subscriber's queue while it sits inside onData. feed in + // batches the fast subscriber can drain, or on a single cpu the feeder + // outruns it and the hub rightly drops both + const batch = 32 + fed := len(keyFrameNAL(1)) + for i := 0; i < avcSubscriberQueue+batch; i += batch { + frames := make([][]byte, batch) + for j := range frames { + frames[j] = deltaFrameNAL(byte(i + j)) + fed += len(frames[j]) + } + source.feed(frames...) + // the fast subscriber kept flowing the whole time + fast.waitForBytes(t, fed) } - source.feed(frames...) - - // the fast subscriber kept flowing the whole time - fast.waitForBytes(t, len(keyFrameNAL(1))+len(frames)*len(frames[0])) releaseSlow() if err := slow.waitUntilFinished(t); !errors.Is(err, errAvcSubscriberTooSlow) { @@ -653,6 +673,7 @@ func TestADroppedSubscriberStopsWithoutDrainingItsStaleQueue(t *testing.T) { firstChunk := concat(spsNAL(), ppsNAL(), keyFrameNAL(1)) source.feed(firstChunk) + slow.waitUntilHandlingAChunk(t) backlog := make([][]byte, 0, avcSubscriberQueue+16) for i := 0; i < avcSubscriberQueue+16; i++ { From 8f01ad14a26a7c4958aa83b0b8fbe8c1ed13fa7f Mon Sep 17 00:00:00 2001 From: gmegidish Date: Sat, 19 Sep 2026 13:07:27 +0200 Subject: [PATCH 3/4] fix(ios): apply requested avc bitrate at start, signal Ready on early stop - --bitrate was never sent to the broadcast extension on iOS; send it over the control channel once the stream conn is up (zero keeps the default) - a stop that lands before the stream went live detaches with a nil error and no OnReady, so Ready was never signaled; answer it with the no-data error --- commands/screenrecord.go | 7 ++++++- devices/ios.go | 8 ++++++++ 2 files changed, 14 insertions(+), 1 deletion(-) diff --git a/commands/screenrecord.go b/commands/screenrecord.go index 6fea65ec..a4b45cfa 100644 --- a/commands/screenrecord.go +++ b/commands/screenrecord.go @@ -273,7 +273,12 @@ func screenRecordAvc(targetDevice devices.ControllableDevice, req ScreenRecordRe } if len(data) == 0 { - return NewErrorResponse(fmt.Errorf("no data captured")) + // a stop that lands before the stream ever went live detaches cleanly + // (nil error) without OnReady firing; whoever waits on Ready still + // needs an answer. no-op if OnReady already signaled. + err := fmt.Errorf("no data captured") + req.signalReady(err) + return NewErrorResponse(err) } outFile, err := os.Create(req.OutputPath) diff --git a/devices/ios.go b/devices/ios.go index f1525a84..2087d2ec 100644 --- a/devices/ios.go +++ b/devices/ios.go @@ -1070,6 +1070,14 @@ func (d *IOSDevice) startAvcStream(config ScreenCaptureConfig, emit func([]byte) d.avcStreamConn = conn d.mu.Unlock() + // zero keeps the extension's default. a failure is not worth losing the + // stream over: the encoder simply stays on its current bitrate. + if config.Bitrate > 0 { + if err := SetAvcBitrate(d, config.Bitrate); err != nil { + utils.Verbose("failed to apply initial avc bitrate %d: %v", config.Bitrate, err) + } + } + // closing must be complete by the time the hub's stop() returns, and by the // time ended() runs: the hub only ever allows one live stream per device, // and a lingering conn would still be the extension's newest client. From 9afdd155249210c3d7096a3a79575fa5853d1f22 Mon Sep 17 00:00:00 2001 From: gmegidish Date: Sat, 19 Sep 2026 14:19:21 +0200 Subject: [PATCH 4/4] fix(server): detach a screencapture as soon as its client goes away Both capture handlers only noticed a gone client inside OnData, which runs when a frame arrives. On a silent stream (static screen) the subscriber stayed attached to the shared avc stream, so the broadcast was never released and the next capture joined a stream nobody was reading. Pass the request context as StopChan so the hub detaches it immediately. --- server/daemon_streams.go | 3 +++ server/server.go | 8 ++++++++ 2 files changed, 11 insertions(+) diff --git a/server/daemon_streams.go b/server/daemon_streams.go index 060ea067..9405e16c 100644 --- a/server/daemon_streams.go +++ b/server/daemon_streams.go @@ -104,6 +104,9 @@ func streamScreenCapture(ctx context.Context, params json.RawMessage, notify fun FPS: fps, Bitrate: req.Bitrate, OnProgress: onProgress, + // OnData only runs when a frame arrives, and a static screen emits none: + // without this a client that disconnected stays subscribed to the stream + StopChan: ctx.Done(), OnData: func(data []byte) bool { if ctx.Err() != nil { return false diff --git a/server/server.go b/server/server.go index 00b4c173..3b49ce8a 100644 --- a/server/server.go +++ b/server/server.go @@ -1851,6 +1851,13 @@ func handleScreenCapture(r *http.Request, w http.ResponseWriter, params json.Raw return fmt.Errorf("error starting agent: %w", err) } + // a write only fails once a frame arrives; on a static screen a viewer that + // went away would otherwise stay subscribed to the stream. r is nil in tests. + var viewerGone <-chan struct{} + if r != nil { + viewerGone = r.Context().Done() + } + // start screen capture and stream to the response writer err = targetDevice.StartScreenCapture(devices.ScreenCaptureConfig{ Format: screenCaptureParams.Format, @@ -1858,6 +1865,7 @@ func handleScreenCapture(r *http.Request, w http.ResponseWriter, params json.Raw Scale: scale, FPS: fps, OnProgress: progressCallback, + StopChan: viewerGone, OnData: func(data []byte) bool { _, writeErr := w.Write(data) if writeErr != nil {