fix(screenrecord): validate the time limit and clean up interrupted recordings

commands/screenrecord.go

@@ -48,8 +48,27 @@ type ScreenRecordResponse struct {
 	EndReason  string `json:"endReason,omitempty"` // "size_limit" when the recording hit the size safety stop
 }
 
+// MaxScreenRecordTimeLimit is the longest time limit, in seconds, a recording
+// accepts: the largest value Android's screenrecord takes (older Android
+// versions refuse anything above 180, and that error is reported as is).
+const MaxScreenRecordTimeLimit = 4294967295
+
+// ValidateScreenRecordTimeLimit rejects a time limit that is negative or
+// above MaxScreenRecordTimeLimit. 0 means no limit.
+func ValidateScreenRecordTimeLimit(timeLimit int) error {
+	if timeLimit < 0 || int64(timeLimit) > MaxScreenRecordTimeLimit {
+		return fmt.Errorf("invalid time limit %d: must be 0 (no limit) or between 1 and %d seconds", timeLimit, MaxScreenRecordTimeLimit)
+	}
+	return nil
+}
+
 // ScreenRecordCommand records the device screen to an MP4 file.
 func ScreenRecordCommand(req ScreenRecordRequest) *CommandResponse {
+	if err := ValidateScreenRecordTimeLimit(req.TimeLimit); err != nil {
+		req.signalReady(err)
+		return NewErrorResponse(err)
+	}
+
 	targetDevice, err := FindDeviceOrAutoSelect(req.DeviceID)
 	if err != nil {
 		req.signalReady(err)

commands/screenrecord_test.go

@@ -1,6 +1,7 @@
 package commands
 
 import (
+	"strings"
 	"testing"
 	"time"
 )
@@ -28,3 +29,39 @@ func TestWithStopChanStopsAfterTimeLimitElapsedSinceFirstFrame(t *testing.T) {
 		t.Fatal("frame after time limit elapsed should stop the recording")
 	}
 }
+
+func TestScreenRecordRejectsNegativeTimeLimit(t *testing.T) {
+	resp := ScreenRecordCommand(ScreenRecordRequest{DeviceID: "no-such-device", TimeLimit: -1})
+
+	if resp.Status != "error" || !strings.Contains(resp.Error, "invalid time limit -1") {
+		t.Fatalf("expected an invalid time limit error, got %+v", resp)
+	}
+}
+
+func TestScreenRecordRejectsTimeLimitAboveMaximum(t *testing.T) {
+	if err := ValidateScreenRecordTimeLimit(MaxScreenRecordTimeLimit + 1); err == nil {
+		t.Fatal("a time limit above the maximum should be rejected")
+	}
+}
+
+func TestScreenRecordAcceptsNoLimitAndMaximumTimeLimit(t *testing.T) {
+	for _, limit := range []int{0, 1, 180, MaxScreenRecordTimeLimit} {
+		if err := ValidateScreenRecordTimeLimit(limit); err != nil {
+			t.Fatalf("time limit %d should be accepted: %v", limit, err)
+		}
+	}
+}
+
+func TestScreenRecordReportsInvalidTimeLimitToReadyWaiter(t *testing.T) {
+	ready := make(chan error, 1)
+	ScreenRecordCommand(ScreenRecordRequest{TimeLimit: -5, Ready: ready})
+
+	select {
+	case err := <-ready:
+		if err == nil {
+			t.Fatal("ready should carry the validation error")
+		}
+	default:
+		t.Fatal("ready was not signaled")
+	}
+}

daemon/client.go

@@ -155,13 +155,13 @@ func Shutdown() error {
 	if _, err := Call(methodShutdown, nil); err != nil {
 		return err
 	}
-	if err := waitForSocketGone(paths.Socket, shutdownGrace); err != nil {
+	if err := waitForSocketGone(paths.Socket, ShutdownGrace); err != nil {
 		return err
 	}
 	if info.PID == 0 || info.PID == os.Getpid() {
 		return nil
 	}
-	return pollUntil(shutdownGrace, func() bool { return !processExists(info.PID) }, "daemon process did not exit")
+	return pollUntil(ShutdownGrace, func() bool { return !processExists(info.PID) }, "daemon process did not exit")
 }
 
 // WaitForSocket polls until something accepts connections on the socket.

daemon/ensure_test.go

@@ -28,7 +28,7 @@ func TestMain(m *testing.M) {
 			Paths:       paths,
 			Version:     os.Getenv(testDaemonVersionEnv),
 			IdleTimeout: time.Minute,
-			Dispatch:    echoDispatcher(make(chan struct{})),
+			Dispatch:    echoDispatcher(make(chan struct{}), make(chan struct{})),
 		})
 		if err != nil {
 			os.Exit(3)

daemon/server.go

@@ -47,10 +47,14 @@ const (
 	errCodeInvalid     = -32600
 
 	idleTick          = 10 * time.Second
-	shutdownGrace     = 10 * time.Second
 	handlerTouchEvery = 30 * time.Second
 )
 
+// ShutdownGrace is how long a stopping daemon waits for in-flight requests
+// to finish before it exits. Work that must complete on shutdown (stopping a
+// recording, cleaning up after it) has to fit inside it.
+const ShutdownGrace = 10 * time.Second
+
 type server struct {
 	opts         Options
 	startedAt    time.Time
@@ -59,6 +63,10 @@ type server struct {
 	wg           sync.WaitGroup
 	shutdown     chan struct{}
 	shutdownOnce sync.Once
+	// handlers is the parent of every request's context; cancelling it on
+	// shutdown stops in-flight streams (a recording finalizes and cleans up)
+	handlers       context.Context
+	cancelHandlers context.CancelFunc
 }
 
 // Run listens on the unix socket and serves until ctx is cancelled, the idle
@@ -85,7 +93,9 @@ func Run(ctx context.Context, opts Options) error {
 		return fmt.Errorf("restrict %s: %w", opts.Paths.Socket, err)
 	}
 
-	s := &server{opts: opts, startedAt: time.Now(), shutdown: make(chan struct{})}
+	handlers, cancelHandlers := context.WithCancel(context.Background())
+	defer cancelHandlers()
+	s := &server{opts: opts, startedAt: time.Now(), shutdown: make(chan struct{}), handlers: handlers, cancelHandlers: cancelHandlers}
 	s.touch()
 
 	if err := writePidFile(opts.Paths.Pid, pidInfo{PID: os.Getpid(), Version: opts.Version, StartedAt: s.startedAt}); err != nil {
@@ -111,6 +121,9 @@ func Run(ctx context.Context, opts Options) error {
 	// join the accept loop first: once it has returned no more wg.Add calls
 	// can race with the wg.Wait inside waitInflight
 	<-accepting
+	// stop in-flight requests as if their clients went away, so a recording
+	// stops and cleans up instead of outliving the daemon on the device
+	s.cancelHandlers()
 	s.waitInflight()
 	if opts.OnShutdown != nil {
 		opts.OnShutdown()
@@ -129,7 +142,7 @@ func (s *server) waitInflight() {
 	go func() { s.wg.Wait(); close(done) }()
 	select {
 	case <-done:
-	case <-time.After(shutdownGrace):
+	case <-time.After(ShutdownGrace):
 		utils.Info("daemon: timed out waiting for in-flight requests")
 	}
 }
@@ -205,7 +218,7 @@ func (s *server) handleConn(conn net.Conn) {
 		return
 	}
 
-	ctx, cancel := context.WithCancel(context.Background())
+	ctx, cancel := context.WithCancel(s.handlers)
 	defer cancel()
 	go cancelOnDisconnect(conn, cancel)
 

daemon/server_test.go

@@ -14,8 +14,9 @@ import (
 )
 
 // echoDispatcher answers "echo" with its params, "fail" with an error,
-// "stream" with three notifications, and "block" by waiting for ctx.
-func echoDispatcher(blocked chan struct{}) Dispatcher {
+// "stream" with three notifications, and "block" by waiting for ctx: it
+// closes entered once it is running and blocked once ctx is cancelled.
+func echoDispatcher(entered chan struct{}, blocked chan struct{}) Dispatcher {
 	return func(ctx context.Context, method string, params json.RawMessage, notify func(any) error) (any, error) {
 		switch method {
 		case "echo":
@@ -30,6 +31,7 @@ func echoDispatcher(blocked chan struct{}) Dispatcher {
 			}
 			return map[string]string{"done": "yes"}, nil
 		case "block":
+			close(entered)
 			<-ctx.Done()
 			close(blocked)
 			return nil, ctx.Err()
@@ -45,6 +47,7 @@ func echoDispatcher(blocked chan struct{}) Dispatcher {
 type runningDaemon struct {
 	paths   FilePaths
 	done    chan error
+	entered chan struct{}
 	blocked chan struct{}
 	cancel  context.CancelFunc
 	result  error
@@ -72,13 +75,13 @@ func startTestDaemon(t *testing.T, idle time.Duration, busy func() bool) *runnin
 	require.NoError(t, err)
 
 	ctx, cancel := context.WithCancel(context.Background())
-	rd := &runningDaemon{paths: paths, done: make(chan error, 1), blocked: make(chan struct{}), cancel: cancel}
+	rd := &runningDaemon{paths: paths, done: make(chan error, 1), entered: make(chan struct{}), blocked: make(chan struct{}), cancel: cancel}
 	go func() {
 		rd.done <- Run(ctx, Options{
 			Paths:       paths,
 			Version:     "test-1",
 			IdleTimeout: idle,
-			Dispatch:    echoDispatcher(rd.blocked),
+			Dispatch:    echoDispatcher(rd.entered, rd.blocked),
 			Busy:        busy,
 		})
 	}()
@@ -249,3 +252,38 @@ func TestOnShutdownRunsAfterRunStopsServing(t *testing.T) {
 		t.Fatal("OnShutdown was not called")
 	}
 }
+
+func TestShutdownCancelsInFlightHandlerContext(t *testing.T) {
+	rd := startTestDaemon(t, 0, nil)
+
+	callDone := make(chan error, 1)
+	go func() {
+		_, err := CallStream(context.Background(), "block", nil, nil)
+		callDone <- err
+	}()
+	rd.waitUntilHandlerIsRunning(t)
+
+	require.NoError(t, Shutdown())
+
+	select {
+	case <-rd.blocked:
+	case <-time.After(3 * time.Second):
+		t.Fatal("in-flight handler was not cancelled when the daemon shut down")
+	}
+	select {
+	case <-callDone:
+	case <-time.After(3 * time.Second):
+		t.Fatal("in-flight call did not return after the daemon shut down")
+	}
+}
+
+// waitUntilHandlerIsRunning waits for a "block" call started in another
+// goroutine to reach its handler.
+func (rd *runningDaemon) waitUntilHandlerIsRunning(t *testing.T) {
+	t.Helper()
+	select {
+	case <-rd.entered:
+	case <-time.After(3 * time.Second):
+		t.Fatal("handler never started")
+	}
+}