JSSE: protect from calling freeSSL() while I/O operations are still in progress

Author: cconlon

src/java/com/wolfssl/provider/jsse/WolfSSLSocket.java

@@ -80,6 +80,10 @@ public class WolfSSLSocket extends SSLSocket {
     private WolfSSLInputStream inStream;
     private WolfSSLOutputStream outStream;
 
+    /* Track active I/O operations to prevent use-after-free */
+    private final java.util.concurrent.atomic.AtomicInteger activeOperations =
+        new java.util.concurrent.atomic.AtomicInteger(0);
+
     private ArrayList<HandshakeCompletedListener> hsListeners = null;
 
     /** TLS handshake initialization called */
@@ -1967,6 +1971,22 @@ public synchronized boolean sessionResumed() throws SSLException {
         return false;
     }
 
+    /**
+     * Helper method to track entry into an I/O operation.
+     * Increments the active operation counter to prevent premature freeSSL().
+     */
+    private void enterIOOperation() {
+        activeOperations.incrementAndGet();
+    }
+
+    /**
+     * Helper method to safely exit an I/O operation.
+     * Must be called for every successful enterIOOperation().
+     */
+    private void exitIOOperation() {
+        activeOperations.decrementAndGet();
+    }
+
     /**
      * Internal private method to check if WolfSSLInputStream
      * and WolfSSLOutputStream are closed.
@@ -2157,11 +2177,11 @@ public synchronized void close() throws IOException {
                     /* Connection is closed, free native WOLFSSL session
                      * to release native memory earlier than garbage
                      * collector might with finalize(), Don't free if we
-                     * have threads still waiting in poll/select or if
+                     * have threads still waiting in poll/select, if
                      * our WolfSSLInputStream or WolfSSLOutputStream are
-                     * still open. */
-                    if (this.ssl.getThreadsBlockedInPoll() == 0 &&
-                        ioStreamsAreClosed()) {
+                     * still open, or if there are active I/O operations. */
+                    if ((this.ssl.getThreadsBlockedInPoll() == 0) &&
+                        ioStreamsAreClosed() && (activeOperations.get() == 0)) {
                         WolfSSLDebug.log(getClass(), WolfSSLDebug.INFO,
                             () -> "calling this.ssl.freeSSL()");
                         this.ssl.freeSSL();
@@ -2171,7 +2191,8 @@ public synchronized void close() throws IOException {
                             () -> "deferring freeing this.ssl, threads " +
                             "blocked in poll: " +
                             this.ssl.getThreadsBlockedInPoll() +
-                            ", or streams not closed");
+                            ", streams not closed, or active operations: " +
+                            activeOperations.get());
                     }
 
                     /* Reset internal WolfSSLEngineHelper to null */
@@ -2710,6 +2731,9 @@ public synchronized int read(byte[] b, int off, int len)
                     "Array index out of bounds");
             }
 
+            /* Enter I/O operation to prevent use-after-free */
+            socket.enterIOOperation();
+
             try {
                 int err;
                 int timeout = socket.getSoTimeout();
@@ -2775,6 +2799,9 @@ public synchronized int read(byte[] b, int off, int len)
                  * IllegalStateException to be thrown from
                  * WolfSSLSession.read(). Return as a SocketException here. */
                 throw new SocketException(e.getMessage());
+            } finally {
+                /* Exit I/O operation */
+                socket.exitIOOperation();
             }
 
             /* return number of bytes read */
@@ -2921,6 +2948,9 @@ public synchronized void write(byte[] b, int off, int len)
                     "Array index out of bounds");
             }
 
+            /* Enter I/O operation to prevent use-after-free */
+            socket.enterIOOperation();
+
             try {
                 int err;
                 int timeout = socket.getSoTimeout();
@@ -2962,6 +2992,9 @@ public synchronized void write(byte[] b, int off, int len)
                                  () -> "got IllegalStateException: " + e +
                                  ", throwing IOException");
                 throw new IOException(e);
+            } finally {
+                /* Exit I/O operation */
+                socket.exitIOOperation();
             }
         }
     } /* end WolfSSLOutputStream inner class */