fix: Close opened BSR files (#3519)

* fix: Close opened BSR files

- Add checks to close a file if it's not nil
- Close files opened during reads for `OpenSession`, `OpenConnection` & `OpenChannel`
- Close opened scanner readers

internal/bsr/bsr.go

@@ -173,7 +173,7 @@ func persistBsrSessionKeys(ctx context.Context, keys *kms.Keys, c *container) er
 // Signature and checksum files will then be verified.
 // Fields on the underlying container will be populated so that the returned Session can be used for BSR
 // playback and conversion to formats such as asciinema
-func OpenSession(ctx context.Context, sessionRecordingId string, f storage.FS, keyUnwrapFn kms.KeyUnwrapCallbackFunc) (*Session, error) {
+func OpenSession(ctx context.Context, sessionRecordingId string, f storage.FS, keyUnwrapFn kms.KeyUnwrapCallbackFunc) (s *Session, err error) {
 	const op = "bsr.OpenSession"
 	switch {
 	case sessionRecordingId == "":
@@ -200,8 +200,15 @@ func OpenSession(ctx context.Context, sessionRecordingId string, f storage.FS, k
 	// Load and verify recording metadata
 	sha256Reader, err := crypto.NewSha256SumReader(ctx, cc.metaFile)
 	if err != nil {
+		cc.metaFile.Close()
 		return nil, fmt.Errorf("%s: %w", op, err)
 	}
+	defer func() {
+		if closeErr := sha256Reader.Close(); closeErr != nil {
+			err = errors.Join(err, fmt.Errorf("%s: %w", op, closeErr))
+		}
+	}()
+
 	meta, err := decodeSessionRecordingMeta(ctx, sha256Reader)
 	if err != nil {
 		return nil, err
@@ -289,7 +296,7 @@ func (s *Session) NewConnection(ctx context.Context, meta *ConnectionRecordingMe
 }
 
 // OpenConnection will open and validate a BSR connection
-func (s *Session) OpenConnection(ctx context.Context, connId string) (*Connection, error) {
+func (s *Session) OpenConnection(ctx context.Context, connId string) (conn *Connection, err error) {
 	const op = "bsr.(Session).OpenConnection"
 	switch {
 	case connId == "":
@@ -319,8 +326,15 @@ func (s *Session) OpenConnection(ctx context.Context, connId string) (*Connectio
 	// Load and verify connection metadata
 	sha256Reader, err := crypto.NewSha256SumReader(ctx, cc.metaFile)
 	if err != nil {
+		cc.metaFile.Close()
 		return nil, fmt.Errorf("%s: %w", op, err)
 	}
+	defer func() {
+		if closeErr := sha256Reader.Close(); closeErr != nil {
+			err = errors.Join(err, fmt.Errorf("%s: %w", op, closeErr))
+		}
+	}()
+
 	sm, err := decodeConnectionRecordingMeta(ctx, sha256Reader)
 	if err != nil {
 		return nil, err
@@ -391,7 +405,7 @@ func (c *Connection) NewChannel(ctx context.Context, meta *ChannelRecordingMeta)
 }
 
 // OpenChannel will open and validate a BSR channel
-func (c *Connection) OpenChannel(ctx context.Context, chanId string) (*Channel, error) {
+func (c *Connection) OpenChannel(ctx context.Context, chanId string) (ch *Channel, err error) {
 	const op = "bsr.OpenChannel"
 	switch {
 	case chanId == "":
@@ -420,8 +434,15 @@ func (c *Connection) OpenChannel(ctx context.Context, chanId string) (*Channel,
 	// Load and verify channel metadata
 	sha256Reader, err := crypto.NewSha256SumReader(ctx, cc.metaFile)
 	if err != nil {
+		cc.metaFile.Close()
 		return nil, fmt.Errorf("%s: %w", op, err)
 	}
+	defer func() {
+		if closeErr := sha256Reader.Close(); closeErr != nil {
+			err = errors.Join(err, fmt.Errorf("%s: %w", op, closeErr))
+		}
+	}()
+
 	sm, err := decodeChannelRecordingMeta(ctx, sha256Reader)
 	if err != nil {
 		return nil, err

internal/bsr/bsr_open_test.go

@@ -469,3 +469,130 @@ func TestOpenBSRMethods_WithoutSummaryAllocFunc(t *testing.T) {
 		})
 	}
 }
+
+func TestCloseBSRMethods(t *testing.T) {
+	ctx := context.Background()
+
+	protocol := Protocol("TEST_CLOSED_FILE")
+	sessionRecordingId := "sr_012344567890"
+	sessionId := "s_012344567890"
+	connectionId := "connection"
+	channelId := "channel"
+
+	keys, err := kms.CreateKeys(ctx, kms.TestWrapper(t), "session")
+	require.NoError(t, err)
+
+	err = RegisterSummaryAllocFunc(protocol, ChannelContainer, func(ctx context.Context) Summary {
+		return &BaseChannelSummary{Id: channelId, ConnectionRecordingId: connectionId}
+	})
+	require.NoError(t, err)
+
+	err = RegisterSummaryAllocFunc(protocol, SessionContainer, func(ctx context.Context) Summary {
+		return &BaseSessionSummary{Id: sessionId, ConnectionCount: 1}
+	})
+	require.NoError(t, err)
+
+	err = RegisterSummaryAllocFunc(protocol, ConnectionContainer, func(ctx context.Context) Summary {
+		return &BaseConnectionSummary{Id: connectionId, ChannelCount: 1}
+	})
+	require.NoError(t, err)
+
+	f := fstest.NewMemFS(fstest.WithOriginalFile())
+
+	srm := &SessionRecordingMeta{
+		Id:       sessionRecordingId,
+		Protocol: protocol,
+	}
+	sessionMeta := TestSessionMeta(sessionId)
+
+	sesh, err := NewSession(ctx, srm, sessionMeta, f, keys, WithSupportsMultiplex(true))
+	require.NoError(t, err)
+	require.NotNil(t, sesh)
+
+	sesh.EncodeSummary(ctx, &BaseChannelSummary{
+		Id:                    "TEST_CHANNEL_ID",
+		ConnectionRecordingId: "TEST_CONNECTION_RECORDING_ID",
+	})
+
+	connMeta := &ConnectionRecordingMeta{Id: connectionId}
+	conn, err := sesh.NewConnection(ctx, connMeta)
+	require.NoError(t, err)
+	require.NotNil(t, conn)
+
+	conn.EncodeSummary(ctx, &BaseConnectionSummary{
+		Id:           "TEST_CONNECTION_ID",
+		ChannelCount: 1,
+	})
+
+	chanMeta := &ChannelRecordingMeta{
+		Id:   channelId,
+		Type: "chan",
+	}
+	ch, err := conn.NewChannel(ctx, chanMeta)
+	require.NoError(t, err)
+	require.NotNil(t, ch)
+
+	ch.EncodeSummary(ctx, &BaseSessionSummary{
+		Id:              "TEST_SESSION_ID",
+		ConnectionCount: 1,
+	})
+
+	ch.Close(ctx)
+	conn.Close(ctx)
+	sesh.Close(ctx)
+
+	keyFn := func(w kms.WrappedKeys) (kms.UnwrappedKeys, error) {
+		u := kms.UnwrappedKeys{
+			BsrKey:  keys.BsrKey,
+			PrivKey: keys.PrivKey,
+		}
+		return u, nil
+	}
+
+	opSesh, err := OpenSession(ctx, srm.Id, f, keyFn)
+	require.NoError(t, err)
+	require.NotNil(t, opSesh)
+
+	opConn, err := opSesh.OpenConnection(ctx, connectionId)
+	require.NoError(t, err)
+	require.NotNil(t, opConn)
+
+	opChan, err := opConn.OpenChannel(ctx, channelId)
+	require.NoError(t, err)
+	require.NotNil(t, opChan)
+
+	// Close all opened containers
+	require.NoError(t, opChan.Close(ctx))
+	require.NoError(t, opConn.Close(ctx))
+	require.NoError(t, opSesh.Close(ctx))
+
+	// Get session container
+	sessionContainer := f.Containers[fmt.Sprintf(bsrFileNameTemplate, sessionRecordingId)]
+	require.NotNil(t, sessionContainer)
+	assert.True(t, sessionContainer.Closed)
+
+	// Ensure all session files are closed
+	for _, file := range sessionContainer.Files {
+		assert.True(t, file.Closed)
+	}
+
+	// Get connection container
+	connectionContainer := sessionContainer.Sub[fmt.Sprintf(connectionFileNameTemplate, connectionId)]
+	require.NotNil(t, connectionContainer)
+	assert.True(t, connectionContainer.Closed)
+
+	// Ensure all connection files are closed
+	for _, file := range connectionContainer.Files {
+		assert.True(t, file.Closed)
+	}
+
+	// Get channel container
+	channelContainer := connectionContainer.Sub[fmt.Sprintf(channelFileNameTemplate, channelId)]
+	require.NotNil(t, channelContainer)
+	assert.True(t, channelContainer.Closed)
+
+	// Ensure all channel files are closed
+	for _, file := range channelContainer.Files {
+		assert.True(t, file.Closed)
+	}
+}

internal/bsr/container.go

@@ -185,7 +185,7 @@ func openContainer(ctx context.Context, t ContainerType, c storage.Container, ke
 	return cc, nil
 }
 
-func (c *container) loadChecksums(ctx context.Context) error {
+func (c *container) loadChecksums(ctx context.Context) (err error) {
 	const op = "bsr.(container).loadChecksums"
 
 	// Open and extract checksum signature
@@ -193,6 +193,12 @@ func (c *container) loadChecksums(ctx context.Context) error {
 	if err != nil {
 		return err
 	}
+	defer func() {
+		if closeErr := checksumsSigFile.Close(); closeErr != nil {
+			err = errors.Join(err, fmt.Errorf("%s: %w", op, closeErr))
+		}
+	}()
+
 	checksumSigBytes := new(bytes.Buffer)
 	_, err = checksumSigBytes.ReadFrom(checksumsSigFile)
 	if err != nil {
@@ -210,6 +216,12 @@ func (c *container) loadChecksums(ctx context.Context) error {
 	if err != nil {
 		return err
 	}
+	defer func() {
+		if closeErr := checksumsFile.Close(); closeErr != nil {
+			err = errors.Join(err, fmt.Errorf("%s: %w", op, closeErr))
+		}
+	}()
+
 	var checksumsBuffer bytes.Buffer
 	cTee := io.TeeReader(checksumsFile, &checksumsBuffer)
 
@@ -237,11 +249,16 @@ func (c *container) loadChecksums(ctx context.Context) error {
 	return nil
 }
 
-func (c *container) loadKey(ctx context.Context, keyFileName string) (*wrapping.KeyInfo, error) {
+func (c *container) loadKey(ctx context.Context, keyFileName string) (k *wrapping.KeyInfo, err error) {
 	keyFile, err := c.container.OpenFile(ctx, keyFileName)
 	if err != nil {
 		return nil, err
 	}
+	defer func() {
+		if closeErr := keyFile.Close(); closeErr != nil {
+			err = errors.Join(err, closeErr)
+		}
+	}()
 
 	keyBytes := new(bytes.Buffer)
 	_, err = keyBytes.ReadFrom(keyFile)
@@ -258,11 +275,16 @@ func (c *container) loadKey(ctx context.Context, keyFileName string) (*wrapping.
 	return key, nil
 }
 
-func (c *container) loadSignature(ctx context.Context, sigFileName string) (*wrapping.SigInfo, error) {
+func (c *container) loadSignature(ctx context.Context, sigFileName string) (s *wrapping.SigInfo, err error) {
 	sigFile, err := c.container.OpenFile(ctx, sigFileName)
 	if err != nil {
 		return nil, err
 	}
+	defer func() {
+		if closeErr := sigFile.Close(); closeErr != nil {
+			err = errors.Join(err, closeErr)
+		}
+	}()
 
 	sigBytes := new(bytes.Buffer)
 	_, err = sigBytes.ReadFrom(sigFile)
@@ -512,28 +534,40 @@ func (c *container) close(_ context.Context) error {
 
 	var closeError error
 
-	if err := c.meta.Close(); err != nil {
-		closeError = errors.Join(closeError, fmt.Errorf("%s: %w", op, err))
+	if !is.Nil(c.meta) {
+		if err := c.meta.Close(); err != nil {
+			closeError = errors.Join(closeError, fmt.Errorf("%s: %w", op, err))
+		}
 	}
 
-	if err := c.sum.Close(); err != nil {
-		closeError = errors.Join(closeError, fmt.Errorf("%s: %w", op, err))
+	if !is.Nil(c.sum) {
+		if err := c.sum.Close(); err != nil {
+			closeError = errors.Join(closeError, fmt.Errorf("%s: %w", op, err))
+		}
 	}
 
-	if err := c.checksums.Close(); err != nil {
-		closeError = errors.Join(closeError, fmt.Errorf("%s: %w", op, err))
+	if !is.Nil(c.checksums) {
+		if err := c.checksums.Close(); err != nil {
+			closeError = errors.Join(closeError, fmt.Errorf("%s: %w", op, err))
+		}
 	}
 
-	if err := c.sigs.Close(); err != nil {
-		closeError = errors.Join(closeError, fmt.Errorf("%s: %w", op, err))
+	if !is.Nil(c.sigs) {
+		if err := c.sigs.Close(); err != nil {
+			closeError = errors.Join(closeError, fmt.Errorf("%s: %w", op, err))
+		}
 	}
 
-	if err := c.journal.Close(); err != nil {
-		closeError = errors.Join(closeError, fmt.Errorf("%s: %w", op, err))
+	if !is.Nil(c.journal) {
+		if err := c.journal.Close(); err != nil {
+			closeError = errors.Join(closeError, fmt.Errorf("%s: %w", op, err))
+		}
 	}
 
-	if err := c.container.Close(); err != nil {
-		closeError = errors.Join(closeError, fmt.Errorf("%s: %w", op, err))
+	if !is.Nil(c.container) {
+		if err := c.container.Close(); err != nil {
+			closeError = errors.Join(closeError, fmt.Errorf("%s: %w", op, err))
+		}
 	}
 
 	return closeError

internal/bsr/convert/convert.go

@@ -49,11 +49,13 @@ func ToAsciicast(ctx context.Context, session *bsr.Session, tmp storage.TempFile
 		if err != nil {
 			return nil, fmt.Errorf("%s: %w", op, err)
 		}
+		defer conn.Close(ctx)
 
 		ch, err := conn.OpenChannel(ctx, chanId)
 		if err != nil {
 			return nil, fmt.Errorf("%s: %w", op, err)
 		}
+		defer ch.Close(ctx)
 
 		switch chs := ch.Summary.(type) {
 		case *ssh.ChannelSummary:
@@ -61,13 +63,21 @@ func ToAsciicast(ctx context.Context, session *bsr.Session, tmp storage.TempFile
 			case ssh.Shell, ssh.Exec:
 				reqScanner, err := ch.OpenRequestScanner(ctx, bsr.Inbound)
 				if err != nil {
+					if !is.Nil(reqScanner) {
+						reqScanner.Close()
+					}
 					return nil, fmt.Errorf("%s: %w", op, err)
 				}
+				defer reqScanner.Close()
 
 				msgScanner, err := ch.OpenMessageScanner(ctx, bsr.Outbound)
 				if err != nil {
+					if !is.Nil(msgScanner) {
+						msgScanner.Close()
+					}
 					return nil, fmt.Errorf("%s: %w", op, err)
 				}
+				defer msgScanner.Close()
 				return sshChannelToAsciicast(ctx, reqScanner, msgScanner, tmp, options...)
 			case "":
 				return nil, fmt.Errorf("%s: session program not set for asciicast conversion", op)

internal/bsr/internal/fstest/fs.go

@@ -35,8 +35,9 @@ const (
 type MemFS struct {
 	Containers map[string]*MemContainer
 
-	newFunc  NewFunc
-	readOnly bool
+	newFunc      NewFunc
+	readOnly     bool
+	originalFile bool
 }
 
 // NewMemFS creates a MemFS. It supports WithNewFunc, WithReadOnly.
@@ -44,9 +45,10 @@ func NewMemFS(options ...Option) *MemFS {
 	opts := getOpts(options...)
 
 	return &MemFS{
-		Containers: make(map[string]*MemContainer),
-		newFunc:    opts.withNewFunc,
-		readOnly:   opts.withReadOnly,
+		Containers:   make(map[string]*MemContainer),
+		newFunc:      opts.withNewFunc,
+		readOnly:     opts.withReadOnly,
+		originalFile: opts.withOriginalFile,
 	}
 }
 
@@ -68,11 +70,12 @@ func (m *MemFS) New(ctx context.Context, n string) (storage.Container, error) {
 	}
 
 	c := &MemContainer{
-		Name:       n,
-		Sub:        make(map[string]*MemContainer),
-		Files:      make(map[string]*MemFile),
-		mode:       defaultContainerPerm,
-		accessMode: storage.ReadWrite,
+		Name:         n,
+		Sub:          make(map[string]*MemContainer),
+		Files:        make(map[string]*MemFile),
+		mode:         defaultContainerPerm,
+		accessMode:   storage.ReadWrite,
+		originalFile: m.originalFile,
 	}
 	m.Containers[n] = c
 	return c, nil
@@ -96,8 +99,9 @@ func (m *MemFS) Open(_ context.Context, n string) (storage.Container, error) {
 type MemContainer struct {
 	Name string
 
-	Sub   map[string]*MemContainer
-	Files map[string]*MemFile
+	Sub          map[string]*MemContainer
+	Files        map[string]*MemFile
+	originalFile bool
 
 	Closed bool
 
@@ -188,6 +192,11 @@ func (m *MemContainer) OpenFile(_ context.Context, n string, option ...storage.O
 			Closed:     src.Closed,
 		}
 	}
+
+	if m.originalFile {
+		return src, nil
+	}
+
 	return dst, nil
 }
 
@@ -211,9 +220,10 @@ func (m *MemContainer) SubContainer(_ context.Context, n string, option ...stora
 			return nil, fmt.Errorf("container %s already exists: %w", n, ErrAlreadyExists)
 		}
 		c = &MemContainer{
-			Name:  n,
-			Sub:   make(map[string]*MemContainer),
-			Files: make(map[string]*MemFile),
+			Name:         n,
+			Sub:          make(map[string]*MemContainer),
+			Files:        make(map[string]*MemFile),
+			originalFile: m.originalFile,
 		}
 	} else {
 		if !exists {

internal/bsr/internal/fstest/options.go

@@ -41,6 +41,7 @@ type options struct {
 	withStatFunc       StatFunc
 	withReadOnly       bool
 	withStorageOptions []storage.Option
+	withOriginalFile   bool
 }
 
 func getDefaultOptions() options {
@@ -50,6 +51,7 @@ func getDefaultOptions() options {
 		withStatFunc:       nil,
 		withReadOnly:       false,
 		withStorageOptions: nil,
+		withOriginalFile:   false,
 	}
 }
 
@@ -87,3 +89,12 @@ func WithStorageOptions(opts []storage.Option) Option {
 		o.withStorageOptions = opts
 	}
 }
+
+// WithOriginalFile is used to decide if the original of a file
+// should be returned instead of a copy of the file during
+// opening a file.
+func WithOriginalFile() Option {
+	return func(o *options) {
+		o.withOriginalFile = true
+	}
+}