Adding SnapshotLock Capabilities

Author: mdzraf

docs/snapshot.md

@@ -3,6 +3,11 @@
 |------------------------------------|-----------------------------------------------------------|
 | fastSnapshotRestoreAvailabilityZones | Comma separated list of availability zones                |
 | outpostArn                         | Arn of the outpost you wish to have the snapshot saved to | 
+| snapshotLockEnabled                | Enable snapshot locking (true/false)                     |
+| snapshotLockMode                   | Lock mode (governance/compliance)                         |
+| snapshotLockDuration               | Lock duration in days                                     |
+| snapshotLockExpirationDate         | Lock expiration date (RFC3339 format)                    |
+| snapshotLockCoolOffPeriod          | Cool-off period in hours (compliance mode only)          | 
 
 The AWS EBS CSI Driver supports [tagging](tagging.md) through `VolumeSnapshotClass.parameters` (in v1.6.0 and later). 
 ## Prerequisites
@@ -44,6 +49,51 @@ parameters:
 
 The driver will attempt to check if the availability zones provided are supported for fast snapshot restore before attempting to create the snapshot. If the `EnableFastSnapshotRestores` API call fails, the driver will hard-fail the request and delete the snapshot. This is to ensure that the snapshot is not left in an inconsistent state.
 
+# Snapshot Lock
+
+The EBS CSI Driver provides support for [EBS Snapshot Lock](https://docs.aws.amazon.com/ebs/latest/userguide/ebs-snapshot-lock.html) via `VolumeSnapshotClass.parameters`. Snapshot locking protects snapshots from accidental or malicious deletion.
+
+Snapshot locking supports two modes:
+- **Governance mode**: Allows privileged users to unlock snapshots
+- **Compliance mode**: Prevents unlocking until expiration (with optional cool-off period)
+
+Lock duration can be specified using either:
+- `snapshotLockDuration`: Duration in days
+- `snapshotLockExpirationDate`: Specific expiration date in RFC3339 format
+
+**Example - Governance Mode with Duration**
+```yaml
+apiVersion: snapshot.storage.k8s.io/v1
+kind: VolumeSnapshotClass
+metadata:
+  name: csi-aws-vsc-locked
+driver: ebs.csi.aws.com
+deletionPolicy: Delete
+parameters:
+  snapshotLockEnabled: "true"
+  snapshotLockMode: "governance"
+  snapshotLockDuration: "7"
+```
+
+**Example - Compliance Mode with Expiration Date**
+```yaml
+apiVersion: snapshot.storage.k8s.io/v1
+kind: VolumeSnapshotClass
+metadata:
+  name: csi-aws-vsc-compliance
+driver: ebs.csi.aws.com
+deletionPolicy: Delete
+parameters:
+  snapshotLockEnabled: "true"
+  snapshotLockMode: "compliance"
+  snapshotLockExpirationDate: "2026-12-31T23:59:59Z"
+  snapshotLockCoolOffPeriod: "24"
+```
+
+## Failure Mode
+
+If the `LockSnapshot` API call fails, the driver will hard-fail the request and delete the snapshot. This ensures that the snapshot is not left in an unlocked state when locking was explicitly requested.
+
 
 # Amazon EBS Local Snapshots on Outposts
 

pkg/cloud/cloud.go

@@ -306,12 +306,18 @@ type ListSnapshotsResponse struct {
 	NextToken string
 }
 
-// SnapshotOptions represents parameters to create an EBS volume.
+// SnapshotOptions represents parameters to create an EBS snapshot.
 type SnapshotOptions struct {
 	Tags       map[string]string
 	OutpostArn string
 }
 
+// SnapshotLockOptions represents the snapshot lock specific parameters for locking en EBS snapshot.
+type SnapshotLockOptions struct {
+	SnapshotLockEnabled bool
+	LockSnapshotInput   ec2.LockSnapshotInput
+}
+
 // ec2ListSnapshotsResponse is a helper struct returned from the AWS API calling function to the main ListSnapshots function.
 type ec2ListSnapshotsResponse struct {
 	Snapshots []types.Snapshot
@@ -1867,6 +1873,15 @@ func (c *cloud) CreateSnapshot(ctx context.Context, volumeID string, snapshotOpt
 	}, nil
 }
 
+func (c *cloud) LockSnapshot(ctx context.Context, lockSnapshotInput ec2.LockSnapshotInput) (*ec2.LockSnapshotOutput, error) {
+	klog.InfoS("Attempting to lock Snapshot", "request parameters: ", lockSnapshotInput)
+	response, err := c.ec2.LockSnapshot(ctx, &lockSnapshotInput)
+	if err != nil {
+		return nil, err
+	}
+	return response, nil
+}
+
 func (c *cloud) DeleteSnapshot(ctx context.Context, snapshotID string) (success bool, err error) {
 	request := &ec2.DeleteSnapshotInput{}
 	request.SnapshotId = aws.String(snapshotID)

pkg/cloud/cloud_test.go

@@ -5320,3 +5320,56 @@ func TestCheckIfIopsIncreaseOnExpansion(t *testing.T) {
 		})
 	}
 }
+
+func TestLockSnapshot(t *testing.T) {
+	testCases := []struct {
+		name      string
+		input     ec2.LockSnapshotInput
+		mockError error
+		expectErr bool
+	}{
+		{
+			name: "success: API call succeeds",
+			input: ec2.LockSnapshotInput{
+				SnapshotId:   aws.String("snap-test-id"),
+				LockMode:     types.LockModeGovernance,
+				LockDuration: aws.Int32(1),
+			},
+			mockError: nil,
+			expectErr: false,
+		},
+		{
+			name: "fail: AWS API error is propagated",
+			input: ec2.LockSnapshotInput{
+				SnapshotId: aws.String("snap-test-id"),
+			},
+			mockError: errors.New("InvalidSnapshot.NotFound"),
+			expectErr: true,
+		},
+	}
+
+	for _, tc := range testCases {
+		t.Run(tc.name, func(t *testing.T) {
+			mockCtrl := gomock.NewController(t)
+			mockEC2 := NewMockEC2API(mockCtrl)
+			c := newCloud(mockEC2)
+
+			ctx := context.Background()
+
+			if tc.mockError != nil {
+				mockEC2.EXPECT().LockSnapshot(ctx, &tc.input).Return(nil, tc.mockError)
+			} else {
+				mockEC2.EXPECT().LockSnapshot(ctx, &tc.input).Return(&ec2.LockSnapshotOutput{}, nil)
+			}
+
+			_, err := c.LockSnapshot(ctx, tc.input)
+
+			if tc.expectErr {
+				require.Error(t, err)
+				require.Equal(t, tc.mockError.Error(), err.Error())
+			} else {
+				assert.NoError(t, err)
+			}
+		})
+	}
+}

pkg/cloud/interface.go

@@ -42,4 +42,5 @@ type Cloud interface {
 	AvailabilityZones(ctx context.Context) (map[string]struct{}, error)
 	DryRun(ctx context.Context) error
 	GetInstancesPatching(ctx context.Context, nodeIDs []string) ([]*types.Instance, error)
+	LockSnapshot(ctx context.Context, lockOptions ec2.LockSnapshotInput) (*ec2.LockSnapshotOutput, error)
 }

pkg/cloud/mock_cloud.go

@@ -126,6 +126,21 @@ const (
 const (
 	// FastSnapshotRestoreAvailabilityZones represents key for fast snapshot restore availability zones.
 	FastSnapshotRestoreAvailabilityZones = "fastsnapshotrestoreavailabilityzones"
+
+	// SnapshotLockEnabled represents a key for indicating whether snapshot lock is enabled or disabled.
+	SnapshotLockEnabled = "snapshotlockenabled"
+
+	// SnapshotLockMode represents a key for indicating whether snapshots are locked in Governance or Compliance mode.
+	SnapshotLockMode = "snapshotlockmode"
+
+	// SnapshotLockDuration is a key for the duration for which to lock the snapshots, specified in days.
+	SnapshotLockDuration = "snapshotlockduration"
+
+	// SnapshotLockExpirationDate is a key for specifying the expiration date for the snapshot lock, specified in the format "YYYY-MM-DDThh:mm:ss.sssZ".
+	SnapshotLockExpirationDate = "snapshotlockexpirationdate"
+
+	// SnapshotLockCoolOffPeriod is a key specifying the cooling-off period for compliance mode, specified in hours.
+	SnapshotLockCoolOffPeriod = "snapshotlockcooloffperiod"
 )
 
 // constants for volume tags and their values.

pkg/cloud/mock_ec2.go

@@ -23,8 +23,11 @@ import (
 	"maps"
 	"strconv"
 	"strings"
+	"time"
 
+	"github.com/aws/aws-sdk-go-v2/aws"
 	"github.com/aws/aws-sdk-go-v2/aws/arn"
+	"github.com/aws/aws-sdk-go-v2/service/ec2/types"
 	"github.com/awslabs/volume-modifier-for-k8s/pkg/rpc"
 	csi "github.com/container-storage-interface/spec/lib/go/csi"
 	"github.com/kubernetes-sigs/aws-ebs-csi-driver/pkg/cloud"
@@ -859,6 +862,7 @@ func (d *ControllerService) CreateSnapshot(ctx context.Context, req *csi.CreateS
 	var vscTags []string
 	var fsrAvailabilityZones []string
 	vsProps := new(template.VolumeSnapshotProps)
+	vsLock := new(cloud.SnapshotLockOptions)
 	for key, value := range req.GetParameters() {
 		switch strings.ToLower(key) {
 		case VolumeSnapshotNameKey:
@@ -876,6 +880,28 @@ func (d *ControllerService) CreateSnapshot(ctx context.Context, req *csi.CreateS
 			} else {
 				return nil, status.Errorf(codes.InvalidArgument, "Invalid parameter value %s is not a valid arn", value)
 			}
+		case SnapshotLockEnabled:
+			vsLock.SnapshotLockEnabled = isTrue(value)
+		case SnapshotLockMode:
+			vsLock.LockSnapshotInput.LockMode = types.LockMode(value)
+		case SnapshotLockDuration:
+			lockDuration, err := strconv.ParseInt(value, 10, 32)
+			if err != nil {
+				return nil, status.Errorf(codes.InvalidArgument, "Could not parse SnapshotLockDuration: %q", value)
+			}
+			vsLock.LockSnapshotInput.LockDuration = aws.Int32(int32(lockDuration))
+		case SnapshotLockExpirationDate:
+			expirationDate, err := time.Parse(time.RFC3339, value)
+			if err != nil {
+				return nil, status.Errorf(codes.InvalidArgument, "Could not parse SnapshotLockExpirationDate: %q", value)
+			}
+			vsLock.LockSnapshotInput.ExpirationDate = &expirationDate
+		case SnapshotLockCoolOffPeriod:
+			lockCoolOffPeriod, err := strconv.ParseInt(value, 10, 32)
+			if err != nil {
+				return nil, status.Errorf(codes.InvalidArgument, "Could not parse SnapshotLockCoolOffPeriod: %q", value)
+			}
+			vsLock.LockSnapshotInput.CoolOffPeriod = aws.Int32(int32(lockCoolOffPeriod))
 		default:
 			if strings.HasPrefix(key, TagKeyPrefix) {
 				vscTags = append(vscTags, value)
@@ -936,12 +962,18 @@ func (d *ControllerService) CreateSnapshot(ctx context.Context, req *csi.CreateS
 	if len(fsrAvailabilityZones) > 0 {
 		_, err := d.cloud.EnableFastSnapshotRestores(ctx, fsrAvailabilityZones, snapshot.SnapshotID)
 		if err != nil {
-			if _, deleteErr := d.cloud.DeleteSnapshot(ctx, snapshot.SnapshotID); deleteErr != nil {
-				return nil, status.Errorf(codes.Internal, "Could not delete snapshot ID %q: %v", snapshotName, deleteErr)
-			}
-			return nil, status.Errorf(codes.Internal, "Failed to create Fast Snapshot Restores for snapshot ID %q: %v", snapshotName, err)
+			return nil, d.cleanupSnapshotOnError(ctx, snapshot.SnapshotID, snapshotName, err, "Failed to create Fast Snapshot Restores")
 		}
 	}
+
+	if vsLock.SnapshotLockEnabled {
+		vsLock.LockSnapshotInput.SnapshotId = &snapshot.SnapshotID
+		_, err := d.cloud.LockSnapshot(ctx, vsLock.LockSnapshotInput)
+		if err != nil {
+			return nil, d.cleanupSnapshotOnError(ctx, snapshot.SnapshotID, snapshotName, err, "Failed to lock snapshot")
+		}
+	}
+
 	return newCreateSnapshotResponse(snapshot), nil
 }
 
@@ -1299,3 +1331,10 @@ func validateFormattingOption(volumeCapabilities []*csi.VolumeCapability, paramN
 func isTrue(value string) bool {
 	return value == trueStr
 }
+
+func (d *ControllerService) cleanupSnapshotOnError(ctx context.Context, snapshotID, snapshotName string, originalErr error, errorMsg string) error {
+	if _, deleteErr := d.cloud.DeleteSnapshot(ctx, snapshotID); deleteErr != nil {
+		return status.Errorf(codes.Internal, "Could not delete snapshot ID %q: %v", snapshotName, deleteErr)
+	}
+	return status.Errorf(codes.Internal, "%s for snapshot ID %q: %v", errorMsg, snapshotName, originalErr)
+}