JWT Token Generation Using AWS Lambda

Introduction

Here at BeamWallet we tend to prefer micro-service architecture approach as opposed to archaic monolithic thinking. Recently we came across a simple/common issue which is generating JWT Token to allow our servers to communicate between each other.

Since all of our systems live in private VPCs  and in the interest of the DRY principle we decided to build a centralised service that would handle all of this. We built the solution using AWS Lambda and also AWS API Gateway.

We have 2 different ways of deploying and for secret storing. 

AWS  Elastic Beanstalk with AWS Secret Manager

Our EBS applications are Java Spring backends which have all the passwords and keys stored in AWS Secret Manager.  On start up of each app all passwords are loaded from AWS Secret Manager and loaded into the System Properties.

Our naming convention for naming secrets stored in the Secret Manager is <system>/<environment>
Eg core/dev  or common/staging (where common has commonly used variables across all systems)

AWS ECS with Ansible

We deploy our other apps to ECS using ansible, where the secrets are stored in our Ansible Vault, and after the passwords are stored in ECS' Task definition as Environment Variables.

Our Solution (In Lambda)

Setting up IAM

Add the following roles

Screen Shot 2018-08-10 at 5.19.15 pm.png

Add the Lambda Function

Screen Shot 2018-08-10 at 5.22.18 pm.png

Configure the Lambda function

Set the handler to org.srini.awslambda.examples.generatejwt.GenerateJWTFunctionHandler::handleRequest

Screen Shot 2018-08-10 at 5.24.47 pm.png

Private Key Retrieval

As mentioned above we store our secrets in 2 different locations so we wanted our function to be smart enough to know where to get the private key from. We have a list which contains which service to get which key from.

The easiest one to get the secrets from was the AWS Secrets using the aws-java-sdk. We have all our private keys with the same name so the naming convention was really easy. 

String endpoint = "secretsmanager.eu-central-1.amazonaws.com";
AwsClientBuilder.EndpointConfiguration config = new AwsClientBuilder.EndpointConfiguration(endpoint, region);
AWSSecretsManagerClientBuilder clientBuilder = AWSSecretsManagerClientBuilder.standard();
clientBuilder.setEndpointConfiguration(config);
AWSSecretsManager client = clientBuilder.build();
GetSecretValueRequest request = new GetSecretValueRequest().withSecretId(service + "/" + environment);
String res = client.getSecretValue(request).getSecretString();
HashMap map = new Gson().fromJson(res, new HashMap<String, String>().getClass());
return map.get(AWS_SECRET_KEYWORD).toString();

The hardest step was  getting the private key from the ECS configuration. To access we needed retrieve the task definition from within the container definition which is inside the service. Again we were following the same naming conventions across all of our services which made this task much easier.

String privateKey = "";
AmazonECSClientBuilder clientBuilder = AmazonECSClientBuilder.standard();
clientBuilder.setRegion(region);
AmazonECS client = clientBuilder.build();
String commonName = environment + "-" + service;
List<Service> services = client.describeServices(new DescribeServicesRequest().withCluster(commonName).withServices(commonName)).getServices();
if (service.isEmpty()) {
    throw new IllegalArgumentException("No cluster/service is found with " + commonName);
}
// will always return 1 or 0
String taskDefinition = services.get(0).getTaskDefinition();
List<ContainerDefinition> containerDefinitions = client.describeTaskDefinition(new DescribeTaskDefinitionRequest().withTaskDefinition(taskDefinition)).getTaskDefinition().getContainerDefinitions();
if (containerDefinitions.isEmpty()) {
    throw new IllegalArgumentException("No container definitions is found with " + commonName);
}
// will always return 1 or 0
for (KeyValuePair keyPair : containerDefinitions.get(0).getEnvironment()) {
    if (keyPair.getName().equals(AWS_ECS_KEYWORD)) {
        privateKey = keyPair.getValue();
    }
}
return privateKey;

Key Generation

We are using the Java Security Library to generate the token, io.jsonwebtoken.jjwt to generate the JWT token, and using Lambda variables to set how long we want the time out to be. 

public class TokenGenerator {
    private final int DEFAULT_TIME_IN_MINUTES = 10;
    private final String timeInMinutes = System.getenv("TIME_IN_MINUTES");
    private final String TOKEN_TYPE = "typ";
    private final String RND = "rnd";

    public String generateToken(String secretKey, String subject) throws NoSuchAlgorithmException, InvalidKeySpecException {
        int TIME_IN_MINUTES = DEFAULT_TIME_IN_MINUTES;
        if (!com.amazonaws.util.StringUtils.isNullOrEmpty(timeInMinutes)) {
            TIME_IN_MINUTES = Integer.parseInt(timeInMinutes);
        }
        String jwtToken = "";
        KeyFactory keyFactory = KeyFactory.getInstance("RSA");
        PrivateKey privateKey = keyFactory.generatePrivate(new PKCS8EncodedKeySpec(Base64.decodeBase64(secretKey)));
        jwtToken = Jwts.builder()
                .claim(TOKEN_TYPE, "service")
                .claim(RND, Date.from(LocalDateTime.now().atZone(ZoneId.systemDefault()).toInstant()))
                .setSubject(subject)
                .setIssuedAt(new Date())
                .setExpiration(Date.from(LocalDateTime.now().plusMinutes(TIME_IN_MINUTES).atZone(ZoneId.systemDefault()).toInstant()))
                .signWith(SignatureAlgorithm.RS512, privateKey)
                .compact();
        return jwtToken;
    }
}

API Gateway

The api gateway was used just as a very simple means to ensure we were able to build/develop and also would be able to use on all environments. To set up the API Gateway

Step 1: Create API

Screen Shot 2018-08-10 at 4.33.02 pm.png

Step 2: Link to Lambda

 

Screen Shot 2018-08-10 at 4.33.45 pm.png

Step 3: Deploy the API

Screen Shot 2018-08-10 at 5.13.00 pm.png