Example: Implementing a Custom Credentials Provider that uses the Post-SAML Workflow Hook

The following Java example implements the AD FS Lambda credentials provider with the post-SAML workflow hook. The AWS Lambda service exchanges the temporary credentials for AWS credentials, which are returned to the connector and are then used to access Athena.

import com.simba.athena.iamsupport.model.CredentialsHolder;
import com.simba.athena.iamsupport.plugin.AdfsCredentialsProvider;

/**
 * The ADFS Lambda Credential Provider.
 */
public class AdfsLambdaCredentialProvider extends AdfsCredentialsProvider
{
    /**
     * Perform custom actions after temporary credentials are available.
     * 
     * @param username         Your AD FS user name.
     * @param samlAssertion    The Base64-encoded SAML assertion.
     * @param credentials      The temporary credentials from the assumeRoleWithSAML 
     *                         request.
     * 
     * @return The CredentialsHolder wrapper object that holds the new credentials. 
     *         Returning null if no post SAML action is performed.
     */
    @Override
    protected CredentialsHolder performPostSAMLAction(
        String username, 
        String samlAssertion, 
        CredentialsHolder credentials) throws SdkClientException 
    {
        // Perform post-SAML work flow here.
        //
        //
        // AWSLambda client;
        try
        {
            // === Example:===
            // === Temporary Credentials from ADFS to BasicSessionCredentials: ===
            // 
            AWSCredentials cred = new BasicSessionCredentials(
                credentials.getAccessKeyId(),
                credentials.getSecretAccessKey(), 
                credentials.getSessionToken());
             
            // === Example===
            // === Example of a custom credentials provider communicating against a
            // === Lambda server with the temporaray credentials acquired from STS 
            // === through the AD FS workflow. Then, it calls AthenaIAMRoleTest, a 
            // === custom Lambda function implemented on the AWS Lambda service, 
            // === using the tempoaray credentials to acquire a new set of 
            // === credentials from the Lambda function. Finally, it returns the new
            // === credentials wrapped with CredentialsHolder class back to the 
            // === AdfsCredentialsProvider to consume. 
            // 
            // // Building a lambda client.
            // AWSCredentialsProvider basic = new AWSStaticCredentialsProvider(cred);
            // AWSLambdaClientBuilder builder = AWSLambdaClientBuilder
            //    .standard()
            //    .withRegion(Regions.US_EAST_1)
            //    .withCredentials(basic);
            
            // client = builder.build();
            // 
            // Result result;
            // 
            // // Serialize a payload object and send it to the Lambda function 
            // // AthenaIAMRoleTest.
            // Payload p = new Payload(username, samlAssertion, true);
            // String payload = s_objectMapper.writeValueAsString(p);
            // System.out.println("PAYLOAD -> " + payload);
            // InvokeRequest invReq = new InvokeRequest()
            //    .withFunctionName("AthenaIAMRoleTest")
            //    .withPayload(payload)
            //    .withRequestCredentialsProvider(basic);
            // 
            // InvokeResult fnret = client.invoke(invReq);
            // String err = fnret.getFunctionError();
            // if (err != null)
            // {
            //     System.err.println("ERROR -> " + err);
            //  // Then we should throw an exception.
            // }
            // else
            // {
            //  // Process new credentials/roles returned by the Lambda function.
            //     ByteBuffer bb = fnret.getPayload();
            //     String data = CHARSET.decode(bb).toString();
            //     System.out.println("DATA -> " + data);
            //     JsonNode actualObj = s_objectMapper.readTree(data);
            //     result = new Result(actualObj);
            //     return CredentialsHolder.newInstance(result.getCredentials());
            // }
            
            
            // === Return the new credentials back to the connector: ===
            // 
            // return CredentialsHolder.newInstance(
            //     new BasicAWSCredentials(
            //     newAccessKey, 
            //     newSecretKey));
            // 
            // 
            // === Or, if the new credentials have expiration information ===
            // === Creates a new instance of the CredentialsHolder. === 
            // === @param credentials    The AWSCredentials. === 
            // === @param expiration     The credential expiration date. === 
            // === @return               The CredentialsHolder. === 
            //
            // return CredentialsHolder.newInstance(
            //    <AWSCredentials> credentials, 
            //    <Date> expiration);
        }
        catch (SdkClientException e)
        {
            // === Exception handling. ===
        }
    }
}