From 2521225a60fdf20ac4d545cbac11fef14fd64483 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Roman=20=C5=A0trobl?= <roman.strobl@wultra.com>
Date: Mon, 11 Dec 2023 20:14:49 +0100
Subject: [PATCH 1/2] Fix #1184: Memory leak in CallbackUrlBehavior for
 RestClient instances (#1185)

(cherry picked from commit 33de1e0eebf42913ebc76cbbb2fbc2c681802c05)
(cherry picked from commit d090724cd1d3e8f81a62671fe84e0781878c8f61)
---
 .../behavior/tasks/CallbackUrlBehavior.java   | 55 +++++++++++++++++--
 1 file changed, 50 insertions(+), 5 deletions(-)

diff --git a/powerauth-java-server/src/main/java/io/getlime/security/powerauth/app/server/service/behavior/tasks/CallbackUrlBehavior.java b/powerauth-java-server/src/main/java/io/getlime/security/powerauth/app/server/service/behavior/tasks/CallbackUrlBehavior.java
index 8f2c50f0f..953f94612 100644
--- a/powerauth-java-server/src/main/java/io/getlime/security/powerauth/app/server/service/behavior/tasks/CallbackUrlBehavior.java
+++ b/powerauth-java-server/src/main/java/io/getlime/security/powerauth/app/server/service/behavior/tasks/CallbackUrlBehavior.java
@@ -49,6 +49,7 @@
 import java.net.MalformedURLException;
 import java.net.URL;
 import java.util.*;
+import java.util.concurrent.ConcurrentHashMap;
 import java.util.function.Consumer;
 import java.util.stream.Collectors;
 
@@ -65,7 +66,10 @@ public class CallbackUrlBehavior {
     private LocalizationProvider localizationProvider;
     private PowerAuthServiceConfiguration configuration;
 
-    private final Map<CallbackUrlEntity, RestClient> restClientCache = new HashMap<>();
+    // Store REST clients in cache with their callback ID as a key
+    private final Map<String, RestClient> restClientCache = new ConcurrentHashMap<>();
+    private final Object restClientCacheLock = new Object();
+
     private final CallbackAuthenticationPublicConverter authenticationPublicConverter = new CallbackAuthenticationPublicConverter();
 
     // Prepare logger
@@ -170,6 +174,8 @@ public UpdateCallbackUrlResponse updateCallbackUrl(UpdateCallbackUrlRequest requ
         }
 
         final CallbackUrlEntity entity = entityOptional.get();
+        evictRestClientFromCache(entity);
+
         entity.setName(request.getName());
         entity.setCallbackUrl(request.getCallbackUrl());
         entity.setAttributes(request.getAttributes());
@@ -244,7 +250,9 @@ public RemoveCallbackUrlResponse removeCallbackUrl(RemoveCallbackUrlRequest requ
         response.setId(request.getId());
         final Optional<CallbackUrlEntity> callbackUrlEntityOptional = callbackUrlRepository.findById(request.getId());
         if (callbackUrlEntityOptional.isPresent()) {
-            callbackUrlRepository.delete(callbackUrlEntityOptional.get());
+            final CallbackUrlEntity callbackEntity = callbackUrlEntityOptional.get();
+            evictRestClientFromCache(callbackEntity);
+            callbackUrlRepository.delete(callbackEntity);
             response.setRemoved(true);
         } else {
             response.setRemoved(false);
@@ -415,15 +423,52 @@ private void notifyCallbackUrl(CallbackUrlEntity callbackUrlEntity, Map<String,
      * @return Rest client.
      * @throws RestClientException Thrown when rest client initialization fails.
      */
-    private synchronized RestClient getRestClient(CallbackUrlEntity callbackUrlEntity) throws RestClientException {
-        RestClient restClient = restClientCache.get(callbackUrlEntity);
+    private RestClient getRestClient(final CallbackUrlEntity callbackUrlEntity) throws RestClientException {
+        final String cacheKey = getRestClientCacheKey(callbackUrlEntity);
+        RestClient restClient = restClientCache.get(cacheKey);
         if (restClient == null) {
+            logger.debug("REST client not found in cache, initializing new REST client, callback cache key: {}", cacheKey);
+            restClient = createRestClientAndStoreInCache(callbackUrlEntity);
+        } else {
+            logger.debug("REST client found in cache, callback cache key: {}", cacheKey);
+        }
+        return restClient;
+    }
+
+    /**
+     * Get a key for the REST client cache from a callback URL entity.
+     * @param callbackUrlEntity Callback URL entity.
+     * @return Cache key.
+     */
+    private String getRestClientCacheKey(final CallbackUrlEntity callbackUrlEntity) {
+        return callbackUrlEntity.getId();
+    }
+
+    /**
+     * Create a new REST client and store it in cache for given callback URL entity.
+     * @param callbackUrlEntity Callback URL entity.
+     * @return Rest client.
+     */
+    public RestClient createRestClientAndStoreInCache(final CallbackUrlEntity callbackUrlEntity) throws RestClientException {
+        final String cacheKey = getRestClientCacheKey(callbackUrlEntity);
+        final RestClient restClient;
+        synchronized (restClientCacheLock) {
             restClient = initializeRestClient(callbackUrlEntity);
-            restClientCache.put(callbackUrlEntity, restClient);
+            restClientCache.put(cacheKey, restClient);
         }
         return restClient;
     }
 
+    /**
+     * Evict an instance from REST client cache for given callback URL entity.
+     * @param callbackUrlEntity Callback URL entity.
+     */
+    private void evictRestClientFromCache(final CallbackUrlEntity callbackUrlEntity) {
+        synchronized (restClientCacheLock) {
+            restClientCache.remove(getRestClientCacheKey(callbackUrlEntity));
+        }
+    }
+
     /**
      * Initialize Rest client instance and configure it based on client configuration.
      * @param callbackUrlEntity Callback URL entity.

From df3c1cc9aab6caffb60d6728f518a6d8652a4c11 Mon Sep 17 00:00:00 2001
From: Roman Strobl <roman.strobl@wultra.com>
Date: Tue, 12 Dec 2023 12:40:31 +0100
Subject: [PATCH 2/2] Fix #1199: Synchronize get access to REST client cache

(cherry picked from commit 683eebf3ec225bebcce9ddbd1939b67a88be6e6d)
(cherry picked from commit 282253718cb9afc899ee3c5849acec638e8280cd)
---
 .../behavior/tasks/CallbackUrlBehavior.java   | 22 +++++++++----------
 1 file changed, 11 insertions(+), 11 deletions(-)

diff --git a/powerauth-java-server/src/main/java/io/getlime/security/powerauth/app/server/service/behavior/tasks/CallbackUrlBehavior.java b/powerauth-java-server/src/main/java/io/getlime/security/powerauth/app/server/service/behavior/tasks/CallbackUrlBehavior.java
index 953f94612..d08be7111 100644
--- a/powerauth-java-server/src/main/java/io/getlime/security/powerauth/app/server/service/behavior/tasks/CallbackUrlBehavior.java
+++ b/powerauth-java-server/src/main/java/io/getlime/security/powerauth/app/server/service/behavior/tasks/CallbackUrlBehavior.java
@@ -425,12 +425,15 @@ private void notifyCallbackUrl(CallbackUrlEntity callbackUrlEntity, Map<String,
      */
     private RestClient getRestClient(final CallbackUrlEntity callbackUrlEntity) throws RestClientException {
         final String cacheKey = getRestClientCacheKey(callbackUrlEntity);
-        RestClient restClient = restClientCache.get(cacheKey);
-        if (restClient == null) {
-            logger.debug("REST client not found in cache, initializing new REST client, callback cache key: {}", cacheKey);
-            restClient = createRestClientAndStoreInCache(callbackUrlEntity);
-        } else {
-            logger.debug("REST client found in cache, callback cache key: {}", cacheKey);
+        RestClient restClient;
+        synchronized (restClientCacheLock) {
+            restClient = restClientCache.get(cacheKey);
+            if (restClient == null) {
+                logger.debug("REST client not found in cache, initializing new REST client, callback cache key: {}", cacheKey);
+                restClient = createRestClientAndStoreInCache(callbackUrlEntity);
+            } else {
+                logger.debug("REST client found in cache, callback cache key: {}", cacheKey);
+            }
         }
         return restClient;
     }
@@ -451,11 +454,8 @@ private String getRestClientCacheKey(final CallbackUrlEntity callbackUrlEntity)
      */
     public RestClient createRestClientAndStoreInCache(final CallbackUrlEntity callbackUrlEntity) throws RestClientException {
         final String cacheKey = getRestClientCacheKey(callbackUrlEntity);
-        final RestClient restClient;
-        synchronized (restClientCacheLock) {
-            restClient = initializeRestClient(callbackUrlEntity);
-            restClientCache.put(cacheKey, restClient);
-        }
+        final RestClient restClient = initializeRestClient(callbackUrlEntity);
+        restClientCache.put(cacheKey, restClient);
         return restClient;
     }