diff --git a/lambda-invoker/pom.xml b/lambda-invoker/pom.xml index 307b8ae..702c2c8 100644 --- a/lambda-invoker/pom.xml +++ b/lambda-invoker/pom.xml @@ -47,6 +47,10 @@ com.networknt metrics + + com.networknt + metrics-config + com.networknt body diff --git a/lambda-invoker/src/main/java/com/networknt/aws/lambda/LambdaFunctionHandler.java b/lambda-invoker/src/main/java/com/networknt/aws/lambda/LambdaFunctionHandler.java index eac7221..1fa8825 100644 --- a/lambda-invoker/src/main/java/com/networknt/aws/lambda/LambdaFunctionHandler.java +++ b/lambda-invoker/src/main/java/com/networknt/aws/lambda/LambdaFunctionHandler.java @@ -13,9 +13,12 @@ import io.undertow.server.HttpHandler; import io.undertow.server.HttpServerExchange; import io.undertow.util.HeaderMap; +import io.undertow.util.Headers; import io.undertow.util.HttpString; import org.slf4j.Logger; import org.slf4j.LoggerFactory; +import software.amazon.awssdk.auth.credentials.AwsCredentialsProvider; +import software.amazon.awssdk.auth.credentials.AwsCredentials; import software.amazon.awssdk.core.SdkBytes; import software.amazon.awssdk.core.client.config.ClientOverrideConfiguration; import software.amazon.awssdk.http.async.SdkAsyncHttpClient; @@ -26,10 +29,16 @@ import software.amazon.awssdk.services.lambda.model.InvokeRequest; import software.amazon.awssdk.services.sts.StsClient; import software.amazon.awssdk.services.sts.auth.StsAssumeRoleCredentialsProvider; +import software.amazon.awssdk.services.sts.auth.StsAssumeRoleWithWebIdentityCredentialsProvider; import software.amazon.awssdk.services.sts.model.AssumeRoleRequest; +import software.amazon.awssdk.services.sts.model.AssumeRoleWithWebIdentityRequest; import java.net.URI; +import java.nio.charset.StandardCharsets; +import java.security.MessageDigest; +import java.security.NoSuchAlgorithmException; import java.time.Duration; +import java.util.Base64; import java.util.Deque; import java.util.HashMap; import java.util.Map; @@ -40,13 +49,81 @@ public class LambdaFunctionHandler implements LightHttpHandler { private static final Logger logger = LoggerFactory.getLogger(LambdaFunctionHandler.class); private static final String MISSING_ENDPOINT_FUNCTION = "ERR10063"; private static final String EMPTY_LAMBDA_RESPONSE = "ERR10064"; + private static final String STS_TYPE_FUNC_USER = "StsFuncUser"; + private static final String STS_TYPE_WEB_IDENTITY = "StsWebIdentity"; + private static final String BEARER_PREFIX = "BEARER"; + private static final String INVALID_WEB_IDENTITY_TOKEN_MESSAGE = "Missing or invalid Bearer token for STS web identity"; private static AbstractMetricsHandler metricsHandler; private LambdaInvokerConfig config; private LambdaAsyncClient client; private StsAssumeRoleCredentialsProvider stsCredentialsProvider; + private MutableStsWebIdentityCredentialsProvider stsWebIdentityCredentialsProvider; private StsClient stsClient; + static final class MutableStsWebIdentityCredentialsProvider implements AwsCredentialsProvider, AutoCloseable { + private final LambdaInvokerConfig config; + private final StsClient stsClient; + private StsAssumeRoleWithWebIdentityCredentialsProvider delegate; + private String tokenFingerprint; + + MutableStsWebIdentityCredentialsProvider(LambdaInvokerConfig config, StsClient stsClient) { + this.config = config; + this.stsClient = stsClient; + } + + synchronized boolean updateToken(String token) { + String nextFingerprint = fingerprintToken(token); + if(nextFingerprint.equals(tokenFingerprint) && delegate != null) { + return false; + } + StsAssumeRoleWithWebIdentityCredentialsProvider nextDelegate = + StsAssumeRoleWithWebIdentityCredentialsProvider.builder() + .stsClient(stsClient) + .refreshRequest(AssumeRoleWithWebIdentityRequest.builder() + .roleArn(config.getRoleArn()) + .roleSessionName(config.getRoleSessionName()) + .durationSeconds(config.getDurationSeconds()) + .webIdentityToken(token) + .build()) + .build(); + StsAssumeRoleWithWebIdentityCredentialsProvider previousDelegate = delegate; + delegate = nextDelegate; + tokenFingerprint = nextFingerprint; + closeDelegate(previousDelegate); + return true; + } + + synchronized String getTokenFingerprint() { + return tokenFingerprint; + } + + @Override + public synchronized AwsCredentials resolveCredentials() { + if(delegate == null) { + throw new IllegalStateException("STS web identity credentials provider has not been initialized with a bearer token"); + } + return delegate.resolveCredentials(); + } + + @Override + public synchronized void close() { + closeDelegate(delegate); + delegate = null; + tokenFingerprint = null; + } + + private void closeDelegate(StsAssumeRoleWithWebIdentityCredentialsProvider provider) { + if(provider != null) { + try { + provider.close(); + } catch (Exception e) { + logger.error("Failed to close the StsAssumeRoleWithWebIdentityCredentialsProvider", e); + } + } + } + } + // Package-private constructor for testing - avoids loading config from file and metrics chain setup LambdaFunctionHandler(LambdaInvokerConfig config) { this.config = config; @@ -70,6 +147,40 @@ public LambdaFunctionHandler() { } private LambdaAsyncClient initClient(LambdaInvokerConfig config) { + AwsCredentialsProvider credentialsProvider = null; + // If any STS type selected, use the matching credentials provider for automatic refresh. + if(STS_TYPE_FUNC_USER.equals(config.getStsType())) { + if(logger.isInfoEnabled()) logger.info("STS AssumeRole is set to " + STS_TYPE_FUNC_USER + " for role: {}", config.getRoleArn()); + stsClient = StsClient.builder() + .region(Region.of(config.getRegion())) + .build(); + stsCredentialsProvider = StsAssumeRoleCredentialsProvider.builder() + .stsClient(stsClient) + .refreshRequest(AssumeRoleRequest.builder() + .roleArn(config.getRoleArn()) + .roleSessionName(config.getRoleSessionName()) + .durationSeconds(config.getDurationSeconds()) + .build()) + .build(); + credentialsProvider = stsCredentialsProvider; + } else if(STS_TYPE_WEB_IDENTITY.equals(config.getStsType())) { + if(logger.isInfoEnabled()) logger.info("STS AssumeRole is set to " + STS_TYPE_WEB_IDENTITY + " for role: {}", config.getRoleArn()); + stsClient = StsClient.builder() + .region(Region.of(config.getRegion())) + .build(); + stsWebIdentityCredentialsProvider = buildMutableStsWebIdentityCredentialsProvider(config, stsClient); + credentialsProvider = stsWebIdentityCredentialsProvider; + } else { + if(logger.isInfoEnabled()) logger.info("No STS AssumeRole is set. Using default credential provider chain for LambdaAsyncClient."); + } + return buildLambdaClient(config, credentialsProvider); + } + + MutableStsWebIdentityCredentialsProvider buildMutableStsWebIdentityCredentialsProvider(LambdaInvokerConfig config, StsClient stsClient) { + return new MutableStsWebIdentityCredentialsProvider(config, stsClient); + } + + LambdaAsyncClient buildLambdaClient(LambdaInvokerConfig config, AwsCredentialsProvider credentialsProvider) { SdkAsyncHttpClient asyncHttpClient = NettyNioAsyncHttpClient.builder() .readTimeout(Duration.ofMillis(config.getApiCallAttemptTimeout())) .writeTimeout(Duration.ofMillis(config.getApiCallAttemptTimeout())) @@ -103,26 +214,24 @@ private LambdaAsyncClient initClient(LambdaInvokerConfig config) { builder.endpointOverride(URI.create(config.getEndpointOverride())); } - // If STS is enabled, use StsAssumeRoleCredentialsProvider for automatic credential refresh - if(config.isStsEnabled()) { - if(logger.isInfoEnabled()) logger.info("STS AssumeRole is enabled. Using StsAssumeRoleCredentialsProvider for role: {}", config.getRoleArn()); - stsClient = StsClient.builder() - .region(Region.of(config.getRegion())) - .build(); - stsCredentialsProvider = StsAssumeRoleCredentialsProvider.builder() - .stsClient(stsClient) - .refreshRequest(AssumeRoleRequest.builder() - .roleArn(config.getRoleArn()) - .roleSessionName(config.getRoleSessionName()) - .durationSeconds(config.getDurationSeconds()) - .build()) - .build(); - builder.credentialsProvider(stsCredentialsProvider); + if(credentialsProvider != null) { + builder.credentialsProvider(credentialsProvider); } return builder.build(); } + boolean updateWebIdentityToken(String token) { + if(stsWebIdentityCredentialsProvider == null) { + throw new IllegalStateException("STS web identity credentials provider is not configured"); + } + return stsWebIdentityCredentialsProvider.updateToken(token); + } + + String currentWebIdentityTokenFingerprint() { + return stsWebIdentityCredentialsProvider == null ? null : stsWebIdentityCredentialsProvider.getTokenFingerprint(); + } + @Override public void handleRequest(HttpServerExchange exchange) throws Exception { LambdaInvokerConfig newConfig = LambdaInvokerConfig.load(); @@ -146,6 +255,14 @@ public void handleRequest(HttpServerExchange exchange) throws Exception { } stsCredentialsProvider = null; } + if(stsWebIdentityCredentialsProvider != null) { + try { + stsWebIdentityCredentialsProvider.close(); + } catch (Exception e) { + logger.error("Failed to close the StsAssumeRoleWithWebIdentityCredentialsProvider", e); + } + stsWebIdentityCredentialsProvider = null; + } if(stsClient != null) { try { stsClient.close(); @@ -187,6 +304,21 @@ public void handleRequest(HttpServerExchange exchange) throws Exception { if(config.isMetricsInjection() && metricsHandler != null) metricsHandler.injectMetrics(exchange, startTime, config.getMetricsName(), endpoint); return; } + if(STS_TYPE_WEB_IDENTITY.equals(config.getStsType())) { + String rawAuthHeader = exchange.getRequestHeaders().getFirst(Headers.AUTHORIZATION); + String token = extractBearerToken(rawAuthHeader); + if(token == null || token.isEmpty()) { + exchange.setStatusCode(401); + exchange.getResponseSender().send(INVALID_WEB_IDENTITY_TOKEN_MESSAGE); + if(config.isMetricsInjection() && metricsHandler != null) metricsHandler.injectMetrics(exchange, startTime, config.getMetricsName(), endpoint); + return; + } + if(updateWebIdentityToken(token)) { + if(logger.isDebugEnabled()) logger.debug("Authorization token changed. Refreshed the shared STS web identity credentials provider."); + } else { + if(logger.isDebugEnabled()) logger.debug("Authorization token unchanged. Reusing the shared STS web identity credentials provider."); + } + } APIGatewayProxyRequestEvent requestEvent = new APIGatewayProxyRequestEvent(); requestEvent.setHttpMethod(httpMethod); requestEvent.setPath(requestPath); @@ -277,4 +409,38 @@ private void setResponseHeaders(HttpServerExchange exchange, Map } } } + + /** + * Extracts the bearer token from a raw Authorization header value. + * Returns the token string if the header starts with "Bearer " (case-insensitive), + * or {@code null} if the header is missing/empty or does not use the Bearer scheme. + */ + static String extractBearerToken(String authorizationHeaderValue) { + if (authorizationHeaderValue == null || authorizationHeaderValue.isEmpty()) { + if(logger.isDebugEnabled()) logger.debug("Missing Authorization header from request. STS AssumeRole with Web Identity may fail"); + return null; + } + if (authorizationHeaderValue.length() > BEARER_PREFIX.length() + 1 && + authorizationHeaderValue.regionMatches(true, 0, BEARER_PREFIX, 0, BEARER_PREFIX.length()) && + authorizationHeaderValue.charAt(BEARER_PREFIX.length()) == ' ') { + String token = authorizationHeaderValue.substring(BEARER_PREFIX.length() + 1).trim(); + if (token.isEmpty()) { + if(logger.isDebugEnabled()) logger.debug("Authorization header contains a blank Bearer token. STS AssumeRole with Web Identity may fail"); + return null; + } + return token; + } + if(logger.isDebugEnabled()) logger.debug("Authorization header does not start with Bearer. STS AssumeRole with Web Identity may fail"); + return null; + } + + static String fingerprintToken(String token) { + try { + MessageDigest digest = MessageDigest.getInstance("SHA-256"); + byte[] hashed = digest.digest(token.getBytes(StandardCharsets.UTF_8)); + return Base64.getEncoder().encodeToString(hashed); + } catch (NoSuchAlgorithmException e) { + throw new IllegalStateException("SHA-256 is not available", e); + } + } } diff --git a/lambda-invoker/src/main/java/com/networknt/aws/lambda/LambdaInvokerConfig.java b/lambda-invoker/src/main/java/com/networknt/aws/lambda/LambdaInvokerConfig.java index 14d2e54..6f0fce7 100644 --- a/lambda-invoker/src/main/java/com/networknt/aws/lambda/LambdaInvokerConfig.java +++ b/lambda-invoker/src/main/java/com/networknt/aws/lambda/LambdaInvokerConfig.java @@ -40,7 +40,7 @@ public class LambdaInvokerConfig { private static final String MAX_CONCURRENCY = "maxConcurrency"; private static final String MAX_PENDING_CONNECTION_ACQUIRES = "maxPendingConnectionAcquires"; private static final String CONNECTION_ACQUISITION_TIMEOUT = "connectionAcquisitionTimeout"; - private static final String STS_ENABLED = "stsEnabled"; + private static final String STS_TYPE = "stsType"; private static final String ROLE_ARN = "roleArn"; private static final String ROLE_SESSION_NAME = "roleSessionName"; private static final String DURATION_SECONDS = "durationSeconds"; @@ -152,21 +152,25 @@ public class LambdaInvokerConfig { ) private String metricsName; - @BooleanField( - configFieldName = STS_ENABLED, - externalizedKeyName = STS_ENABLED, + @StringField( + configFieldName = STS_TYPE, + externalizedKeyName = STS_TYPE, description = "Enable STS AssumeRole to obtain temporary credentials for Lambda invocation instead of using the\n" + - "permanent IAM credentials. When set to true, the handler will call STS AssumeRole with the configured\n" + - "roleArn, roleSessionName, and durationSeconds to get short-lived credentials. This is the recommended\n" + - "approach for production environments to follow the principle of least privilege.\n", - defaultValue = "false" + "permanent IAM credentials. Only 2 STS types supported: StsFuncUser and StsWebIdentity.\n" + + "If STS is not to be used set this property as empty. When StsFuncUser is set the handler will\n" + + "use the configured AWS IAM User to assume the given RoleARN. When StsWebIdentity is set the handler will\n" + + "use the request bearer token as the WEB_IDENTITY_TOKEN to be exchanged for STS token. Regardless of the\n" + + "selected type, the handler will call STS AssumeRole with the configured roleArn, roleSessionName, and\n" + + "durationSeconds to get short-lived credentials. Using one of the STS types is the recommended approach\n" + + "for production environments to follow the principle of least privilege.\n", + defaultValue = "" ) - private boolean stsEnabled; + private String stsType; @StringField( configFieldName = ROLE_ARN, externalizedKeyName = ROLE_ARN, - description = "The ARN of the IAM role to assume when stsEnabled is true. For example,\n" + + description = "The ARN of the IAM role to assume when stsType is not empty. For example,\n" + "arn:aws:iam::123456789012:role/LambdaInvokerRole\n" ) private String roleArn; @@ -329,12 +333,12 @@ public void setConnectionAcquisitionTimeout(int connectionAcquisitionTimeout) { this.connectionAcquisitionTimeout = connectionAcquisitionTimeout; } - public boolean isStsEnabled() { - return stsEnabled; + public String getStsType() { + return stsType; } - public void setStsEnabled(boolean stsEnabled) { - this.stsEnabled = stsEnabled; + public void setStsType(String stsType) { + this.stsType = stsType; } public String getRoleArn() { @@ -398,8 +402,8 @@ private void setConfigData() { object = mappedConfig.get(CONNECTION_ACQUISITION_TIMEOUT); if (object != null) connectionAcquisitionTimeout = Config.loadIntegerValue(CONNECTION_ACQUISITION_TIMEOUT, object); - object = mappedConfig.get(STS_ENABLED); - if(object != null) stsEnabled = Config.loadBooleanValue(STS_ENABLED, object); + object = mappedConfig.get(STS_TYPE); + if (object != null) stsType = (String) object; object = mappedConfig.get(ROLE_ARN); if(object != null) roleArn = (String) object; @@ -456,8 +460,17 @@ private void setConfigMap() { } private void validate() { - if (stsEnabled && (roleArn == null || roleArn.trim().isEmpty())) { - throw new ConfigException(ROLE_ARN + " must be configured when " + STS_ENABLED + " is true."); + String normalizedStsType = stsType == null ? null : stsType.trim(); + // Write normalized value back so downstream equals() comparisons work correctly + // even when the config value has leading/trailing whitespace (e.g. "StsWebIdentity "). + stsType = normalizedStsType; + if (normalizedStsType != null && !normalizedStsType.isEmpty()) { + if (!"StsFuncUser".equals(normalizedStsType) && !"StsWebIdentity".equals(normalizedStsType)) { + throw new ConfigException(STS_TYPE + " must be one of [StsFuncUser, StsWebIdentity], but was: " + normalizedStsType); + } + if (roleArn == null || roleArn.trim().isEmpty()) { + throw new ConfigException(ROLE_ARN + " must be configured when " + STS_TYPE + " is not empty."); + } } } } diff --git a/lambda-invoker/src/main/resources/config/lambda-invoker-schema.json b/lambda-invoker/src/main/resources/config/lambda-invoker-schema.json index 0fc97ad..b88a3e1 100644 --- a/lambda-invoker/src/main/resources/config/lambda-invoker-schema.json +++ b/lambda-invoker/src/main/resources/config/lambda-invoker-schema.json @@ -1,7 +1,7 @@ { "$schema" : "http://json-schema.org/draft-07/schema#", "type" : "object", - "required" : [ "region", "endpointOverride", "apiCallTimeout", "apiCallAttemptTimeout", "maxRetries", "maxConcurrency", "maxPendingConnectionAcquires", "connectionAcquisitionTimeout", "logType", "functions", "metricsInjection", "metricsName", "stsEnabled", "roleArn", "roleSessionName", "durationSeconds" ], + "required" : [ "region", "endpointOverride", "apiCallTimeout", "apiCallAttemptTimeout", "maxRetries", "maxConcurrency", "maxPendingConnectionAcquires", "connectionAcquisitionTimeout", "logType", "functions", "metricsInjection", "metricsName", "stsType", "roleArn", "roleSessionName", "durationSeconds" ], "properties" : { "region" : { "type" : "string", @@ -70,14 +70,14 @@ "description" : "When the metrics info is injected into the metrics handler, we need to pass a metric name to it so that\nthe metrics info can be categorized in a tree structure under the name. By default, it is lambda-response,\nand users can change it.\n", "default" : "lambda-response" }, - "stsEnabled" : { - "type" : "boolean", - "description" : "Enable STS AssumeRole to obtain temporary credentials for Lambda invocation instead of using the\npermanent IAM credentials. When set to true, the handler will call STS AssumeRole with the configured\nroleArn, roleSessionName, and durationSeconds to get short-lived credentials. This is the recommended\napproach for production environments to follow the principle of least privilege.\n", - "default" : false + "stsType" : { + "type" : "string", + "description" : "Enable STS AssumeRole to obtain temporary credentials for Lambda invocation instead of using the\npermanent IAM credentials. Only 2 STS types supported: StsFuncUser and StsWebIdentity.\nIf STS is not to be used set this property as empty. When StsFuncUser is set the handler will\nuse the configured AWS IAM User to assume the given RoleARN. When StsWebIdentity is set the handler will\nuse the request bearer token as the WEB_IDENTITY_TOKEN to be exchanged for STS token. Regardless of the\nselected type, the handler will call STS AssumeRole with the configured roleArn, roleSessionName, and\ndurationSeconds to get short-lived credentials. Using one of the STS types is the recommended approach\nfor production environments to follow the principle of least privilege.", + "default" : "" }, "roleArn" : { "type" : "string", - "description" : "The ARN of the IAM role to assume when stsEnabled is true. For example,\narn:aws:iam::123456789012:role/LambdaInvokerRole\n" + "description" : "The ARN of the IAM role to assume when stsType is set. For example,\narn:aws:iam::123456789012:role/LambdaInvokerRole\n" }, "roleSessionName" : { "type" : "string", diff --git a/lambda-invoker/src/main/resources/config/lambda-invoker.yaml b/lambda-invoker/src/main/resources/config/lambda-invoker.yaml index 8b7422a..79452ed 100644 --- a/lambda-invoker/src/main/resources/config/lambda-invoker.yaml +++ b/lambda-invoker/src/main/resources/config/lambda-invoker.yaml @@ -32,11 +32,15 @@ metricsInjection: ${lambda-invoker.metricsInjection:false} # and users can change it. metricsName: ${lambda-invoker.metricsName:lambda-response} # Enable STS AssumeRole to obtain temporary credentials for Lambda invocation instead of using the -# permanent IAM credentials. When set to true, the handler will call STS AssumeRole with the configured -# roleArn, roleSessionName, and durationSeconds to get short-lived credentials. This is the recommended -# approach for production environments to follow the principle of least privilege. -stsEnabled: ${lambda-invoker.stsEnabled:false} -# The ARN of the IAM role to assume when stsEnabled is true. For example, +# permanent IAM credentials. Only 2 STS types supported: StsFuncUser and StsWebIdentity. +# If STS is not to be used set this property as empty. When StsFuncUser is set the handler will +# use the configured AWS IAM User to assume the given RoleARN. When StsWebIdentity is set the handler will +# use the request bearer token as the WEB_IDENTITY_TOKEN to be exchanged for STS token. Regardless of the +# selected type, the handler will call STS AssumeRole with the configured roleArn, roleSessionName, and +# durationSeconds to get short-lived credentials. Using one of the STS types is the recommended approach +# for production environments to follow the principle of least privilege. +stsType: ${lambda-invoker.stsType:} +# The ARN of the IAM role to assume when stsType is set. For example, # arn:aws:iam::123456789012:role/LambdaInvokerRole roleArn: ${lambda-invoker.roleArn:} # The session name to use when assuming the role. This is used for auditing and tracking in CloudTrail. diff --git a/lambda-invoker/src/main/resources/config/lambda-invoker.yml b/lambda-invoker/src/main/resources/config/lambda-invoker.yml index 8b7422a..79452ed 100644 --- a/lambda-invoker/src/main/resources/config/lambda-invoker.yml +++ b/lambda-invoker/src/main/resources/config/lambda-invoker.yml @@ -32,11 +32,15 @@ metricsInjection: ${lambda-invoker.metricsInjection:false} # and users can change it. metricsName: ${lambda-invoker.metricsName:lambda-response} # Enable STS AssumeRole to obtain temporary credentials for Lambda invocation instead of using the -# permanent IAM credentials. When set to true, the handler will call STS AssumeRole with the configured -# roleArn, roleSessionName, and durationSeconds to get short-lived credentials. This is the recommended -# approach for production environments to follow the principle of least privilege. -stsEnabled: ${lambda-invoker.stsEnabled:false} -# The ARN of the IAM role to assume when stsEnabled is true. For example, +# permanent IAM credentials. Only 2 STS types supported: StsFuncUser and StsWebIdentity. +# If STS is not to be used set this property as empty. When StsFuncUser is set the handler will +# use the configured AWS IAM User to assume the given RoleARN. When StsWebIdentity is set the handler will +# use the request bearer token as the WEB_IDENTITY_TOKEN to be exchanged for STS token. Regardless of the +# selected type, the handler will call STS AssumeRole with the configured roleArn, roleSessionName, and +# durationSeconds to get short-lived credentials. Using one of the STS types is the recommended approach +# for production environments to follow the principle of least privilege. +stsType: ${lambda-invoker.stsType:} +# The ARN of the IAM role to assume when stsType is set. For example, # arn:aws:iam::123456789012:role/LambdaInvokerRole roleArn: ${lambda-invoker.roleArn:} # The session name to use when assuming the role. This is used for auditing and tracking in CloudTrail. diff --git a/lambda-invoker/src/test/java/com/networknt/aws/lambda/LambdaFunctionHandlerTest.java b/lambda-invoker/src/test/java/com/networknt/aws/lambda/LambdaFunctionHandlerTest.java index 857f630..02f5242 100644 --- a/lambda-invoker/src/test/java/com/networknt/aws/lambda/LambdaFunctionHandlerTest.java +++ b/lambda-invoker/src/test/java/com/networknt/aws/lambda/LambdaFunctionHandlerTest.java @@ -5,12 +5,63 @@ import org.junit.jupiter.api.*; import org.slf4j.Logger; import org.slf4j.LoggerFactory; +import software.amazon.awssdk.auth.credentials.AwsCredentialsProvider; +import software.amazon.awssdk.services.lambda.LambdaAsyncClient; + +import java.lang.reflect.Proxy; import java.util.HashMap; import java.util.Map; +import static org.junit.jupiter.api.Assertions.*; + class LambdaFunctionHandlerTest { static final Logger logger = LoggerFactory.getLogger(LambdaFunctionHandlerTest.class); + private static class TestLambdaFunctionHandler extends LambdaFunctionHandler { + int buildLambdaClientCalls; + + TestLambdaFunctionHandler(LambdaInvokerConfig config) { + super(config); + } + + @Override + LambdaAsyncClient buildLambdaClient(LambdaInvokerConfig config, AwsCredentialsProvider credentialsProvider) { + buildLambdaClientCalls++; + return (LambdaAsyncClient) Proxy.newProxyInstance( + LambdaAsyncClient.class.getClassLoader(), + new Class[]{LambdaAsyncClient.class}, + (proxy, method, args) -> { + if("close".equals(method.getName())) { + return null; + } + if(method.getReturnType().equals(boolean.class)) { + return false; + } + if(method.getReturnType().equals(int.class)) { + return 0; + } + return null; + } + ); + } + } + + private LambdaInvokerConfig webIdentityConfig() { + LambdaInvokerConfig config = new LambdaInvokerConfig(); + config.setRegion("us-east-1"); + config.setApiCallTimeout(60000); + config.setApiCallAttemptTimeout(20000); + config.setMaxConcurrency(10); + config.setMaxPendingConnectionAcquires(100); + config.setConnectionAcquisitionTimeout(10); + config.setMaxRetries(0); + config.setRoleArn("arn:aws:iam::123456789012:role/LambdaInvokerRole"); + config.setRoleSessionName("test-session"); + config.setDurationSeconds(900); + config.setStsType("StsWebIdentity"); + return config; + } + @Test void testAPIGatewayProxyRequestEvent() throws Exception { APIGatewayProxyRequestEvent requestEvent = new APIGatewayProxyRequestEvent(); @@ -26,4 +77,86 @@ void testAPIGatewayProxyRequestEvent() throws Exception { requestEvent.setBody(null); System.out.println(JsonMapper.objectMapper.writeValueAsString(requestEvent)); } + + // --- extractBearerToken tests --- + + @Test + void testExtractBearerToken_nullHeader_returnsNull() { + assertNull(LambdaFunctionHandler.extractBearerToken(null)); + } + + @Test + void testExtractBearerToken_emptyHeader_returnsNull() { + assertNull(LambdaFunctionHandler.extractBearerToken("")); + } + + @Test + void testExtractBearerToken_validBearerMixedCase_extractsToken() { + assertEquals("mytoken123", LambdaFunctionHandler.extractBearerToken("Bearer mytoken123")); + } + + @Test + void testExtractBearerToken_validBearerLowercase_extractsToken() { + assertEquals("mytoken123", LambdaFunctionHandler.extractBearerToken("bearer mytoken123")); + } + + @Test + void testExtractBearerToken_validBearerUppercase_extractsToken() { + assertEquals("mytoken123", LambdaFunctionHandler.extractBearerToken("BEARER mytoken123")); + } + + @Test + void testExtractBearerToken_bearerWithJwtToken_extractsFullToken() { + String jwt = "eyJhbGciOiJSUzI1NiJ9.eyJzdWIiOiJ1c2VyIn0.signature"; + assertEquals(jwt, LambdaFunctionHandler.extractBearerToken("Bearer " + jwt)); + } + + @Test + void testExtractBearerToken_nonBearerScheme_returnsNull() { + assertNull(LambdaFunctionHandler.extractBearerToken("Basic dXNlcjpwYXNz")); + } + + @Test + void testExtractBearerToken_headerIsBearerOnly_returnsNull() { + // "BEARER" alone is exactly BEARER_PREFIX.length() characters, not greater, so returns null + assertNull(LambdaFunctionHandler.extractBearerToken("BEARER")); + } + + @Test + void testExtractBearerToken_tokenIsEmptyAfterBearer_returnsNull() { + // Empty bearer tokens are treated as invalid and return null so the STS path skips refresh. + assertNull(LambdaFunctionHandler.extractBearerToken("Bearer ")); + } + + @Test + void testUpdateWebIdentityToken_refreshesProviderWithoutRebuildingClient() { + TestLambdaFunctionHandler handler = new TestLambdaFunctionHandler(webIdentityConfig()); + + assertTrue(handler.updateWebIdentityToken("token-1")); + assertEquals(1, handler.buildLambdaClientCalls); + assertNotNull(handler.currentWebIdentityTokenFingerprint()); + assertNotEquals("token-1", handler.currentWebIdentityTokenFingerprint()); + } + + @Test + void testUpdateWebIdentityToken_reusesProviderWhenTokenUnchanged() { + TestLambdaFunctionHandler handler = new TestLambdaFunctionHandler(webIdentityConfig()); + assertTrue(handler.updateWebIdentityToken("token-1")); + String fingerprint = handler.currentWebIdentityTokenFingerprint(); + + assertFalse(handler.updateWebIdentityToken("token-1")); + assertEquals(fingerprint, handler.currentWebIdentityTokenFingerprint()); + assertEquals(1, handler.buildLambdaClientCalls); + } + + @Test + void testUpdateWebIdentityToken_refreshesProviderWhenTokenChanges() { + TestLambdaFunctionHandler handler = new TestLambdaFunctionHandler(webIdentityConfig()); + assertTrue(handler.updateWebIdentityToken("token-1")); + String firstFingerprint = handler.currentWebIdentityTokenFingerprint(); + + assertTrue(handler.updateWebIdentityToken("token-2")); + assertNotEquals(firstFingerprint, handler.currentWebIdentityTokenFingerprint()); + assertEquals(1, handler.buildLambdaClientCalls); + } } diff --git a/lambda-invoker/src/test/java/com/networknt/aws/lambda/LambdaInvokerConfigTest.java b/lambda-invoker/src/test/java/com/networknt/aws/lambda/LambdaInvokerConfigTest.java index 7ca6e87..37ace75 100644 --- a/lambda-invoker/src/test/java/com/networknt/aws/lambda/LambdaInvokerConfigTest.java +++ b/lambda-invoker/src/test/java/com/networknt/aws/lambda/LambdaInvokerConfigTest.java @@ -1,35 +1,108 @@ package com.networknt.aws.lambda; -import com.networknt.config.Config; import com.networknt.config.ConfigException; import com.networknt.config.JsonMapper; import org.junit.jupiter.api.Test; import java.util.Map; +import java.util.concurrent.atomic.AtomicReference; -import static org.junit.jupiter.api.Assertions.assertThrows; -import static org.junit.jupiter.api.Assertions.assertNotNull; -import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.*; public class LambdaInvokerConfigTest { - private static LambdaInvokerConfig config = (LambdaInvokerConfig) Config.getInstance().getJsonObjectConfig(LambdaInvokerConfig.CONFIG_NAME, LambdaInvokerConfig.class); - @Test public void testFunctionMapping() { + LambdaInvokerConfig config = LambdaInvokerConfig.load("lambda-invoker-sts-with-role"); Map functions = config.getFunctions(); + assertNotNull(functions); System.out.println(JsonMapper.toJson(functions)); } @Test - public void testStsEnabledWithoutRoleArnThrowsConfigException() { + public void testStsTypeWithoutRoleArnThrowsConfigException() { assertThrows(ConfigException.class, () -> LambdaInvokerConfig.load("lambda-invoker-sts-no-role"), - "ConfigException was not thrown despite stsEnabled=true with blank roleArn"); + "ConfigException was not thrown despite stsType defined with blank roleArn"); } @Test - public void testStsEnabledWithRoleArnSucceeds() { + public void testStsTypeWithRoleArnSucceeds() { LambdaInvokerConfig stsConfig = LambdaInvokerConfig.load("lambda-invoker-sts-with-role"); assertNotNull(stsConfig); assertEquals("arn:aws:iam::123456789012:role/TestRole", stsConfig.getRoleArn()); } + + @Test + public void testStsTypeFuncWithRoleArn() { + LambdaInvokerConfig stsConfig = LambdaInvokerConfig.load("lambda-invoker-sts-func"); + assertNotNull(stsConfig); + assertEquals("arn:aws:iam::123456789012:role/TestRole", stsConfig.getRoleArn()); + } + + @Test + public void testInvalidStsTypeThrowsConfigException() { + ConfigException ex = assertThrows(ConfigException.class, + () -> LambdaInvokerConfig.load("lambda-invoker-sts-invalid-type"), + "ConfigException was not thrown for unsupported stsType value"); + assertTrue(ex.getMessage().contains("stsType"), "Exception message should mention stsType"); + } + + @Test + public void testStsTypeWithTrailingSpaceIsNormalized() { + // "StsWebIdentity " (trailing space) must be trimmed so getStsType() equals "StsWebIdentity" + LambdaInvokerConfig stsConfig = LambdaInvokerConfig.load("lambda-invoker-sts-trailing-space"); + assertNotNull(stsConfig); + assertEquals("StsWebIdentity", stsConfig.getStsType(), + "validate() should strip trailing whitespace and write normalized value back to stsType"); + } + + @Test + public void testTokenRefreshDecisionWhenAuthorizationMissing() { + AtomicReference cache = new AtomicReference<>("token1"); + + assertFalse(shouldRefreshStsWebIdentityToken(cache, null), + "Missing Authorization should not trigger a refresh"); + assertEquals("token1", cache.get(), "Cached token should remain unchanged when Authorization is missing"); + } + + @Test + public void testTokenRefreshDecisionWhenTokenUnchanged() { + AtomicReference cache = new AtomicReference<>("token2"); + + assertFalse(shouldRefreshStsWebIdentityToken(cache, "Bearer token2"), + "Same token should use the cached fast-path and avoid refresh"); + assertEquals("token2", cache.get(), "Cached token should remain unchanged when the token matches"); + } + + @Test + public void testTokenRefreshDecisionWhenTokenChanges() { + AtomicReference cache = new AtomicReference<>("token2"); + + assertTrue(shouldRefreshStsWebIdentityToken(cache, "Bearer token3"), + "Different token should trigger a refresh"); + assertEquals("token3", cache.get(), "Cached token should be updated after a token change"); + } + + private static boolean shouldRefreshStsWebIdentityToken(AtomicReference cache, String authorization) { + String incomingToken = extractAuthorizationToken(authorization); + if (incomingToken == null || incomingToken.isBlank()) { + return false; + } + + String cachedToken = cache.get(); + if (incomingToken.equals(cachedToken)) { + return false; + } + + cache.set(incomingToken); + return true; + } + + private static String extractAuthorizationToken(String authorization) { + if (authorization == null) { + return null; + } + + String prefix = "Bearer "; + return authorization.startsWith(prefix) ? authorization.substring(prefix.length()) : authorization; + } } diff --git a/lambda-invoker/src/test/resources/config/lambda-invoker-sts-func.yml b/lambda-invoker/src/test/resources/config/lambda-invoker-sts-func.yml new file mode 100644 index 0000000..5a9b194 --- /dev/null +++ b/lambda-invoker/src/test/resources/config/lambda-invoker-sts-func.yml @@ -0,0 +1,2 @@ +stsType: StsFuncUser +roleArn: arn:aws:iam::123456789012:role/TestRole diff --git a/lambda-invoker/src/test/resources/config/lambda-invoker-sts-invalid-type.yml b/lambda-invoker/src/test/resources/config/lambda-invoker-sts-invalid-type.yml new file mode 100644 index 0000000..78bf639 --- /dev/null +++ b/lambda-invoker/src/test/resources/config/lambda-invoker-sts-invalid-type.yml @@ -0,0 +1,2 @@ +stsType: StsWebIdentiy +roleArn: arn:aws:iam::123456789012:role/TestRole diff --git a/lambda-invoker/src/test/resources/config/lambda-invoker-sts-no-role.yml b/lambda-invoker/src/test/resources/config/lambda-invoker-sts-no-role.yml index 9841484..3508542 100644 --- a/lambda-invoker/src/test/resources/config/lambda-invoker-sts-no-role.yml +++ b/lambda-invoker/src/test/resources/config/lambda-invoker-sts-no-role.yml @@ -1,2 +1,2 @@ -stsEnabled: true +stsType: StsWebIdentity roleArn: diff --git a/lambda-invoker/src/test/resources/config/lambda-invoker-sts-trailing-space.yml b/lambda-invoker/src/test/resources/config/lambda-invoker-sts-trailing-space.yml new file mode 100644 index 0000000..bd258e5 --- /dev/null +++ b/lambda-invoker/src/test/resources/config/lambda-invoker-sts-trailing-space.yml @@ -0,0 +1,2 @@ +stsType: "StsWebIdentity " +roleArn: arn:aws:iam::123456789012:role/TestRole diff --git a/lambda-invoker/src/test/resources/config/lambda-invoker-sts-with-role.yml b/lambda-invoker/src/test/resources/config/lambda-invoker-sts-with-role.yml index 7ddc658..d286c92 100644 --- a/lambda-invoker/src/test/resources/config/lambda-invoker-sts-with-role.yml +++ b/lambda-invoker/src/test/resources/config/lambda-invoker-sts-with-role.yml @@ -1,2 +1,2 @@ -stsEnabled: true +stsType: StsWebIdentity roleArn: arn:aws:iam::123456789012:role/TestRole diff --git a/pom.xml b/pom.xml index ae15638..d2bad0a 100644 --- a/pom.xml +++ b/pom.xml @@ -245,6 +245,11 @@ metrics ${project.version} + + com.networknt + metrics-config + ${project.version} + com.networknt status