STSAssumeRoleAWSCredentialsProvider.java

package qty.aws;

import java.lang.reflect.InvocationHandler;
import java.lang.reflect.Method;
import java.lang.reflect.Proxy;

import com.amazonaws.auth.AWSCredentialsProvider;
import com.amazonaws.auth.AWSSessionCredentials;
import com.amazonaws.auth.BasicAWSCredentials;
import com.amazonaws.services.securitytoken.AWSSecurityTokenServiceClient;
import com.amazonaws.services.securitytoken.model.AssumeRoleRequest;
import com.amazonaws.services.securitytoken.model.AssumeRoleResult;
import com.amazonaws.services.securitytoken.model.Credentials;

public class STSAssumeRoleAWSCredentialsProvider implements AWSCredentialsProvider {

    private static final int SESSION_DURATION_SECONDS = 900;
    private String awsAccessKey;
    private String awsSecretKey;
    private String roleArn;
    private AWSSecurityTokenServiceClient client;

    public STSAssumeRoleAWSCredentialsProvider(AWSCredentialsProvider provider, String roleArn) {
        this(provider.getCredentials().getAWSAccessKeyId(), provider.getCredentials().getAWSSecretKey(), roleArn);
    }

    public STSAssumeRoleAWSCredentialsProvider(String accessKey, String secretKey, String roleArn) {
        this.awsAccessKey = accessKey;
        this.awsSecretKey = secretKey;
        this.roleArn = roleArn;
        client = new AWSSecurityTokenServiceClient(new BasicAWSCredentials(awsAccessKey, awsSecretKey));
        generateCredentials();
    }

    protected AWSCredentialsAdapter generateCredentials() {
        AssumeRoleResult assumeRoleResult = client.assumeRole(new AssumeRoleRequest()
                .withDurationSeconds(SESSION_DURATION_SECONDS)
                .withRoleSessionName("session-" + System.currentTimeMillis()).withRoleArn(this.roleArn));
        return new AWSCredentialsAdapter(assumeRoleResult.getCredentials());
    }

    @Override
    public AWSSessionCredentials getCredentials() {
        AWSSessionCredentials awsCredentials = (AWSSessionCredentials) Proxy.newProxyInstance(
                AWSSessionCredentials.class.getClassLoader(), new Class[] { AWSSessionCredentials.class },
                new AutoRenewCredentials(this));

        return awsCredentials;
    }

    @Override
    public void refresh() {
    }

    static class AutoRenewCredentials implements InvocationHandler {

        private STSAssumeRoleAWSCredentialsProvider provider;
        private AWSCredentialsAdapter credentials;

        public AutoRenewCredentials(STSAssumeRoleAWSCredentialsProvider provider) {
            this.provider = provider;
            this.credentials = provider.generateCredentials();
        }

        @Override
        public Object invoke(Object proxy, Method method, Object[] args) throws Throwable {
            synchronized (this) {
                if (!credentials.validate()) {
                    credentials = provider.generateCredentials();
                }
            }

            return method.invoke(credentials, args);
        }
    }

    static class AWSCredentialsAdapter implements AWSSessionCredentials {
        private static final int MARGIN_OF_SAFTY = 60 * 1000;
        private Credentials credentials;

        public AWSCredentialsAdapter(Credentials credentials) {
            this.credentials = credentials;
        }

        @Override
        public String getAWSAccessKeyId() {
            return credentials.getAccessKeyId();
        }

        @Override
        public String getAWSSecretKey() {
            return credentials.getSecretAccessKey();
        }

        @Override
        public String getSessionToken() {
            return credentials.getSessionToken();
        }

        public boolean validate() {
            return System.currentTimeMillis() + MARGIN_OF_SAFTY < credentials.getExpiration().getTime();
        }

    }

}