From 09cf85466ceff890f2981a1fe7f7d48d81af872c Mon Sep 17 00:00:00 2001 From: Zac Bergquist Date: Fri, 20 Oct 2023 16:57:26 -0600 Subject: [PATCH] tsh: split large desktop recordings into multiple files AVI files contain 32-bit pointers, which means we cant just add frames to a single file until the end of time. This change properly detects when a recording export has reached the maximum size of an AVI file and splits the recording into multiple files. Closes #33109 Closes #33110 --- tool/tsh/recording_export.go | 77 +++++++++++++++++++++++-------- tool/tsh/recording_export_test.go | 10 ++-- 2 files changed, 64 insertions(+), 23 deletions(-) diff --git a/tool/tsh/recording_export.go b/tool/tsh/recording_export.go index 3922d4601bf2f..246e68531790b 100644 --- a/tool/tsh/recording_export.go +++ b/tool/tsh/recording_export.go @@ -19,11 +19,13 @@ package main import ( "bytes" "context" + "errors" "fmt" "image" "image/draw" "image/jpeg" "image/png" + "strings" "github.com/gravitational/trace" "github.com/icza/mjpeg" @@ -54,30 +56,39 @@ func onExportRecording(cf *CLIConf) error { authClient := proxyClient.CurrentCluster() defer authClient.Close() - fname := cf.OutFile - if fname == "" { - fname = cf.SessionID + ".avi" + filenamePrefix := cf.SessionID + if cf.OutFile != "" { + // trim the extension if it was provided (we'll add it later on) + filenamePrefix = strings.TrimSuffix( + strings.TrimSuffix(cf.OutFile, ".avi"), ".AVI") } - frames, err := writeMovie(cf.Context, authClient, session.ID(cf.SessionID), fname) - // there may be a partial file, even if we encountered an error, - // so we indicate to the user when we wrote something - if frames > 0 { - fmt.Printf("wrote recording to %v\n", fname) + _, err = writeMovie(cf.Context, authClient, session.ID(cf.SessionID), filenamePrefix, fmt.Printf) + return trace.Wrap(err) +} + +func makeAVIFileName(prefix string, currentFile int) string { + if currentFile == 0 { + return prefix + ".avi" } - return trace.Wrap(err) + return fmt.Sprintf("%v-%d.avi", prefix, currentFile) } -// writeMovie writes streams the events for the specified session into a movie file -// identified by fname. It returns the number of frames that were written and an error. -func writeMovie(ctx context.Context, ss events.SessionStreamer, sid session.ID, fname string) (int, error) { +// writeMovie writes the events for the specified session into one or more movie files +// beginning with the specified prefix. It returns the number of frames that were written and an error. +func writeMovie(ctx context.Context, ss events.SessionStreamer, sid session.ID, prefix string, + write func(format string, args ...any) (int, error)) (frames int, err error) { + var screen *image.NRGBA var movie mjpeg.AviWriter lastEmitted := int64(-1) buf := new(bytes.Buffer) - frameCount := 0 + + var frameCount, fileCount int + var width, height int32 + currentFilename := makeAVIFileName(prefix, fileCount) evts, errs := ss.StreamSessionEvents(ctx, sid, 0) loop: @@ -85,12 +96,16 @@ loop: select { case err := <-errs: if movie != nil { - movie.Close() + if err := movie.Close(); err == nil && write != nil && frames > 0 { + write("wrote %v\n", currentFilename) + } } return frameCount, trace.Wrap(err) case <-ctx.Done(): if movie != nil { - movie.Close() + if err := movie.Close(); err == nil && write != nil && frames > 0 { + write("wrote %v\n", currentFilename) + } } return frameCount, ctx.Err() case evt, more := <-evts: @@ -123,12 +138,13 @@ loop: // the window during a session. If this changes, we'd have to // find the maximum window size first. log.Debugf("allocating %dx%d screen", msg.Width, msg.Height) + width, height = int32(msg.Width), int32(msg.Height) screen = image.NewNRGBA(image.Rectangle{ Min: image.Pt(0, 0), Max: image.Pt(int(msg.Width), int(msg.Height)), }) - movie, err = mjpeg.New(fname, int32(msg.Width), int32(msg.Height), framesPerSecond) + movie, err = mjpeg.New(currentFilename, width, height, framesPerSecond) if err != nil { return frameCount, trace.Wrap(err) } @@ -170,7 +186,27 @@ loop: return frameCount, trace.Wrap(err) } for i := 0; i < int(framesToEmit); i++ { - if err := movie.AddFrame(buf.Bytes()); err != nil { + err := movie.AddFrame(buf.Bytes()) + if errors.Is(err, mjpeg.ErrTooLarge) { + // this file can't get any larger - time to open a new file + if err := movie.Close(); err != nil { + return frameCount, trace.WrapWithMessage(err, "failed to write partial recording") + } + if write != nil { + write("wrote %v\n", currentFilename) + } + fileCount++ + currentFilename = makeAVIFileName(prefix, fileCount) + movie, err = mjpeg.New(currentFilename, width, height, framesPerSecond) + if err != nil { + return frameCount, trace.Wrap(err) + } + + // write the frame to the new file + if err := movie.AddFrame(buf.Bytes()); err != nil { + return frameCount, trace.Wrap(err) + } + } else if err != nil { return frameCount, trace.Wrap(err) } frameCount++ @@ -190,7 +226,12 @@ loop: return 0, trace.BadParameter("operation canceled") } - return frameCount, trace.Wrap(movie.Close()) + err = movie.Close() + if err == nil && write != nil { + write("wrote %v\n", currentFilename) + } + + return frameCount, trace.Wrap(err) } func imgFromPNGMessage(msg tdp.Message) (image.Image, error) { diff --git a/tool/tsh/recording_export_test.go b/tool/tsh/recording_export_test.go index 82acdc53abcf3..d8cf2c2cbdfac 100644 --- a/tool/tsh/recording_export_test.go +++ b/tool/tsh/recording_export_test.go @@ -55,7 +55,7 @@ func TestWriteMovieCanBeCanceled(t *testing.T) { ctx, cancel := context.WithCancel(context.Background()) cancel() - frames, err := writeMovie(ctx, fs, "test", "test.avi") + frames, err := writeMovie(ctx, fs, "test", "test.avi", nil) require.Equal(t, context.Canceled, err) require.Equal(t, 0, frames) } @@ -68,7 +68,7 @@ func TestWriteMovieDoesNotSupportSSH(t *testing.T) { } fs := eventstest.NewFakeStreamer(events, 0) - frames, err := writeMovie(context.Background(), fs, "test", "test.avi") + frames, err := writeMovie(context.Background(), fs, "test", "test.avi", nil) require.True(t, trace.IsBadParameter(err), "expected bad paramater error, got %v", err) require.Equal(t, 0, frames) } @@ -88,7 +88,7 @@ func TestWriteMovieMultipleScreenSpecs(t *testing.T) { fs := eventstest.NewFakeStreamer(events, 0) t.Cleanup(func() { os.RemoveAll("test.avi") }) - frames, err := writeMovie(context.Background(), fs, session.ID("test"), "test.avi") + frames, err := writeMovie(context.Background(), fs, session.ID("test"), "test.avi", nil) require.True(t, trace.IsBadParameter(err), "expected bad paramater error, got %v", err) require.Equal(t, 0, frames) } @@ -105,7 +105,7 @@ func TestWriteMovieWritesOneFrame(t *testing.T) { } fs := eventstest.NewFakeStreamer(events, 0) t.Cleanup(func() { os.RemoveAll("test.avi") }) - frames, err := writeMovie(context.Background(), fs, session.ID("test"), "test.avi") + frames, err := writeMovie(context.Background(), fs, session.ID("test"), "test.avi", nil) require.NoError(t, err) require.Equal(t, 1, frames) } @@ -122,7 +122,7 @@ func TestWriteMovieWritesManyFrames(t *testing.T) { } fs := eventstest.NewFakeStreamer(events, 0) t.Cleanup(func() { os.RemoveAll("test.avi") }) - frames, err := writeMovie(context.Background(), fs, session.ID("test"), "test.avi") + frames, err := writeMovie(context.Background(), fs, session.ID("test"), "test.avi", nil) require.NoError(t, err) require.Equal(t, framesPerSecond, frames) }