Dropwizard Deep Dive – Part 3: Multitenancy

Hello once again! This is Part 3 of a three-part series on extending Dropwizard to have custom authentication, authorization, and multitenancy. In Part 1, we set up custom authentication in Dropwizard, and in Part 2, we extended that to have role-based authorization. For this final part, we are going to diverge slightly and tackle the related but different concept of multitenancy.

Just to get our definitions straight, I want to explain what “multitenancy” means in this post. It refers to extending our API in such a way that users can only access data from a particular high-level organization they have access to within an API. This problem comes up often when providing software as a service; each customer needs to use the same API, but they should only have access to their own data.

In Parts 1 and 2, you saw that Java and Dropwizard have some built-in annotations that allow you to declaratively add auth to each of your endpoints. In this part, we are going to build a custom solution on top of what we’ve already created to enforce secure multitenancy in our API. To accomplish this, we will leverage several built-in abstractions in Jersey and Hibernate. In fact, pretty much all of the things we are going to do here are not specific to Dropwizard and could be leveraged in any app using Hibernate and Jersey.

Speaking of Hibernate, this tutorial assumes you are using that—and only that—to manage database access. The entire tutorial is based on using the filter feature built into Hibernate. I know Dropwizard has first-class support for JDBI and libraries for other data stores, but if you are using any of those, you are on your own as far as adapting and extending this solution goes.

Adding Tenant Relations to DB

The first step to adding multitenancy is to add the DB relations. We need to provide a TenantModel to represent a company. Each entity restricted to a company (UserModel and WidgetModel) also needs to reference a TenantModel. These relations are easy to add with Hibernate.


@Entity
@Table(name = "Tenant")
public class TenantModel {
  @Id
  @GeneratedValue(strategy = GenerationType.IDENTITY)
  private Long id;

  private String name;

  public Long getId() {
    return id;
  }

  public void setId(Long id) {
    this.id = id;
  }

  public String getName() {
    return name;
  }

  public void setName(String name) {
    this.name = name;
  }
}

@Entity
@Table(name = "Widget")
public class WidgetModel {
  @Id
  @GeneratedValue(strategy = GenerationType.IDENTITY)
  private Integer id;

  private String name;

  @Enumerated(EnumType.STRING)
  private WidgetScope scope;

  @ManyToOne
  @JoinColumn(name = "tenantId")
  private TenantModel tenant;

Basing Multitenancy On Path Params

Now that we have a tenant model and all the necessary associations, we can start up our app again. While reading our data will work as before, you’ll notice that all writes now get an error. We need to associate a tenant model when we write users or widgets, and we aren’t doing that.

In order to associate a tenant, we need to get a tenant ID from somewhere. In our case, that somewhere is going to be the URL path. The path is a nice place, because unlike a body, it is a part of every request–and unlike a query param, it is not optional. This makes it easy to send and gives us an automatic 404 if we forget it.

Adding it to our resource paths is easy:


@Path("/tenants/{tenantId}/widgets")
@Produces(MediaType.APPLICATION_JSON)
public class WidgetResource {
  ...

Using Thread Locals to Store Tenant

Now that we have a tenant ID on every request, we need to use it to fetch and store our tenant model so that we can write that back out to all associated entities. Normally, we would use a Jersey filter to do that. In this case, however, we are going to use another Jersey feature called a RequestEventListener. I’ll explain why we use this approach in a little bit.

To implement a RequestEventListener, we first need to implement an ApplicationEventListener. This outer listener will set up our other listener to listen for each incoming request.


@Provider
@Priority(10000) // This needs to be the last listener to run
public class MultitenancyApplicationListener implements ApplicationEventListener {
  private TenantDAO tenantDAO;

  public MultitenancyApplicationListener(TenantDAO tenantDAO) {
    this.tenantDAO = tenantDAO;
  }

  @Override
  public void onEvent(ApplicationEvent event) {}

  @Override
  public RequestEventListener onRequest(RequestEvent requestEvent) {
    return new MultitenancyRequestListener(tenantDAO);
  }
}

With that in place, we can implement our actual RequestEventListener which will do the real work.


public class MultitenancyRequestListener implements RequestEventListener {
  private TenantDAO tenantDAO;
  private Cache<Long, TenantModel> tenants;

  public MultitenancyRequestListener(TenantDAO tenantDAO) {
    this.tenantDAO = tenantDAO;

    tenants = CacheBuilder.newBuilder()
      .maximumSize(1000)
      .expireAfterWrite(1, TimeUnit.HOURS)
      .build();
  }

  @Override
  public void onEvent(RequestEvent event) {
    if (event.getType() == RequestEvent.Type.RESOURCE_METHOD_START) {
      try {
        Long tenantId = Long.valueOf(event.getContainerRequest().getUriInfo().getPathParameters().getFirst("tenantId"));

        TenantModel tenant = tenants.get(tenantId, () -> tenantDAO.getTenant(tenantId).get());

        TenantRequestData.tenant.set(tenant);
      } catch (Exception e) {
        throw new WebApplicationException("Tenant not found", FORBIDDEN);
      }
    }
  }
}

There is a lot going on here, so let’s break it down.

In our onEvent method, we first check the event we are getting notified of in the RESOURCE_METHOD_START event. There are a number of different events that get fired in the event listener lifecycle, but right now, this is the only one in which we are interested. Once we get a RESOURCE_METHOD_START event, we extract the tenantId from the URL path, and use that to do a lookup of the tenant data. All of these lookups are wrapped in a cache, because it will be running on every single request and get lots of duplicate data. Assuming we load a tenant successfully, we set a global ThreadLocal with the tenant data for this request so we can access it from anywhere in the request.

Now that we have our tenant in our global ThreadLocal, doing writes is easy. We simply pull the tenant out in any resource methods that need it, and we are guaranteed that we have the right one for the request.


public WidgetModel createWidget(Widget widget) {
    WidgetModel widgetModel = new WidgetModel();
    widgetModel.setName(widget.getName());
    widgetModel.setScope(widget.getScope());
    widgetModel.setTenant(TenantRequestData.tenant.get());

    return persist(widgetModel);
  }

Using Hibernate Filters to Enforce Multitenancy in DB Reads

Now that we can write entities with the correct tenant, we need to handle reads as well.

We’ll use Hibernate filters as the basis of our read multitenancy solution. Filters allow us to inject restrictions into the “where” clause of any query on a particular object in a particular session. Once we add a tenant relation to the WidgetModel from our example (we can assume a tenant in this case is a company), then it is easy to add such a filter.


@Entity
@Table(name = "Widget")
@FilterDef(
  name = "restrictToTenant",
  defaultCondition = "tenantId = :tenantId",
  parameters = @ParamDef(name = "tenantId", type = "long")
)
@Filter(name = "restrictToTenant")
public class WidgetModel {
  ...
}

We will also want to add a tenant to the UserModel, which means it is easiest to move the @FilterDef annotation to a package-info.java file. Combined with the default condition, this package level filter def is now easy to add to any entity.


@FilterDef(
  name = "restrictToTenant",
  defaultCondition = "tenantId = :tenantId",
  parameters = @ParamDef(name = "tenantId", type = "long")
)
package dao.entities;

import org.hibernate.annotations.FilterDef;
import org.hibernate.annotations.ParamDef;

We’ll also have to configure the Hibernate bundle in our app class to load our whole package and see package level filters.


private final ScanningHibernateBundle<ExampleConfig> hibernate = new ScanningHibernateBundle<ExampleConfig>("dao.entities") {
    @Override
    protected void configure(Configuration configuration) {
      // Register package so global filters in package-info.java get seen.
      configuration.addPackage("dao.entities");
      super.configure(configuration);
    }

    @Override
    public PooledDataSourceFactory getDataSourceFactory(ExampleConfig config) {
      return config.getDatabaseConfig();
    }
  };

Now we have some Hibernate filters on all the appropriate entities, but if you run the app, you will see nothing has changed. This is because Hibernate filters are disabled by default and need to be manually enabled for every individual Hibernate session.

Since we are using Dropwizard’s @UnitOfWork annotation, we already have a separate session that gets set up for each request. So we need to add some code that will run before each request, after the unit of work annotation, to enable our multitenancy filter. Fortunately, the RequestEventListener we set up earlier meets all of these criteria. In particular, we can use the @Priority annotation to guarantee it runs after @UnitOfWork has set up a database session, which is why we used an event listener instead of the simpler filter.

Since we already have a tenant ID in our listener, is is easy to amend it to enable our Hibernate filter.


public class MultitenancyRequestListener implements RequestEventListener {
  private TenantDAO tenantDAO;
  private SessionFactory sessionFactory;
  private Cache<Long, TenantModel> tenants;

  public MultitenancyRequestListener(TenantDAO tenantDAO, SessionFactory sessionFactory) {
    this.tenantDAO = tenantDAO;
    this.sessionFactory = sessionFactory;

    tenants = CacheBuilder.newBuilder()
      .maximumSize(1000)
      .expireAfterWrite(1, TimeUnit.HOURS)
      .build();
  }

  @Override
  public void onEvent(RequestEvent event) {
    if (event.getType() == RequestEvent.Type.RESOURCE_METHOD_START) {
      try {
        Long tenantId = Long.valueOf(event.getContainerRequest().getUriInfo().getPathParameters().getFirst("tenantId"));
        TenantModel tenant = tenants.get(tenantId, () -> tenantDAO.getTenant(tenantId).get());

        TenantRequestData.tenant.set(tenant);
        sessionFactory.getCurrentSession().enableFilter("restrictToTenant").setParameter("tenantId", tenantId);
      } catch (Exception e) {
        throw new WebApplicationException("Tenant not found", FORBIDDEN);
      }
    }
  }
}

And just like that, we have read multitenancy. You’ll see that if you hit the API endpoints with different tenant IDs now, you will get different responses.

Locking Down Between Tenants

Our multitenancy is working great, but we now have a problem: The roles we added in the last post are no longer suitable. If you browse the API, you’ll see that a MANAGER user of Tenant 1 can browse the top-secret widgets of Tenant 2 without a problem. This is because our roles themselves have no concept of different tenants. We need to lock this down so that our tenants are isolated and only accessible to their own users.

Fortunately, doing this will be pretty easy. Since we already have a tenant ID in our URL, and it’s attached to the user associated with the request, we simply need to compare the two when doing role checks. This is really easy to add to our security context.


public class CustomSecurityContext implements SecurityContext {
  private final CustomAuthUser principal;
  private final Long tenantId;
  private final SecurityContext securityContext;

  public CustomSecurityContext(CustomAuthUser principal, Long tenantId, SecurityContext securityContext) {
    this.principal = principal;
    this.tenantId = tenantId;
    this.securityContext = securityContext;
  }

  ...

  @Override
  public boolean isUserInRole(String role) {
    return principal.getTenantId().equals(tenantId) && role.equals(principal.getRole().name());
  }

Then, we just need to provide the security context the tenant ID from the path in CustomAuthFilter:


  @Override
  public void filter(ContainerRequestContext requestContext) throws IOException {
    Optional<CustomAuthUser> authenticatedUser;

    try {
      CustomCredentials credentials = getCredentials(requestContext);
      authenticatedUser = authenticator.authenticate(credentials);
    } catch (AuthenticationException e) {
      throw new WebApplicationException("Unable to validate credentials", Response.Status.UNAUTHORIZED);
    }

    if (authenticatedUser.isPresent()) {
      // Provide tenant ID from path to security context
      Long tenantId = parseTenantId(requestContext);
      SecurityContext securityContext = new CustomSecurityContext(authenticatedUser.get(), tenantId, requestContext.getSecurityContext());
      requestContext.setSecurityContext(securityContext);
    } else {
      throw new WebApplicationException("Credentials not valid", Response.Status.UNAUTHORIZED);
    }
  }

And the tenant ID from the current user ID to the authenticated user in CustomAuthenticator:


  @Override
  @UnitOfWork
  public Optional<CustomAuthUser> authenticate(CustomCredentials credentials) throws AuthenticationException {
    CustomAuthUser authenticatedUser = null;
    Optional<UserModel> user = userDAO.getUser(credentials.getUserId());

    if (user.isPresent()) {
      UserModel userModel = user.get();
      Optional<TokenModel> token = tokenDAO.findTokenForUser(userModel);

      if (token.isPresent()) {
        TokenModel tokenModel = token.get();

        if (tokenModel.getId().equals(credentials.getToken())) {
          // Pass the user's tenant ID
          authenticatedUser = new CustomAuthUser(userModel.getName(), userModel.getTenant().getId(), userModel.getRole());
        }
      }
    }

    return Optional.fromNullable(authenticatedUser);
  }

We now have secure, isolated multitenancy for both reads and writes in our API, just by having a logged-in user and a path parameter.

If you’ve stuck around for this entire thing (or even just this part), thanks a lot! I hope this guide is helpful. You can see the code for the entire series here.