NIFI-7385 Provided reverse-indexed TokenCache implementation.

Cleaned up code style.
Unit test was failing on Windows 1.8 GitHub Actions build but no other environment. Increased artificial delay to avoid timing issues.

Co-authored-by: Andy LoPresto <alopresto@apache.org>

This closes #4271.

Signed-off-by: Andy LoPresto <alopresto@apache.org>

nifi-nar-bundles/nifi-framework-bundle/nifi-framework/nifi-web/nifi-web-security/src/main/java/org/apache/nifi/web/security/otp/OtpAuthenticationFilter.java

@@ -16,15 +16,18 @@
  */
 package org.apache.nifi.web.security.otp;
 
+import java.util.regex.Pattern;
+import javax.servlet.http.HttpServletRequest;
 import org.apache.nifi.web.security.NiFiAuthenticationFilter;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 import org.springframework.security.core.Authentication;
 
-import javax.servlet.http.HttpServletRequest;
-import java.util.regex.Pattern;
-
 /**
+ * This filter is used to capture one time passwords (OTP) from requests made to download files through the browser.
+ * It's required because when we initiate a download in the browser, it must be opened in a new tab. The new tab
+ * cannot be initialized with authentication headers, so we must add a token as a query parameter instead. As
+ * tokens in URL strings are visible in various places, this must only be used once - hence our OTP.
  */
 public class OtpAuthenticationFilter extends NiFiAuthenticationFilter {
 

nifi-nar-bundles/nifi-framework-bundle/nifi-framework/nifi-web/nifi-web-security/src/main/java/org/apache/nifi/web/security/otp/OtpAuthenticationProvider.java

@@ -28,7 +28,7 @@ import org.springframework.security.core.Authentication;
 import org.springframework.security.core.AuthenticationException;
 
 /**
- *
+ * This provider will be used when the request is attempting to authenticate with a download or ui extension OTP/token.
  */
 public class OtpAuthenticationProvider extends NiFiAuthenticationProvider {
 

nifi-nar-bundles/nifi-framework-bundle/nifi-framework/nifi-web/nifi-web-security/src/main/java/org/apache/nifi/web/security/otp/OtpService.java

@@ -16,23 +16,19 @@
  */
 package org.apache.nifi.web.security.otp;
 
-import com.google.common.cache.Cache;
+import java.nio.charset.StandardCharsets;
-import com.google.common.cache.CacheBuilder;
+import java.security.InvalidKeyException;
+import java.security.NoSuchAlgorithmException;
+import java.security.SecureRandom;
+import java.util.concurrent.TimeUnit;
+import javax.crypto.Mac;
+import javax.crypto.spec.SecretKeySpec;
 import org.apache.commons.codec.binary.Base64;
 import org.apache.nifi.web.security.token.OtpAuthenticationToken;
 import org.apache.nifi.web.security.util.CacheKey;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
-import javax.crypto.Mac;
-import javax.crypto.spec.SecretKeySpec;
-import java.nio.charset.StandardCharsets;
-import java.security.InvalidKeyException;
-import java.security.NoSuchAlgorithmException;
-import java.security.SecureRandom;
-import java.util.concurrent.ConcurrentMap;
-import java.util.concurrent.TimeUnit;
-
 /**
  * OtpService is a service for generating and verifying one time password tokens.
  */
@@ -45,8 +41,8 @@ public class OtpService {
     // protected for testing purposes
     protected static final int MAX_CACHE_SOFT_LIMIT = 100;
 
-    private final Cache<CacheKey, String> downloadTokenCache;
+    private final TokenCache downloadTokens;
-    private final Cache<CacheKey, String> uiExtensionCache;
+    private final TokenCache uiExtensionTokens;
 
     /**
      * Creates a new OtpService with an expiration of 5 minutes.
@@ -64,8 +60,8 @@ public class OtpService {
      * @throws IllegalArgumentException If duration is negative
      */
     public OtpService(final int duration, final TimeUnit units) {
-        downloadTokenCache = CacheBuilder.newBuilder().expireAfterWrite(duration, units).build();
+        downloadTokens = new TokenCache("download tokens", duration, units);
-        uiExtensionCache = CacheBuilder.newBuilder().expireAfterWrite(duration, units).build();
+        uiExtensionTokens = new TokenCache("UI extension tokens", duration, units);
     }
 
     /**
@@ -75,7 +71,7 @@ public class OtpService {
      * @return                          The one time use download token
      */
     public String generateDownloadToken(final OtpAuthenticationToken authenticationToken) {
-        return generateToken(downloadTokenCache.asMap(), authenticationToken);
+        return generateToken(downloadTokens, authenticationToken);
     }
 
     /**
@@ -86,7 +82,7 @@ public class OtpService {
      * @throws OtpAuthenticationException   When the specified token does not correspond to an authenticated identity
      */
     public String getAuthenticationFromDownloadToken(final String token) throws OtpAuthenticationException {
-        return getAuthenticationFromToken(downloadTokenCache.asMap(), token);
+        return getAuthenticationFromToken(downloadTokens, token);
     }
 
     /**
@@ -96,7 +92,7 @@ public class OtpService {
      * @return                          The one time use UI extension token
      */
     public String generateUiExtensionToken(final OtpAuthenticationToken authenticationToken) {
-        return generateToken(uiExtensionCache.asMap(), authenticationToken);
+        return generateToken(uiExtensionTokens, authenticationToken);
     }
 
     /**
@@ -107,42 +103,55 @@ public class OtpService {
      * @throws OtpAuthenticationException   When the specified token does not correspond to an authenticated identity
      */
     public String getAuthenticationFromUiExtensionToken(final String token) throws OtpAuthenticationException {
-        return getAuthenticationFromToken(uiExtensionCache.asMap(), token);
+        return getAuthenticationFromToken(uiExtensionTokens, token);
     }
 
     /**
      * Generates a token and stores it in the specified cache.
      *
-     * @param cache                     The cache
+     * @param tokenCache                A cache that maps tokens to users
      * @param authenticationToken       The authentication
      * @return                          The one time use token
      */
-    private String generateToken(final ConcurrentMap<CacheKey, String> cache, final OtpAuthenticationToken authenticationToken) {
+    private String generateToken(final TokenCache tokenCache, final OtpAuthenticationToken authenticationToken) {
-        if (cache.size() >= MAX_CACHE_SOFT_LIMIT) {
+        final String userId = (String) authenticationToken.getPrincipal();
-            throw new IllegalStateException("The maximum number of single use tokens have been issued.");
+
+        // If the user has a token already, return it
+        if(tokenCache.containsValue(userId)) {
+            return (tokenCache.getKeyForValue(userId)).getKey();
+        } else {
+            // Otherwise, generate a token
+            if (tokenCache.size() >= MAX_CACHE_SOFT_LIMIT) {
+                throw new IllegalStateException("The maximum number of single use tokens have been issued.");
+            }
+
+            // Hash the authentication and build a cache key
+            final CacheKey cacheKey = new CacheKey(hash(authenticationToken));
+
+            // Store the token and user in the cache
+            tokenCache.put(cacheKey, userId);
+
+            // Return the token
+            return cacheKey.getKey();
         }
-
-        // hash the authentication and build a cache key
-        final CacheKey cacheKey = new CacheKey(hash(authenticationToken));
-
-        // store the token unless the token is already stored which should not update it's original timestamp
-        cache.putIfAbsent(cacheKey, authenticationToken.getName());
-
-        // return the token
-        return cacheKey.getKey();
     }
 
     /**
-     * Gets the corresponding authentication for the specified one time use token. The specified token will be removed.
+     * Gets the corresponding authentication for the specified one time use token. The specified token will be removed
+     * from the token cache.
      *
-     * @param cache                     The cache
+     * @param tokenCache                A cache that maps tokens to users
      * @param token                     The one time use token
      * @return                          The authenticated identity
      */
-    private String getAuthenticationFromToken(final ConcurrentMap<CacheKey, String> cache, final String token) throws OtpAuthenticationException {
+    private String getAuthenticationFromToken(final TokenCache tokenCache, final String token) throws OtpAuthenticationException {
-        final String authenticatedUser = cache.remove(new CacheKey(token));
+        final CacheKey cacheKey = new CacheKey(token);
+        final String authenticatedUser = (String) tokenCache.getIfPresent(cacheKey);
+
         if (authenticatedUser == null) {
             throw new OtpAuthenticationException("Unable to validate the access token.");
+        } else {
+            tokenCache.invalidate(cacheKey);
         }
 
         return authenticatedUser;

nifi-nar-bundles/nifi-framework-bundle/nifi-framework/nifi-web/nifi-web-security/src/main/java/org/apache/nifi/web/security/otp/TokenCache.java

@@ -0,0 +1,144 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.nifi.web.security.otp;
+
+import com.google.common.cache.AbstractCache;
+import com.google.common.cache.Cache;
+import com.google.common.cache.CacheBuilder;
+import com.google.common.cache.CacheStats;
+import java.util.Map;
+import java.util.concurrent.ConcurrentMap;
+import java.util.concurrent.TimeUnit;
+import org.apache.nifi.web.security.util.CacheKey;
+import org.checkerframework.checker.nullness.qual.Nullable;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+/**
+ * This class provides a specific wrapper implementation based on the Guava {@link Cache} but with
+ * reverse-index capability because of the special use case (a user [the cache value] can only have
+ * one active token [the cache key] at a time). This allows reverse lookup semantics.
+ */
+public class TokenCache extends AbstractCache<CacheKey, String> {
+    private static final Logger logger = LoggerFactory.getLogger(TokenCache.class);
+
+    private final String contentsDescription;
+    private final Cache<CacheKey, String> internalCache;
+
+    public TokenCache(String contentsDescription, final int duration, final TimeUnit units) {
+        this.contentsDescription = contentsDescription;
+        internalCache = CacheBuilder.newBuilder().expireAfterWrite(duration, units).build();
+    }
+
+    /**
+     * Returns the value associated with {@code key} in this cache, or {@code null} if there is no
+     * cached value for {@code key}.
+     *
+     * @param key the (wrapped) {@code token}
+     * @since 11.0
+     * @return the retrieved value ({@code user})
+     */
+    @Override
+    public @Nullable String getIfPresent(Object key) {
+        return internalCache.getIfPresent(key);
+    }
+
+    /**
+     * Puts the provided value ({@code user}) in the cache at the provided key (wrapped {@code token}).
+     *
+     * @param key the cache key
+     * @param value the value to insert
+     * @since 11.0
+     */
+    @Override
+    public void put(CacheKey key, String value) {
+            internalCache.put(key, value);
+    }
+
+    /**
+     * Returns {@code true} if the cache contains the provided value.
+     *
+     * @param value the value ({@code user}) to look for
+     * @return true if the user exists in the cache
+     */
+    public boolean containsValue(String value) {
+        return internalCache.asMap().containsValue(value);
+    }
+
+    /**
+     * Returns the {@link CacheKey} representing the key ({@code token}) associated with the provided value ({@code user}).
+     *
+     * @param value the value ({@code user}) to look for
+     * @return the CacheKey ({@code token}) associated with this user, or {@code null} if the user has no tokens in this cache
+     */
+    @Nullable
+    public CacheKey getKeyForValue(String value) {
+        if (containsValue(value)) {
+            Map<CacheKey, String> cacheMap = internalCache.asMap();
+            for (Map.Entry<CacheKey, String> e : cacheMap.entrySet()) {
+                if (e.getValue().equals(value)) {
+                    return e.getKey();
+                }
+            }
+            throw new IllegalStateException("The value existed in the cache but expired during retrieval");
+        } else {
+            return null;
+        }
+    }
+
+    // Override the unsupported abstract methods from the parent
+
+    @Override
+    public void invalidate(Object key) {
+        internalCache.invalidate(key);
+    }
+
+    @Override
+    public void invalidateAll() {
+        internalCache.invalidateAll(internalCache.asMap().keySet());
+    }
+
+    @Override
+    public long size() {
+        return internalCache.size();
+    }
+
+    @Override
+    public CacheStats stats() {
+        return internalCache.stats();
+    }
+
+    @Override
+    public ConcurrentMap<CacheKey, String> asMap() {
+        return internalCache.asMap();
+    }
+
+    /**
+     * Returns a string representation of the cache.
+     *
+     * @return a string representation of the cache
+     */
+    @Override
+    public String toString() {
+        return new StringBuilder("TokenCache for ")
+                .append(contentsDescription)
+                .append(" with ")
+                .append(internalCache.size())
+                .append(" elements")
+                .toString();
+    }
+}

nifi-nar-bundles/nifi-framework-bundle/nifi-framework/nifi-web/nifi-web-security/src/test/groovy/org/apache/nifi/web/security/otp/TokenCacheTest.groovy

@@ -0,0 +1,249 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.nifi.web.security.otp
+
+
+import org.apache.nifi.web.security.token.OtpAuthenticationToken
+import org.apache.nifi.web.security.util.CacheKey
+import org.bouncycastle.jce.provider.BouncyCastleProvider
+import org.junit.After
+import org.junit.Before
+import org.junit.BeforeClass
+import org.junit.Test
+import org.junit.runner.RunWith
+import org.junit.runners.JUnit4
+import org.slf4j.Logger
+import org.slf4j.LoggerFactory
+
+import java.security.Security
+import java.util.concurrent.TimeUnit
+
+@RunWith(JUnit4.class)
+class TokenCacheTest extends GroovyTestCase {
+    private static final Logger logger = LoggerFactory.getLogger(TokenCache.class)
+
+    private static final String andy = "alopresto"
+    private static final String nathan = "ngough"
+    private static final String matt = "mgilman"
+
+    private static final int LONG_CACHE_EXPIRATION = 10
+    private static final int SHORT_CACHE_EXPIRATION = 1
+
+    @BeforeClass
+    static void setUpOnce() throws Exception {
+        Security.addProvider(new BouncyCastleProvider())
+
+        logger.metaClass.methodMissing = { String name, args ->
+            logger.info("[${name?.toUpperCase()}] ${(args as List).join(" ")}")
+        }
+    }
+
+    @Before
+    void setUp() throws Exception {
+
+    }
+
+    @After
+    void tearDown() throws Exception {
+
+    }
+
+    /**
+     * Returns a simple "hash" of the provided principal (for test purposes, simply reverses the String).
+     *
+     * @param principal the token principal
+     * @return the hashed token output
+     */
+    private static String hash(def principal) {
+        principal.toString().reverse()
+    }
+
+    /**
+     * Returns the {@link CacheKey} constructed from the provided token.
+     *
+     * @param token the authentication token
+     * @return the cache key
+     */
+    private static CacheKey buildCacheKey(OtpAuthenticationToken token) {
+        new CacheKey(hash(token.principal))
+    }
+
+    @Test
+    void testShouldCheckIfContainsValue() throws Exception {
+        // Arrange
+        TokenCache tokenCache = new TokenCache("test tokens", LONG_CACHE_EXPIRATION, TimeUnit.SECONDS)
+
+        OtpAuthenticationToken andyToken = new OtpAuthenticationToken(andy)
+        OtpAuthenticationToken nathanToken = new OtpAuthenticationToken(nathan)
+
+        tokenCache.put(buildCacheKey(andyToken), andy)
+        tokenCache.put(buildCacheKey(nathanToken), nathan)
+
+        logger.info(tokenCache.toString())
+
+        // Act
+        boolean containsAndyToken = tokenCache.containsValue(andy)
+        boolean containsNathanToken = tokenCache.containsValue(nathan)
+        boolean containsMattToken = tokenCache.containsValue(matt)
+
+        // Assert
+        assert containsAndyToken
+        assert containsNathanToken
+        assert !containsMattToken
+    }
+
+    @Test
+    void testShouldGetKeyByValue() throws Exception {
+        // Arrange
+        TokenCache tokenCache = new TokenCache("test tokens", LONG_CACHE_EXPIRATION, TimeUnit.SECONDS)
+
+        OtpAuthenticationToken andyToken = new OtpAuthenticationToken(andy)
+        OtpAuthenticationToken nathanToken = new OtpAuthenticationToken(nathan)
+
+        tokenCache.put(buildCacheKey(andyToken), andy)
+        tokenCache.put(buildCacheKey(nathanToken), nathan)
+
+        logger.info(tokenCache.toString())
+
+        // Act
+        CacheKey keyForAndyToken = tokenCache.getKeyForValue(andy)
+        CacheKey keyForNathanToken = tokenCache.getKeyForValue(nathan)
+        CacheKey keyForMattToken = tokenCache.getKeyForValue(matt)
+
+        def tokens = [keyForAndyToken, keyForNathanToken, keyForMattToken]
+        logger.info("Retrieved tokens: ${tokens}")
+
+        // Assert
+        assert keyForAndyToken.getKey() == hash(andyToken.principal)
+        assert keyForNathanToken.getKey() == hash(nathanToken.principal)
+        assert !keyForMattToken
+    }
+
+    @Test
+    void testShouldNotGetKeyByValueAfterExpiration() throws Exception {
+        // Arrange
+        TokenCache tokenCache = new TokenCache("test tokens", SHORT_CACHE_EXPIRATION, TimeUnit.SECONDS)
+
+        OtpAuthenticationToken andyToken = new OtpAuthenticationToken(andy)
+        OtpAuthenticationToken nathanToken = new OtpAuthenticationToken(nathan)
+
+        tokenCache.put(buildCacheKey(andyToken), andy)
+        tokenCache.put(buildCacheKey(nathanToken), nathan)
+
+        logger.info(tokenCache.toString())
+
+        // Sleep to allow the cache entries to expire (was failing on Windows JDK 8 when only sleeping for 1 second)
+        sleep(SHORT_CACHE_EXPIRATION * 2 * 1000)
+
+        // Act
+        CacheKey keyForAndyToken = tokenCache.getKeyForValue(andy)
+        CacheKey keyForNathanToken = tokenCache.getKeyForValue(nathan)
+        CacheKey keyForMattToken = tokenCache.getKeyForValue(matt)
+
+        def tokens = [keyForAndyToken, keyForNathanToken, keyForMattToken]
+        logger.info("Retrieved tokens: ${tokens}")
+
+        // Assert
+        assert !keyForAndyToken
+        assert !keyForNathanToken
+        assert !keyForMattToken
+    }
+
+    @Test
+    void testShouldInvalidateSingleKey() throws Exception {
+        // Arrange
+        TokenCache tokenCache = new TokenCache("test tokens", LONG_CACHE_EXPIRATION, TimeUnit.SECONDS)
+
+        OtpAuthenticationToken andyToken = new OtpAuthenticationToken(andy)
+        OtpAuthenticationToken nathanToken = new OtpAuthenticationToken(nathan)
+        OtpAuthenticationToken mattToken = new OtpAuthenticationToken(matt)
+
+        CacheKey andyKey = buildCacheKey(andyToken)
+        CacheKey nathanKey = buildCacheKey(nathanToken)
+        CacheKey mattKey = buildCacheKey(mattToken)
+
+        tokenCache.put(andyKey, andy)
+        tokenCache.put(nathanKey, nathan)
+        tokenCache.put(mattKey, matt)
+
+        logger.info(tokenCache.toString())
+
+        // Act
+        tokenCache.invalidate(andyKey)
+
+        // Assert
+        assert !tokenCache.containsValue(andy)
+        assert tokenCache.containsValue(nathan)
+        assert tokenCache.containsValue(matt)
+    }
+
+    @Test
+    void testShouldInvalidateMultipleKeys() throws Exception {
+        // Arrange
+        TokenCache tokenCache = new TokenCache("test tokens", LONG_CACHE_EXPIRATION, TimeUnit.SECONDS)
+
+        OtpAuthenticationToken andyToken = new OtpAuthenticationToken(andy)
+        OtpAuthenticationToken nathanToken = new OtpAuthenticationToken(nathan)
+        OtpAuthenticationToken mattToken = new OtpAuthenticationToken(matt)
+
+        CacheKey andyKey = buildCacheKey(andyToken)
+        CacheKey nathanKey = buildCacheKey(nathanToken)
+        CacheKey mattKey = buildCacheKey(mattToken)
+
+        tokenCache.put(andyKey, andy)
+        tokenCache.put(nathanKey, nathan)
+        tokenCache.put(mattKey, matt)
+
+        logger.info(tokenCache.toString())
+
+        // Act
+        tokenCache.invalidateAll([andyKey, nathanKey])
+
+        // Assert
+        assert !tokenCache.containsValue(andy)
+        assert !tokenCache.containsValue(nathan)
+        assert tokenCache.containsValue(matt)
+    }
+
+    @Test
+    void testShouldInvalidateAll() throws Exception {
+        // Arrange
+        TokenCache tokenCache = new TokenCache("test tokens", LONG_CACHE_EXPIRATION, TimeUnit.SECONDS)
+
+        OtpAuthenticationToken andyToken = new OtpAuthenticationToken(andy)
+        OtpAuthenticationToken nathanToken = new OtpAuthenticationToken(nathan)
+        OtpAuthenticationToken mattToken = new OtpAuthenticationToken(matt)
+
+        CacheKey andyKey = buildCacheKey(andyToken)
+        CacheKey nathanKey = buildCacheKey(nathanToken)
+        CacheKey mattKey = buildCacheKey(mattToken)
+
+        tokenCache.put(andyKey, andy)
+        tokenCache.put(nathanKey, nathan)
+        tokenCache.put(mattKey, matt)
+
+        logger.info(tokenCache.toString())
+
+        // Act
+        tokenCache.invalidateAll()
+
+        // Assert
+        assert !tokenCache.containsValue(andy)
+        assert !tokenCache.containsValue(nathan)
+        assert !tokenCache.containsValue(matt)
+    }
+}

nifi-nar-bundles/nifi-framework-bundle/nifi-framework/nifi-web/nifi-web-security/src/test/java/org/apache/nifi/web/security/otp/OtpServiceTest.java

@@ -16,19 +16,21 @@
  */
 package org.apache.nifi.web.security.otp;
 
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertNotEquals;
+import static org.junit.Assert.assertNotNull;
+import static org.junit.Assert.fail;
+
+import java.util.concurrent.TimeUnit;
 import org.apache.nifi.web.security.token.OtpAuthenticationToken;
 import org.junit.Before;
 import org.junit.Test;
 
-import java.util.concurrent.TimeUnit;
-
-import static org.junit.Assert.assertEquals;
-import static org.junit.Assert.assertNotNull;
-import static org.junit.Assert.fail;
-
 public class OtpServiceTest {
 
     private final static String USER_1 = "user-identity-1";
+    private final static int CACHE_EXPIRY_TIME = 1;
+    private final static int WAIT_TIME = 2000;
 
     private OtpService otpService;
 
@@ -87,7 +89,7 @@ public class OtpServiceTest {
                 final OtpAuthenticationToken authenticationToken = new OtpAuthenticationToken("user-identity-" + i);
                 otpService.generateDownloadToken(authenticationToken);
             } catch (final IllegalStateException iae) {
-                // ensure we failed when we've past the limit
+                // ensure we failed when we've passed the limit
                 assertEquals(OtpService.MAX_CACHE_SOFT_LIMIT + 1, i);
                 throw iae;
             }
@@ -102,7 +104,7 @@ public class OtpServiceTest {
                 final OtpAuthenticationToken authenticationToken = new OtpAuthenticationToken("user-identity-" + i);
                 otpService.generateUiExtensionToken(authenticationToken);
             } catch (final IllegalStateException iae) {
-                // ensure we failed when we've past the limit
+                // ensure we failed when we've passed the limit
                 assertEquals(OtpService.MAX_CACHE_SOFT_LIMIT + 1, i);
                 throw iae;
             }
@@ -121,29 +123,169 @@ public class OtpServiceTest {
 
     @Test(expected = OtpAuthenticationException.class)
     public void testUiExtensionTokenExpiration() throws Exception {
-        final OtpService otpServiceWithTightExpiration = new OtpService(2, TimeUnit.SECONDS);
+        final OtpService otpServiceWithTightExpiration = new OtpService(CACHE_EXPIRY_TIME, TimeUnit.SECONDS);
 
         final OtpAuthenticationToken authenticationToken = new OtpAuthenticationToken(USER_1);
         final String downloadToken = otpServiceWithTightExpiration.generateUiExtensionToken(authenticationToken);
 
-        // sleep for 4 seconds which should sufficiently expire the valid token
+        // sleep for 2 seconds which should sufficiently expire the valid token
-        Thread.sleep(4 * 1000);
+        Thread.sleep(WAIT_TIME);
 
-        // attempt to get the token now that its expired
+        // attempt to get the token now that it's expired
         otpServiceWithTightExpiration.getAuthenticationFromUiExtensionToken(downloadToken);
     }
 
     @Test(expected = OtpAuthenticationException.class)
     public void testDownloadTokenExpiration() throws Exception {
-        final OtpService otpServiceWithTightExpiration = new OtpService(2, TimeUnit.SECONDS);
+        final OtpService otpServiceWithTightExpiration = new OtpService(CACHE_EXPIRY_TIME, TimeUnit.SECONDS);
 
         final OtpAuthenticationToken authenticationToken = new OtpAuthenticationToken(USER_1);
         final String downloadToken = otpServiceWithTightExpiration.generateDownloadToken(authenticationToken);
 
-        // sleep for 4 seconds which should sufficiently expire the valid token
+        // sleep for 2 seconds which should sufficiently expire the valid token
-        Thread.sleep(4 * 1000);
+        Thread.sleep(WAIT_TIME);
 
-        // attempt to get the token now that its expired
+        // attempt to get the token now that it's expired
         otpServiceWithTightExpiration.getAuthenticationFromDownloadToken(downloadToken);
     }
+
+    @Test
+    public void testDownloadTokenIsTheSameForSubsequentRequests() {
+        final OtpAuthenticationToken authenticationToken = new OtpAuthenticationToken(USER_1);
+        final String downloadToken = otpService.generateDownloadToken(authenticationToken);
+        final String secondDownloadToken = otpService.generateDownloadToken(authenticationToken);
+
+        assertEquals(downloadToken, secondDownloadToken);
+    }
+
+    @Test
+    public void testDownloadTokenIsTheSameForSubsequentRequestsUntilUsed() {
+        final OtpAuthenticationToken authenticationToken = new OtpAuthenticationToken(USER_1);
+
+        // generate two tokens
+        final String downloadToken = otpService.generateDownloadToken(authenticationToken);
+        final String secondDownloadToken = otpService.generateDownloadToken(authenticationToken);
+
+        assertEquals(downloadToken, secondDownloadToken);
+
+        // use the token
+        otpService.getAuthenticationFromDownloadToken(downloadToken);
+
+        // make sure the next token is now different
+        final String thirdDownloadToken = otpService.generateDownloadToken(authenticationToken);
+        assertNotEquals(downloadToken, thirdDownloadToken);
+    }
+
+    @Test
+    public void testDownloadTokenIsValidForSubsequentGenerateAndUse() {
+        final OtpAuthenticationToken authenticationToken = new OtpAuthenticationToken(USER_1);
+
+        // generate a token
+        final String downloadToken = otpService.generateDownloadToken(authenticationToken);
+
+        // use the token
+        final String auth = otpService.getAuthenticationFromDownloadToken(downloadToken);
+        assertEquals(USER_1, auth);
+
+        // generate a new token, make sure it's different, then authenticate with it
+        final String secondDownloadToken = otpService.generateDownloadToken(authenticationToken);
+        assertNotEquals(downloadToken, secondDownloadToken);
+        final String secondAuth = otpService.getAuthenticationFromDownloadToken(secondDownloadToken);
+        assertEquals(USER_1, secondAuth);
+    }
+
+    @Test
+    public void testSingleUserCannotGenerateTooManyUIExtensionTokens() throws Exception {
+        // ensure we'll try to loop past the limit
+        for (int i = 1; i < OtpService.MAX_CACHE_SOFT_LIMIT + 10; i++) {
+            final OtpAuthenticationToken authenticationToken = new OtpAuthenticationToken("user-identity-1");
+            otpService.generateUiExtensionToken(authenticationToken);
+
+        }
+
+        // make sure other users can still generate tokens
+        final OtpAuthenticationToken anotherAuthenticationToken = new OtpAuthenticationToken("user-identity-2");
+        final String auth = otpService.generateUiExtensionToken(anotherAuthenticationToken);
+        assertNotNull(auth);
+    }
+
+    @Test
+    public void testSingleUserCannotGenerateTooManyDownloadTokens() throws Exception {
+        // ensure we'll try to loop past the limit
+        for (int i = 1; i < OtpService.MAX_CACHE_SOFT_LIMIT + 10; i++) {
+            final OtpAuthenticationToken authenticationToken = new OtpAuthenticationToken("user-identity-1");
+            otpService.generateDownloadToken(authenticationToken);
+
+        }
+
+        // make sure other users can still generate tokens
+        final OtpAuthenticationToken anotherAuthenticationToken = new OtpAuthenticationToken("user-identity-2");
+        final String auth = otpService.generateDownloadToken(anotherAuthenticationToken);
+        assertNotNull(auth);
+    }
+
+    @Test(expected = OtpAuthenticationException.class)
+    public void testDownloadTokenNotValidAfterUse() throws Exception {
+        final OtpAuthenticationToken authenticationToken = new OtpAuthenticationToken(USER_1);
+        final String downloadToken = otpService.generateDownloadToken(authenticationToken);
+
+        // use the token
+        final String authenticatedUser = otpService.getAuthenticationFromDownloadToken(downloadToken);
+
+        // check we authenticated successfully
+        assertNotNull(authenticatedUser);
+        assertEquals(USER_1, authenticatedUser);
+
+        // check authentication fails with the used token
+        final String failedAuthentication = otpService.getAuthenticationFromDownloadToken(downloadToken);
+    }
+
+    @Test(expected = OtpAuthenticationException.class)
+    public void testUIExtensionTokenNotValidAfterUse() throws Exception {
+        final OtpAuthenticationToken authenticationToken = new OtpAuthenticationToken(USER_1);
+        final String downloadToken = otpService.generateDownloadToken(authenticationToken);
+
+        // use the token
+        final String authenticatedUser = otpService.getAuthenticationFromUiExtensionToken(downloadToken);
+
+        // check we authenticated successfully
+        assertNotNull(authenticatedUser);
+        assertEquals(USER_1, authenticatedUser);
+
+        // check authentication fails with the used token
+        final String failedAuthentication = otpService.getAuthenticationFromUiExtensionToken(downloadToken);
+    }
+
+    @Test
+    public void testShouldGenerateNewDownloadTokenAfterExpiration() throws Exception {
+        final OtpService otpServiceWithTightExpiration = new OtpService(CACHE_EXPIRY_TIME, TimeUnit.SECONDS);
+
+        final OtpAuthenticationToken authenticationToken = new OtpAuthenticationToken(USER_1);
+        final String downloadToken = otpServiceWithTightExpiration.generateDownloadToken(authenticationToken);
+
+        // sleep for 2 seconds which should sufficiently expire the valid token
+        Thread.sleep(WAIT_TIME);
+
+        // get a new token and make sure the previous one had expired
+        final String secondDownloadToken = otpServiceWithTightExpiration.generateDownloadToken(authenticationToken);
+        assertNotEquals(downloadToken, secondDownloadToken);
+    }
+
+    @Test
+    public void testDownloadTokenRemainsTheSameBeforeExpirationButNotAfter() throws Exception {
+        final OtpService otpServiceWithTightExpiration = new OtpService(CACHE_EXPIRY_TIME, TimeUnit.SECONDS);
+
+        final OtpAuthenticationToken authenticationToken = new OtpAuthenticationToken(USER_1);
+        final String downloadToken = otpServiceWithTightExpiration.generateDownloadToken(authenticationToken);
+        final String secondDownloadToken = otpServiceWithTightExpiration.generateDownloadToken(authenticationToken);
+
+        assertEquals(downloadToken, secondDownloadToken);
+
+        // sleep for 2 seconds which should sufficiently expire the valid token
+        Thread.sleep(WAIT_TIME);
+
+        // get a new token and make sure the previous one had expired
+        final String thirdDownloadToken = otpServiceWithTightExpiration.generateDownloadToken(authenticationToken);
+        assertNotEquals(downloadToken, thirdDownloadToken);
+    }
 }