Guard Executor/Scheduler unwrapping against implementations returning…

… themselves. If unwrap isn't able to unwrap then it would return itself causing a cast exception. [resolves #220] Signed-off-by: Mark Paluch <[email protected]>
r2dbc · Nov 28, 2024 · 55f8f30 · 55f8f30
1 parent bf6540e
commit 55f8f30
Show file tree

Hide file tree

Showing 2 changed files with 212 additions and 151 deletions.
diff --git a/src/main/java/io/r2dbc/pool/ConnectionPool.java b/src/main/java/io/r2dbc/pool/ConnectionPool.java
@@ -108,22 +108,10 @@ public ConnectionPool(ConnectionPoolConfiguration configuration) {
 
                         Connection connection = ref.poolable();
                         Scheduler scheduler = null;
-                        Executor executor = null;
                         Mono<Connection> conn;
 
                         if (connection instanceof Wrapped<?>) {
-
-                            Wrapped<?> wrapped = (Wrapped<?>) connection;
-
-                            scheduler = wrapped.unwrap(Scheduler.class);
-
-                            if (scheduler == null) {
-                                executor = wrapped.unwrap(Executor.class);
-                            }
-
-                            if (executor != null) {
-                                scheduler = Schedulers.fromExecutor(executor);
-                            }
+                            scheduler = findScheduler((Wrapped<?>) connection);
                         }
 
                         if (scheduler != null) {
@@ -167,6 +155,27 @@ public ConnectionPool(ConnectionPoolConfiguration configuration) {
         this.create = configuration.getAcquireRetry() > 0 ? create.retry(configuration.getAcquireRetry()) : create;
     }
 
+    @Nullable
+    @SuppressWarnings("DataFlowIssue")
+    private static Scheduler findScheduler(Wrapped<?> connection) {
+
+        Object unwrapped = connection.unwrap(Scheduler.class);
+
+        if (!(unwrapped instanceof Scheduler)) {
+            unwrapped = connection.unwrap(Executor.class);
+        }
+
+        if (unwrapped instanceof Executor) {
+            unwrapped = Schedulers.fromExecutor((Executor) unwrapped);
+        }
+
+        if (unwrapped instanceof Scheduler) {
+            return (Scheduler) unwrapped;
+        }
+
+        return null;
+    }
+
     private Mono<Connection> prepareConnection(ConnectionPoolConfiguration configuration, PooledRef<Connection> ref, Connection connection, Function<Connection, Mono<Void>> allocateValidation) {
 
         Mono<Void> prepare = null;