OAuth 2.0, AWS Cognito, Spring Boot, Spring Security, Java, Testing

AWS Cognito - Mock users in Spring Boot integration tests

A tutorial about how to mock the security context of an AWS Cognito user in Spring Boot integration tests.

Author Martin
Martin24th March 2021

In this tutorial we will show you, how you can mock different AWS Cognito users in your Spring Boot integration tests. This tutorial assumes that you have a Spring Boot service which is configured as an OAuth 2.0 client, with AWS Cognito as the OAuth 2.0 provider. Moreover, we assume that you use Spring Security to annotate your controllers to manage access control on your routes. If you are looking for an article that shows you how to setup Spring Boot with AWS Cognito, you might want to consider this one.

Example

Lets consider this example controller with two protected routes.

1@RestController
2public class CommentController {
3
4  @PreAuthorize("hasRole('COMMENTER')")
5  @GetMapping("/comments")
6  public String getComments() {
7    return "comments";
8  }
9
10  @PreAuthorize("hasRole('MODERATOR')")
11  @DELETE("/comments")
12  public void deleteComments() {
13    // delete comments
14  }
15
16}

In the example above a user with the role "COMMENTER" is allowed to read comments and a user with role "MODERATOR" is allowed to delete comments. Let's see how we can use Spring Boot integration tests to verify, that the access control rules are working as expected.

Integration Tests with Spring Boot

Spring Boot provides an annotation called @SpringBootTest that creates an application context for our test cases. This makes it very convenient to simulate the running application and test its behaviour.

To properly test the two endpoints of the CommentController we need to call the endpoints with different users, one with the COMMENTER role and the other one with the MODERATOR role. For this purpose Spring Boot provides us another annotation called @WithMockUser. By annotating your test method with @WithMockUser, your test case will run with a SecurityContext that has been populated with a UsernamePasswordAuthenticationToken. This is very handy if your application relies on UsernamePasswordAuthenticationToken, which is the default authentication implementation in a Spring Boot application.

With AWS Cognito we have a security context that stores the current user as an OAuth2AuthenticationToken. Additionally, the roles that you configure for your users on AWS Cognito are encoded in an attribute called "cognito:groups". To simulate such a logged in user, we can create our own test annotation similar to @WithMockUser.

Custom security context factory

First, we need a custom security context factory, that allows us to inject our mocked OAuth2AuthenticationToken into the test security context.

1package com.example.support.security;
2
3import java.util.Collections;
4import java.util.HashMap;
5import java.util.List;
6import java.util.Map;
7import org.springframework.security.core.GrantedAuthority;
8import org.springframework.security.core.context.SecurityContext;
9import org.springframework.security.core.context.SecurityContextHolder;
10import org.springframework.security.oauth2.client.authentication.OAuth2AuthenticationToken;
11import org.springframework.security.oauth2.core.user.DefaultOAuth2User;
12import org.springframework.security.oauth2.core.user.OAuth2User;
13import org.springframework.security.oauth2.core.user.OAuth2UserAuthority;
14import org.springframework.security.test.context.support.WithSecurityContextFactory;
15
16public class WithMockAwsCognitoUserSecurityContextFactory
17    implements WithSecurityContextFactory<WithAwsCognitoUser> {
18
19    @Override
20    public SecurityContext createSecurityContext(final WithAwsCognitoUser withAwsCognitoUser) {
21        Map<String, Object> attributes = new HashMap<>();
22        attributes.put("sub", withAwsCognitoUser.username());
23        attributes.put("cognito:username", withAwsCognitoUser.username());
24        attributes.put("email", withAwsCognitoUser.email());
25        attributes.put("cognito:groups", withAwsCognitoUser.roles());
26
27        List<GrantedAuthority> authorities = Collections.singletonList(
28            new OAuth2UserAuthority("ROLE_USER", attributes));
29        OAuth2User user = new DefaultOAuth2User(authorities, attributes, "sub");
30
31        SecurityContext context = SecurityContextHolder.createEmptyContext();
32        context.setAuthentication(
33            new OAuth2AuthenticationToken(user, authorities, "client-registration-id"));
34        return context;
35    }
36}

Lines 21:25

We take the values from the annotation e.g. @WithAwsCognitoUser(username = "moderator", roles = "[MODERATOR]", email = "moderator@test.com") and map them into the attributes object, which looks exactly like the attributes that you get from AWS Cognito OpenID connect tokens.

Lines 27:29

Here we create a new OAuth2User with the attributes object.

Lines 31:33

We initialize an empty security context for the test case and set the currently authenticated user with an OAuth2AuthenticationToken object.

Custom security annotation

This annotation allows us to define the currently logged in user on each test case.

1package com.example.support.security;
2
3import java.lang.annotation.Documented;
4import java.lang.annotation.ElementType;
5import java.lang.annotation.Inherited;
6import java.lang.annotation.Retention;
7import java.lang.annotation.RetentionPolicy;
8import java.lang.annotation.Target;
9import org.springframework.security.test.context.support.TestExecutionEvent;
10import org.springframework.security.test.context.support.WithSecurityContext;
11
12@Target({ElementType.METHOD, ElementType.TYPE})
13@Retention(RetentionPolicy.RUNTIME)
14@Inherited
15@Documented
16@WithSecurityContext(
17    factory = WithMockAwsCognitoUserSecurityContextFactory.class
18)
19public @interface WithAwsCognitoUser {
20    String username() default "";
21
22    String roles() default "";
23
24    String email() default "test@test.com";
25
26    TestExecutionEvent setupBefore() default TestExecutionEvent.TEST_METHOD;
27}

How to use the annotion?

This example test case shows how you can define the logged in user on each test case and assert the response from the controller.

1@SpringBootTest
2class CommentControllerTest {
3
4    @Autowired
5    protected MockMvc mockMvc;
6
7    @Test
8    @WithAwsCognitoUser(roles = "[COMMENTER]")
9    void givenRequestFromCommenter_getComments_shouldBeSuccessful() {
10        mockMvc.perform(MockMvcRequestBuilders.get("/comments"))
11            .andExpect(status().isOk())
12    }
13
14    @Test
15    @WithAwsCognitoUser(roles = "[MODERATOR]")
16    void givenRequestFromModerator_deleteComments_shouldBeSuccessful() {
17        mockMvc.perform(MockMvcRequestBuilders.delete("/comments"))
18            .andExpect(status().isOk())
19    }
20
21    @Test
22    @WithAwsCognitoUser(roles = "[COMMENTER]")
23    void givenRequestFromCommenter_deleteComments_shouldBeForbidden() {
24        mockMvc.perform(MockMvcRequestBuilders.delete("/comments"))
25            .andExpect(status().is4xxClientError())
26    }
27}

Tipp: Create domain specific annotations

In order to reduce duplication and improve readability of your test cases, you can create domain specific annotations.

1package com.example.support.security;
2
3@Target({ElementType.METHOD, ElementType.TYPE})
4@Retention(RetentionPolicy.RUNTIME)
5@Inherited
6@WithAwsCognitoUser(username = "moderator", roles = "[MODERATOR]", email = "moderator@test.com")
7public @interface WithModerator {
8}
1@Test
2@WithModerator
3void givenRequestFromModerator_deleteComments_shouldBeSuccessful() {
4    mockMvc.perform(MockMvcRequestBuilders.delete("/comments"))
5        .andExpect(status().isOk())
6}

Conclusion

With only a few lines of code Spring Boot allows you to mock an authorized AWS Cognito user in your test cases. This should give you additional confidence in your integration tests and security annotations. If you want to know how you can map AWS Cognito roles to Spring Security granted authorities, then check out this blog post.