KNOX-2627 - Limiting the number of managed tokens per user (#463)

diff --git a/gateway-server/src/main/java/org/apache/knox/gateway/config/impl/GatewayConfigImpl.java b/gateway-server/src/main/java/org/apache/knox/gateway/config/impl/GatewayConfigImpl.java
index 43895e9..22ff816 100644
--- a/gateway-server/src/main/java/org/apache/knox/gateway/config/impl/GatewayConfigImpl.java
+++ b/gateway-server/src/main/java/org/apache/knox/gateway/config/impl/GatewayConfigImpl.java
@@ -262,9 +262,11 @@
   private static final String KNOX_TOKEN_ALIAS_PERSISTENCE_INTERVAL = GATEWAY_CONFIG_FILE_PREFIX + ".knox.token.state.alias.persistence.interval";
   private static final String KNOX_TOKEN_PERMISSIVE_VALIDATION_ENABLED = GATEWAY_CONFIG_FILE_PREFIX + ".knox.token.permissive.validation";
   private static final String KNOX_TOKEN_HASH_ALGORITHM = GATEWAY_CONFIG_FILE_PREFIX + ".knox.token.hash.algorithm";
+  public static final String KNOX_TOKEN_USER_LIMIT = GATEWAY_CONFIG_FILE_PREFIX + ".knox.token.limit.per.user";
   private static final long KNOX_TOKEN_EVICTION_INTERVAL_DEFAULT = TimeUnit.MINUTES.toSeconds(5);
   private static final long KNOX_TOKEN_EVICTION_GRACE_PERIOD_DEFAULT = TimeUnit.HOURS.toSeconds(24);
   private static final long KNOX_TOKEN_ALIAS_PERSISTENCE_INTERVAL_DEFAULT = TimeUnit.SECONDS.toSeconds(15);
+  public static final int KNOX_TOKEN_USER_LIMIT_DEFAULT = 10;
   private static final boolean KNOX_TOKEN_PERMISSIVE_VALIDATION_ENABLED_DEFAULT = false;
 
   private static final String KNOX_HOMEPAGE_PROFILE_PREFIX =  "knox.homepage.profile.";
@@ -1186,6 +1188,11 @@
   }
 
   @Override
+  public int getMaximumNumberOfTokensPerUser() {
+    return getInt(KNOX_TOKEN_USER_LIMIT, KNOX_TOKEN_USER_LIMIT_DEFAULT);
+  }
+
+  @Override
   public Set<String> getHiddenTopologiesOnHomepage() {
     final Set<String> hiddenTopologies = new HashSet<>(getTrimmedStringCollection(KNOX_HOMEPAGE_HIDDEN_TOPOLOGIES));
     return hiddenTopologies == null || hiddenTopologies.isEmpty() ? KNOX_HOMEPAGE_HIDDEN_TOPOLOGIES_DEFAULT : hiddenTopologies;
diff --git a/gateway-service-knoxtoken/src/main/java/org/apache/knox/gateway/service/knoxtoken/TokenResource.java b/gateway-service-knoxtoken/src/main/java/org/apache/knox/gateway/service/knoxtoken/TokenResource.java
index aa1f2cf..61b8f43 100644
--- a/gateway-service-knoxtoken/src/main/java/org/apache/knox/gateway/service/knoxtoken/TokenResource.java
+++ b/gateway-service-knoxtoken/src/main/java/org/apache/knox/gateway/service/knoxtoken/TokenResource.java
@@ -141,6 +141,8 @@
 
   private Optional<Long> maxTokenLifetime = Optional.empty();
 
+  private int tokenLimitPerUser;
+
   private List<String> allowedRenewers;
 
   @Context
@@ -212,6 +214,8 @@
       final AliasService aliasService = services.getService(ServiceType.ALIAS_SERVICE);
       tokenMAC = new TokenMAC(gatewayConfig.getKnoxTokenHashAlgorithm(), aliasService.getPasswordFromAliasForGateway(TokenMAC.KNOX_TOKEN_HASH_KEY_ALIAS_NAME));
 
+      tokenLimitPerUser = gatewayConfig.getMaximumNumberOfTokensPerUser();
+
       String renewIntervalValue = context.getInitParameter(TOKEN_EXP_RENEWAL_INTERVAL);
       if (renewIntervalValue != null && !renewIntervalValue.isEmpty()) {
         try {
@@ -585,6 +589,15 @@
       jku = request.getRequestURL().substring(0, idx) + JWKSResource.JWKS_PATH;
     }
 
+    if (tokenStateService != null) {
+      if (tokenLimitPerUser != -1) { // if -1 => unlimited tokens for all users
+        if (tokenStateService.getTokens(p.getName()).size() == tokenLimitPerUser) {
+          log.tokenLimitExceeded(p.getName());
+          return Response.status(Response.Status.FORBIDDEN).entity("{ \"Unable to get token - token limit exceeded.\" }").build();
+        }
+      }
+    }
+
     try {
       final boolean managedToken = tokenStateService != null;
       JWT token;
diff --git a/gateway-service-knoxtoken/src/main/java/org/apache/knox/gateway/service/knoxtoken/TokenServiceMessages.java b/gateway-service-knoxtoken/src/main/java/org/apache/knox/gateway/service/knoxtoken/TokenServiceMessages.java
index f2a65e5..8cae4b6 100644
--- a/gateway-service-knoxtoken/src/main/java/org/apache/knox/gateway/service/knoxtoken/TokenServiceMessages.java
+++ b/gateway-service-knoxtoken/src/main/java/org/apache/knox/gateway/service/knoxtoken/TokenServiceMessages.java
@@ -83,4 +83,7 @@
 
   @Message( level = MessageLevel.WARN, text = "Invalid duration used for JWT token lifespan ({0}) using the configured TTL for KnoxToken service")
   void invalidLifetimeValue(String lifetimeStr);
+
+  @Message( level = MessageLevel.ERROR, text = "Unable to get token for user {0}: token limit exceeded")
+  void tokenLimitExceeded(String userName);
 }
diff --git a/gateway-service-knoxtoken/src/test/java/org/apache/knox/gateway/service/knoxtoken/TokenServiceResourceTest.java b/gateway-service-knoxtoken/src/test/java/org/apache/knox/gateway/service/knoxtoken/TokenServiceResourceTest.java
index ff333df..1b84a4a 100644
--- a/gateway-service-knoxtoken/src/test/java/org/apache/knox/gateway/service/knoxtoken/TokenServiceResourceTest.java
+++ b/gateway-service-knoxtoken/src/test/java/org/apache/knox/gateway/service/knoxtoken/TokenServiceResourceTest.java
@@ -17,11 +17,14 @@
  */
 package org.apache.knox.gateway.service.knoxtoken;
 
+import static org.apache.knox.gateway.config.impl.GatewayConfigImpl.KNOX_TOKEN_USER_LIMIT;
+import static org.apache.knox.gateway.config.impl.GatewayConfigImpl.KNOX_TOKEN_USER_LIMIT_DEFAULT;
 import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertFalse;
 import static org.junit.Assert.assertNotEquals;
 import static org.junit.Assert.assertNotNull;
 import static org.junit.Assert.assertTrue;
+import static org.junit.Assert.fail;
 
 import com.fasterxml.jackson.core.type.TypeReference;
 import com.fasterxml.jackson.databind.ObjectMapper;
@@ -81,7 +84,9 @@
 import java.util.Map;
 import java.util.Set;
 import java.util.TimeZone;
+import java.util.TreeSet;
 import java.util.concurrent.Callable;
+import java.util.concurrent.ConcurrentHashMap;
 import java.util.concurrent.ExecutorService;
 import java.util.concurrent.Executors;
 import java.util.concurrent.Future;
@@ -94,9 +99,10 @@
   private static RSAPublicKey publicKey;
   private static RSAPrivateKey privateKey;
 
-  private static String TOKEN_API_PATH = "https://gateway-host:8443/gateway/sandbox/knoxtoken/api/v1";
-  private static String TOKEN_PATH = "/token";
-  private static String JKWS_PATH = "/jwks.json";
+  private static final String TOKEN_API_PATH = "https://gateway-host:8443/gateway/sandbox/knoxtoken/api/v1";
+  private static final String TOKEN_PATH = "/token";
+  private static final String JKWS_PATH = "/jwks.json";
+  private static final String USER_NAME = "alice";
 
   private ServletContext context;
   private HttpServletRequest request;
@@ -136,7 +142,7 @@
     request = EasyMock.createNiceMock(HttpServletRequest.class);
     EasyMock.expect(request.getServletContext()).andReturn(context).anyTimes();
     Principal principal = EasyMock.createNiceMock(Principal.class);
-    EasyMock.expect(principal.getName()).andReturn("alice").anyTimes();
+    EasyMock.expect(principal.getName()).andReturn(USER_NAME).anyTimes();
     EasyMock.expect(request.getUserPrincipal()).andReturn(principal).anyTimes();
     EasyMock.expect(request.getRequestURL()).andReturn(new StringBuffer(TOKEN_API_PATH+TOKEN_PATH)).anyTimes();
     if (contextExpectations.containsKey(TokenResource.LIFESPAN)) {
@@ -159,6 +165,8 @@
       EasyMock.expect(config.getServiceParameter(tokenStateServiceType, "impl")).andReturn(contextExpectations.get(tokenStateServiceType)).anyTimes();
     }
     EasyMock.expect(config.getKnoxTokenHashAlgorithm()).andReturn(HmacAlgorithms.HMAC_SHA_256.getName()).anyTimes();
+    EasyMock.expect(config.getMaximumNumberOfTokensPerUser())
+        .andReturn(contextExpectations.containsKey(KNOX_TOKEN_USER_LIMIT) ? Integer.parseInt(contextExpectations.get(KNOX_TOKEN_USER_LIMIT)) : -1).anyTimes();
     tss = new TestTokenStateService();
     EasyMock.expect(services.getService(ServiceType.TOKEN_STATE_SERVICE)).andReturn(tss).anyTimes();
 
@@ -937,6 +945,48 @@
     assertTrue((expiresDate.getTime() - now.getTime()) < oneMinute); // the configured TTL was used even if lifespan was supplied
   }
 
+  @Test
+  public void testConfiguredTokenLimitPerUser() throws Exception {
+    testLimitingTokensPerUser(String.valueOf(KNOX_TOKEN_USER_LIMIT_DEFAULT), KNOX_TOKEN_USER_LIMIT_DEFAULT);
+  }
+
+  @Test
+  public void testUnlimitedTokensPerUser() throws Exception {
+    testLimitingTokensPerUser(String.valueOf("-1"), 100);
+  }
+
+  @Test
+  public void tesTokenLimitPerUserExceeded() throws Exception {
+    try {
+      testLimitingTokensPerUser(String.valueOf("10"), 11);
+      fail("Exception should have been thrown");
+    } catch (Exception e) {
+      assertTrue(e.getMessage().contains("Unable to get token - token limit exceeded."));
+    }
+  }
+
+  private void testLimitingTokensPerUser(String configuredLimit, int numberOfTokens) throws Exception {
+    final Map<String, String> contextExpectations = new HashMap<>();
+    contextExpectations.put(KNOX_TOKEN_USER_LIMIT, configuredLimit);
+    configureCommonExpectations(contextExpectations, Boolean.TRUE);
+
+    final TokenResource tr = new TokenResource();
+    tr.request = request;
+    tr.context = context;
+    tr.init();
+
+    for (int i = 0; i < numberOfTokens; i++) {
+      final Response getTokenResponse = tr.doGet();
+      if (getTokenResponse.getStatus() != Response.Status.OK.getStatusCode()) {
+        throw new Exception(getTokenResponse.getEntity().toString());
+      }
+    }
+    final Response getKnoxTokensResponse = tr.getUserTokens(USER_NAME);
+    final Collection<String> tokens = ((Map<String, Collection<String>>) JsonUtils.getObjectFromJsonString(getKnoxTokensResponse.getEntity().toString()))
+        .get("tokens");
+    assertEquals(tokens.size(), numberOfTokens);
+  }
+
   /**
    *
    * @param isTokenStateServerManaged true, if server-side token state management should be enabled; Otherwise, false or null.
@@ -1218,6 +1268,7 @@
     private Map<String, Long> expirationData = new HashMap<>();
     private Map<String, Long> issueTimes = new HashMap<>();
     private Map<String, Long> maxLifetimes = new HashMap<>();
+    private final Map<String, TokenMetadata> tokenMetadata = new ConcurrentHashMap<>();
 
     long getIssueTime(final String token) {
       return issueTimes.get(token);
@@ -1310,16 +1361,26 @@
 
     @Override
     public void addMetadata(String tokenId, TokenMetadata metadata) {
+      tokenMetadata.put(tokenId, metadata);
     }
 
     @Override
     public TokenMetadata getTokenMetadata(String tokenId) throws UnknownTokenException {
-      return null;
+      return tokenMetadata.get(tokenId);
     }
 
     @Override
     public Collection<KnoxToken> getTokens(String userName) {
-      return null;
+      final Collection<KnoxToken> tokens = new TreeSet<>();
+      tokenMetadata.entrySet().stream().filter(entry -> entry.getValue().getUserName().equals(userName)).forEach(metadata -> {
+        String tokenId = metadata.getKey();
+        try {
+          tokens.add(new KnoxToken(tokenId, getTokenIssueTime(tokenId), getTokenExpiration(tokenId), 0L, metadata.getValue()));
+        } catch (UnknownTokenException e) {
+          // NOP
+        }
+      });
+      return tokens;
     }
 
     @Override
diff --git a/gateway-spi/src/main/java/org/apache/knox/gateway/config/GatewayConfig.java b/gateway-spi/src/main/java/org/apache/knox/gateway/config/GatewayConfig.java
index 4474c33..3599bc9 100644
--- a/gateway-spi/src/main/java/org/apache/knox/gateway/config/GatewayConfig.java
+++ b/gateway-spi/src/main/java/org/apache/knox/gateway/config/GatewayConfig.java
@@ -703,6 +703,14 @@
   String getKnoxTokenHashAlgorithm();
 
   /**
+   * @return the maximum number of tokens a user can manage at the same time. -1
+   *         means that users are allowed to create/manage as many tokens as they
+   *         want. This configuration only applies when server-managed token state
+   *         is enabled either in gateway-site or at the topology level.
+   */
+  int getMaximumNumberOfTokensPerUser();
+
+  /**
    * @return the list of topologies that should be hidden on Knox homepage
    */
   Set<String> getHiddenTopologiesOnHomepage();
diff --git a/gateway-test-release-utils/src/main/java/org/apache/knox/gateway/GatewayTestConfig.java b/gateway-test-release-utils/src/main/java/org/apache/knox/gateway/GatewayTestConfig.java
index c65fb66..50b2046 100644
--- a/gateway-test-release-utils/src/main/java/org/apache/knox/gateway/GatewayTestConfig.java
+++ b/gateway-test-release-utils/src/main/java/org/apache/knox/gateway/GatewayTestConfig.java
@@ -812,6 +812,11 @@
   }
 
   @Override
+  public int getMaximumNumberOfTokensPerUser() {
+    return 0;
+  }
+
+  @Override
   public Set<String> getHiddenTopologiesOnHomepage() {
     return Collections.emptySet();
   }
diff --git a/gateway-util-common/src/main/java/org/apache/knox/gateway/util/JsonUtils.java b/gateway-util-common/src/main/java/org/apache/knox/gateway/util/JsonUtils.java
index 9d2e613..bb0d4e3 100644
--- a/gateway-util-common/src/main/java/org/apache/knox/gateway/util/JsonUtils.java
+++ b/gateway-util-common/src/main/java/org/apache/knox/gateway/util/JsonUtils.java
@@ -73,10 +73,10 @@
   }
 
   public static Object getObjectFromJsonString(String json) {
-    Map<String, String> obj = null;
+    Map<String, Object> obj = null;
     JsonFactory factory = new JsonFactory();
     ObjectMapper mapper = new ObjectMapper(factory);
-    TypeReference<Map<String, String>> typeRef = new TypeReference<Map<String, String>>() {};
+    TypeReference<Map<String, Object>> typeRef = new TypeReference<Map<String, Object>>() {};
     try {
       obj = mapper.readValue(json, typeRef);
     } catch (IOException e) {