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. === } } }