From 1548d48f004f5eac8e27ad10784de50f474346d6 Mon Sep 17 00:00:00 2001 From: Elim Tsiagbey Date: Mon, 31 Jul 2023 17:13:24 -0400 Subject: [PATCH] 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 | 27 ++++- internal/bsr/bsr_open_test.go | 127 ++++++++++++++++++++++++ internal/bsr/container.go | 64 +++++++++--- internal/bsr/convert/convert.go | 10 ++ internal/bsr/internal/fstest/fs.go | 40 +++++--- internal/bsr/internal/fstest/options.go | 11 ++ 6 files changed, 246 insertions(+), 33 deletions(-) diff --git a/internal/bsr/bsr.go b/internal/bsr/bsr.go index 546fa8a66..86f7e3c9b 100644 --- a/internal/bsr/bsr.go +++ b/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 diff --git a/internal/bsr/bsr_open_test.go b/internal/bsr/bsr_open_test.go index e3ca0dc7a..4f350fad1 100644 --- a/internal/bsr/bsr_open_test.go +++ b/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) + } +} diff --git a/internal/bsr/container.go b/internal/bsr/container.go index 5f0914ade..c239b2f93 100644 --- a/internal/bsr/container.go +++ b/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 diff --git a/internal/bsr/convert/convert.go b/internal/bsr/convert/convert.go index d6847a434..277e75d9c 100644 --- a/internal/bsr/convert/convert.go +++ b/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) diff --git a/internal/bsr/internal/fstest/fs.go b/internal/bsr/internal/fstest/fs.go index ef702300d..ef6ef920e 100644 --- a/internal/bsr/internal/fstest/fs.go +++ b/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 { diff --git a/internal/bsr/internal/fstest/options.go b/internal/bsr/internal/fstest/options.go index a9b7def30..d3500f428 100644 --- a/internal/bsr/internal/fstest/options.go +++ b/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 + } +}