VAULT-36947: Support force unloading a snapshot (#8740) (#9036)

* portion of changes for autoloading

* add test checking for panic

* add endpoint for force unloading

* separate method for force unload

* changelog

* don't redefine constants

Co-authored-by: miagilepner <[email protected]>

Author: hc-github-team-secure-vault-core

api/sys_raft.go

@@ -503,10 +503,27 @@ func (c *Sys) RaftUnloadSnapshot(snapID string) (*Secret, error) {
 // RaftUnloadSnapshotWithContext unloads a snapshot from the raft cluster.
 // It accepts a snapshot ID to identify the snapshot to be unloaded.
 func (c *Sys) RaftUnloadSnapshotWithContext(ctx context.Context, snapID string) (*Secret, error) {
+	return c.raftUnloadSnapshotWithContext(ctx, snapID, false)
+}
+
+// RaftForceUnloadSnapshot wraps RaftForceUnloadSnapshotWithContext using context.Background.
+func (c *Sys) RaftForceUnloadSnapshot(snapID string) (*Secret, error) {
+	return c.RaftForceUnloadSnapshotWithContext(context.Background(), snapID)
+}
+
+// RaftForceUnloadSnapshotWithContext forcefully unloads the given snapshot
+func (c *Sys) RaftForceUnloadSnapshotWithContext(ctx context.Context, snapID string) (*Secret, error) {
+	return c.raftUnloadSnapshotWithContext(ctx, snapID, true)
+}
+
+func (c *Sys) raftUnloadSnapshotWithContext(ctx context.Context, snapID string, force bool) (*Secret, error) {
 	ctx, cancelFunc := c.c.withConfiguredTimeout(ctx)
 	defer cancelFunc()
 
 	r := c.c.NewRequest(http.MethodDelete, "/v1/sys/storage/raft/snapshot-load/"+snapID)
+	if force {
+		r.Params.Set("force", "true")
+	}
 
 	resp, err := c.c.rawRequestWithContext(ctx, r)
 	if err != nil {

changelog/_8740.txt

@@ -0,0 +1,6 @@
+```release-note:improvement
+core/snapshot-load (enterprise): Add a `force` query parameter to the `DELETE sys/storage/raft/snapshot-load/{snapshot_id}` endpoint to allow for forced deletion of snapshots. This is useful when the snapshot is in a state that prevents normal deletion, such as being in the process of loading.
+```
+```release-note:improvement
+cli (enterprise): Add a `-force` flag to `vault operator raft snapshot unload` command to force deletion of a loaded snapshot.
+```

physical/raft/snapshot.go

@@ -364,7 +364,7 @@ func (s *BoltSnapshotSink) writeBoltDBFile() error {
 		defer close(s.doneWritingCh)
 		defer boltDB.Close()
 
-		err := loadSnapshot(boltDB, s.logger, reader, nil, false)
+		err := loadSnapshot(context.Background(), boltDB, s.logger, reader, nil, false)
 		if err != nil {
 			s.writeError = err
 			return
@@ -503,8 +503,8 @@ func snapshotName(term, index uint64) string {
 // FSM. The caller is responsible for closing the reader.
 // If pathsToFilter is not nil, the function will filter out any keys that are
 // found in the pathsToFilter tree.
-func LoadReadOnlySnapshot(fsm *FSM, snapshotFile io.ReadCloser, filterKey func(key string) bool, logger log.Logger) error {
-	return loadSnapshot(fsm.db, logger, snapshotFile, filterKey, true)
+func LoadReadOnlySnapshot(ctx context.Context, fsm *FSM, snapshotFile io.ReadCloser, filterKey func(key string) bool, logger log.Logger) error {
+	return loadSnapshot(ctx, fsm.db, logger, snapshotFile, filterKey, true)
 }
 
 // loadSnapshot loads a snapshot from a file into the supplied boltDB database.
@@ -514,7 +514,7 @@ func LoadReadOnlySnapshot(fsm *FSM, snapshotFile io.ReadCloser, filterKey func(k
 // to 1.0.
 // If pathsToFilter is not nil, the function will filter out any keys that are
 // found in the pathsToFilter tree.
-func loadSnapshot(db *bolt.DB, logger log.Logger, snapshotFile io.ReadCloser, filterKey func(key string) bool, readOnly bool) error {
+func loadSnapshot(ctx context.Context, db *bolt.DB, logger log.Logger, snapshotFile io.ReadCloser, filterKey func(key string) bool, readOnly bool) error {
 	// The delimited reader will parse full proto messages from the snapshot data.
 	protoReader := NewDelimitedReader(snapshotFile, math.MaxInt32)
 	defer protoReader.Close()
@@ -524,6 +524,11 @@ func loadSnapshot(db *bolt.DB, logger log.Logger, snapshotFile io.ReadCloser, fi
 	entry := new(pb.StorageEntry)
 	for !done {
 		err := db.Update(func(tx *bolt.Tx) error {
+			select {
+			case <-ctx.Done():
+				return fmt.Errorf("context canceled while loading snapshot: %w", ctx.Err())
+			default:
+			}
 			b, err := tx.CreateBucketIfNotExists(dataBucketName)
 			if readOnly {
 				b.FillPercent = 1.0

physical/raft/snapshot_test.go

@@ -994,7 +994,7 @@ func TestLoadReadOnlySnapshot(t *testing.T) {
 	// Create an FSM to load the snapshot data into.
 	fsm, err := NewFSM(dir, "test-fsm", logger)
 
-	err = LoadReadOnlySnapshot(fsm, snapshotFile, toExclude, logger)
+	err = LoadReadOnlySnapshot(context.Background(), fsm, snapshotFile, toExclude, logger)
 	require.NoError(t, err)
 	value, err := fsm.Get(context.Background(), "/path/to/exclude/1")
 	require.NoError(t, err)

vault/snapshots/source.go

@@ -30,6 +30,18 @@ func (m *manualUploadSource) Type(_ context.Context) string {
 	return "manual"
 }
 
-func (m *manualUploadSource) ReadCloser(_ context.Context) (io.ReadCloser, error) {
-	return m.r, nil
+func (m *manualUploadSource) ReadCloser(ctx context.Context) (io.ReadCloser, error) {
+	return &ctxAwareReadCloser{ctx, m.r}, nil
+}
+
+type ctxAwareReadCloser struct {
+	ctx context.Context
+	io.ReadCloser
+}
+
+func (c *ctxAwareReadCloser) Read(p []byte) (n int, err error) {
+	if c.ctx.Err() != nil {
+		return 0, c.ctx.Err()
+	}
+	return c.ReadCloser.Read(p)
 }