ScottLogic / prompt-injection

Application which investigates defensive measures against prompt injection attacks on an LLM, with a focus on the exposure of external tools.
MIT License
15 stars 10 forks source link

Connect cognito userpool with Application Load Balancer #808

Closed chriswilty closed 4 months ago

chriswilty commented 7 months ago

Here is a SO post on integrating cognito with Application Load Balancer (ALB). It is straightforward, but requires the ALB to be secure, hence this issue is blocked by #806.

Also dependent on #807

Update - 02/04/2024

This turned out to be a heap more difficult than expected. ALB does not provide any way to verify an auth token in a request header or cookie, rather, ALB wants to provide the connection to your Identity Provider. This is in contrast to API Gateway which provides an explicit way to verify a provided auth token.

In theory we could allow ALB to provide the auth connector, however, that would mean using the Cognito hosted UI, which does not provide enough flexibility to allow matching our SpyLogic style. In contrast, AWS Amplify provides its own UI with greater ability to customise, plus, if that wasn't good enough we could provide our own UI and hook it up to the auth endpoints provided by Cognito. So, we have no choice but to initiate auth through the UI and simply verify the token at the load balancer...

... which is not possible! There is, however, another potential solution. We could use a CloudFront distribution as a (non-caching) proxy to the ALB with an Edge Lambda function to verify the token, and if it is valid, add a custom header to the request, containing a UUID generated during stack deployment which the ALB will therefore know about. The ALB can be configured with a simple rule to forward traffic to the target group (i.e. fargate service) only if it contains this header with the expected UUID value, and the default rule can be amended to return 401 Unauthorized. As this traffic is entirely within the AWS network, the token value should never leak to the outside world. The end result is that traffic not routed through cloudfront will always return 401, and traffic that is routed through cloudfront must have a valid token - which Amplify will take care of in the UI.