Allow CMKs and bucket keys

README.md

@@ -67,18 +67,21 @@ terraform apply \
 
 ## Terraform Variables
 
-| Variable                  | Required | Default                      | Description                                            |
-| ------------------------- | -------- | ---------------------------- | ------------------------------------------------------ |
-| `aws_s3_bucket_name`      | Yes      | -                            | Name of the S3 bucket to create for storing audit logs |
-| `render_api_key`          | Yes      | -                            | Render API key for accessing audit logs                |
-| `render_workspace_ids`    | No       | `[]`                         | List of workspace IDs to fetch audit logs from         |
-| `render_organization_id`  | No       | `""`                         | Organization ID for Enterprise audit logs              |
-| `aws_iam_user_name`       | No       | `render-audit-log-processor` | Name of the IAM user created for S3 access             |
-| `render_cronjob_name`     | No       | `render-auditlogs`           | Name of the Render Cron Job                            |
-| `render_cronjob_schedule` | No       | `1/15 * * * *`               | Cron schedule (default: every 15 minutes)              |
-| `render_cronjob_plan`     | No       | `starter`                    | Render plan for the Cron Job                           |
-| `render_cronjob_region`   | No       | `oregon`                     | Region to deploy the Cron Job                          |
-| `render_project_name`     | No       | `audit-logs`                 | Name of the Render project                             |
+| Variable                    | Required | Default                      | Description                                            |
+| --------------------------- | -------- | ---------------------------- | ------------------------------------------------------ |
+| `aws_s3_bucket_name`        | Yes      | -                            | Name of the S3 bucket to create for storing audit logs |
+| `render_api_key`            | Yes      | -                            | Render API key for accessing audit logs                |
+| `render_workspace_ids`      | No       | `[]`                         | List of workspace IDs to fetch audit logs from         |
+| `render_organization_id`    | No       | `""`                         | Organization ID for Enterprise audit logs              |
+| `aws_iam_user_name`         | No       | `render-audit-log-processor` | Name of the IAM user created for S3 access             |
+| `aws_s3_bucket_key_enabled` | No       | `false`                      | Enable S3 bucket key to reduce KMS calls               |
+| `aws_s3_kms_key_id`         | No       | `""`                         | ARN for KMS key to use for encryption                  |
+| `aws_s3_use_kms`            | No       | `false`                      | Use KMS for encryption (instead of SSE-S3)             |
+| `render_cronjob_name`       | No       | `render-auditlogs`           | Name of the Render Cron Job                            |
+| `render_cronjob_schedule`   | No       | `1/15 * * * *`               | Cron schedule (default: every 15 minutes)              |
+| `render_cronjob_plan`       | No       | `starter`                    | Render plan for the Cron Job                           |
+| `render_cronjob_region`     | No       | `oregon`                     | Region to deploy the Cron Job                          |
+| `render_project_name`       | No       | `audit-logs`                 | Name of the Render project                             |
 
 ## Architecture
 
@@ -108,6 +111,11 @@ RENDER_API_KEY=your-api-key
 AWS_ACCESS_KEY_ID=your-aws-key
 AWS_SECRET_ACCESS_KEY=your-aws-secret
 AWS_REGION=us-west-2
+
+# Optional: KMS encryption settings (defaults to SSE-S3 if not set)
+S3_USE_KMS=true
+S3_KMS_KEY_ID=arn:aws:kms:us-west-2:123456789012:key/your-key-id  # Optional
+S3_BUCKET_KEY_ENABLED=true  # Optional
 ```
 
 2. Run the application:

main.go

@@ -30,7 +30,11 @@ func main() {
 	}
 
 	// Create S3 uploader
-	uploader, err := aws.NewUploader(ctx, s3.NewFromConfig(cfg.AWSConfig), cfg.S3Bucket, cfg.AWSRegion)
+	uploader, err := aws.NewUploaderWithOptions(ctx, s3.NewFromConfig(cfg.AWSConfig), cfg.S3Bucket, cfg.AWSRegion, aws.UploaderOptions{
+		UseKMS:           cfg.S3UseKMS,
+		KMSKeyID:         cfg.S3KMSKeyID,
+		BucketKeyEnabled: cfg.S3BucketKeyEnabled,
+	})
 	if err != nil {
 		log.Fatal("Error creating S3 uploader:", err)
 	}

pkg/aws/checkpoint.go

@@ -60,13 +60,27 @@ func (u *Uploader) SaveCheckpoint(ctx context.Context, cp *Checkpoint, logType a
 		return fmt.Errorf("error marshaling checkpoint: %w", err)
 	}
 
-	_, err = u.client.PutObject(ctx, &s3.PutObjectInput{
-		Bucket:               aws.String(u.bucket),
-		Key:                  aws.String(fmt.Sprintf("%s=%s/%s", logType, workspace, checkpointKey)),
-		Body:                 bytes.NewReader(data),
-		ContentType:          aws.String("application/json"),
-		ServerSideEncryption: types.ServerSideEncryptionAes256,
-	})
+	putInput := &s3.PutObjectInput{
+		Bucket:      aws.String(u.bucket),
+		Key:         aws.String(fmt.Sprintf("%s=%s/%s", logType, workspace, checkpointKey)),
+		Body:        bytes.NewReader(data),
+		ContentType: aws.String("application/json"),
+	}
+
+	// Configure server-side encryption
+	if u.opts.UseKMS {
+		putInput.ServerSideEncryption = types.ServerSideEncryptionAwsKms
+		if u.opts.KMSKeyID != "" {
+			putInput.SSEKMSKeyId = aws.String(u.opts.KMSKeyID)
+		}
+		if u.opts.BucketKeyEnabled {
+			putInput.BucketKeyEnabled = aws.Bool(true)
+		}
+	} else {
+		putInput.ServerSideEncryption = types.ServerSideEncryptionAes256
+	}
+
+	_, err = u.client.PutObject(ctx, putInput)
 	if err != nil {
 		return fmt.Errorf("error writing checkpoint to S3: %w", err)
 	}

pkg/aws/checkpoint_test.go

@@ -158,6 +158,38 @@ func TestSaveCheckpoint(t *testing.T) {
 		require.NoError(t, err)
 	})
 
+	t.Run("uses KMS with key ID and bucket key enabled", func(t *testing.T) {
+		checkpoint := &awspkg.Checkpoint{
+			LastCursor:    "kms-cursor",
+			LastTimestamp: testTime,
+		}
+
+		const kmsKey = "arn:aws:kms:us-west-2:123456789012:key/abcdefab-1234-5678-9abc-def012345678"
+
+		s3Client := &mockS3Client{
+			putObjectFunc: func(ctx context.Context, params *s3.PutObjectInput, optFns ...func(*s3.Options)) (*s3.PutObjectOutput, error) {
+				require.Equal(t, "test-bucket", *params.Bucket)
+				require.Equal(t, "workspace=test-workspace/checkpoint.json", *params.Key)
+				require.Equal(t, types.ServerSideEncryptionAwsKms, params.ServerSideEncryption)
+				require.NotNil(t, params.SSEKMSKeyId)
+				require.Equal(t, kmsKey, *params.SSEKMSKeyId)
+				require.NotNil(t, params.BucketKeyEnabled)
+				require.True(t, *params.BucketKeyEnabled)
+				return &s3.PutObjectOutput{}, nil
+			},
+		}
+
+		uploader, err := awspkg.NewUploaderWithOptions(ctx, s3Client, "test-bucket", "test-region", awspkg.UploaderOptions{
+			UseKMS:           true,
+			KMSKeyID:         kmsKey,
+			BucketKeyEnabled: true,
+		})
+		require.NoError(t, err)
+
+		err = uploader.SaveCheckpoint(ctx, checkpoint, auditlogs.WorkspaceAuditLog, "test-workspace")
+		require.NoError(t, err)
+	})
+
 	t.Run("returns error on S3 error", func(t *testing.T) {
 		checkpoint := &awspkg.Checkpoint{
 			LastCursor:    "test-cursor",

pkg/aws/uploader.go

@@ -21,15 +21,27 @@ type S3Client interface {
 	PutObject(ctx context.Context, params *s3.PutObjectInput, optFns ...func(*s3.Options)) (*s3.PutObjectOutput, error)
 }
 
+type UploaderOptions struct {
+	UseKMS           bool
+	KMSKeyID         string
+	BucketKeyEnabled bool
+}
+
 type Uploader struct {
 	client S3Client
 	bucket string
+	opts   UploaderOptions
 }
 
 func NewUploader(ctx context.Context, client S3Client, bucket, region string) (*Uploader, error) {
+	return NewUploaderWithOptions(ctx, client, bucket, region, UploaderOptions{})
+}
+
+func NewUploaderWithOptions(ctx context.Context, client S3Client, bucket, region string, opts UploaderOptions) (*Uploader, error) {
 	return &Uploader{
 		client: client,
 		bucket: bucket,
+		opts:   opts,
 	}, nil
 }
 
@@ -56,13 +68,28 @@ func (u *Uploader) UploadAuditLogs(ctx context.Context, auditLogType auditlogs.L
 	key := generateS3Key(auditLogType, id, data[0].AuditLog.Timestamp)
 
 	// Upload to S3
-	_, err = u.client.PutObject(ctx, &s3.PutObjectInput{
-		Bucket:               aws.String(u.bucket),
-		Key:                  aws.String(key),
-		Body:                 bytes.NewReader(compressedData.Bytes()),
-		ContentType:          aws.String("application/gzip"),
-		ServerSideEncryption: types.ServerSideEncryptionAes256,
-	})
+	putInput := &s3.PutObjectInput{
+		Bucket:      aws.String(u.bucket),
+		Key:         aws.String(key),
+		Body:        bytes.NewReader(compressedData.Bytes()),
+		ContentType: aws.String("application/gzip"),
+	}
+
+	// Configure server-side encryption
+	if u.opts.UseKMS {
+		putInput.ServerSideEncryption = types.ServerSideEncryptionAwsKms
+		if u.opts.KMSKeyID != "" {
+			putInput.SSEKMSKeyId = aws.String(u.opts.KMSKeyID)
+		}
+		if u.opts.BucketKeyEnabled {
+			putInput.BucketKeyEnabled = aws.Bool(true)
+		}
+	} else {
+		// Default to SSE-S3 (AES256)
+		putInput.ServerSideEncryption = types.ServerSideEncryptionAes256
+	}
+
+	_, err = u.client.PutObject(ctx, putInput)
 	if err != nil {
 		return "", fmt.Errorf("error uploading to S3: %w", err)
 	}

pkg/aws/uploader_test.go

@@ -12,6 +12,7 @@ import (
 	"time"
 
 	"github.com/aws/aws-sdk-go-v2/service/s3"
+	"github.com/aws/aws-sdk-go-v2/service/s3/types"
 	"github.com/stretchr/testify/require"
 
 	"github.com/renderinc/render-auditlogs/pkg/auditlogs"
@@ -108,4 +109,100 @@ func TestUploadAuditLogs(t *testing.T) {
 		require.Contains(t, err.Error(), "error uploading to S3")
 		require.Empty(t, s3URI)
 	})
+
+	t.Run("uses default SSE-S3 encryption when KMS not enabled", func(t *testing.T) {
+		t.Parallel()
+		s3Client := &mockS3Client{
+			putObjectFunc: func(ctx context.Context, params *s3.PutObjectInput, optFns ...func(*s3.Options)) (*s3.PutObjectOutput, error) {
+				require.Equal(t, "test-bucket", *params.Bucket)
+				require.Equal(t, types.ServerSideEncryptionAes256, params.ServerSideEncryption)
+				require.Nil(t, params.SSEKMSKeyId)
+				require.Nil(t, params.BucketKeyEnabled)
+				return &s3.PutObjectOutput{}, nil
+			},
+		}
+
+		uploader, err := aws.NewUploaderWithOptions(ctx, s3Client, "test-bucket", "test-region", aws.UploaderOptions{
+			UseKMS: false,
+		})
+		require.NoError(t, err)
+
+		s3URI, err := uploader.UploadAuditLogs(ctx, auditlogs.WorkspaceAuditLog, "workspace-123", testData)
+
+		require.NoError(t, err)
+		require.NotEmpty(t, s3URI)
+	})
+
+	t.Run("uses KMS encryption without specific key ID", func(t *testing.T) {
+		t.Parallel()
+		s3Client := &mockS3Client{
+			putObjectFunc: func(ctx context.Context, params *s3.PutObjectInput, optFns ...func(*s3.Options)) (*s3.PutObjectOutput, error) {
+				require.Equal(t, "test-bucket", *params.Bucket)
+				require.Equal(t, types.ServerSideEncryptionAwsKms, params.ServerSideEncryption)
+				require.Nil(t, params.SSEKMSKeyId)
+				require.Nil(t, params.BucketKeyEnabled)
+				return &s3.PutObjectOutput{}, nil
+			},
+		}
+
+		uploader, err := aws.NewUploaderWithOptions(ctx, s3Client, "test-bucket", "test-region", aws.UploaderOptions{
+			UseKMS: true,
+		})
+		require.NoError(t, err)
+
+		s3URI, err := uploader.UploadAuditLogs(ctx, auditlogs.WorkspaceAuditLog, "workspace-123", testData)
+
+		require.NoError(t, err)
+		require.NotEmpty(t, s3URI)
+	})
+
+	t.Run("uses KMS encryption with specific key ID", func(t *testing.T) {
+		t.Parallel()
+		s3Client := &mockS3Client{
+			putObjectFunc: func(ctx context.Context, params *s3.PutObjectInput, optFns ...func(*s3.Options)) (*s3.PutObjectOutput, error) {
+				require.Equal(t, "test-bucket", *params.Bucket)
+				require.Equal(t, types.ServerSideEncryptionAwsKms, params.ServerSideEncryption)
+				require.NotNil(t, params.SSEKMSKeyId)
+				require.Equal(t, "arn:aws:kms:us-west-2:123456789012:key/12345678-1234-1234-1234-123456789012", *params.SSEKMSKeyId)
+				require.Nil(t, params.BucketKeyEnabled)
+				return &s3.PutObjectOutput{}, nil
+			},
+		}
+
+		uploader, err := aws.NewUploaderWithOptions(ctx, s3Client, "test-bucket", "test-region", aws.UploaderOptions{
+			UseKMS:   true,
+			KMSKeyID: "arn:aws:kms:us-west-2:123456789012:key/12345678-1234-1234-1234-123456789012",
+		})
+		require.NoError(t, err)
+
+		s3URI, err := uploader.UploadAuditLogs(ctx, auditlogs.WorkspaceAuditLog, "workspace-123", testData)
+
+		require.NoError(t, err)
+		require.NotEmpty(t, s3URI)
+	})
+
+	t.Run("uses KMS encryption with bucket key enabled", func(t *testing.T) {
+		t.Parallel()
+		s3Client := &mockS3Client{
+			putObjectFunc: func(ctx context.Context, params *s3.PutObjectInput, optFns ...func(*s3.Options)) (*s3.PutObjectOutput, error) {
+				require.Equal(t, "test-bucket", *params.Bucket)
+				require.Equal(t, types.ServerSideEncryptionAwsKms, params.ServerSideEncryption)
+				require.Nil(t, params.SSEKMSKeyId)
+				require.NotNil(t, params.BucketKeyEnabled)
+				require.True(t, *params.BucketKeyEnabled)
+				return &s3.PutObjectOutput{}, nil
+			},
+		}
+
+		uploader, err := aws.NewUploaderWithOptions(ctx, s3Client, "test-bucket", "test-region", aws.UploaderOptions{
+			UseKMS:           true,
+			BucketKeyEnabled: true,
+		})
+		require.NoError(t, err)
+
+		s3URI, err := uploader.UploadAuditLogs(ctx, auditlogs.WorkspaceAuditLog, "workspace-123", testData)
+
+		require.NoError(t, err)
+		require.NotEmpty(t, s3URI)
+	})
 }

pkg/env/env.go

@@ -16,6 +16,9 @@ type Config struct {
 	WorkspaceIDS       []string `required:"true" split_words:"true"`
 	OrganizationID     string   `required:"false" split_words:"true"`
 	S3Bucket           string   `required:"true" split_words:"true"`
+	S3BucketKeyEnabled bool     `required:"false" split_words:"true"`
+	S3KMSKeyID         string   `required:"false" split_words:"true"`
+	S3UseKMS           bool     `required:"false" split_words:"true"`
 	RenderAPIKey       string   `required:"true" split_words:"true"`
 	AWSAccessKeyID     string   `required:"true" split_words:"true"`
 	AWSSecretAccessKey string   `required:"true" split_words:"true"`

terraform/main.tf

@@ -4,6 +4,7 @@ module "aws" {
 
   aws_s3_bucket_name = var.aws_s3_bucket_name
   aws_iam_user_name = var.aws_iam_user_name
+  aws_s3_use_kms = var.aws_s3_use_kms
 }
 
 module "render" {

terraform/modules/aws/s3.tf

@@ -49,7 +49,14 @@ resource "aws_s3_bucket_policy" "render_audit_logs" {
           Principal = "*",
           Action = "s3:PutObject",
           Resource = "arn:aws:s3:::${aws_s3_bucket.render_audit_logs.id}/*",
-          Condition = {
+          Condition = var.aws_s3_use_kms ? {
+            StringEquals = {
+              "s3:x-amz-server-side-encryption" = "aws:kms"
+            },
+            Null = {
+              "s3:x-amz-server-side-encryption-aws-kms-key-id" = "true"
+            }
+          } : {
             StringNotEquals = {
               "s3:x-amz-server-side-encryption" = "AES256"
             }

terraform/modules/aws/variables.tf

@@ -7,3 +7,8 @@ variable "aws_iam_user_name" {
   type = string
   default = "render-audit-log-processor"
 }
+
+variable "aws_s3_use_kms" {
+  type    = bool
+  default = false
+}

terraform/modules/render-audit-logs/render.tf

@@ -27,6 +27,9 @@ resource "render_cron_job" "render-audit-logs" {
     "WORKSPACE_IDS" = { value = join(",", var.render_workspace_ids) }
     "RENDER_API_KEY" = { value = var.render_api_key }
     "S3_BUCKET" = { value = var.aws_s3_bucket_name }
+    "S3_BUCKET_KEY_ENABLED" = { value = var.aws_s3_bucket_key_enabled }
+    "S3_KMS_KEY_ID" = { value = var.aws_s3_kms_key_id }
+    "S3_USE_KMS" = { value = var.aws_s3_use_kms }
   }
 }
 

terraform/modules/render-audit-logs/variables.tf

@@ -2,6 +2,21 @@ variable "aws_s3_bucket_name" {
   type = string
 }
 
+variable "aws_s3_bucket_key_enabled" {
+  type = bool
+  default = false
+}
+
+variable "aws_s3_kms_key_id" {
+  type = string
+  default = ""
+}
+
+variable "aws_s3_use_kms" {
+  type = bool
+  default = false
+}
+
 variable "aws_access_key" {
   type = string
   sensitive = true