Aws Sdk Java V2 Kms

AWS SDK for Java 2.x - AWS KMS (Key Management Service)

Overview

Provides AWS KMS patterns using AWS SDK for Java 2.x. Covers key management, encryption/decryption, envelope encryption, digital signatures, and Spring Boot integration.

Instructions

1. Set Up IAM Permissions - Grant kms: actions with least privilege 2. Create KMS Client - Instantiate KmsClient with region and credentials 3. Create Keys - Use createKey() → Verify key state is ENABLED before proceeding 4. Set Key Policies - Define key usage permissions → Test access before production 5. Encrypt Data - Use encrypt() for data <4KB; Verify ciphertext is not empty 6. Envelope Encryption - For larger data, use generateDataKey() → Verify data key generation succeeded 7. Digital Signatures - Create signing keys → Verify signatureValid=true after sign/verify 8. Key Rotation - Enable auto-rotation → Confirm rotation schedule is active*

When to Use

• Creating/managing symmetric encryption keys for data protection
• Implementing envelope encryption for large data
• Generating data keys for local encryption with KMS-managed keys
• Setting up digital signatures with asymmetric keys
• Integrating encryption into Spring Boot applications

Dependencies

Maven

<dependency>
    <groupId>software.amazon.awssdk</groupId>
    <artifactId>kms</artifactId>
</dependency>

Gradle

implementation 'software.amazon.awssdk:kms:2.x.x'

Client Setup

Basic Synchronous Client

import software.amazon.awssdk.regions.Region;
import software.amazon.awssdk.services.kms.KmsClient;

KmsClient kmsClient = KmsClient.builder()
    .region(Region.US_EAST_1)
    .build();

Basic Asynchronous Client

import software.amazon.awssdk.services.kms.KmsAsyncClient;

KmsAsyncClient kmsAsyncClient = KmsAsyncClient.builder()
    .region(Region.US_EAST_1)
    .build();

Advanced Client Configuration

KmsClient kmsClient = KmsClient.builder()
    .region(Region.of(System.getenv("AWS_REGION")))
    .credentialsProvider(DefaultCredentialsProvider.create())
    .overrideConfiguration(c -> c.retryPolicy(RetryPolicy.builder()
        .numRetries(3)
        .build()))
    .build();

Basic Key Management

Create Encryption Key

public String createEncryptionKey(KmsClient kmsClient, String description) {
    CreateKeyRequest request = CreateKeyRequest.builder()
        .description(description)
        .keyUsage(KeyUsageType.ENCRYPT_DECRYPT)
        .build();

    CreateKeyResponse response = kmsClient.createKey(request);
    return response.keyMetadata().keyId();
}

Describe Key

public KeyMetadata getKeyMetadata(KmsClient kmsClient, String keyId) {
    DescribeKeyRequest request = DescribeKeyRequest.builder()
        .keyId(keyId)
        .build();

    return kmsClient.describeKey(request).keyMetadata();
}

Enable/Disable Key

public void toggleKeyState(KmsClient kmsClient, String keyId, boolean enable) {
    if (enable) {
        kmsClient.enableKey(EnableKeyRequest.builder().keyId(keyId).build());
    } else {
        kmsClient.disableKey(DisableKeyRequest.builder().keyId(keyId).build());
    }
}

Basic Encryption and Decryption

Encrypt Data

public String encryptData(KmsClient kmsClient, String keyId, String plaintext) {
    SdkBytes plaintextBytes = SdkBytes.fromString(plaintext, StandardCharsets.UTF_8);

    EncryptRequest request = EncryptRequest.builder()
        .keyId(keyId)
        .plaintext(plaintextBytes)
        .build();

    EncryptResponse response = kmsClient.encrypt(request);
    return Base64.getEncoder().encodeToString(
        response.ciphertextBlob().asByteArray());
}

Decrypt Data

public String decryptData(KmsClient kmsClient, String ciphertextBase64) {
    byte[] ciphertext = Base64.getDecoder().decode(ciphertextBase64);
    SdkBytes ciphertextBytes = SdkBytes.fromByteArray(ciphertext);

    DecryptRequest request = DecryptRequest.builder()
        .ciphertextBlob(ciphertextBytes)
        .build();

    DecryptResponse response = kmsClient.decrypt(request);
    return response.plaintext().asString(StandardCharsets.UTF_8);
}

Envelope Encryption Pattern

Generate and Use Data Key

public DataKeyResult encryptWithEnvelope(KmsClient kmsClient, String masterKeyId, byte[] data) {
    try {
        GenerateDataKeyRequest keyRequest = GenerateDataKeyRequest.builder()
            .keyId(masterKeyId)
            .keySpec(DataKeySpec.AES_256)
            .build();

        GenerateDataKeyResponse keyResponse = kmsClient.generateDataKey(keyRequest);

        // Validate response
        if (keyResponse.plaintext() == null || keyResponse.ciphertextBlob() == null) {
            throw new IllegalStateException("Data key generation returned null");
        }

        byte[] encryptedData = encryptWithAES(data, keyResponse.plaintext().asByteArray());

        // Clear plaintext key from memory
        Arrays.fill(keyResponse.plaintext().asByteArray(), (byte) 0);

        return new DataKeyResult(encryptedData, keyResponse.ciphertextBlob().asByteArray());

    } catch (KmsException e) {
        throw new RuntimeException("Envelope encryption failed: " + e.awsErrorDetails().errorCode(), e);
    }
}

public byte[] decryptWithEnvelope(KmsClient kmsClient, DataKeyResult encryptedEnvelope) {
    try {
        DecryptRequest keyDecryptRequest = DecryptRequest.builder()
            .ciphertextBlob(SdkBytes.fromByteArray(encryptedEnvelope.encryptedKey()))
            .build();

        DecryptResponse keyDecryptResponse = kmsClient.decrypt(keyDecryptRequest);

        // Validate response
        if (keyDecryptResponse.plaintext() == null) {
            throw new IllegalStateException("Key decryption returned null");
        }

        byte[] decryptedData = decryptWithAES(
            encryptedEnvelope.encryptedData(),
            keyDecryptResponse.plaintext().asByteArray());

        // Clear plaintext key from memory
        Arrays.fill(keyDecryptResponse.plaintext().asByteArray(), (byte) 0);

        return decryptedData;

    } catch (KmsException e) {
        throw new RuntimeException("Envelope decryption failed: " + e.awsErrorDetails().errorCode(), e);
    }
}

Digital Signatures

Create Signing Key and Sign Data

public String createAndSignData(KmsClient kmsClient, String description, String message) {
    // Create signing key
    CreateKeyRequest keyRequest = CreateKeyRequest.builder()
        .description(description)
        .keySpec(KeySpec.RSA_2048)
        .keyUsage(KeyUsageType.SIGN_VERIFY)
        .build();

    CreateKeyResponse keyResponse = kmsClient.createKey(keyRequest);
    String keyId = keyResponse.keyMetadata().keyId();

    // Sign data
    SignRequest signRequest = SignRequest.builder()
        .keyId(keyId)
        .message(SdkBytes.fromString(message, StandardCharsets.UTF_8))
        .signingAlgorithm(SigningAlgorithmSpec.RSASSA_PSS_SHA_256)
        .build();

    SignResponse signResponse = kmsClient.sign(signRequest);
    return Base64.getEncoder().encodeToString(
        signResponse.signature().asByteArray());
}

Verify Signature

public boolean verifySignature(KmsClient kmsClient,
                             String keyId,
                             String message,
                             String signatureBase64) {
    byte[] signature = Base64.getDecoder().decode(signatureBase64);

    VerifyRequest verifyRequest = VerifyRequest.builder()
        .keyId(keyId)
        .message(SdkBytes.fromString(message, StandardCharsets.UTF_8))
        .signature(SdkBytes.fromByteArray(signature))
        .signingAlgorithm(SigningAlgorithmSpec.RSASSA_PSS_SHA_256)
        .build();

    VerifyResponse verifyResponse = kmsClient.verify(verifyRequest);
    return verifyResponse.signatureValid();
}

Spring Boot Integration

Configuration Class

@Configuration
public class KmsConfiguration {

    @Bean
    public KmsClient kmsClient() {
        return KmsClient.builder()
            .region(Region.US_EAST_1)
            .build();
    }

    @Bean
    public KmsAsyncClient kmsAsyncClient() {
        return KmsAsyncClient.builder()
            .region(Region.US_EAST_1)
            .build();
    }
}

Encryption Service

@Service
@RequiredArgsConstructor
public class KmsEncryptionService {

    private final KmsClient kmsClient;

    @Value("${kms.encryption-key-id}")
    private String keyId;

    public String encrypt(String plaintext) {
        try {
            EncryptRequest request = EncryptRequest.builder()
                .keyId(keyId)
                .plaintext(SdkBytes.fromString(plaintext, StandardCharsets.UTF_8))
                .build();

            EncryptResponse response = kmsClient.encrypt(request);
            return Base64.getEncoder().encodeToString(
                response.ciphertextBlob().asByteArray());

        } catch (KmsException e) {
            throw new RuntimeException("Encryption failed", e);
        }
    }

    public String decrypt(String ciphertextBase64) {
        try {
            byte[] ciphertext = Base64.getDecoder().decode(ciphertextBase64);

            DecryptRequest request = DecryptRequest.builder()
                .ciphertextBlob(SdkBytes.fromByteArray(ciphertext))
                .build();

            DecryptResponse response = kmsClient.decrypt(request);
            return response.plaintext().asString(StandardCharsets.UTF_8);

        } catch (KmsException e) {
            throw new RuntimeException("Decryption failed", e);
        }
    }
}

Examples

Basic Encryption Example

public class BasicEncryptionExample {
    public static void main(String[] args) {
        KmsClient kmsClient = KmsClient.builder()
            .region(Region.US_EAST_1)
            .build();

        // Create key
        String keyId = createEncryptionKey(kmsClient, "Example encryption key");
        System.out.println("Created key: " + keyId);

        // Encrypt and decrypt
        String plaintext = "Hello, World!";
        String encrypted = encryptData(kmsClient, keyId, plaintext);
        String decrypted = decryptData(kmsClient, encrypted);

        System.out.println("Original: " + plaintext);
        System.out.println("Decrypted: " + decrypted);
    }
}

Envelope Encryption Example

public class EnvelopeEncryptionExample {
    public static void main(String[] args) {
        KmsClient kmsClient = KmsClient.builder()
            .region(Region.US_EAST_1)
            .build();

        String masterKeyId = "alias/your-master-key";
        String largeData = "This is a large amount of data that needs encryption...";
        byte[] data = largeData.getBytes(StandardCharsets.UTF_8);

        // Encrypt using envelope pattern
        DataKeyResult encryptedEnvelope = encryptWithEnvelope(
            kmsClient, masterKeyId, data);

        // Decrypt
        byte[] decryptedData = decryptWithEnvelope(
            kmsClient, encryptedEnvelope);

        String result = new String(decryptedData, StandardCharsets.UTF_8);
        System.out.println("Decrypted: " + result);
    }
}

Best Practices

Security

• Always use envelope encryption for large data - Encrypt data locally and only encrypt the data key with KMS
• Use encryption context - Add contextual information to track and audit usage
• Never log sensitive data - Avoid logging plaintext or encryption keys
• Implement proper key lifecycle - Enable automatic rotation and set deletion policies
• Use separate keys for different purposes - Don't reuse keys across multiple applications

Performance

• Cache encrypted data keys - Reduce KMS API calls by caching data keys
• Use async operations - Leverage async clients for non-blocking I/O
• Reuse client instances - Don't create new clients for each operation
• Implement connection pooling - Configure proper connection pooling settings

Error Handling

• Implement retry logic - Handle throttling exceptions with exponential backoff
• Check key states - Verify key is enabled before performing operations
• Use circuit breakers - Prevent cascading failures during KMS outages
• Log errors comprehensively - Include KMS error codes and context

References

For detailed implementation patterns, advanced techniques, and comprehensive examples:

• @references/technical-guide.md - Complete technical implementation patterns
• @references/spring-boot-integration.md - Spring Boot integration patterns
• @references/testing.md - Testing strategies and examples
• @references/best-practices.md - Security and operational best practices

Related Skills

• @aws-sdk-java-v2-core - Core AWS SDK patterns and configuration
• @aws-sdk-java-v2-dynamodb - DynamoDB integration patterns
• @aws-sdk-java-v2-secrets-manager - Secrets management patterns
• @spring-boot-dependency-injection - Spring dependency injection patterns

External References

• AWS KMS Developer Guide
• AWS SDK for Java 2.x Documentation
• KMS Best Practices

Constraints and Warnings

• Data Size Limit: Direct encryption limited to 4KB; use envelope encryption for larger data
• Key Usage Limits: KMS has quotas on API calls per second
• Key Material: Imported key material cannot be managed by AWS for rotation
• Key Deletion: Key deletion requires 7-30 day waiting period
• Regional Boundaries: KMS keys cannot be used across regions
• Cost Considerations: KMS charges per API call and for key storage
• Asymmetric Keys: Not all regions support asymmetric key types
• Key Policies: Changes to key policies require careful IAM review
• Envelope Encryption: Proper implementation required for data key security
• Logging: Enable CloudTrail to audit all KMS API usage

AWS KMS Best Practices

Security Best Practices

Key Management

1. Use Separate Keys for Different Purposes

• Create unique keys for different applications or data types
• Avoid reusing keys across multiple purposes
• Use aliases instead of raw key IDs for references

   // Good: Create specific keys
   String encryptionKey = kms.createKey("Database encryption key");
   String signingKey = kms.createSigningKey("Document signing key");

   // Bad: Use the same key for everything

2. Enable Automatic Key Rotation

• Enable automatic key rotation for enhanced security
• Review rotation schedules based on compliance requirements

   public void enableKeyRotation(KmsClient kmsClient, String keyId) {
       EnableKeyRotationRequest request = EnableKeyRotationRequest.builder()
           .keyId(keyId)
           .build();
       kmsClient.enableKeyRotation(request);
   }

3. Implement Key Lifecycle Policies

• Set key expiration dates based on data retention policies
• Schedule key deletion when no longer needed
• Use key policies to enforce lifecycle rules

4. Use Key Aliases

• Always use aliases instead of raw key IDs
• Create meaningful aliases following naming conventions
• Regularly review and update aliases

   public void createKeyWithAlias(KmsClient kmsClient, String alias, String description) {
       // Create key
       CreateKeyResponse response = kmsClient.createKey(
           CreateKeyRequest.builder()
               .description(description)
               .build());

       // Create alias
       CreateAliasRequest aliasRequest = CreateAliasRequest.builder()
           .aliasName(alias)
           .targetKeyId(response.keyMetadata().keyId())
           .build();
       kmsClient.createAlias(aliasRequest);
   }

Encryption Security

1. Never Log Plaintext or Encryption Keys

• Avoid logging sensitive data in any form
• Ensure proper logging configuration to prevent accidental exposure

   // Bad: Logging sensitive data
   logger.info("Encrypted data: {}", encryptedData);

   // Good: Log only metadata
   logger.info("Encryption completed for user: {}", userId);

2. Use Encryption Context

• Always include encryption context for additional security
• Use contextual information to verify data integrity

   public Map<String, String> createEncryptionContext(String userId, String dataType) {
       return Map.of(
           "userId", userId,
           "dataType", dataType,
           "timestamp", Instant.now().toString()
       );
   }

3. Implement Least Privilege IAM Policies

• Grant minimal required permissions to KMS keys
• Use IAM policies to restrict access to specific resources

   {
       "Version": "2012-10-17",
       "Statement": [
           {
               "Effect": "Allow",
               "Principal": {"AWS": "arn:aws:iam::123456789012:role/app-role"},
               "Action": [
                   "kms:Encrypt",
                   "kms:Decrypt",
                   "kms:DescribeKey"
               ],
               "Resource": "arn:aws:kms:us-east-1:123456789012:key/your-key-id",
               "Condition": {
                   "StringEquals": {
                       "kms:EncryptionContext:userId": "${aws:userid}"
                   }
               }
           }
       ]
   }

4. Clear Sensitive Data from Memory

• Explicitly clear sensitive data from memory after use
• Use secure memory management practices

   public void secureMemoryExample() {
       byte[] sensitiveKey = new byte[32];
       // ... use the key ...

       // Clear sensitive data
       Arrays.fill(sensitiveKey, (byte) 0);
   }

Performance Best Practices

1. Cache Data Keys for Envelope Encryption

• Cache encrypted data keys to avoid repeated KMS calls
• Use appropriate cache eviction policies
• Monitor cache hit rates

   public class DataKeyCache {
       private final Cache<String, byte[]> keyCache;

       public DataKeyCache() {
           this.keyCache = Caffeine.newBuilder()
               .expireAfterWrite(1, TimeUnit.HOURS)
               .maximumSize(1000)
               .build();
       }

       public byte[] getCachedDataKey(String keyId, KmsClient kmsClient) {
           return keyCache.get(keyId, k -> {
               GenerateDataKeyResponse response = kmsClient.generateDataKey(
                   GenerateDataKeyRequest.builder()
                       .keyId(keyId)
                       .keySpec(DataKeySpec.AES_256)
                       .build());
               return response.ciphertextBlob().asByteArray();
           });
       }
   }

2. Use Async Operations for Non-Blocking I/O

• Leverage async clients for parallel processing
• Use CompletableFuture for chaining operations

   public CompletableFuture<Void> processMultipleAsync(List<String> dataItems) {
       List<CompletableFuture<Void>> futures = dataItems.stream()
           .map(item -> CompletableFuture.runAsync(() ->
               encryptAndStoreItem(item)))
           .collect(Collectors.toList());

       return CompletableFuture.allOf(futures.toArray(new CompletableFuture[0]));
   }

3. Implement Connection Pooling

• Configure connection pooling for better resource utilization
• Set appropriate pool sizes based on load

   public KmsClient createPooledClient() {
       return KmsClient.builder()
           .region(Region.US_EAST_1)
           .httpClientBuilder(ApacheHttpClient.builder()
               .maxConnections(100)
               .connectionTimeToLive(Duration.ofSeconds(30))
               .build())
           .build();
   }

4. Reuse KMS Client Instances

• Create and reuse client instances rather than creating new ones
• Use dependency injection for client management

   @Service
   @RequiredArgsConstructor
   public class KmsService {
       private final KmsClient kmsClient; // Inject and reuse

       public void performOperation() {
           // Use the same client instance
           kmsClient.someOperation();
       }
   }

Cost Optimization

1. Use Envelope Encryption for Large Data

• Generate data keys for encrypting large datasets
• Only use KMS for encrypting the data key, not the entire dataset

   public class EnvelopeEncryption {
       private final KmsClient kmsClient;

       public byte[] encryptLargeData(byte[] largeData) {
           // Generate data key
           GenerateDataKeyResponse response = kmsClient.generateDataKey(
               GenerateDataKeyRequest.builder()
                   .keyId("master-key-id")
                   .keySpec(DataKeySpec.AES_256)
                   .build());

           byte[] encryptedKey = response.ciphertextBlob().asByteArray();
           byte[] plaintextKey = response.plaintext().asByteArray();

           // Encrypt data with local key
           byte[] encryptedData = localEncrypt(largeData, plaintextKey);

           // Return both encrypted data and encrypted key
           return combine(encryptedKey, encryptedData);
       }
   }

2. Cache Encrypted Data Keys

• Cache encrypted data keys to avoid repeated KMS calls
• Use time-based cache expiration

3. Monitor API Usage

• Track KMS API calls for billing and optimization
• Set up CloudWatch alarms for unexpected usage

   public class KmsUsageMonitor {
       private final MeterRegistry meterRegistry;

       public void recordEncryption() {
           meterRegistry.counter("kms.encryption.count").increment();
           meterRegistry.timer("kms.encryption.time").record(() -> {
               // Perform encryption
           });
       }
   }

4. Use Data Key Caching Libraries

• Implement proper caching strategies
• Consider using dedicated caching solutions for data keys

Error Handling Best Practices

1. Implement Retry Logic for Throttling

• Add retry logic for throttling exceptions
• Use exponential backoff for retries

   public class KmsRetryHandler {
       private static final int MAX_RETRIES = 3;
       private static final long INITIAL_DELAY = 1000; // 1 second

       public <T> T executeWithRetry(Supplier<T> operation) {
           int attempt = 0;
           while (attempt < MAX_RETRIES) {
               try {
                   return operation.get();
               } catch (KmsException e) {
                   if (!isRetryable(e) || attempt == MAX_RETRIES - 1) {
                       throw e;
                   }
                   attempt++;
                   try {
                       Thread.sleep(INITIAL_DELAY * (long) Math.pow(2, attempt));
                   } catch (InterruptedException ie) {
                       Thread.currentThread().interrupt();
                       throw new RuntimeException("Retry interrupted", ie);
                   }
               }
           }
           throw new IllegalStateException("Should not reach here");
       }

       private boolean isRetryable(KmsException e) {
           return "ThrottlingException".equals(e.awsErrorDetails().errorCode());
       }
   }

2. Handle Key State Errors Gracefully

• Check key state before performing operations
• Handle key states like PendingDeletion, Disabled, etc.

   public void performOperationWithKeyStateCheck(KmsClient kmsClient, String keyId) {
       KeyMetadata metadata = describeKey(kmsClient, keyId);

       switch (metadata.keyState()) {
           case ENABLED:
               // Perform operation
               break;
           case DISABLED:
               throw new IllegalStateException("Key is disabled");
           case PENDING_DELETION:
               throw new IllegalStateException("Key is scheduled for deletion");
           default:
               throw new IllegalStateException("Unknown key state: " + metadata.keyState());
       }
   }

3. Log KMS-Specific Error Codes

• Implement comprehensive error logging
• Map KMS error codes to meaningful application errors

   public class KmsErrorHandler {
       public String mapKmsErrorToAppError(KmsException e) {
           String errorCode = e.awsErrorDetails().errorCode();
           switch (errorCode) {
               case "NotFoundException":
                   return "Key not found";
               case "DisabledException":
                   return "Key is disabled";
               case "AccessDeniedException":
                   return "Access denied";
               case "InvalidKeyUsageException":
                   return "Invalid key usage";
               default:
                   return "KMS error: " + errorCode;
           }
       }
   }

4. Implement Circuit Breakers

• Use circuit breakers to handle KMS unavailability
• Prevent cascading failures during outages

   public class KmsCircuitBreaker {
       private final CircuitBreaker circuitBreaker;

       public KmsCircuitBreaker() {
           this.circuitBreaker = CircuitBreaker.builder()
               .name("kmsService")
               .failureRateThreshold(50)
               .waitDurationInOpenState(Duration.ofSeconds(30))
               .ringBufferSizeInHalfOpenState(2)
               .ringBufferSizeInClosedState(2)
               .build();
       }

       public <T> T executeWithCircuitBreaker(Callable<T> operation) {
           return circuitBreaker.executeCallable(() -> {
               try {
                   return operation.call();
               } catch (KmsException e) {
                   if (isFailure(e)) {
                       throw new CircuitBreakerOpenException("KMS service unavailable");
                   }
                   throw e;
               }
           });
       }

       private boolean isFailure(KmsException e) {
           return "KMSDisabledException".equals(e.awsErrorDetails().errorCode());
       }
   }

Testing Best Practices

1. Test with Mock KMS Client

• Use mock clients for unit tests
• Verify all expected interactions

   @Test
   void shouldEncryptWithProperEncryptionContext() {
       // Arrange
       when(kmsClient.encrypt(any(EncryptRequest.class))).thenReturn(...);

       // Act
       String result = encryptionService.encrypt("test", "user123");

       // Assert
       verify(kmsClient).encrypt(argThat(request ->
           request.encryptionContext().containsKey("userId") &&
           request.encryptionContext().get("userId").equals("user123")));
   }

2. Test Error Scenarios

• Test various error conditions
• Verify proper error handling and recovery

3. Performance Testing

• Test performance under load
• Measure latency and throughput

4. Integration Testing with Local KMS

• Test with local KMS when possible
• Verify integration with real AWS services

Monitoring and Observability

1. Implement Comprehensive Logging

• Log all KMS operations with appropriate levels
• Include correlation IDs for tracing

   public class KmsLoggingAspect {
       private static final Logger logger = LoggerFactory.getLogger(KmsService.class);

       @Around("execution(* com.yourcompany.kms..*.*(..))")
       public Object logKmsOperation(ProceedingJoinPoint joinPoint) throws Throwable {
           String operation = joinPoint.getSignature().getName();
           logger.info("Starting KMS operation: {}", operation);

           long startTime = System.currentTimeMillis();
           try {
               Object result = joinPoint.proceed();
               long duration = System.currentTimeMillis() - startTime;
               logger.info("Completed KMS operation: {} in {}ms", operation, duration);
               return result;
           } catch (Exception e) {
               long duration = System.currentTimeMillis() - startTime;
               logger.error("KMS operation {} failed in {}ms: {}", operation, duration, e.getMessage());
               throw e;
           }
       }
   }

2. Set Up CloudWatch Alarms

• Monitor API call rates
• Set up alarms for error rates
• Track key usage patterns

3. Use Distributed Tracing

• Implement tracing for KMS operations
• Correlate KMS calls with application operations

4. Monitor Key Usage Metrics

• Track key usage patterns
• Monitor for unusual usage patterns

Compliance and Auditing

1. Enable KMS Key Usage Logging

• Configure CloudTrail to log KMS operations
• Enable detailed logging for compliance

2. Regular Security Audits

• Conduct regular audits of KMS key usage
• Review access policies periodically

3. Comprehensive Backup Strategy

• Implement key backup and recovery procedures
• Test backup restoration processes

4. Comprehensive Access Reviews

• Regularly review IAM policies for KMS access
• Remove unnecessary permissions

Advanced Security Considerations

1. Multi-Region KMS Keys

• Consider multi-region keys for disaster recovery
• Test failover scenarios

2. Cross-Account Access

• Implement proper cross-account access controls
• Use resource-based policies for account sharing

3. Custom Key Stores

• Consider custom key stores for enhanced security
• Implement proper key management in custom stores

4. Key Material External

• Use imported key material for enhanced control
• Implement proper key rotation for imported keys

Development Best Practices

1. Use Dependency Injection

• Inject KMS clients rather than creating them directly
• Use proper configuration management

   @Configuration
   @ConfigurationProperties(prefix = "aws.kms")
   public class KmsProperties {
       private String region = "us-east-1";
       private String encryptionKeyId;
       private int maxRetries = 3;

       // Getters and setters
   }

2. Proper Configuration Management

• Use environment-specific configurations
• Secure sensitive configuration values

3. Version Control and Documentation

• Keep KMS-related code well documented
• Track key usage patterns in version control

4. Code Reviews

• Conduct thorough code reviews for KMS-related code
• Focus on security and error handling

Implementation Checklists

Key Setup Checklist

• [ ] Create appropriate KMS keys for different purposes
• [ ] Enable automatic key rotation
• [ ] Set up key aliases
• [ ] Configure IAM policies with least privilege
• [ ] Set up CloudTrail logging

Implementation Checklist

• [ ] Use envelope encryption for large data
• [ ] Implement proper error handling
• [ ] Add comprehensive logging
• [ ] Set up monitoring and alarms
• [ ] Write comprehensive tests

Security Checklist

• [ ] Never log sensitive data
• [ ] Use encryption context
• [ ] Implement proper access controls
• [ ] Clear sensitive data from memory
• [ ] Regularly audit access patterns

By following these best practices, you can ensure that your AWS KMS implementation is secure, performant, cost-effective, and maintainable.

Spring Boot Integration with AWS KMS

Configuration

Basic Configuration

@Configuration
public class KmsConfiguration {

    @Bean
    public KmsClient kmsClient() {
        return KmsClient.builder()
            .region(Region.US_EAST_1)
            .build();
    }

    @Bean
    public KmsAsyncClient kmsAsyncClient() {
        return KmsAsyncClient.builder()
            .region(Region.US_EAST_1)
            .build();
    }
}

Configuration with Custom Settings

@Configuration
@ConfigurationProperties(prefix = "aws.kms")
public class KmsAdvancedConfiguration {

    private Region region = Region.US_EAST_1;
    private String endpoint;
    private Duration timeout = Duration.ofSeconds(10);
    private String accessKeyId;
    private String secretAccessKey;

    @Bean
    public KmsClient kmsClient() {
        KmsClientBuilder builder = KmsClient.builder()
            .region(region)
            .overrideConfiguration(c -> c.retryPolicy(RetryPolicy.builder()
                .numRetries(3)
                .build()));

        if (endpoint != null) {
            builder.endpointOverride(URI.create(endpoint));
        }

        // Add credentials if provided
        if (accessKeyId != null && secretAccessKey != null) {
            builder.credentialsProvider(StaticCredentialsProvider.create(
                AwsBasicCredentials.create(accessKeyId, secretAccessKey)));
        }

        return builder.build();
    }

    // Getters and Setters
    public Region getRegion() { return region; }
    public void setRegion(Region region) { this.region = region; }
    public String getEndpoint() { return endpoint; }
    public void setEndpoint(String endpoint) { this.endpoint = endpoint; }
    public Duration getTimeout() { return timeout; }
    public void setTimeout(Duration timeout) { this.timeout = timeout; }
    public String getAccessKeyId() { return accessKeyId; }
    public void setAccessKeyId(String accessKeyId) { this.accessKeyId = accessKeyId; }
    public String getSecretAccessKey() { return secretAccessKey; }
    public void setSecretAccessKey(String secretAccessKey) { this.secretAccessKey = secretAccessKey; }
}

Application Properties

# AWS KMS Configuration
aws.kms.region=us-east-1
aws.kms.endpoint=
aws.kms.timeout=10s
aws.kms.access-key-id=
aws.kms.secret-access-key=

# KMS Key Configuration
kms.encryption-key-id=alias/your-encryption-key
kms.signing-key-id=alias/your-signing-key

Encryption Service

Basic Encryption Service

@Service
public class KmsEncryptionService {

    private final KmsClient kmsClient;

    @Value("${kms.encryption-key-id}")
    private String keyId;

    public KmsEncryptionService(KmsClient kmsClient) {
        this.kmsClient = kmsClient;
    }

    public String encrypt(String plaintext) {
        try {
            EncryptRequest request = EncryptRequest.builder()
                .keyId(keyId)
                .plaintext(SdkBytes.fromString(plaintext, StandardCharsets.UTF_8))
                .build();

            EncryptResponse response = kmsClient.encrypt(request);

            // Return Base64-encoded ciphertext
            return Base64.getEncoder()
                .encodeToString(response.ciphertextBlob().asByteArray());

        } catch (KmsException e) {
            throw new RuntimeException("Encryption failed", e);
        }
    }

    public String decrypt(String ciphertextBase64) {
        try {
            byte[] ciphertext = Base64.getDecoder().decode(ciphertextBase64);

            DecryptRequest request = DecryptRequest.builder()
                .ciphertextBlob(SdkBytes.fromByteArray(ciphertext))
                .build();

            DecryptResponse response = kmsClient.decrypt(request);

            return response.plaintext().asString(StandardCharsets.UTF_8);

        } catch (KmsException e) {
            throw new RuntimeException("Decryption failed", e);
        }
    }
}

Secure Data Repository

@Repository
public class SecureDataRepository {

    private final KmsEncryptionService encryptionService;
    private final JdbcTemplate jdbcTemplate;

    public SecureDataRepository(KmsEncryptionService encryptionService,
                                JdbcTemplate jdbcTemplate) {
        this.encryptionService = encryptionService;
        this.jdbcTemplate = jdbcTemplate;
    }

    public void saveSecureData(String id, String sensitiveData) {
        String encryptedData = encryptionService.encrypt(sensitiveData);

        jdbcTemplate.update(
            "INSERT INTO secure_data (id, encrypted_value) VALUES (?, ?)",
            id, encryptedData);
    }

    public String getSecureData(String id) {
        String encryptedData = jdbcTemplate.queryForObject(
            "SELECT encrypted_value FROM secure_data WHERE id = ?",
            String.class, id);

        return encryptionService.decrypt(encryptedData);
    }
}

Advanced Envelope Encryption Service

@Service
public class EnvelopeEncryptionService {

    private final KmsClient kmsClient;

    @Value("${kms.master-key-id}")
    private String masterKeyId;

    private final Cache<String, DataKeyPair> keyCache =
        Caffeine.newBuilder()
            .expireAfterWrite(1, TimeUnit.HOURS)
            .maximumSize(100)
            .build();

    public EnvelopeEncryptionService(KmsClient kmsClient) {
        this.kmsClient = kmsClient;
    }

    public EncryptedEnvelope encryptLargeData(byte[] data) {
        // Check cache for existing key
        DataKeyPair dataKeyPair = keyCache.getIfPresent(masterKeyId);

        if (dataKeyPair == null) {
            // Generate new data key
            GenerateDataKeyResponse dataKeyResponse = kmsClient.generateDataKey(
                GenerateDataKeyRequest.builder()
                    .keyId(masterKeyId)
                    .keySpec(DataKeySpec.AES_256)
                    .build());

            dataKeyPair = new DataKeyPair(
                dataKeyResponse.plaintext().asByteArray(),
                dataKeyResponse.ciphertextBlob().asByteArray());

            // Cache the encrypted key (not plaintext)
            keyCache.put(masterKeyId, dataKeyPair);
        }

        try {
            // Encrypt data with plaintext data key
            byte[] encryptedData = encryptWithAES(data, dataKeyPair.plaintext());

            // Clear plaintext key from memory immediately after use
            Arrays.fill(dataKeyPair.plaintext(), (byte) 0);

            return new EncryptedEnvelope(encryptedData, dataKeyPair.encrypted());

        } catch (Exception e) {
            throw new RuntimeException("Envelope encryption failed", e);
        }
    }

    public byte[] decryptLargeData(EncryptedEnvelope envelope) {
        // Get data key from cache or decrypt from KMS
        DataKeyPair dataKeyPair = keyCache.getIfPresent(masterKeyId);

        if (dataKeyPair == null || !Arrays.equals(dataKeyPair.encrypted(), envelope.encryptedKey())) {
            // Decrypt data key from KMS
            DecryptResponse decryptResponse = kmsClient.decrypt(
                DecryptRequest.builder()
                    .ciphertextBlob(SdkBytes.fromByteArray(envelope.encryptedKey()))
                    .build());

            dataKeyPair = new DataKeyPair(
                decryptResponse.plaintext().asByteArray(),
                envelope.encryptedKey());

            // Cache for future use
            keyCache.put(masterKeyId, dataKeyPair);
        }

        try {
            // Decrypt data with plaintext data key
            byte[] decryptedData = decryptWithAES(envelope.encryptedData(), dataKeyPair.plaintext());

            // Clear plaintext key from memory
            Arrays.fill(dataKeyPair.plaintext(), (byte) 0);

            return decryptedData;

        } catch (Exception e) {
            throw new RuntimeException("Envelope decryption failed", e);
        }
    }

    private byte[] encryptWithAES(byte[] data, byte[] key) throws Exception {
        SecretKeySpec keySpec = new SecretKeySpec(key, "AES");
        Cipher cipher = Cipher.getInstance("AES/GCM/NoPadding");
        GCMParameterSpec spec = new GCMParameterSpec(128, key, key.length - 16);
        cipher.init(Cipher.ENCRYPT_MODE, keySpec, spec);
        return cipher.doFinal(data);
    }

    private byte[] decryptWithAES(byte[] data, byte[] key) throws Exception {
        SecretKeySpec keySpec = new SecretKeySpec(key, "AES");
        Cipher cipher = Cipher.getInstance("AES/GCM/NoPadding");
        GCMParameterSpec spec = new GCMParameterSpec(128, key, key.length - 16);
        cipher.init(Cipher.DECRYPT_MODE, keySpec, spec);
        return cipher.doFinal(data);
    }

    public record DataKeyPair(byte[] plaintext, byte[] encrypted) {}
    public record EncryptedEnvelope(byte[] encryptedData, byte[] encryptedKey) {}
}

Data Encryption Interceptor

SQL Encryption Interceptor

public class KmsDataEncryptInterceptor implements StatementInterceptor {

    private final KmsEncryptionService encryptionService;

    public KmsDataEncryptInterceptor(KmsEncryptionService encryptionService) {
        this.encryptionService = encryptionService;
    }

    @Override
    public ResultSet intercept(ResultSet rs, Statement statement, Connection connection) throws SQLException {
        return new EncryptingResultSetWrapper(rs, encryptionService);
    }

    @Override
    public void interceptAfterExecution(Statement statement) {
        // No-op
    }
}

class EncryptingResultSetWrapper implements ResultSet {

    private final ResultSet delegate;
    private final KmsEncryptionService encryptionService;

    public EncryptingResultSetWrapper(ResultSet delegate, KmsEncryptionService encryptionService) {
        this.delegate = delegate;
        this.encryptionService = encryptionService;
    }

    @Override
    public String getString(String columnLabel) throws SQLException {
        String value = delegate.getString(columnLabel);
        if (value == null) return null;

        // Check if this is an encrypted column
        if (isEncryptedColumn(columnLabel)) {
            return encryptionService.decrypt(value);
        }

        return value;
    }

    private boolean isEncryptedColumn(String columnLabel) {
        // Implement logic to identify encrypted columns
        return columnLabel.contains("encrypted") || columnLabel.contains("secure");
    }

    // Delegate other methods to original ResultSet
    @Override
    public boolean next() throws SQLException {
        return delegate.next();
    }

    // ... other ResultSet method implementations
}

Configuration Profiles

Development Profile

# src/main/resources/application-dev.properties
aws.kms.region=us-east-1
kms.encryption-key-id=alias/dev-encryption-key
logging.level.com.yourcompany=DEBUG

Production Profile

# src/main/resources/application-prod.properties
aws.kms.region=${AWS_REGION:us-east-1}
kms.encryption-key-id=${KMS_ENCRYPTION_KEY_ID:alias/production-encryption-key}
logging.level.com.yourcompany=WARN
spring.cloud.aws.credentials.access-key=${AWS_ACCESS_KEY_ID}
spring.cloud.aws.credentials.secret-key=${AWS_SECRET_ACCESS_KEY}

Test Configuration

@Configuration
@Profile("test")
public class KmsTestConfiguration {

    @Bean
    @Primary
    public KmsClient testKmsClient() {
        // Return a mock or test-specific KMS client
        return mock(KmsClient.class);
    }

    @Bean
    public KmsEncryptionService testKmsEncryptionService() {
        return new KmsEncryptionService(testKmsClient());
    }
}

Health Checks and Monitoring

KMS Health Indicator

@Component
public class KmsHealthIndicator implements HealthIndicator {

    private final KmsClient kmsClient;
    private final String keyId;

    public KmsHealthIndicator(KmsClient kmsClient,
                             @Value("${kms.encryption-key-id}") String keyId) {
        this.kmsClient = kmsClient;
        this.keyId = keyId;
    }

    @Override
    public Health health() {
        try {
            // Test KMS connectivity by describing the key
            DescribeKeyRequest request = DescribeKeyRequest.builder()
                .keyId(keyId)
                .build();

            DescribeKeyResponse response = kmsClient.describeKey(request);

            // Check if key is in a healthy state
            KeyState keyState = response.keyMetadata().keyState();
            boolean isHealthy = keyState == KeyState.ENABLED;

            if (isHealthy) {
                return Health.up()
                    .withDetail("keyId", keyId)
                    .withDetail("keyState", keyState)
                    .withDetail("keyArn", response.keyMetadata().arn())
                    .build();
            } else {
                return Health.down()
                    .withDetail("keyId", keyId)
                    .withDetail("keyState", keyState)
                    .withDetail("message", "KMS key is not in ENABLED state")
                    .build();
            }

        } catch (KmsException e) {
            return Health.down()
                .withDetail("keyId", keyId)
                .withDetail("error", e.awsErrorDetails().errorMessage())
                .withDetail("errorCode", e.awsErrorDetails().errorCode())
                .build();
        }
    }
}

Metrics Collection

@Service
public class KmsMetricsCollector {

    private final MeterRegistry meterRegistry;
    private final KmsClient kmsClient;

    private final Counter encryptionCounter;
    private final Counter decryptionCounter;
    private final Timer encryptionTimer;
    private final Timer decryptionTimer;

    public KmsMetricsCollector(MeterRegistry meterRegistry, KmsClient kmsClient) {
        this.meterRegistry = meterRegistry;
        this.kmsClient = kmsClient;

        this.encryptionCounter = Counter.builder("kms.encryption.count")
            .description("Number of encryption operations")
            .register(meterRegistry);

        this.decryptionCounter = Counter.builder("kms.decryption.count")
            .description("Number of decryption operations")
            .register(meterRegistry);

        this.encryptionTimer = Timer.builder("kms.encryption.time")
            .description("Time taken for encryption operations")
            .register(meterRegistry);

        this.decryptionTimer = Timer.builder("kms.decryption.time")
            .description("Time taken for decryption operations")
            .register(meterRegistry);
    }

    public String encryptWithMetrics(String plaintext) {
        encryptionCounter.increment();

        return encryptionTimer.record(() -> {
            try {
                EncryptRequest request = EncryptRequest.builder()
                    .keyId("your-key-id")
                    .plaintext(SdkBytes.fromString(plaintext, StandardCharsets.UTF_8))
                    .build();

                EncryptResponse response = kmsClient.encrypt(request);
                return Base64.getEncoder().encodeToString(
                    response.ciphertextBlob().asByteArray());

            } catch (KmsException e) {
                meterRegistry.counter("kms.encryption.errors")
                    .increment();
                throw e;
            }
        });
    }
}

AWS KMS Technical Guide

Key Management Operations

Create KMS Key

import software.amazon.awssdk.services.kms.model.*;
import java.util.stream.Collectors;

public String createKey(KmsClient kmsClient, String description) {
    try {
        CreateKeyRequest request = CreateKeyRequest.builder()
            .description(description)
            .keyUsage(KeyUsageType.ENCRYPT_DECRYPT)
            .origin(OriginType.AWS_KMS)
            .build();

        CreateKeyResponse response = kmsClient.createKey(request);

        String keyId = response.keyMetadata().keyId();
        System.out.println("Created key: " + keyId);

        return keyId;

    } catch (KmsException e) {
        System.err.println("Error creating key: " + e.awsErrorDetails().errorMessage());
        throw e;
    }
}

Create Key with Custom Key Store

public String createKeyWithCustomStore(KmsClient kmsClient,
                                       String description,
                                       String customKeyStoreId) {
    CreateKeyRequest request = CreateKeyRequest.builder()
        .description(description)
        .keyUsage(KeyUsageType.ENCRYPT_DECRYPT)
        .origin(OriginType.AWS_CLOUDHSM)
        .customKeyStoreId(customKeyStoreId)
        .build();

    CreateKeyResponse response = kmsClient.createKey(request);

    return response.keyMetadata().keyId();
}

List Keys

import java.util.List;

public List<KeyListEntry> listKeys(KmsClient kmsClient) {
    try {
        ListKeysRequest request = ListKeysRequest.builder()
            .limit(100)
            .build();

        ListKeysResponse response = kmsClient.listKeys(request);

        response.keys().forEach(key -> {
            System.out.println("Key ARN: " + key.keyArn());
            System.out.println("Key ID: " + key.keyId());
            System.out.println();
        });

        return response.keys();

    } catch (KmsException e) {
        System.err.println("Error listing keys: " + e.awsErrorDetails().errorMessage());
        throw e;
    }
}

List Keys with Pagination (Async)

import software.amazon.awssdk.services.kms.paginators.ListKeysPublisher;
import java.util.concurrent.CompletableFuture;

public CompletableFuture<Void> listAllKeysAsync(KmsAsyncClient kmsAsyncClient) {
    ListKeysRequest request = ListKeysRequest.builder()
        .limit(15)
        .build();

    ListKeysPublisher keysPublisher = kmsAsyncClient.listKeysPaginator(request);

    return keysPublisher
        .subscribe(r -> r.keys().forEach(key ->
            System.out.println("Key ARN: " + key.keyArn())))
        .whenComplete((result, exception) -> {
            if (exception != null) {
                System.err.println("Error: " + exception.getMessage());
            } else {
                System.out.println("Successfully listed all keys");
            }
        });
}

Describe Key

public KeyMetadata describeKey(KmsClient kmsClient, String keyId) {
    try {
        DescribeKeyRequest request = DescribeKeyRequest.builder()
            .keyId(keyId)
            .build();

        DescribeKeyResponse response = kmsClient.describeKey(request);
        KeyMetadata metadata = response.keyMetadata();

        System.out.println("Key ID: " + metadata.keyId());
        System.out.println("Key ARN: " + metadata.arn());
        System.out.println("Key State: " + metadata.keyState());
        System.out.println("Creation Date: " + metadata.creationDate());
        System.out.println("Enabled: " + metadata.enabled());

        return metadata;

    } catch (KmsException e) {
        System.err.println("Error describing key: " + e.awsErrorDetails().errorMessage());
        throw e;
    }
}

Enable/Disable Key

public void enableKey(KmsClient kmsClient, String keyId) {
    try {
        EnableKeyRequest request = EnableKeyRequest.builder()
            .keyId(keyId)
            .build();

        kmsClient.enableKey(request);
        System.out.println("Key enabled: " + keyId);

    } catch (KmsException e) {
        System.err.println("Error enabling key: " + e.awsErrorDetails().errorMessage());
        throw e;
    }
}

public void disableKey(KmsClient kmsClient, String keyId) {
    try {
        DisableKeyRequest request = DisableKeyRequest.builder()
            .keyId(keyId)
            .build();

        kmsClient.disableKey(request);
        System.out.println("Key disabled: " + keyId);

    } catch (KmsException e) {
        System.err.println("Error disabling key: " + e.awsErrorDetails().errorMessage());
        throw e;
    }
}

Encryption and Decryption

Encrypt Data

import software.amazon.awssdk.core.SdkBytes;
import java.nio.charset.StandardCharsets;

public byte[] encryptData(KmsClient kmsClient, String keyId, String plaintext) {
    try {
        SdkBytes plaintextBytes = SdkBytes.fromString(plaintext, StandardCharsets.UTF_8);

        EncryptRequest request = EncryptRequest.builder()
            .keyId(keyId)
            .plaintext(plaintextBytes)
            .build();

        EncryptResponse response = kmsClient.encrypt(request);

        byte[] encryptedData = response.ciphertextBlob().asByteArray();
        System.out.println("Data encrypted successfully");

        return encryptedData;

    } catch (KmsException e) {
        System.err.println("Error encrypting data: " + e.awsErrorDetails().errorMessage());
        throw e;
    }
}

Decrypt Data

public String decryptData(KmsClient kmsClient, byte[] ciphertext) {
    try {
        SdkBytes ciphertextBytes = SdkBytes.fromByteArray(ciphertext);

        DecryptRequest request = DecryptRequest.builder()
            .ciphertextBlob(ciphertextBytes)
            .build();

        DecryptResponse response = kmsClient.decrypt(request);

        String decryptedText = response.plaintext().asString(StandardCharsets.UTF_8);
        System.out.println("Data decrypted successfully");

        return decryptedText;

    } catch (KmsException e) {
        System.err.println("Error decrypting data: " + e.awsErrorDetails().errorMessage());
        throw e;
    }
}

Encrypt with Encryption Context

import java.util.Map;

public byte[] encryptWithContext(KmsClient kmsClient,
                                 String keyId,
                                 String plaintext,
                                 Map<String, String> encryptionContext) {
    try {
        EncryptRequest request = EncryptRequest.builder()
            .keyId(keyId)
            .plaintext(SdkBytes.fromString(plaintext, StandardCharsets.UTF_8))
            .encryptionContext(encryptionContext)
            .build();

        EncryptResponse response = kmsClient.encrypt(request);

        return response.ciphertextBlob().asByteArray();

    } catch (KmsException e) {
        System.err.println("Error encrypting with context: " + e.awsErrorDetails().errorMessage());
        throw e;
    }
}

Data Key Generation (Envelope Encryption)

Generate Data Key

public record DataKeyPair(byte[] plaintext, byte[] encrypted) {}

public DataKeyPair generateDataKey(KmsClient kmsClient, String keyId) {
    try {
        GenerateDataKeyRequest request = GenerateDataKeyRequest.builder()
            .keyId(keyId)
            .keySpec(DataKeySpec.AES_256)
            .build();

        GenerateDataKeyResponse response = kmsClient.generateDataKey(request);

        byte[] plaintextKey = response.plaintext().asByteArray();
        byte[] encryptedKey = response.ciphertextBlob().asByteArray();

        System.out.println("Data key generated");

        return new DataKeyPair(plaintextKey, encryptedKey);

    } catch (KmsException e) {
        System.err.println("Error generating data key: " + e.awsErrorDetails().errorMessage());
        throw e;
    }
}

Generate Data Key Without Plaintext

public byte[] generateDataKeyWithoutPlaintext(KmsClient kmsClient, String keyId) {
    try {
        GenerateDataKeyWithoutPlaintextRequest request =
            GenerateDataKeyWithoutPlaintextRequest.builder()
                .keyId(keyId)
                .keySpec(DataKeySpec.AES_256)
                .build();

        GenerateDataKeyWithoutPlaintextResponse response =
            kmsClient.generateDataKeyWithoutPlaintext(request);

        return response.ciphertextBlob().asByteArray();

    } catch (KmsException e) {
        System.err.println("Error generating data key: " + e.awsErrorDetails().errorMessage());
        throw e;
    }
}

Digital Signing

Create Signing Key

public String createSigningKey(KmsClient kmsClient, String description) {
    try {
        CreateKeyRequest request = CreateKeyRequest.builder()
            .description(description)
            .keySpec(KeySpec.RSA_2048)
            .keyUsage(KeyUsageType.SIGN_VERIFY)
            .origin(OriginType.AWS_KMS)
            .build();

        CreateKeyResponse response = kmsClient.createKey(request);

        return response.keyMetadata().keyId();

    } catch (KmsException e) {
        System.err.println("Error creating signing key: " + e.awsErrorDetails().errorMessage());
        throw e;
    }
}

Sign Data

public byte[] signData(KmsClient kmsClient, String keyId, String message) {
    try {
        SdkBytes messageBytes = SdkBytes.fromString(message, StandardCharsets.UTF_8);

        SignRequest request = SignRequest.builder()
            .keyId(keyId)
            .message(messageBytes)
            .signingAlgorithm(SigningAlgorithmSpec.RSASSA_PSS_SHA_256)
            .build();

        SignResponse response = kmsClient.sign(request);

        byte[] signature = response.signature().asByteArray();
        System.out.println("Data signed successfully");

        return signature;

    } catch (KmsException e) {
        System.err.println("Error signing data: " + e.awsErrorDetails().errorMessage());
        throw e;
    }
}

Verify Signature

public boolean verifySignature(KmsClient kmsClient,
                                String keyId,
                                String message,
                                byte[] signature) {
    try {
        VerifyRequest request = VerifyRequest.builder()
            .keyId(keyId)
            .message(SdkBytes.fromString(message, StandardCharsets.UTF_8))
            .signature(SdkBytes.fromByteArray(signature))
            .signingAlgorithm(SigningAlgorithmSpec.RSASSA_PSS_SHA_256)
            .build();

        VerifyResponse response = kmsClient.verify(request);

        boolean isValid = response.signatureValid();
        System.out.println("Signature valid: " + isValid);

        return isValid;

    } catch (KmsException e) {
        System.err.println("Error verifying signature: " + e.awsErrorDetails().errorMessage());
        throw e;
    }
}

Sign and Verify (Async)

public CompletableFuture<Boolean> signAndVerifyAsync(KmsAsyncClient kmsAsyncClient,
                                                      String message) {
    String signMessage = message;

    // Create signing key
    CreateKeyRequest createKeyRequest = CreateKeyRequest.builder()
        .keySpec(KeySpec.RSA_2048)
        .keyUsage(KeyUsageType.SIGN_VERIFY)
        .origin(OriginType.AWS_KMS)
        .build();

    return kmsAsyncClient.createKey(createKeyRequest)
        .thenCompose(createKeyResponse -> {
            String keyId = createKeyResponse.keyMetadata().keyId();

            SdkBytes messageBytes = SdkBytes.fromString(signMessage, StandardCharsets.UTF_8);
            SignRequest signRequest = SignRequest.builder()
                .keyId(keyId)
                .message(messageBytes)
                .signingAlgorithm(SigningAlgorithmSpec.RSASSA_PSS_SHA_256)
                .build();

            return kmsAsyncClient.sign(signRequest)
                .thenCompose(signResponse -> {
                    byte[] signedBytes = signResponse.signature().asByteArray();

                    VerifyRequest verifyRequest = VerifyRequest.builder()
                        .keyId(keyId)
                        .message(messageBytes)
                        .signature(SdkBytes.fromByteArray(signedBytes))
                        .signingAlgorithm(SigningAlgorithmSpec.RSASSA_PSS_SHA_256)
                        .build();

                    return kmsAsyncClient.verify(verifyRequest)
                        .thenApply(VerifyResponse::signatureValid);
                });
        })
        .exceptionally(throwable -> {
            throw new RuntimeException("Failed to sign or verify", throwable);
        });
}

Key Tagging

Tag Key

public void tagKey(KmsClient kmsClient, String keyId, Map<String, String> tags) {
    try {
        List<Tag> tagList = tags.entrySet().stream()
            .map(entry -> Tag.builder()
                .tagKey(entry.getKey())
                .tagValue(entry.getValue())
                .build())
            .collect(Collectors.toList());

        TagResourceRequest request = TagResourceRequest.builder()
            .keyId(keyId)
            .tags(tagList)
            .build();

        kmsClient.tagResource(request);
        System.out.println("Key tagged successfully");

    } catch (KmsException e) {
        System.err.println("Error tagging key: " + e.awsErrorDetails().errorMessage());
        throw e;
    }
}

List Tags

public Map<String, String> listTags(KmsClient kmsClient, String keyId) {
    try {
        ListResourceTagsRequest request = ListResourceTagsRequest.builder()
            .keyId(keyId)
            .build();

        ListResourceTagsResponse response = kmsClient.listResourceTags(request);

        return response.tags().stream()
            .collect(Collectors.toMap(Tag::tagKey, Tag::tagValue));

    } catch (KmsException e) {
        System.err.println("Error listing tags: " + e.awsErrorDetails().errorMessage());
        throw e;
    }
}

Advanced Techniques

Envelope Encryption Service

@Service
public class EnvelopeEncryptionService {

    private final KmsClient kmsClient;

    @Value("${kms.master-key-id}")
    private String masterKeyId;

    public EnvelopeEncryptionService(KmsClient kmsClient) {
        this.kmsClient = kmsClient;
    }

    public EncryptedEnvelope encryptLargeData(byte[] data) {
        // Generate data key
        GenerateDataKeyResponse dataKeyResponse = kmsClient.generateDataKey(
            GenerateDataKeyRequest.builder()
                .keyId(masterKeyId)
                .keySpec(DataKeySpec.AES_256)
                .build());

        byte[] plaintextKey = dataKeyResponse.plaintext().asByteArray();
        byte[] encryptedKey = dataKeyResponse.ciphertextBlob().asByteArray();

        try {
            // Encrypt data with plaintext data key
            byte[] encryptedData = encryptWithAES(data, plaintextKey);

            // Clear plaintext key from memory
            Arrays.fill(plaintextKey, (byte) 0);

            return new EncryptedEnvelope(encryptedData, encryptedKey);

        } catch (Exception e) {
            throw new RuntimeException("Envelope encryption failed", e);
        }
    }

    public byte[] decryptLargeData(EncryptedEnvelope envelope) {
        // Decrypt data key
        DecryptResponse decryptResponse = kmsClient.decrypt(
            DecryptRequest.builder()
                .ciphertextBlob(SdkBytes.fromByteArray(envelope.encryptedKey()))
                .build());

        byte[] plaintextKey = decryptResponse.plaintext().asByteArray();

        try {
            // Decrypt data with plaintext data key
            byte[] decryptedData = decryptWithAES(envelope.encryptedData(), plaintextKey);

            // Clear plaintext key from memory
            Arrays.fill(plaintextKey, (byte) 0);

            return decryptedData;

        } catch (Exception e) {
            throw new RuntimeException("Envelope decryption failed", e);
        }
    }

    private byte[] encryptWithAES(byte[] data, byte[] key) throws Exception {
        SecretKeySpec keySpec = new SecretKeySpec(key, "AES");
        Cipher cipher = Cipher.getInstance("AES/GCM/NoPadding");
        cipher.init(Cipher.ENCRYPT_MODE, keySpec);
        return cipher.doFinal(data);
    }

    private byte[] decryptWithAES(byte[] data, byte[] key) throws Exception {
        SecretKeySpec keySpec = new SecretKeySpec(key, "AES");
        Cipher cipher = Cipher.getInstance("AES/GCM/NoPadding");
        cipher.init(Cipher.DECRYPT_MODE, keySpec);
        return cipher.doFinal(data);
    }

    public record EncryptedEnvelope(byte[] encryptedData, byte[] encryptedKey) {}
}

Error Handling Strategies

public class KmsErrorHandler {

    private static final int MAX_RETRIES = 3;
    private static final long RETRY_DELAY_MS = 1000;

    public <T> T executeWithRetry(Supplier<T> operation, String operationName) {
        int attempt = 0;
        KmsException lastException = null;

        while (attempt < MAX_RETRIES) {
            try {
                return operation.get();
            } catch (KmsException e) {
                lastException = e;
                attempt++;

                // Check if it's a throttling error and retryable
                if (e.awsErrorDetails().errorCode().equals("ThrottlingException") && attempt < MAX_RETRIES) {
                    try {
                        Thread.sleep(RETRY_DELAY_MS);
                    } catch (InterruptedException ie) {
                        Thread.currentThread().interrupt();
                        throw new RuntimeException("Retry interrupted", ie);
                    }
                } else {
                    // Non-retryable error or max retries exceeded
                    throw e;
                }
            }
        }

        throw new RuntimeException(String.format("Failed to execute %s after %d attempts", operationName, MAX_RETRIES), lastException);
    }

    public boolean isRetryableError(KmsException e) {
        String errorCode = e.awsErrorDetails().errorCode();
        return "ThrottlingException".equals(errorCode)
            || "TooManyRequestsException".equals(errorCode)
            || "LimitExceededException".equals(errorCode);
    }
}

Connection Pooling Configuration

import software.amazon.awssdk.http.apache.ApacheHttpClient;
import org.apache.http.impl.client.CloseableHttpClient;
import org.apache.http.impl.client.HttpClients;
import org.apache.http.impl.conn.PoolingHttpClientConnectionManager;

public class KmsConnectionPool {

    public static KmsClient createPooledClient() {
        // Configure connection pool
        PoolingHttpClientConnectionManager connectionManager =
            new PoolingHttpClientConnectionManager();
        connectionManager.setMaxTotal(100);
        connectionManager.setDefaultMaxPerRoute(20);

        CloseableHttpClient httpClient = HttpClients.custom()
            .setConnectionManager(connectionManager)
            .build();

        ApacheHttpClient.Builder httpClientBuilder = ApacheHttpClient.builder()
            .httpClient(httpClient);

        return KmsClient.builder()
            .region(Region.US_EAST_1)
            .httpClientBuilder(httpClientBuilder)
            .build();
    }
}

Testing AWS KMS Integration

Unit Testing with Mocked Client

Basic Unit Test

import software.amazon.awssdk.core.SdkBytes;
import software.amazon.awssdk.services.kms.KmsClient;
import software.amazon.awssdk.services.kms.model.*;
import org.junit.jupiter.api.Test;
import org.junit.jupiter.api.extension.ExtendWith;
import org.mockito.InjectMocks;
import org.mockito.Mock;
import org.mockito.junit.jupiter.MockitoExtension;

import java.nio.charset.StandardCharsets;

import static org.assertj.core.api.Assertions.assertThat;
import static org.mockito.ArgumentMatchers.any;
import static org.mockito.Mockito.when;

@ExtendWith(MockitoExtension.class)
class KmsEncryptionServiceTest {

    @Mock
    private KmsClient kmsClient;

    @InjectMocks
    private KmsEncryptionService encryptionService;

    @Test
    void shouldEncryptData() {
        // Arrange
        String plaintext = "sensitive data";
        byte[] ciphertext = "encrypted".getBytes();

        when(kmsClient.encrypt(any(EncryptRequest.class)))
            .thenReturn(EncryptResponse.builder()
                .ciphertextBlob(SdkBytes.fromByteArray(ciphertext))
                .build());

        // Act
        String result = encryptionService.encrypt(plaintext);

        // Assert
        assertThat(result).isNotEmpty();
        verify(kmsClient).encrypt(any(EncryptRequest.class));
    }

    @Test
    void shouldDecryptData() {
        // Arrange
        String encryptedText = "ciphertext";
        String expectedPlaintext = "sensitive data";

        when(kmsClient.decrypt(any(DecryptRequest.class)))
            .thenReturn(DecryptResponse.builder()
                .plaintext(SdkBytes.fromString(expectedPlaintext, StandardCharsets.UTF_8))
                .build());

        // Act
        String result = encryptionService.decrypt(encryptedText);

        // Assert
        assertThat(result).isEqualTo(expectedPlaintext);
        verify(kmsClient).decrypt(any(DecryptRequest.class));
    }

    @Test
    void shouldThrowExceptionOnEncryptionFailure() {
        // Arrange
        when(kmsClient.encrypt(any(EncryptRequest.class)))
            .thenThrow(KmsException.builder()
                .awsErrorDetails(AwsErrorDetails.builder()
                    .errorCode("KMSDisabledException")
                    .errorMessage("KMS is disabled")
                    .build())
                .build());

        // Act & Assert
        assertThatThrownBy(() -> encryptionService.encrypt("test"))
            .isInstanceOf(RuntimeException.class)
            .hasMessageContaining("Encryption failed");
    }
}

Parameterized Tests

import org.junit.jupiter.params.ParameterizedTest;
import org.junit.jupiter.params.provider.CsvSource;

class KmsEncryptionParameterizedTest {

    @Mock
    private KmsClient kmsClient;

    @InjectMocks
    private KmsEncryptionService encryptionService;

    @ParameterizedTest
    @CsvSource({
        "hello, world",
        "12345, 67890",
        "special@chars, normal",
        "very long string with multiple words, another string",
        "", // empty string
        "null test, null test"
    })
    void shouldEncryptAndDecrypt(String plaintext, String testIdentifier) {
        // Arrange
        byte[] ciphertext = "encrypted".getBytes();

        when(kmsClient.encrypt(any(EncryptRequest.class)))
            .thenReturn(EncryptResponse.builder()
                .ciphertextBlob(SdkBytes.fromByteArray(ciphertext))
                .build());

        when(kmsClient.decrypt(any(DecryptRequest.class)))
            .thenReturn(DecryptResponse.builder()
                .plaintext(SdkBytes.fromString(plaintext, StandardCharsets.UTF_8))
                .build());

        // Act
        String encrypted = encryptionService.encrypt(plaintext);
        String decrypted = encryptionService.decrypt(encrypted);

        // Assert
        assertThat(decrypted).isEqualTo(plaintext);
    }
}

Integration Testing with Testcontainers

Local KMS Mock Setup

import org.testcontainers.containers.localstack.LocalStackContainer;
import org.testcontainers.junit.jupiter.Container;
import org.testcontainers.junit.jupiter.Testcontainers;
import software.amazon.awssdk.auth.credentials.AwsBasicCredentials;
import software.amazon.awssdk.auth.credentials.StaticCredentialsProvider;
import software.amazon.awssdk.services.kms.KmsClient;
import software.amazon.awssdk.regions.Region;
import org.junit.jupiter.api.BeforeAll;
import org.junit.jupiter.api.TestInstance;

import static org.testcontainers.containers.localstack.LocalStackContainer.Service.KMS;

@Testcontainers
@TestInstance(TestInstance.Lifecycle.PER_CLASS)
class KmsIntegrationTest {

    @Container
    private static final LocalStackContainer localStack =
        new LocalStackContainer(DockerImageName.parse("localstack/localstack:latest"))
            .withServices(KMS);

    private KmsClient kmsClient;

    @BeforeAll
    void setup() {
        kmsClient = KmsClient.builder()
            .region(Region.of(localStack.getRegion()))
            .endpointOverride(localStack.getEndpointOverride(KMS))
            .credentialsProvider(StaticCredentialsProvider.create(
                AwsBasicCredentials.create(localStack.getAccessKey(), localStack.getSecretKey())))
            .build();
    }

    @Test
    void shouldCreateAndManageKeysWithLocalKms() {
        // Create a key
        String keyId = createTestKey(kmsClient, "test-key");
        assertThat(keyId).isNotEmpty();

        // Describe the key
        KeyMetadata metadata = describeKey(kmsClient, keyId);
        assertThat(metadata.keyState()).isEqualTo(KeyState.ENABLED);

        // List keys
        List<KeyListEntry> keys = listKeys(kmsClient);
        assertThat(keys).hasSizeGreaterThan(0);
    }
}

Testing with Spring Boot Test Slices

KmsServiceSlice Test

import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.boot.test.autoconfigure.web.servlet.AutoConfigureMockMvc;
import org.springframework.boot.test.context.SpringBootTest;
import org.springframework.boot.test.mock.mockito.MockBean;
import org.springframework.test.context.ActiveProfiles;
import org.springframework.test.web.servlet.MockMvc;

@SpringBootTest
@AutoConfigureMockMvc
@ActiveProfiles("test")
class KmsControllerIntegrationTest {

    @Autowired
    private MockMvc mockMvc;

    @MockBean
    private KmsEncryptionService kmsEncryptionService;

    @Test
    void shouldEncryptData() throws Exception {
        String plaintext = "test data";
        String encrypted = "encrypted-data";

        when(kmsEncryptionService.encrypt(plaintext)).thenReturn(encrypted);

        mockMvc.perform(post("/api/kms/encrypt")
                .contentType(MediaType.APPLICATION_JSON)
                .content("{\"data\":\"" + plaintext + "\"}"))
            .andExpect(status().isOk())
            .andExpect(jsonPath("$.data").value(encrypted));

        verify(kmsEncryptionService).encrypt(plaintext);
    }

    @Test
    void shouldHandleEncryptionErrors() throws Exception {
        when(kmsEncryptionService.encrypt(any()))
            .thenThrow(new RuntimeException("KMS error"));

        mockMvc.perform(post("/api/kms/encrypt")
                .contentType(MediaType.APPLICATION_JSON)
                .content("{\"data\":\"test\"}"))
            .andExpect(status().isInternalServerError());
    }
}

Testing with SpringBootTest and Configuration

import org.springframework.boot.test.context.TestConfiguration;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Primary;

@TestConfiguration
class KmsTestConfiguration {

    @Bean
    @Primary
    public KmsClient testKmsClient() {
        // Create a mock KMS client for testing
        KmsClient mockClient = mock(KmsClient.class);

        // Mock key creation
        when(mockClient.createKey(any(CreateKeyRequest.class)))
            .thenReturn(CreateKeyResponse.builder()
                .keyMetadata(KeyMetadata.builder()
                    .keyId("test-key-id")
                    .keyArn("arn:aws:kms:us-east-1:123456789012:key/test-key-id")
                    .keyState(KeyState.ENABLED)
                    .build())
                .build());

        // Mock encryption
        when(mockClient.encrypt(any(EncryptRequest.class)))
            .thenReturn(EncryptResponse.builder()
                .ciphertextBlob(SdkBytes.fromString("encrypted-data", StandardCharsets.UTF_8))
                .build());

        // Mock decryption
        when(mockClient.decrypt(any(DecryptRequest.class)))
            .thenReturn(DecryptResponse.builder()
                .plaintext(SdkBytes.fromString("decrypted-data", StandardCharsets.UTF_8))
                .build());

        return mockClient;
    }
}

@SpringBootTest(classes = {Application.class, KmsTestConfiguration.class})
class KmsServiceWithTestConfigIntegrationTest {

    @Autowired
    private KmsEncryptionService encryptionService;

    @Test
    void shouldUseTestConfiguration() {
        String result = encryptionService.encrypt("test");
        assertThat(result).isNotEmpty();
    }
}

Testing Envelope Encryption

Envelope Encryption Test

import static org.mockito.ArgumentMatchers.any;
import static org.mockito.Mockito.*;

class EnvelopeEncryptionServiceTest {

    @Mock
    private KmsClient kmsClient;

    @InjectMocks
    private EnvelopeEncryptionService envelopeEncryptionService;

    @Test
    void shouldEncryptAndDecryptLargeData() {
        // Arrange
        byte[] testData = "large test data".getBytes();
        byte[] encryptedDataKey = "encrypted-data-key".getBytes();

        // Mock data key generation
        when(kmsClient.generateDataKey(any(GenerateDataKeyRequest.class)))
            .thenReturn(GenerateDataKeyResponse.builder()
                .plaintext(SdkBytes.fromByteArray("data-key".getBytes()))
                .ciphertextBlob(SdkBytes.fromByteArray(encryptedDataKey))
                .build());

        // Mock data key decryption
        when(kmsClient.decrypt(any(DecryptRequest.class)))
            .thenReturn(DecryptResponse.builder()
                .plaintext(SdkBytes.fromByteArray("data-key".getBytes()))
                .build());

        // Act
        EncryptedEnvelope encryptedEnvelope = envelopeEncryptionService.encryptLargeData(testData);
        byte[] decryptedData = envelopeEncryptionService.decryptLargeData(encryptedEnvelope);

        // Assert
        assertThat(encryptedEnvelope.encryptedData()).isNotEmpty();
        assertThat(encryptedEnvelope.encryptedKey()).isEqualTo(encryptedDataKey);
        assertThat(decryptedData).isEqualTo(testData);

        // Verify interactions
        verify(kmsClient).generateDataKey(any(GenerateDataKeyRequest.class));
        verify(kmsClient).decrypt(any(DecryptRequest.class));
    }

    @Test
    void shouldClearSensitiveDataFromMemory() {
        // Arrange
        byte[] testData = "test data".getBytes();
        byte[] encryptedDataKey = "encrypted-key".getBytes();

        when(kmsClient.generateDataKey(any(GenerateDataKeyRequest.class)))
            .thenReturn(GenerateDataKeyResponse.builder()
                .plaintext(SdkBytes.fromByteArray("sensitive-data-key".getBytes()))
                .ciphertextBlob(SdkBytes.fromByteArray(encryptedDataKey))
                .build());

        when(kmsClient.decrypt(any(DecryptRequest.class)))
            .thenReturn(DecryptResponse.builder()
                .plaintext(SdkBytes.fromByteArray("sensitive-data-key".getBytes()))
                .build());

        // Act
        envelopeEncryptionService.encryptLargeData(testData);
        envelopeEncryptionService.decryptLargeData(new EncryptedEnvelope(testData, encryptedDataKey));

        // Note: Memory clearing is difficult to test directly
        // In real tests, you would verify no sensitive data remains in memory traces
    }
}

Testing Digital Signatures

Digital Signature Tests

class DigitalSignatureServiceTest {

    @Mock
    private KmsClient kmsClient;

    @InjectMocks
    private DigitalSignatureService signatureService;

    @Test
    void shouldSignAndVerifyData() {
        // Arrange
        String message = "test message";
        byte[] signature = "signature-data".getBytes();

        when(kmsClient.sign(any(SignRequest.class)))
            .thenReturn(SignResponse.builder()
                .signature(SdkBytes.fromByteArray(signature))
                .build());

        when(kmsClient.verify(any(VerifyRequest.class)))
            .thenReturn(VerifyResponse.builder()
                .signatureValid(true)
                .build());

        // Act
        byte[] signedSignature = signatureService.signData(message);
        boolean isValid = signatureService.verifySignature(message, signedSignature);

        // Assert
        assertThat(signedSignature).isEqualTo(signature);
        assertThat(isValid).isTrue();
    }

    @Test
    void shouldDetectInvalidSignature() {
        // Arrange
        String message = "test message";
        byte[] signature = "invalid-signature".getBytes();

        when(kmsClient.verify(any(VerifyRequest.class)))
            .thenReturn(VerifyResponse.builder()
                .signatureValid(false)
                .build());

        // Act & Assert
        assertThatThrownBy(() ->
            signatureService.verifySignature(message, signature))
            .isInstanceOf(SecurityException.class)
            .hasMessageContaining("Invalid signature");
    }
}

Performance Testing

Performance Test with JMH

import org.openjdk.jmh.annotations.*;
import org.openjdk.jmh.infra.Blackhole;
import java.util.concurrent.TimeUnit;

@BenchmarkMode(Mode.AverageTime)
@OutputTimeUnit(TimeUnit.MILLISECONDS)
@Warmup(iterations = 3, time = 1)
@Measurement(iterations = 5, time = 1)
@Fork(1)
class KmsPerformanceTest {

    @MockBean
    private KmsClient kmsClient;

    @Autowired
    private KmsEncryptionService encryptionService;

    @Benchmark
    public void testEncryptionPerformance(Blackhole bh) {
        String testData = "performance test data with some content";
        when(kmsClient.encrypt(any(EncryptRequest.class)))
            .thenReturn(EncryptResponse.builder()
                .ciphertextBlob(SdkBytes.fromString("encrypted", StandardCharsets.UTF_8))
                .build());

        String result = encryptionService.encrypt(testData);
        bh.consume(result);
    }

    @Benchmark
    public void testDecryptionPerformance(Blackhole bh) {
        String encryptedData = "encrypted-performance-data";
        when(kmsClient.decrypt(any(DecryptRequest.class)))
            .thenReturn(DecryptResponse.builder()
                .plaintext(SdkBytes.fromString("decrypted", StandardCharsets.UTF_8))
                .build());

        String result = encryptionService.decrypt(encryptedData);
        bh.consume(result);
    }
}

Testing Error Scenarios

Error Handling Tests

class KmsErrorHandlingTest {

    @Mock
    private KmsClient kmsClient;

    @InjectMocks
    private KmsEncryptionService encryptionService;

    @Test
    void shouldHandleThrottlingException() {
        // Arrange
        when(kmsClient.encrypt(any(EncryptRequest.class)))
            .thenThrow(KmsException.builder()
                .awsErrorDetails(AwsErrorDetails.builder()
                    .errorCode("ThrottlingException")
                    .errorMessage("Rate exceeded")
                    .build())
                .build());

        // Act & Assert
        assertThatThrownBy(() -> encryptionService.encrypt("test"))
            .isInstanceOf(RuntimeException.class)
            .hasMessageContaining("Rate limit exceeded");
    }

    @Test
    void shouldHandleDisabledKey() {
        // Arrange
        when(kmsClient.encrypt(any(EncryptRequest.class)))
            .thenThrow(KmsException.builder()
                .awsErrorDetails(AwsErrorDetails.builder()
                    .errorCode("DisabledException")
                    .errorMessage("Key is disabled")
                    .build())
                .build());

        // Act & Assert
        assertThatThrownBy(() -> encryptionService.encrypt("test"))
            .isInstanceOf(RuntimeException.class)
            .hasMessageContaining("Key is disabled");
    }

    @Test
    void shouldHandleNotFoundException() {
        // Arrange
        when(kmsClient.encrypt(any(EncryptRequest.class)))
            .thenThrow(KmsException.builder()
                .awsErrorDetails(AwsErrorDetails.builder()
                    .errorCode("NotFoundException")
                    .errorMessage("Key not found")
                    .build())
                .build());

        // Act & Assert
        assertThatThrownBy(() -> encryptionService.encrypt("test"))
            .isInstanceOf(RuntimeException.class)
            .hasMessageContaining("Key not found");
    }
}

Integration Testing with AWS Local

Testcontainers KMS Setup

import org.testcontainers.containers.localstack.LocalStackContainer;
import software.amazon.awssdk.services.kms.KmsClient;
import static org.testcontainers.containers.localstack.LocalStackContainer.Service.KMS;

@SpringBootTest
class KmsAwsLocalIntegrationTest {

    @Container
    private static final LocalStackContainer localStack =
        new LocalStackContainer(DockerImageName.parse("localstack/localstack:latest"))
            .withServices(KMS)
            .withEnv("DEFAULT_REGION", "us-east-1");

    private KmsClient kmsClient;

    @BeforeEach
    void setup() {
        kmsClient = KmsClient.builder()
            .region(Region.AWS_GLOBAL)
            .endpointOverride(localStack.getEndpointOverride(KMS))
            .credentialsProvider(StaticCredentialsProvider.create(
                AwsBasicCredentials.create(localStack.getAccessKey(), localStack.getSecretKey())))
            .build();
    }

    @Test
    void shouldCreateKeyInLocalKms() {
        // This test creates a real key in the local KMS instance
        CreateKeyRequest request = CreateKeyRequest.builder()
            .description("Test key")
            .keyUsage(KeyUsageType.ENCRYPT_DECRYPT)
            .build();

        CreateKeyResponse response = kmsClient.createKey(request);
        assertThat(response.keyMetadata().keyId()).isNotEmpty();
    }
}