Skip to content

Commit d3f444a

Browse files
msm1992msm1992
committed
Migrate Lambda mediator to AWS SDK 2.x
1 parent 522d8ed commit d3f444a

File tree

4 files changed

+157
-153
lines changed

4 files changed

+157
-153
lines changed

‎components/apimgt/org.wso2.carbon.apimgt.gateway/src/main/java/org/wso2/carbon/apimgt/gateway/mediators/AWSLambdaMediator.java

+79-78
Original file line numberDiff line numberDiff line change
@@ -17,25 +17,6 @@
1717
*/
1818
package org.wso2.carbon.apimgt.gateway.mediators;
1919

20-
import com.amazonaws.ClientConfiguration;
21-
import com.amazonaws.SdkClientException;
22-
import com.amazonaws.auth.AWSCredentialsProvider;
23-
import com.amazonaws.auth.AWSStaticCredentialsProvider;
24-
import com.amazonaws.auth.BasicAWSCredentials;
25-
import com.amazonaws.auth.BasicSessionCredentials;
26-
import com.amazonaws.auth.DefaultAWSCredentialsProviderChain;
27-
import com.amazonaws.client.builder.AwsClientBuilder.EndpointConfiguration;
28-
import com.amazonaws.regions.Regions;
29-
import com.amazonaws.services.lambda.AWSLambda;
30-
import com.amazonaws.services.lambda.AWSLambdaClientBuilder;
31-
import com.amazonaws.services.lambda.model.InvocationType;
32-
import com.amazonaws.services.lambda.model.InvokeRequest;
33-
import com.amazonaws.services.lambda.model.InvokeResult;
34-
import com.amazonaws.services.securitytoken.AWSSecurityTokenService;
35-
import com.amazonaws.services.securitytoken.AWSSecurityTokenServiceClientBuilder;
36-
import com.amazonaws.services.securitytoken.model.AssumeRoleRequest;
37-
import com.amazonaws.services.securitytoken.model.AssumeRoleResult;
38-
import com.amazonaws.services.securitytoken.model.Credentials;
3920
import com.google.gson.JsonObject;
4021
import com.google.gson.JsonParser;
4122
import org.apache.axiom.om.OMElement;
@@ -58,10 +39,29 @@
5839
import org.wso2.carbon.apimgt.gateway.internal.ServiceReferenceHolder;
5940
import org.wso2.carbon.apimgt.gateway.utils.redis.RedisCacheUtils;
6041
import org.wso2.carbon.apimgt.impl.APIConstants;
42+
import software.amazon.awssdk.auth.credentials.*;
43+
import software.amazon.awssdk.core.SdkBytes;
44+
import software.amazon.awssdk.core.client.config.ClientOverrideConfiguration;
45+
import software.amazon.awssdk.core.exception.SdkClientException;
46+
import software.amazon.awssdk.http.SdkHttpClient;
47+
import software.amazon.awssdk.http.apache.ApacheHttpClient;
48+
import software.amazon.awssdk.regions.Region;
49+
import software.amazon.awssdk.regions.providers.DefaultAwsRegionProviderChain;
50+
import software.amazon.awssdk.services.lambda.LambdaClient;
51+
import software.amazon.awssdk.services.lambda.model.InvocationType;
52+
import software.amazon.awssdk.services.lambda.model.InvokeRequest;
53+
import software.amazon.awssdk.services.lambda.model.InvokeResponse;
54+
import software.amazon.awssdk.services.sts.StsClient;
55+
import software.amazon.awssdk.services.sts.model.AssumeRoleRequest;
56+
import software.amazon.awssdk.services.sts.model.AssumeRoleResponse;
57+
import software.amazon.awssdk.services.sts.model.Credentials;
6158

6259
import java.io.ByteArrayOutputStream;
6360
import java.io.IOException;
61+
import java.net.URI;
62+
import java.net.URISyntaxException;
6463
import java.nio.charset.StandardCharsets;
64+
import java.time.Duration;
6565
import java.util.Iterator;
6666
import java.util.Set;
6767
import java.util.TreeMap;
@@ -79,7 +79,7 @@ public class AWSLambdaMediator extends AbstractMediator {
7979
private String roleArn = "";
8080
private String roleSessionName = "";
8181
private String roleRegion = "";
82-
private int resourceTimeout = APIConstants.AWS_DEFAULT_CONNECTION_TIMEOUT;
82+
private Duration resourceTimeout = Duration.ofMillis(APIConstants.AWS_DEFAULT_CONNECTION_TIMEOUT);
8383
private boolean isContentEncodingEnabled = false;
8484
private static final String PATH_PARAMETERS = "pathParameters";
8585
private static final String QUERY_STRING_PARAMETERS = "queryStringParameters";
@@ -165,15 +165,15 @@ public boolean mediate(MessageContext messageContext) {
165165
log.debug("Passing the payload " + payload.toString() + " to AWS Lambda function with resource name "
166166
+ resourceName);
167167
}
168-
InvokeResult invokeResult = invokeLambda(payload.toString());
168+
InvokeResponse invokeResult = invokeLambda(payload.toString());
169169

170170
if (invokeResult != null) {
171171
if (log.isDebugEnabled()) {
172172
log.debug("AWS Lambda function: " + resourceName + " is invoked successfully.");
173173
}
174-
JsonUtil.getNewJsonPayload(axis2MessageContext, new String(invokeResult.getPayload().array()),
174+
JsonUtil.getNewJsonPayload(axis2MessageContext, new String(invokeResult.payload().asByteArray(), StandardCharsets.UTF_8),
175175
true, true);
176-
axis2MessageContext.setProperty(APIMgtGatewayConstants.HTTP_SC, invokeResult.getStatusCode());
176+
axis2MessageContext.setProperty(APIMgtGatewayConstants.HTTP_SC, invokeResult.statusCode());
177177
axis2MessageContext.setProperty(APIMgtGatewayConstants.REST_MESSAGE_TYPE, APIConstants.APPLICATION_JSON_MEDIA_TYPE);
178178
axis2MessageContext.setProperty(APIMgtGatewayConstants.REST_CONTENT_TYPE, APIConstants.APPLICATION_JSON_MEDIA_TYPE);
179179
axis2MessageContext.removeProperty(APIConstants.NO_ENTITY_BODY);
@@ -198,39 +198,39 @@ public boolean mediate(MessageContext messageContext) {
198198
* @param payload - input parameters to pass to AWS Lambda function as a JSONString
199199
* @return InvokeResult
200200
*/
201-
private InvokeResult invokeLambda(String payload) {
201+
private InvokeResponse invokeLambda(String payload) {
202202
try {
203203
// Validate resource timeout and set client configuration
204-
if (resourceTimeout < 1000 || resourceTimeout > 900000) {
204+
if (resourceTimeout.toMillis() < 1000 || resourceTimeout.toMillis() > 900000) {
205205
setResourceTimeout(APIConstants.AWS_DEFAULT_CONNECTION_TIMEOUT);
206206
}
207-
ClientConfiguration clientConfig = new ClientConfiguration();
208-
clientConfig.setSocketTimeout(resourceTimeout);
207+
ClientOverrideConfiguration clientConfig = ClientOverrideConfiguration.builder()
208+
.apiCallTimeout(resourceTimeout).build();
209209

210-
AWSLambda awsLambdaClient;
210+
LambdaClient awsLambdaClient;
211211
if (StringUtils.isEmpty(accessKey) && StringUtils.isEmpty(secretKey)) {
212212
if (log.isDebugEnabled()) {
213213
log.debug("Using temporary credentials supplied by the IAM role attached to AWS instance");
214214
}
215215
if (StringUtils.isEmpty(roleArn) && StringUtils.isEmpty(roleSessionName)
216216
&& StringUtils.isEmpty(roleRegion)) {
217-
awsLambdaClient = AWSLambdaClientBuilder.standard()
218-
.withCredentials(DefaultAWSCredentialsProviderChain.getInstance())
219-
.withClientConfiguration(clientConfig)
217+
awsLambdaClient = LambdaClient.builder()
218+
.credentialsProvider(DefaultCredentialsProvider.create())
219+
.httpClientBuilder(ApacheHttpClient.builder())
220+
.overrideConfiguration(clientConfig)
220221
.build();
221222
} else if (StringUtils.isNotEmpty(roleArn) && StringUtils.isNotEmpty(roleSessionName)
222223
&& StringUtils.isNotEmpty(roleRegion)) {
224+
Region region = new DefaultAwsRegionProviderChain().getRegion();
223225
Credentials sessionCredentials = getSessionCredentials(
224-
DefaultAWSCredentialsProviderChain.getInstance(), roleArn, roleSessionName,
225-
String.valueOf(Regions.getCurrentRegion()));
226-
BasicSessionCredentials basicSessionCredentials = new BasicSessionCredentials(
227-
sessionCredentials.getAccessKeyId(),
228-
sessionCredentials.getSecretAccessKey(),
229-
sessionCredentials.getSessionToken());
230-
awsLambdaClient = AWSLambdaClientBuilder.standard()
231-
.withCredentials(new AWSStaticCredentialsProvider(basicSessionCredentials))
232-
.withClientConfiguration(clientConfig)
233-
.withRegion(roleRegion)
226+
DefaultCredentialsProvider.create(), roleArn, roleSessionName,
227+
String.valueOf(region));
228+
AwsSessionCredentials basicSessionCredentials = AwsSessionCredentials.create(sessionCredentials.accessKeyId(), sessionCredentials.secretAccessKey(), sessionCredentials.sessionToken());
229+
awsLambdaClient = LambdaClient.builder()
230+
.credentialsProvider(StaticCredentialsProvider.create(basicSessionCredentials))
231+
.httpClientBuilder(ApacheHttpClient.builder())
232+
.overrideConfiguration(clientConfig)
233+
.region(Region.of(roleRegion))
234234
.build();
235235
} else {
236236
log.error("Missing AWS STS configurations");
@@ -241,26 +241,25 @@ private InvokeResult invokeLambda(String payload) {
241241
if (log.isDebugEnabled()) {
242242
log.debug("Using user given stored credentials");
243243
}
244-
BasicAWSCredentials awsCredentials = new BasicAWSCredentials(accessKey, secretKey);
244+
AwsBasicCredentials awsCredentials = AwsBasicCredentials.create(accessKey, secretKey);
245245
if (StringUtils.isEmpty(roleArn) && StringUtils.isEmpty(roleSessionName)
246246
&& StringUtils.isEmpty(roleRegion)) {
247-
awsLambdaClient = AWSLambdaClientBuilder.standard()
248-
.withCredentials(new AWSStaticCredentialsProvider(awsCredentials))
249-
.withClientConfiguration(clientConfig)
250-
.withRegion(region)
247+
awsLambdaClient = LambdaClient.builder()
248+
.credentialsProvider(StaticCredentialsProvider.create(awsCredentials))
249+
.httpClientBuilder(ApacheHttpClient.builder())
250+
.overrideConfiguration(clientConfig)
251+
.region(Region.of(region))
251252
.build();
252253
} else if (StringUtils.isNotEmpty(roleArn) && StringUtils.isNotEmpty(roleSessionName)
253254
&& StringUtils.isNotEmpty(roleRegion)) {
254255
Credentials sessionCredentials = getSessionCredentials(
255-
new AWSStaticCredentialsProvider(awsCredentials), roleArn, roleSessionName, region);
256-
BasicSessionCredentials basicSessionCredentials = new BasicSessionCredentials(
257-
sessionCredentials.getAccessKeyId(),
258-
sessionCredentials.getSecretAccessKey(),
259-
sessionCredentials.getSessionToken());
260-
awsLambdaClient = AWSLambdaClientBuilder.standard()
261-
.withCredentials(new AWSStaticCredentialsProvider(basicSessionCredentials))
262-
.withClientConfiguration(clientConfig)
263-
.withRegion(roleRegion)
256+
StaticCredentialsProvider.create(awsCredentials), roleArn, roleSessionName, region);
257+
AwsSessionCredentials basicSessionCredentials = AwsSessionCredentials.create(sessionCredentials.accessKeyId(), sessionCredentials.secretAccessKey(), sessionCredentials.sessionToken());
258+
awsLambdaClient = LambdaClient.builder()
259+
.credentialsProvider(StaticCredentialsProvider.create(basicSessionCredentials))
260+
.httpClientBuilder(ApacheHttpClient.builder())
261+
.overrideConfiguration(clientConfig)
262+
.region(Region.of(roleRegion))
264263
.build();
265264
} else {
266265
log.error("Missing AWS STS configurations");
@@ -270,20 +269,22 @@ private InvokeResult invokeLambda(String payload) {
270269
log.error("Missing AWS Credentials");
271270
return null;
272271
}
273-
InvokeRequest invokeRequest = new InvokeRequest()
274-
.withFunctionName(resourceName)
275-
.withPayload(payload)
276-
.withInvocationType(InvocationType.RequestResponse)
277-
.withSdkClientExecutionTimeout(resourceTimeout);
272+
273+
SdkBytes payloadBytes = SdkBytes.fromUtf8String(payload);
274+
InvokeRequest invokeRequest = InvokeRequest.builder()
275+
.functionName(resourceName)
276+
.payload(payloadBytes)
277+
.invocationType(InvocationType.REQUEST_RESPONSE)
278+
.build();
278279
return awsLambdaClient.invoke(invokeRequest);
279-
} catch (SdkClientException e) {
280+
} catch (SdkClientException | URISyntaxException e) {
280281
log.error("Error while invoking the lambda function", e);
281282
}
282283
return null;
283284
}
284285

285-
private Credentials getSessionCredentials(AWSCredentialsProvider credentialsProvider, String roleArn,
286-
String roleSessionName, String region) {
286+
private Credentials getSessionCredentials(AwsCredentialsProvider credentialsProvider, String roleArn,
287+
String roleSessionName, String region) throws URISyntaxException {
287288
Credentials sessionCredentials = null;
288289
if (ServiceReferenceHolder.getInstance().isRedisEnabled()) {
289290
Object previousCredentialsObject = new RedisCacheUtils(ServiceReferenceHolder.getInstance().getRedisPool())
@@ -295,30 +296,30 @@ private Credentials getSessionCredentials(AWSCredentialsProvider credentialsProv
295296
sessionCredentials = CredentialsCache.getInstance().getCredentialsMap().get(roleSessionName);
296297
}
297298
if (sessionCredentials != null) {
298-
long expirationTime = sessionCredentials.getExpiration().getTime();
299+
long expirationTime = sessionCredentials.expiration().toEpochMilli();
299300
long currentTime = System.currentTimeMillis();
300301
long timeDifference = expirationTime - currentTime;
301302
if (timeDifference > 1000) {
302303
return sessionCredentials;
303304
}
304305
}
305-
AWSSecurityTokenService awsSTSClient;
306+
StsClient awsSTSClient;
306307
if (StringUtils.isEmpty(region)) {
307-
awsSTSClient = AWSSecurityTokenServiceClientBuilder.standard()
308-
.withCredentials(credentialsProvider)
308+
awsSTSClient = StsClient.builder()
309+
.credentialsProvider(credentialsProvider)
309310
.build();
310311
} else {
311-
awsSTSClient = AWSSecurityTokenServiceClientBuilder.standard()
312-
.withCredentials(credentialsProvider)
313-
.withEndpointConfiguration(new EndpointConfiguration("https://sts." + region + ".amazonaws.com",
314-
region))
312+
awsSTSClient = StsClient.builder()
313+
.credentialsProvider(credentialsProvider)
314+
.endpointOverride(new URI("https://sts." + region + ".amazonaws.com"))
315315
.build();
316316
}
317-
AssumeRoleRequest roleRequest = new AssumeRoleRequest()
318-
.withRoleArn(roleArn)
319-
.withRoleSessionName(roleSessionName);
320-
AssumeRoleResult assumeRoleResult = awsSTSClient.assumeRole(roleRequest);
321-
sessionCredentials = assumeRoleResult.getCredentials();
317+
AssumeRoleRequest roleRequest = AssumeRoleRequest.builder()
318+
.roleArn(roleArn)
319+
.roleSessionName(roleSessionName)
320+
.build();
321+
AssumeRoleResponse assumeRoleResult = awsSTSClient.assumeRole(roleRequest);
322+
sessionCredentials = assumeRoleResult.credentials();
322323
if (ServiceReferenceHolder.getInstance().isRedisEnabled()) {
323324
new RedisCacheUtils(ServiceReferenceHolder.getInstance().getRedisPool())
324325
.addObject(roleSessionName, sessionCredentials);
@@ -416,7 +417,7 @@ public String getResourceName() {
416417
}
417418

418419
public int getResourceTimeout() {
419-
return resourceTimeout;
420+
return (int) resourceTimeout.toMillis();
420421
}
421422

422423
public void setAccessKey(String accessKey) {
@@ -448,7 +449,7 @@ public void setResourceName(String resourceName) {
448449
}
449450

450451
public void setResourceTimeout(int resourceTimeout) {
451-
this.resourceTimeout = resourceTimeout;
452+
this.resourceTimeout = Duration.ofMillis(resourceTimeout);
452453
}
453454

454455
public void setIsContentEncodingEnabled(boolean isContentEncodingEnabled) {

‎components/apimgt/org.wso2.carbon.apimgt.gateway/src/main/java/org/wso2/carbon/apimgt/gateway/mediators/CredentialsCache.java

+1-1
Original file line numberDiff line numberDiff line change
@@ -17,7 +17,7 @@
1717
*/
1818
package org.wso2.carbon.apimgt.gateway.mediators;
1919

20-
import com.amazonaws.services.securitytoken.model.Credentials;
20+
import software.amazon.awssdk.services.sts.model.Credentials;
2121

2222
import java.util.HashMap;
2323
import java.util.Map;

0 commit comments

Comments
 (0)