You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@nifi.apache.org by al...@apache.org on 2020/06/09 02:49:25 UTC

[nifi] branch master updated: 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.

This is an automated email from the ASF dual-hosted git repository.

alopresto pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/nifi.git


The following commit(s) were added to refs/heads/master by this push:
     new 01e42df  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.
01e42df is described below

commit 01e42dfb3291c3a3549023edadafd2d8023f3042
Author: Nathan Gough <th...@gmail.com>
AuthorDate: Mon May 4 12:07:36 2020 -0400

    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 <al...@apache.org>
    
    This closes #4271.
    
    Signed-off-by: Andy LoPresto <al...@apache.org>
---
 .../web/security/otp/OtpAuthenticationFilter.java  |   9 +-
 .../security/otp/OtpAuthenticationProvider.java    |   2 +-
 .../apache/nifi/web/security/otp/OtpService.java   |  81 ++++---
 .../apache/nifi/web/security/otp/TokenCache.java   | 144 ++++++++++++
 .../nifi/web/security/otp/TokenCacheTest.groovy    | 249 +++++++++++++++++++++
 .../nifi/web/security/otp/OtpServiceTest.java      | 176 +++++++++++++--
 6 files changed, 604 insertions(+), 57 deletions(-)

diff --git a/nifi-nar-bundles/nifi-framework-bundle/nifi-framework/nifi-web/nifi-web-security/src/main/java/org/apache/nifi/web/security/otp/OtpAuthenticationFilter.java b/nifi-nar-bundles/nifi-framework-bundle/nifi-framework/nifi-web/nifi-web-security/src/main/java/org/apache/nifi/web/security/otp/OtpAuthenticationFilter.java
index 98d8e68..c09a4bb 100644
--- a/nifi-nar-bundles/nifi-framework-bundle/nifi-framework/nifi-web/nifi-web-security/src/main/java/org/apache/nifi/web/security/otp/OtpAuthenticationFilter.java
+++ b/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 {
 
diff --git a/nifi-nar-bundles/nifi-framework-bundle/nifi-framework/nifi-web/nifi-web-security/src/main/java/org/apache/nifi/web/security/otp/OtpAuthenticationProvider.java b/nifi-nar-bundles/nifi-framework-bundle/nifi-framework/nifi-web/nifi-web-security/src/main/java/org/apache/nifi/web/security/otp/OtpAuthenticationProvider.java
index f375df2..bcc42cd 100644
--- a/nifi-nar-bundles/nifi-framework-bundle/nifi-framework/nifi-web/nifi-web-security/src/main/java/org/apache/nifi/web/security/otp/OtpAuthenticationProvider.java
+++ b/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 {
 
diff --git a/nifi-nar-bundles/nifi-framework-bundle/nifi-framework/nifi-web/nifi-web-security/src/main/java/org/apache/nifi/web/security/otp/OtpService.java b/nifi-nar-bundles/nifi-framework-bundle/nifi-framework/nifi-web/nifi-web-security/src/main/java/org/apache/nifi/web/security/otp/OtpService.java
index bcd26a4..e8d96ae 100644
--- a/nifi-nar-bundles/nifi-framework-bundle/nifi-framework/nifi-web/nifi-web-security/src/main/java/org/apache/nifi/web/security/otp/OtpService.java
+++ b/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,22 +16,18 @@
  */
 package org.apache.nifi.web.security.otp;
 
-import com.google.common.cache.Cache;
-import com.google.common.cache.CacheBuilder;
-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;
+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;
 
 /**
  * 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 Cache<CacheKey, String> uiExtensionCache;
+    private final TokenCache downloadTokens;
+    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();
-        uiExtensionCache = CacheBuilder.newBuilder().expireAfterWrite(duration, units).build();
+        downloadTokens = new TokenCache("download tokens", duration, units);
+        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) {
-        if (cache.size() >= MAX_CACHE_SOFT_LIMIT) {
-            throw new IllegalStateException("The maximum number of single use tokens have been issued.");
+    private String generateToken(final TokenCache tokenCache, final OtpAuthenticationToken authenticationToken) {
+        final String userId = (String) authenticationToken.getPrincipal();
+
+        // 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 {
-        final String authenticatedUser = cache.remove(new CacheKey(token));
+    private String getAuthenticationFromToken(final TokenCache tokenCache, final String token) throws OtpAuthenticationException {
+        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;
diff --git a/nifi-nar-bundles/nifi-framework-bundle/nifi-framework/nifi-web/nifi-web-security/src/main/java/org/apache/nifi/web/security/otp/TokenCache.java b/nifi-nar-bundles/nifi-framework-bundle/nifi-framework/nifi-web/nifi-web-security/src/main/java/org/apache/nifi/web/security/otp/TokenCache.java
new file mode 100644
index 0000000..a19efe2
--- /dev/null
+++ b/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();
+    }
+}
diff --git a/nifi-nar-bundles/nifi-framework-bundle/nifi-framework/nifi-web/nifi-web-security/src/test/groovy/org/apache/nifi/web/security/otp/TokenCacheTest.groovy b/nifi-nar-bundles/nifi-framework-bundle/nifi-framework/nifi-web/nifi-web-security/src/test/groovy/org/apache/nifi/web/security/otp/TokenCacheTest.groovy
new file mode 100644
index 0000000..64f8b6e
--- /dev/null
+++ b/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)
+    }
+}
diff --git a/nifi-nar-bundles/nifi-framework-bundle/nifi-framework/nifi-web/nifi-web-security/src/test/java/org/apache/nifi/web/security/otp/OtpServiceTest.java b/nifi-nar-bundles/nifi-framework-bundle/nifi-framework/nifi-web/nifi-web-security/src/test/java/org/apache/nifi/web/security/otp/OtpServiceTest.java
index 56fc9b1..6ff920d 100644
--- a/nifi-nar-bundles/nifi-framework-bundle/nifi-framework/nifi-web/nifi-web-security/src/test/java/org/apache/nifi/web/security/otp/OtpServiceTest.java
+++ b/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 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.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;
+
 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
-        Thread.sleep(4 * 1000);
+        // sleep for 2 seconds which should sufficiently expire the valid token
+        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
-        Thread.sleep(4 * 1000);
+        // sleep for 2 seconds which should sufficiently expire the valid token
+        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);
+    }
+}
\ No newline at end of file