Implementing Ouath2 BFF Pattern with AWS Amplify and Cognito

Question

I am using AWS Cognito as my Authentication server and using Amplify UI library with react to handle authentication in my client. I am mainly using Amplify since it allows me to do more customisations in login windows that hosted UI.

I also want to have OAuth2 Backend For Frontend (BFF) with this setup, since it allows :

• Using Cognito confidential oauth client instead of public client
• Store tokens in backend instead of browser

So far I have been managed to redirect auth calls via my BFF, append SECURITY_HASH to requests, get tokens and send them back to frontend.

My Amplify config is as below:

      Amplify.configure({
        Auth: {
          Cognito: {
            userPoolClientId: 'clientID',
            userPoolId: 'ap-southeast-xxxxx',
            userPoolEndpoint: 'http://localhost:7081/cognito'
          }
        }
      });

And my BFF looks like this:

@PostMapping(
        path = "/cognito",
        consumes = "application/x-amz-json-1.1"
)
public Object postCognitoResponses(
        @RequestBody String data,
        @RequestHeader("X-Amz-Target") String target
) throws URISyntaxException, JsonProcessingException {
    Map<String, Object> requestMap = objectMapper.readValue(data, Map.class);
    CognitoIdentityProviderClient cognitoClient = CognitoIdentityProviderClient.builder()
            .region(Region.of("ap-southeast-2"))
            .build();
    if (target.equals("AWSCognitoIdentityProviderService.InitiateAuth")) {
        Map<String, String> authParameters = (Map<String, String>) requestMap.get("AuthParameters");
        String username = authParameters.get("USERNAME");  //Not available during "REFRESH_TOKEN_AUTH" flow, need a workaround
        String secretHash = getPasswordHash(username, requestMap.get("ClientId").toString(),"");
        authParameters.put("SECRET_HASH", secretHash);
        InitiateAuthRequest initiateAuthRequest = InitiateAuthRequest.builder()
                .authFlow(requestMap.get("AuthFlow").toString())
                .clientId(requestMap.get("ClientId").toString())
                .authParameters( authParameters)
                .build();
        InitiateAuthResponse initiateAuthResponse = cognitoClient.initiateAuth(initiateAuthRequest);
        AuthResponseDTO responseDTO = new AuthResponseDTO(initiateAuthResponse);
        return responseDTO;
    } else if (target.equals("AWSCognitoIdentityProviderService.RespondToAuthChallenge")) {
        String challengeName = (String) requestMap.get("ChallengeName");
        Map<String, String> challengeResponses = (Map<String, String>) requestMap.get("ChallengeResponses");

        String clientId = (String) requestMap.get("ClientId");
        String username = challengeResponses.get("USERNAME");
        String secretHash = getPasswordHash(username, clientId,"");
        challengeResponses.put("SECRET_HASH", secretHash);
        RespondToAuthChallengeRequest.Builder respondToAuthChallengeRequestBuilder = RespondToAuthChallengeRequest.builder()
                .challengeName(challengeName)
                .clientId(clientId)
                .challengeResponses(challengeResponses);
        if (requestMap.containsKey("Session")) {
            respondToAuthChallengeRequestBuilder.session(requestMap.get("Session").toString());
        }
        RespondToAuthChallengeRequest respondToAuthChallengeRequest = respondToAuthChallengeRequestBuilder.build();
        RespondToAuthChallengeResponse respondToAuthChallengeResponse = cognitoClient.respondToAuthChallenge(respondToAuthChallengeRequest);
        return new CognitoChallengeResponseDTO(respondToAuthChallengeResponse);
    } else if (target.equals("AWSCognitoIdentityProviderService.GetUser")) {

    } else if (target.equals("AWSCognitoIdentityProviderService.ForgotPassword")) {

    } else if (target.equals("AWSCognitoIdentityProviderService.RevokeToken")) {

    }
    return "test";

}

The response I send from BFF to frontend at the end of the flow will have tokens send from cognito. In my next step, I want to store the tokens in BFF and send a session to frontend via Set-Cookie instead of sending tokens. When making calls to APIs, I will proxy everything via BFF which will append the token after validating the session ID.

Currently, when I remove the token from response, Amplify returns an error and I don't seem to have any control over this behaviour. Does anyone know if there is a way to achieve this?

I am aware that I can use a custom store and save the tokens in backend. However, it will still send token to frontend before storing it, which I'd like to avoid if possible.