Adding SnapshotLock Capabilities

Author: mdzraf

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/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"
@@ -856,9 +859,11 @@ func (d *ControllerService) CreateSnapshot(ctx context.Context, req *csi.CreateS
 		cloud.AwsEbsDriverTagKey: isManagedByDriver,
 	}
 
+	// Get Parameters Here
 	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 +881,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 +963,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 +1332,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)
+}

pkg/driver/constants.go

@@ -42,4 +42,5 @@ type EC2API interface {
 	CreateTags(ctx context.Context, params *ec2.CreateTagsInput, optFns ...func(*ec2.Options)) (*ec2.CreateTagsOutput, error)
 	DeleteTags(ctx context.Context, params *ec2.DeleteTagsInput, optFns ...func(*ec2.Options)) (*ec2.DeleteTagsOutput, error)
 	EnableFastSnapshotRestores(ctx context.Context, params *ec2.EnableFastSnapshotRestoresInput, optFns ...func(*ec2.Options)) (*ec2.EnableFastSnapshotRestoresOutput, error)
+	LockSnapshot(ctx context.Context, params *ec2.LockSnapshotInput, optFns ...func(*ec2.Options)) (*ec2.LockSnapshotOutput, error)
 }

pkg/driver/controller.go

@@ -210,6 +210,10 @@ func (d *fakeCloud) EnableFastSnapshotRestores(ctx context.Context, availability
 	return &ec2.EnableFastSnapshotRestoresOutput{}, nil
 }
 
+func (d *fakeCloud) LockSnapshot(ctx context.Context, lockOptions ec2.LockSnapshotInput) (*ec2.LockSnapshotOutput, error) {
+	return &ec2.LockSnapshotOutput{}, nil
+}
+
 func (d *fakeCloud) GetDiskByName(ctx context.Context, name string, capacityBytes int64) (*cloud.Disk, error) {
 	return &cloud.Disk{}, nil
 }