[TT-14473] - support for encrypted aws kinesis

README.md

@@ -1358,7 +1358,9 @@
 
 `region` - The AWS region your Kinesis stream is located - i.e. eu-west-2
 
-`batch_size` - Optional. The maximum size of the records in a batch is 5MiB. If your records are larger in size setting this batch size paramter can guarantee you don't have failed delivery due to too large a batch. Default size if unset is 100.
+`batch_size` - The maximum size of the records in a batch is 5MiB. If your records are larger in size setting this batch size paramter can guarantee you don't have failed delivery due to too large a batch. Default size if unset is 100 (optional)
+
+`kms_key_id` - The AWS KMS key ID used to encrypt the records in the Kinesis stream. If not provided, the records will not be encrypted (optional)
 
 ###### JSON / Conf File
 
@@ -1368,7 +1370,8 @@
       "meta": {
         "stream_name": "my-stream",
         "region": "eu-west-2",
-        "batch_size": 100
+        "batch_size": 100,
+        "kms_key_id": "your-kms-key-id"
       }
     },
 ```
@@ -1381,6 +1384,7 @@
 TYK_PMP_PUMPS_KINESIs_META_STREAMNAME=my-stream
 TYK_PMP_PUMPS_KINESIS_META_REGION=eu-west-2
 TYK_PMP_PUMPS_KINESIS_META_BATCHSIZE=100
+TYK_PMP_PUMPS_KINESIS_META_KMSKEYID=your-kms-key-id
 ```
 
 # Base Pump Configurations

pumps/kinesis.go

@@ -4,8 +4,10 @@
 	"context"
 	"crypto/rand"
 	"encoding/json"
+	"errors"
 	"fmt"
 	"math/big"
+	"time"
 
 	"github.com/TykTechnologies/tyk-pump/analytics"
 	"github.com/mitchellh/mapstructure"
@@ -41,6 +43,9 @@
 	// Each record in the request can be as large as 1 MiB, up to a limit of 5 MiB for the entire request, including partition keys.
 	// Each shard can support writes up to 1,000 records per second, up to a maximum data write total of 1 MiB per second.
 	BatchSize int `mapstructure:"batch_size"`
+	// The KMS Key ID used for server-side encryption of the Kinesis stream.
+	// Defaults to an empty string if not provided.
+	KMSKeyID string `mapstructure:"kms_key_id" default:""`
 }
 
 var (
@@ -85,6 +90,53 @@
 
 	// Create Kinesis client
 	p.client = kinesis.NewFromConfig(cfg)
+
+	// Check if KMSKeyID is provided and enable server-side encryption
+	if p.kinesisConf.KMSKeyID != "" {
+		ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second)
+		defer cancel()
+
+		// First, check if encryption is already enabled
+		describeOutput, err := p.client.DescribeStream(ctx, &kinesis.DescribeStreamInput{
+			StreamName: aws.String(p.kinesisConf.StreamName),
+		})
+
+		switch {
+		case err != nil:
+			p.log.Fatalf("Failed to describe Kinesis stream: %v", err)
+		case describeOutput.StreamDescription.EncryptionType == types.EncryptionTypeKms:
+			currentKeyID := aws.ToString(describeOutput.StreamDescription.KeyId)
+			if currentKeyID == p.kinesisConf.KMSKeyID {
+				p.log.Info("Server-side encryption is already enabled with the specified KMS Key ID")
+			} else {
+				p.log.Fatal("Server-side encryption is already enabled with a different KMS Key ID")
+			}
+		default:
+			// Encryption not enabled, proceed to enable it
+			_, err := p.client.StartStreamEncryption(ctx, &kinesis.StartStreamEncryptionInput{
+				StreamName:     aws.String(p.kinesisConf.StreamName),
+				EncryptionType: types.EncryptionTypeKms,
+				KeyId:          aws.String(p.kinesisConf.KMSKeyID),
+			})
+
+			if err != nil {
+				var resourceInUseErr *types.ResourceInUseException
+				if errors.As(err, &resourceInUseErr) {
+					p.log.Info("Server-side encryption is already enabled for the Kinesis stream.")
+				} else {
+					p.log.Fatalf("Failed to enable server-side encryption for Kinesis stream: %v", err)
+				}
+			} else {
+				kmsKeyID := p.kinesisConf.KMSKeyID
+				loggableKeyID := "***" // Default to a fully masked key
+				if len(kmsKeyID) >= 8 {
+					loggableKeyID = fmt.Sprintf("%s***%s", kmsKeyID[:4], kmsKeyID[len(kmsKeyID)-4:])
+				}
+				p.log.Info("Server-side encryption enabled for Kinesis stream with KMS Key ID: ", loggableKeyID)
+			}
+		}
+	}
+
 	p.log.Info(p.GetName() + " Initialized")
 
 	return nil

pumps/kinesis_test.go

@@ -0,0 +1,196 @@
+package pumps
+
+import (
+	"testing"
+
+	"github.com/mitchellh/mapstructure"
+	"github.com/stretchr/testify/assert"
+)
+
+func TestKinesisPump_New(t *testing.T) {
+	pump := &KinesisPump{}
+	newPump := pump.New()
+	assert.IsType(t, &KinesisPump{}, newPump)
+}
+
+func TestKinesisPump_GetName(t *testing.T) {
+	pump := &KinesisPump{}
+	assert.Equal(t, "Kinesis Pump", pump.GetName())
+}
+
+func TestKinesisConf_KMSKeyID_Configuration(t *testing.T) {
+	tests := []struct {
+		name     string
+		kmsKeyID string
+		expected string
+	}{
+		{
+			name:     "KMSKeyID provided",
+			kmsKeyID: "arn:aws:kms:us-east-1:123456789012:key/12345678-1234-1234-1234-123456789012",
+			expected: "arn:aws:kms:us-east-1:123456789012:key/12345678-1234-1234-1234-123456789012",
+		},
+		{
+			name:     "KMSKeyID empty",
+			kmsKeyID: "",
+			expected: "",
+		},
+		{
+			name:     "KMSKeyID alias format",
+			kmsKeyID: "alias/my-key",
+			expected: "alias/my-key",
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			conf := &KinesisConf{
+				KMSKeyID: tt.kmsKeyID,
+			}
+			assert.Equal(t, tt.expected, conf.KMSKeyID)
+		})
+	}
+}
+
+func TestKinesisPump_KMSKeyID_DefaultValue(t *testing.T) {
+	conf := &KinesisConf{}
+	assert.Equal(t, "", conf.KMSKeyID)
+}
+
+func TestSplitIntoBatches(t *testing.T) {
+	records := make([]interface{}, 25)
+	batches := splitIntoBatches(records, 10)
+	assert.Len(t, batches, 3)
+	assert.Len(t, batches[0], 10)
+	assert.Len(t, batches[1], 10)
+	assert.Len(t, batches[2], 5)
+}
+
+func TestKinesisPump_KMSKeyID_LogMasking(t *testing.T) {
+	kmsKeyID := "arn:aws:kms:us-east-1:123456789012:key/12345678-1234-1234-1234-123456789012"
+	if len(kmsKeyID) >= 8 {
+		masked := kmsKeyID[:4] + "***" + kmsKeyID[len(kmsKeyID)-4:]
+		assert.Equal(t, "arn:***9012", masked)
+	}
+}
+
+// Tests for the new describe stream encryption logic
+func TestKinesisPump_EncryptionConfig_SameKey(t *testing.T) {
+	kmsKeyID := "arn:aws:kms:us-east-1:123456789012:key/test-key-id"
+
+	config := map[string]interface{}{
+		"stream_name": "test-stream",
+		"region":      "us-east-1",
+		"kms_key_id":  kmsKeyID,
+	}
+
+	kinesisConf := &KinesisConf{}
+	err := mapstructure.Decode(config, kinesisConf)
+	assert.NoError(t, err)
+	assert.Equal(t, kmsKeyID, kinesisConf.KMSKeyID)
+	assert.Equal(t, "test-stream", kinesisConf.StreamName)
+	assert.Equal(t, "us-east-1", kinesisConf.Region)
+}
+
+func TestKinesisPump_EncryptionConfig_DifferentKey(t *testing.T) {
+	currentKeyID := "arn:aws:kms:us-east-1:123456789012:key/current-key"
+	newKeyID := "arn:aws:kms:us-east-1:123456789012:key/new-key"
+
+	config := map[string]interface{}{
+		"stream_name": "test-stream",
+		"region":      "us-east-1",
+		"kms_key_id":  newKeyID,
+	}
+
+	kinesisConf := &KinesisConf{}
+	err := mapstructure.Decode(config, kinesisConf)
+	assert.NoError(t, err)
+	assert.Equal(t, newKeyID, kinesisConf.KMSKeyID)
+	assert.Equal(t, "test-stream", kinesisConf.StreamName)
+	assert.Equal(t, "us-east-1", kinesisConf.Region)
+
+	// Verify the keys are different (simulating the scenario)
+	assert.NotEqual(t, currentKeyID, newKeyID)
+}
+
+func TestKinesisPump_EncryptionConfig_NotEncrypted(t *testing.T) {
+	kmsKeyID := "arn:aws:kms:us-east-1:123456789012:key/test-key-id"
+
+	config := map[string]interface{}{
+		"stream_name": "test-stream",
+		"region":      "us-east-1",
+		"kms_key_id":  kmsKeyID,
+	}
+
+	kinesisConf := &KinesisConf{}
+	err := mapstructure.Decode(config, kinesisConf)
+	assert.NoError(t, err)
+	assert.Equal(t, kmsKeyID, kinesisConf.KMSKeyID)
+	assert.Equal(t, "test-stream", kinesisConf.StreamName)
+	assert.Equal(t, "us-east-1", kinesisConf.Region)
+}
+
+func TestKinesisPump_EncryptionConfig_NoKMSKeyID(t *testing.T) {
+	config := map[string]interface{}{
+		"stream_name": "test-stream",
+		"region":      "us-east-1",
+		// No kms_key_id provided - should skip encryption
+	}
+
+	kinesisConf := &KinesisConf{}
+	err := mapstructure.Decode(config, kinesisConf)
+	assert.NoError(t, err)
+	assert.Equal(t, "", kinesisConf.KMSKeyID)
+	assert.Equal(t, "test-stream", kinesisConf.StreamName)
+	assert.Equal(t, "us-east-1", kinesisConf.Region)
+}
+
+func TestKinesisPump_BatchSize_Configuration(t *testing.T) {
+	//nolint:govet
+	tests := []struct {
+		name          string
+		batchSize     interface{}
+		expectedValue int
+	}{
+		{
+			name:          "Default batch size (not provided)",
+			batchSize:     nil,
+			expectedValue: 0, // Will be set to 100 in Init()
+		},
+		{
+			name:          "Custom batch size",
+			batchSize:     250,
+			expectedValue: 250,
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			config := map[string]interface{}{
+				"stream_name": "test-stream",
+				"region":      "us-east-1",
+			}
+
+			if tt.batchSize != nil {
+				config["batch_size"] = tt.batchSize
+			}
+
+			kinesisConf := &KinesisConf{}
+			err := mapstructure.Decode(config, kinesisConf)
+			assert.NoError(t, err)
+			assert.Equal(t, tt.expectedValue, kinesisConf.BatchSize)
+		})
+	}
+}
+
+func TestKinesisPump_StreamName_Required(t *testing.T) {
+	config := map[string]interface{}{
+		"region": "us-east-1",
+		// Missing stream_name
+	}
+
+	kinesisConf := &KinesisConf{}
+	err := mapstructure.Decode(config, kinesisConf)
+	assert.NoError(t, err)
+	assert.Equal(t, "", kinesisConf.StreamName) // Should be empty when not provided
+	assert.Equal(t, "us-east-1", kinesisConf.Region)
+}