From 4b2a71e31d407748281e732dc7bb487d9b177953 Mon Sep 17 00:00:00 2001
From: Andrea Di Cesare <andrea@softinstigate.com>
Date: Wed, 13 Nov 2024 14:05:07 +0100
Subject: [PATCH] :sparkles: Add authDb to the MongoAccount to track the db
 used during authentication

---
 .../MongoRealmAuthenticator.java              | 30 +++++++------------
 .../security/tokens/JwtTokenManager.java      |  4 +--
 2 files changed, 12 insertions(+), 22 deletions(-)

diff --git a/security/src/main/java/org/restheart/security/authenticators/MongoRealmAuthenticator.java b/security/src/main/java/org/restheart/security/authenticators/MongoRealmAuthenticator.java
index 1a4b62911..8b4b8f8cd 100644
--- a/security/src/main/java/org/restheart/security/authenticators/MongoRealmAuthenticator.java
+++ b/security/src/main/java/org/restheart/security/authenticators/MongoRealmAuthenticator.java
@@ -27,6 +27,7 @@
 import java.util.Map;
 
 import org.bson.BsonDocument;
+import org.bson.BsonString;
 import org.mindrot.jbcrypt.BCrypt;
 import org.restheart.cache.Cache;
 import org.restheart.cache.CacheFactory;
@@ -146,7 +147,7 @@ public void init() {
         if (cacheEnabled) {
             final int cacheSize = arg(config, "cache-size");
             final int cacheTTL = arg(config, "cache-ttl");
-            this.USERS_CACHE = CacheFactory.createLocalLoadingCache(cacheSize, this.cacheExpirePolicy, cacheTTL, key -> findAccount(accountIdTrasformer(key)));
+            this.USERS_CACHE = CacheFactory.createLocalLoadingCache(cacheSize, this.cacheExpirePolicy, cacheTTL, key -> findAccount(key));
         }
 
         try {
@@ -354,7 +355,7 @@ private MongoRealmAccount getAccount(final String usersDb, final String id) {
         final var cacheKey = new CacheKey(id, usersDb);
 
         if (USERS_CACHE == null) {
-            return findAccount(this.accountIdTrasformer(cacheKey));
+            return findAccount(cacheKey);
         } else {
             final var _account = USERS_CACHE.getLoading(cacheKey);
 
@@ -366,18 +367,6 @@ private MongoRealmAccount getAccount(final String usersDb, final String id) {
         }
     }
 
-    /**
-     * Override this method to trasform the account id. By default it returns
-     * the id without any transformation. For example, it could be overridden to
-     * force the id to be lowercase.
-     *
-     * @param key the cache key
-     * @return the trasformed cache key(default is identity)
-     */
-    protected CacheKey accountIdTrasformer(final CacheKey key) {
-        return key;
-    }
-
     /**
      * if client authenticates passing the real credentials, update the account
      * in the auth-token cache, otherwise the client authenticating with the
@@ -448,7 +437,7 @@ public boolean checkUserCollection() throws IllegalStateException {
         return true;
     }
 
-    public long countAccounts() {
+    private long countAccounts() {
         try {
             return mclient.getDatabase(this.getUsersDb()).getCollection(this.getUsersCollection()).estimatedDocumentCount();
         } catch (final Throwable t) {
@@ -457,7 +446,7 @@ public long countAccounts() {
         }
     }
 
-    public void createDefaultAccount() {
+    private void createDefaultAccount() {
         if (this.mclient == null) {
             LOGGER.error("Cannot find account: mongo service is not enabled.");
             return;
@@ -472,7 +461,7 @@ public void createDefaultAccount() {
         }
     }
 
-    public MongoRealmAccount findAccount(final CacheKey key) {
+    private MongoRealmAccount findAccount(final CacheKey key) {
         final var coll = mclient.getDatabase(key.db()).getCollection(this.getUsersCollection()).withDocumentClass(BsonDocument.class);
 
         BsonDocument _account;
@@ -542,9 +531,10 @@ public MongoRealmAccount findAccount(final CacheKey key) {
             }
         });
 
-        return new MongoRealmAccount(key.db(), key.id(), _password.getAsJsonPrimitive().getAsString().toCharArray(), roles,
-                // used this because password has been removed from account
-                BsonDocument.parse(account.toString()));
+        var properties = BsonDocument.parse(account.toString()); // used this because password has been removed from account
+        properties.put("authDb", new BsonString(key.db()));
+
+        return new MongoRealmAccount(key.db(), key.id(), _password.getAsJsonPrimitive().getAsString().toCharArray(), roles, properties);
     }
 
     /**
diff --git a/security/src/main/java/org/restheart/security/tokens/JwtTokenManager.java b/security/src/main/java/org/restheart/security/tokens/JwtTokenManager.java
index 4a88c0206..8f1f0dede 100644
--- a/security/src/main/java/org/restheart/security/tokens/JwtTokenManager.java
+++ b/security/src/main/java/org/restheart/security/tokens/JwtTokenManager.java
@@ -231,7 +231,7 @@ private Token newToken(Account account, Date expires) {
             .withSubject(account.getPrincipal().getName())
             .withExpiresAt(expires)
             .withIssuer(issuer)
-            .withArrayClaim("roles", (String[]) account.getRoles().toArray());
+            .withArrayClaim("roles", account.getRoles().toArray(new String[account.getRoles().size()]));
 
         Builder[] builder = { _builder };
 
@@ -246,7 +246,7 @@ private Token newToken(Account account, Date expires) {
 
         var raw = builder[0].sign(algo);
 
-        return new Token(raw.toCharArray(), expires, (String[]) account.getRoles().toArray(), properties);
+        return new Token(raw.toCharArray(), expires, account.getRoles().toArray(new String[account.getRoles().size()]), properties);
     }