Store MFA authentication status directly on tokens

  Move MFA status from being computed from user identities to being stored
  as a boolean field on tokens themselves. This provides per-token MFA
  tracking and eliminates the need for complex identity lookups.

Author: dauglyon

src/main/java/us/kbase/auth2/lib/Authentication.java

@@ -745,13 +745,19 @@ public void forceResetAllPasswords(final IncomingToken token)
 
 	private NewToken login(final UserName userName, final TokenCreationContext tokenCtx)
 			throws AuthStorageException {
+		return login(userName, tokenCtx, false);
+	}
+	
+	private NewToken login(final UserName userName, final TokenCreationContext tokenCtx, 
+			final Boolean mfaAuthenticated) throws AuthStorageException {
 		final NewToken nt = new NewToken(StoredToken.getBuilder(
-					TokenType.LOGIN, randGen.randomUUID(), userName)
-				.withLifeTime(clock.instant(),
-						cfg.getAppConfig().getTokenLifetimeMS(TokenLifetimeType.LOGIN))
-				.withContext(tokenCtx)
-				.build(),
-				randGen.getToken());
+				TokenType.LOGIN, randGen.randomUUID(), userName)
+			.withLifeTime(clock.instant(),
+					cfg.getAppConfig().getTokenLifetimeMS(TokenLifetimeType.LOGIN))
+			.withContext(tokenCtx)
+			.withMfaAuthenticated(mfaAuthenticated)
+			.build(),
+			randGen.getToken());
 		storage.storeToken(nt.getStoredToken(), nt.getTokenHash());
 		setLastLogin(userName);
 		logInfo("Logged in user {} with token {}",
@@ -905,7 +911,9 @@ public NewToken createToken(
 		final NewToken nt = new NewToken(StoredToken.getBuilder(tokenType, id, au.getUserName())
 				.withLifeTime(clock.instant(), life)
 				.withContext(tokenCtx)
-				.withTokenName(tokenName).build(),
+				.withTokenName(tokenName)
+				.withMfaAuthenticated(null) // Agent/Dev/Serv tokens don't have MFA status
+				.build(),
 				randGen.getToken());
 		storage.storeToken(nt.getStoredToken(), nt.getTokenHash());
 		logInfo("User {} created {} token {}", au.getUserName().getName(), tokenType, id);
@@ -2043,6 +2051,7 @@ public NewToken testModeCreateToken(
 		final NewToken nt = new NewToken(StoredToken.getBuilder(tokenType, id, userName)
 				.withLifeTime(clock.instant(), TEST_MODE_DATA_LIFETIME_MS)
 				.withNullableTokenName(tokenName)
+				.withMfaAuthenticated(null) // Test mode tokens don't have MFA status
 				.build(),
 				randGen.getToken());
 		storage.testModeStoreToken(nt.getStoredToken(), nt.getTokenHash());
@@ -2339,7 +2348,9 @@ public NewToken login(
 						linked, u.get().getUserName().getName());
 			}
 		}
-		return login(u.get().getUserName(), tokenCtx);
+		final Boolean mfaStatus = ri.get().getDetails() != null ? 
+				ri.get().getDetails().getMfaAuthenticated() : null;
+		return login(u.get().getUserName(), tokenCtx, mfaStatus);
 	}
 
 	private Optional<RemoteIdentity> getIdentity(
@@ -3142,52 +3153,6 @@ public long getSuggestedTokenCacheTime() throws AuthStorageException {
 		return cfg.getAppConfig().getTokenLifetimeMS(TokenLifetimeType.EXT_CACHE);
 	}
 
-	/** Get MFA status for a token by checking the token user's identities.
-	 * @param token the token to check.
-	 * @return true if MFA was used, false if password only, null if unknown or not supported.
-	 * @throws AuthStorageException if an error occurred accessing the storage system.
-	 * @throws InvalidTokenException if the token is invalid.
-	 */
-	public Boolean getMfaStatus(final IncomingToken token) 
-			throws AuthStorageException, InvalidTokenException {
-		if (token == null) {
-			return null;
-		}
-		
-		// Get the stored token to find the username
-		final StoredToken storedToken = getToken(token);
-		final UserName userName = storedToken.getUserName();
-		
-		// Check if the user exists before trying to get user details
-		try {
-			storage.getUser(userName);
-		} catch (NoSuchUserException e) {
-			// User doesn't exist, return null for MFA status
-			return null;
-		}
-		
-		try {
-			final AuthUser user = getUser(token);
-			final Set<us.kbase.auth2.lib.identity.RemoteIdentity> identities = user.getIdentities();
-			
-			// Check for identities with MFA information from supported providers
-			if (identities != null) {
-				for (final us.kbase.auth2.lib.identity.RemoteIdentity identity : identities) {
-					if (identity != null && identity.getDetails() != null) {
-						final Boolean mfaStatus = identity.getDetails().getMfaAuthenticated();
-						if (mfaStatus != null) {
-							return mfaStatus;
-						}
-					}
-				}
-			}
-			
-			return null; // No MFA information available
-		} catch (DisabledUserException e) {
-			// Return null for disabled users
-			return null;
-		}
-	}
 
 	/** Get the external configuration without providing any credentials.
 	 * 

src/main/java/us/kbase/auth2/lib/storage/mongo/Fields.java

@@ -137,6 +137,8 @@ public class Fields {
 	public static final String TOKEN_CUSTOM_KEY = "k";
 	/** A value for a custom context key / value pair. */
 	public static final String TOKEN_CUSTOM_VALUE = "v";
+	/** Whether the token was created with multi-factor authentication. */
+	public static final String TOKEN_MFA_AUTHENTICATED = "mfaauth";
 
 	/* ************************
 	 * temporary session data fields

src/main/java/us/kbase/auth2/lib/storage/mongo/MongoStorage.java

@@ -871,7 +871,8 @@ private void storeToken(final String collection, final StoredToken token, final
 				.append(Fields.TOKEN_DEVICE, ctx.getDevice().orElse(null))
 				.append(Fields.TOKEN_IP, ctx.getIpAddress().isPresent() ?
 						ctx.getIpAddress().get().getHostAddress() : null)
-				.append(Fields.TOKEN_CUSTOM_CONTEXT, toCustomContextList(ctx.getCustomContext()));
+				.append(Fields.TOKEN_CUSTOM_CONTEXT, toCustomContextList(ctx.getCustomContext()))
+				.append(Fields.TOKEN_MFA_AUTHENTICATED, token.getMfaAuthenticated());
 		try {
 			db.getCollection(collection).insertOne(td);
 		} catch (MongoWriteException mwe) {
@@ -969,6 +970,7 @@ private StoredToken getToken(final Document t) throws AuthStorageException {
 						t.getDate(Fields.TOKEN_EXPIRY).toInstant())
 				.withNullableTokenName(getTokenName(t.getString(Fields.TOKEN_NAME)))
 				.withContext(toTokenCreationContext(t))
+				.withMfaAuthenticated(t.getBoolean(Fields.TOKEN_MFA_AUTHENTICATED))
 				.build();
 	}
 

src/main/java/us/kbase/auth2/lib/token/StoredToken.java

@@ -23,6 +23,7 @@ public class StoredToken {
 	private final UserName userName;
 	private final Instant creationDate;
 	private final Instant expirationDate;
+	private final Boolean mfaAuthenticated;
 
 	private StoredToken(
 			final UUID id,
@@ -31,7 +32,8 @@ private StoredToken(
 			final UserName userName,
 			final TokenCreationContext context,
 			final Instant creationDate,
-			final Instant expirationDate) {
+			final Instant expirationDate,
+			final Boolean mfaAuthenticated) {
 		// this stuff is here just in case naughty users use casting to skip a builder step
 		requireNonNull(creationDate, "created");
 		// no way to test this one
@@ -43,6 +45,7 @@ private StoredToken(
 		this.expirationDate = expirationDate;
 		this.creationDate = creationDate;
 		this.id = id;
+		this.mfaAuthenticated = mfaAuthenticated;
 	}
 
 	/** Get the type of the token.
@@ -93,6 +96,13 @@ public Instant getCreationDate() {
 	public Instant getExpirationDate() {
 		return expirationDate;
 	}
+
+	/** Get whether the token was created with multi-factor authentication.
+	 * @return true if the token was created with MFA, false if not, null if unknown.
+	 */
+	public Boolean getMfaAuthenticated() {
+		return mfaAuthenticated;
+	}
 
 	@Override
 	public int hashCode() {
@@ -105,6 +115,7 @@ public int hashCode() {
 		result = prime * result + ((tokenName == null) ? 0 : tokenName.hashCode());
 		result = prime * result + ((type == null) ? 0 : type.hashCode());
 		result = prime * result + ((userName == null) ? 0 : userName.hashCode());
+		result = prime * result + ((mfaAuthenticated == null) ? 0 : mfaAuthenticated.hashCode());
 		return result;
 	}
 
@@ -165,6 +176,13 @@ public boolean equals(Object obj) {
 		} else if (!userName.equals(other.userName)) {
 			return false;
 		}
+		if (mfaAuthenticated == null) {
+			if (other.mfaAuthenticated != null) {
+				return false;
+			}
+		} else if (!mfaAuthenticated.equals(other.mfaAuthenticated)) {
+			return false;
+		}
 		return true;
 	}
 
@@ -224,6 +242,12 @@ public interface OptionalsStep {
 		 */
 		OptionalsStep withContext(TokenCreationContext context);
 
+		/** Specify whether the token was created with multi-factor authentication.
+		 * @param mfaAuthenticated true if MFA was used, false if not, null if unknown.
+		 * @return this builder.
+		 */
+		OptionalsStep withMfaAuthenticated(Boolean mfaAuthenticated);
+		
 		/** Build the token.
 		 * @return a new StoredToken.
 		 */
@@ -239,6 +263,7 @@ private static class Builder implements LifeStep, OptionalsStep {
 		private final UserName userName;
 		private Instant creationDate;
 		private Instant expirationDate;
+		private Boolean mfaAuthenticated;
 
 		private Builder(final TokenType type, final UUID id, final UserName userName) {
 			requireNonNull(type, "type");
@@ -269,10 +294,16 @@ public OptionalsStep withContext(final TokenCreationContext context) {
 			return this;
 		}
 
+		@Override
+		public OptionalsStep withMfaAuthenticated(final Boolean mfaAuthenticated) {
+			this.mfaAuthenticated = mfaAuthenticated;
+			return this;
+		}
+
 		@Override
 		public StoredToken build() {
 			return new StoredToken(id, type, tokenName, userName, context,
-					creationDate, expirationDate);
+					creationDate, expirationDate, mfaAuthenticated);
 		}
 
 		@Override

src/main/java/us/kbase/auth2/service/api/APIToken.java

@@ -10,31 +10,10 @@ public class APIToken extends ExternalToken {
 	private final long cachefor;
 	private final Boolean mfaAuthenticated;
 
-	/**
-	 * Constructor without MFA status calculation.
-	 * MFA status will be set to null.
-	 * 
-	 * @param token the stored token
-	 * @param tokenCacheTimeMillis the token cache time in milliseconds
-	 */
 	public APIToken(final StoredToken token, final long tokenCacheTimeMillis) {
 		super(token);
 		cachefor = tokenCacheTimeMillis;
-		mfaAuthenticated = null;
-	}
-	
-	/**
-	 * Constructor with MFA status provided.
-	 * 
-	 * @param token the stored token
-	 * @param tokenCacheTimeMillis the token cache time in milliseconds
-	 * @param mfaAuthenticated the MFA authentication status
-	 */
-	public APIToken(final StoredToken token, final long tokenCacheTimeMillis, 
-			final Boolean mfaAuthenticated) {
-		super(token);
-		cachefor = tokenCacheTimeMillis;
-		this.mfaAuthenticated = mfaAuthenticated;
+		mfaAuthenticated = token.getMfaAuthenticated();
 	}
 
 	public long getCachefor() {

src/main/java/us/kbase/auth2/service/api/Token.java

@@ -58,7 +58,7 @@ public APIToken viewToken(@HeaderParam(APIConstants.HEADER_TOKEN) final String t
 			throws NoTokenProvidedException, InvalidTokenException, AuthStorageException {
 		final IncomingToken it = getToken(token);
 		final StoredToken ht = auth.getToken(it);
-		return new APIToken(ht, auth.getSuggestedTokenCacheTime(), auth.getMfaStatus(it));
+		return new APIToken(ht, auth.getSuggestedTokenCacheTime());
 	}
 
 	private static class CreateToken extends IncomingJSON {

src/test/java/us/kbase/test/auth2/service/api/TokenEndpointTest.java

@@ -180,6 +180,60 @@ TokenType.AGENT, id, new UserName("foo"))
 		assertThat("incorrect token", response, is(expected));
 	}
 
+	@Test
+	public void getTokenMfaIsolation() throws Exception {
+		// Create user
+		final UserName userName = new UserName("testuser");
+		final UUID userUuid = UUID.randomUUID();
+		manager.storage.createLocalUser(LocalUser.getLocalUserBuilder(
+				userName, userUuid, new DisplayName("Test User"), inst(10000))
+				.withEmailAddress(new EmailAddress("[email protected]")).build(),
+				new PasswordHashAndSalt("password".getBytes(), "salt".getBytes()));
+		
+		// Create two tokens with different MFA status
+		final UUID token1Id = UUID.randomUUID();
+		final UUID token2Id = UUID.randomUUID();
+		
+		final IncomingToken token1 = new IncomingToken("token1hash");
+		final IncomingToken token2 = new IncomingToken("token2hash");
+		
+		// Token 1 with MFA=true
+		manager.storage.storeToken(StoredToken.getBuilder(
+				TokenType.LOGIN, token1Id, userName)
+				.withLifeTime(Instant.ofEpochMilli(10000), Instant.ofEpochMilli(1000000000000000L))
+				.withMfaAuthenticated(true)
+				.build(), token1.getHashedToken().getTokenHash());
+		
+		// Token 2 with MFA=false
+		manager.storage.storeToken(StoredToken.getBuilder(
+				TokenType.LOGIN, token2Id, userName)
+				.withLifeTime(Instant.ofEpochMilli(10000), Instant.ofEpochMilli(1000000000000000L))
+				.withMfaAuthenticated(false)
+				.build(), token2.getHashedToken().getTokenHash());
+		
+		// Test token1 returns MFA=true
+		final URI target1 = UriBuilder.fromUri(host).path("/api/V2/token").build();
+		final WebTarget wt1 = CLI.target(target1);
+		final Builder req1 = wt1.request().header("authorization", token1.getToken());
+		final Response res1 = req1.get();
+		
+		assertThat("incorrect response code for token1", res1.getStatus(), is(200));
+		@SuppressWarnings("unchecked")
+		final Map<String, Object> response1 = res1.readEntity(Map.class);
+		assertThat("token1 should have MFA=true", response1.get("mfaAuthenticated"), is(true));
+		
+		// Test token2 returns MFA=false
+		final URI target2 = UriBuilder.fromUri(host).path("/api/V2/token").build();
+		final WebTarget wt2 = CLI.target(target2);
+		final Builder req2 = wt2.request().header("authorization", token2.getToken());
+		final Response res2 = req2.get();
+		
+		assertThat("incorrect response code for token2", res2.getStatus(), is(200));
+		@SuppressWarnings("unchecked")
+		final Map<String, Object> response2 = res2.readEntity(Map.class);
+		assertThat("token2 should have MFA=false", response2.get("mfaAuthenticated"), is(false));
+	}
+	
 	@Test
 	public void getTokenFailNoToken() throws Exception {
 		final URI target = UriBuilder.fromUri(host).path("/api/V2/token").build();