#HOWTO: JAX-RS user-based rate-limiting with JSR-375

Recently I had the requirement to limit the access to specific JAX-RS endpoints and to keep track of the user’s current amount of API calls. To solve this problem I asked Adam Bien (@AdamBien) in his monthly Airhacks Q&A about this requirement and he gave me a hint for a possible solution while using the ContainerRequestFilter interface for filtering the access to JAX-RS resources.

In this blog post I’ll show you how to implement a simple user-based rate-limiting for a HTTP REST endpoint. I’ll deploy the application to Payara and make use of Payara’s in-memory H2 database to store the available users and their current/max API budget.  To secure the resource I’ll use the JSR-375 (Java EE Security API) and define an IdentityStore based on a database.

Let’s start with the JAX-RS resource we want to secure. In our example, I am defining one resource which is available under /resources/stocks and returns a hard-coded stock price for Google’s Alphabet share:

@Path("stocks")
public class StockResource {

    @GET
    @RolesAllowed("USER")
    public Response getAllStocks() {
        JsonObject json = Json.createObjectBuilder()
          .add("name", "Alphabet Inc.")
          .add("price", 1220.5)
          .build();
        return Response.ok(json.toString()).build();
    }
}
@ApplicationPath("/resources")
public class JAXRSApplication extends Application {
}

I annotated the method with @RolesAllowed to only allow users with the role “USER” to access this endpoint. Unkown users will therefor get a 401 Unauthorized HTTP Status and won’t have access to this endpoint.

To define the IdentityStore and the authentication mechanism I make use of some annotations from the new Java EE Security API and define a configuration class:

@DatabaseIdentityStoreDefinition(
        dataSourceLookup = "jdbc/__default",
        callerQuery = "SELECT password FROM user WHERE username = ?",
        groupsQuery = "SELECT role FROM user_roles where username = ?",
        hashAlgorithm = Pbkdf2PasswordHash.class
)
@BasicAuthenticationMechanismDefinition
@DeclareRoles({"USER", "ADMIN"})
@ApplicationScoped
public class ApplicationSecurityConfig {
}

The first annotation is responsible for setting up the IdentityStore based on the provided configurations. I am using Payara’s default datasource and define the required SQL statements for validating an incoming user and for retrieving its roles. In addition, I am selecting the hash algorithm which should be used for hashing the passwords (Pbkdf2PasswordHash is the default algorithm).

The second annotation is for configuring the authentication mechanism. For simplification is decided to use a @BasicAuthenticationMechanismDefinition  but there is also a form-based or a custom mechanism available.

The third annotation is used for defining the available roles for the application’s context.

For storing the user information in the database I modeled two JPA entities: User and UserRoles (the table design is not a best practice as I denormalized the database structure, so don’t use this at home):

@Entity
public class User {

    @Id
    @GeneratedValue(strategy = GenerationType.IDENTITY)
    private Long id;
    private String username;
    private String password;
    private long amountOfApiCalls = 0;
    private long maxApiCallsPerMinute = 10;

    // getter & setter

}
@Entity
@Table(name = "user_roles")
public class UserRoles {

    @Id
    @GeneratedValue(strategy = GenerationType.IDENTITY)
    private Long id;
    private String username;
    private String role;

    // getter & setter

}

The User table contains two columns for storing the current amount of API calls and the maximum amount of API calls within a minute.

For a quick setup of pre-defined users and roles I am inserting some users on the application’s startup:

@Singleton
@Startup
public class DatabaseSetup {

    @Inject
    Pbkdf2PasswordHash pbkdf2PasswordHash;

    @Resource(lookup = "java:comp/DefaultDataSource")
    DataSource dataSource;

    @PostConstruct
    public void initDefaultUser() {

        executeUpdate("INSERT INTO user (id, username, password, amountOfApiCalls, maxApiCallsPerMinute) " 
           + "VALUES (1, 'rieckpil', '" 
           + this.pbkdf2PasswordHash.generate("HelloWorld".toCharArray()) + "', 0, 10)");

         // some more user

        executeUpdate("INSERT INTO user_roles (id, username, role) VALUES (1, 'rieckpil', 'USER')");
        executeUpdate("INSERT INTO user_roles (id, username, role) VALUES (2, 'rieckpil', 'ADMIN')");
      
        // some more roles

        System.out.println("Successfully initialized database with default user");
    }

    private void executeUpdate(String query) {
        try {
            this.dataSource.getConnection().createStatement().executeUpdate(query);
        } catch (SQLException e) {
            e.printStackTrace();
            throw new RuntimeException("unable to setup database");
        }
    }
}

The Pbkdf2PasswordHash is available through CDI an can be used anywhere in your code to create a new hash or verify an incoming hash.

For re-setting the API budget for every user, I am using an EJB timer to schedule this task every minute:

@Singleton
@Startup
public class ApiBudgetRefresher {

    @PersistenceContext
    EntityManager entityManager;

    @Schedule(minute = "*", hour = "*", persistent = false)
    public void updateApiBudget() {
        System.out.println("-- refreshing API budget for all users");

        List<User> userList = entityManager.createQuery("SELECT u FROM User u", User.class).getResultList();
        for (User user : userList) {
            user.setAmountOfApiCalls(0);
        }

        System.out.println("-- successfully refreshed API budget for all users");
    }
}

Now comes the interesting part for the implementation of the ContainerRequestFilter interface:

@Provider
public class RateLimitingFilter implements ContainerRequestFilter {

    @PersistenceContext
    EntityManager entityManager;

    @Transactional
    @Override
    public void filter(ContainerRequestContext requestContext) throws IOException {

        SecurityContext securityContext = requestContext.getSecurityContext();
        String username = securityContext.getUserPrincipal().getName();

        User user = entityManager.createQuery("SELECT u FROM User u WHERE u.username=:username", User.class).setParameter(
                "username", username).getSingleResult();

        if (user.getAmountOfApiCalls() >= user.getMaxApiCallsPerMinute()) {
            requestContext.abortWith(Response.status(Response.Status.TOO_MANY_REQUESTS).build());
        }

        user.setAmountOfApiCalls(user.getAmountOfApiCalls() + 1);
        System.out.println(user);
    }
}

The class provides an implementation for the filter() method and is recognized from JAX-RS through the @Provider annotation. To update all current users I am injecting the EntityManager to this filter. As this is not an EJB we need the  @Transactional annotation here to explicitly wrap this code in a transaction. The code will first retrieve the username from the Principal an will then check if the current budget of the user allows a new API call. If not the filter will return with a HTTP Status code 429 – Too Many Requests.

Accessing the endpoint with the credentials for rieckpil as a base64 encoded string less than eleven times will result in a valid result:

$ curl -i -H 'Authorization: Basic cmllY2twaWw6SGVsbG9Xb3JsZA==' http://localhost:8080/api-rate-limiting/resources/stocks
HTTP/1.1 200 OK
Server: Payara Server  5.182 #badassfish
X-Powered-By: Servlet/4.0 JSP/2.3 (Payara Server  5.182 #badassfish Java/Oracle Corporation/1.8)
Content-Type: text/plain
Content-Length: 39
X-Frame-Options: SAMEORIGIN

{"name":"Alphabet Inc.","price":1220.5}

After the tenth API call, the user will get the following result and has to wait one minute:

$ curl -i -H 'Authorization: Basic cmllY2twaWw6SGVsbG9Xb3JsZA==' http://localhost:8080/api-rate-limiting/resources/stocks
HTTP/1.1 429 Too Many Requests
Server: Payara Server  5.182 #badassfish
X-Powered-By: Servlet/4.0 JSP/2.3 (Payara Server  5.182 #badassfish Java/Oracle Corporation/1.8)
Content-Language:
Content-Type: text/html
Content-Length: 1100
X-Frame-Options: SAMEORIGIN

// error HTML page

You can try this example on your local machine as I provided a Docker image in the GitHub repository.

Happy rate-limiting,

Phil

Leave a comment

Your email address will not be published. Required fields are marked *