Merge "Wrap cached sessions before returning" into nyc-dev

src/main/java/org/conscrypt/AbstractSessionContext.java

@@ -241,7 +241,7 @@ abstract class AbstractSessionContext implements SSLSessionContext {
      *
      * @return a session or null if the session can't be converted
      */
-    SSLSession toSession(byte[] data, String host, int port) {
+    OpenSSLSessionImpl toSession(byte[] data, String host, int port) {
         ByteArrayInputStream bais = new ByteArrayInputStream(data);
         DataInputStream dais = new DataInputStream(bais);
         try {
@@ -271,6 +271,14 @@ abstract class AbstractSessionContext implements SSLSessionContext {
         }
     }
 
+    protected SSLSession wrapSSLSessionIfNeeded(SSLSession session) {
+        if (session instanceof OpenSSLSessionImpl) {
+            return Platform.wrapSSLSession((OpenSSLSessionImpl) session);
+        } else {
+            return session;
+        }
+    }
+
     @Override
     public SSLSession getSession(byte[] sessionId) {
         if (sessionId == null) {
@@ -282,7 +290,11 @@ abstract class AbstractSessionContext implements SSLSessionContext {
             session = sessions.get(key);
         }
         if (session != null && session.isValid()) {
-            return session;
+            if (session instanceof OpenSSLSessionImpl) {
+                return Platform.wrapSSLSession((OpenSSLSessionImpl) session);
+            } else {
+                return session;
+            }
         }
         return null;
     }

src/main/java/org/conscrypt/ClientSessionContext.java

@@ -30,8 +30,7 @@ public class ClientSessionContext extends AbstractSessionContext {
      * Sessions indexed by host and port. Protect from concurrent
      * access by holding a lock on sessionsByHostAndPort.
      */
-    final Map<HostAndPort, SSLSession> sessionsByHostAndPort
-        = new HashMap<HostAndPort, SSLSession>();
+    private final HashMap<HostAndPort, SSLSession> sessionsByHostAndPort = new HashMap<>();
 
     private SSLClientSessionCache persistentCache;
 
@@ -77,7 +76,7 @@ public class ClientSessionContext extends AbstractSessionContext {
             session = sessionsByHostAndPort.get(hostAndPortKey);
         }
         if (session != null && session.isValid()) {
-            return session;
+            return wrapSSLSessionIfNeeded(session);
         }
 
         // Look in persistent cache.
@@ -90,7 +89,7 @@ public class ClientSessionContext extends AbstractSessionContext {
                     synchronized (sessionsByHostAndPort) {
                         sessionsByHostAndPort.put(hostAndPortKey, session);
                     }
-                    return session;
+                    return wrapSSLSessionIfNeeded(session);
                 }
             }
         }

src/main/java/org/conscrypt/OpenSSLExtendedSessionImpl.java

@@ -37,6 +37,10 @@ public class OpenSSLExtendedSessionImpl extends ExtendedSSLSession {
         this.delegate = delegate;
     }
 
+    public OpenSSLSessionImpl getDelegate() {
+        return delegate;
+    }
+
     public String[] getLocalSupportedSignatureAlgorithms() {
         throw new UnsupportedOperationException();
     }

src/main/java/org/conscrypt/SSLParametersImpl.java

@@ -43,6 +43,7 @@ import javax.crypto.SecretKey;
 import javax.net.ssl.KeyManager;
 import javax.net.ssl.KeyManagerFactory;
 import javax.net.ssl.SSLException;
+import javax.net.ssl.SSLSession;
 import javax.net.ssl.SSLHandshakeException;
 import javax.net.ssl.TrustManager;
 import javax.net.ssl.TrustManagerFactory;
@@ -422,17 +423,25 @@ public class SSLParametersImpl implements Cloneable {
 
     OpenSSLSessionImpl getSessionToReuse(long sslNativePointer, String hostname, int port)
             throws SSLException {
-        final OpenSSLSessionImpl sessionToReuse;
+        OpenSSLSessionImpl sessionToReuse = null;
+
         if (client_mode) {
             // look for client session to reuse
-            sessionToReuse = getCachedClientSession(clientSessionContext, hostname, port);
-            if (sessionToReuse != null) {
-                NativeCrypto.SSL_set_session(sslNativePointer,
-                        sessionToReuse.sslSessionNativePointer);
+            SSLSession cachedSession = getCachedClientSession(clientSessionContext, hostname, port);
+            if (cachedSession != null) {
+                if (cachedSession instanceof OpenSSLSessionImpl) {
+                    sessionToReuse = (OpenSSLSessionImpl) cachedSession;
+                } else if (cachedSession instanceof OpenSSLExtendedSessionImpl) {
+                    sessionToReuse = ((OpenSSLExtendedSessionImpl) cachedSession).getDelegate();
+                }
+
+                if (sessionToReuse != null) {
+                    NativeCrypto.SSL_set_session(sslNativePointer,
+                            sessionToReuse.sslSessionNativePointer);
+                }
             }
-        } else {
-            sessionToReuse = null;
         }
+
         return sessionToReuse;
     }
 
@@ -766,12 +775,13 @@ public class SSLParametersImpl implements Cloneable {
     /**
      * Gets the suitable session reference from the session cache container.
      */
-    OpenSSLSessionImpl getCachedClientSession(ClientSessionContext sessionContext, String hostName,
+    SSLSession getCachedClientSession(ClientSessionContext sessionContext, String hostName,
             int port) {
         if (hostName == null) {
             return null;
         }
-        OpenSSLSessionImpl session = (OpenSSLSessionImpl) sessionContext.getSession(hostName, port);
+
+        SSLSession session = sessionContext.getSession(hostName, port);
         if (session == null) {
             return null;
         }

src/main/java/org/conscrypt/ServerSessionContext.java

@@ -54,19 +54,21 @@ public class ServerSessionContext extends AbstractSessionContext {
 
     @Override
     public SSLSession getSession(byte[] sessionId) {
-        SSLSession session = super.getSession(sessionId);
-        if (session != null) {
-            return session;
+        // First see if AbstractSessionContext can satisfy the request.
+        SSLSession cachedSession = super.getSession(sessionId);
+        if (cachedSession != null) {
+            // This will already have gone through Platform#wrapSSLSession
+            return cachedSession;
         }
 
-        // Check persistent cache.
+        // Then check the persistent cache.
         if (persistentCache != null) {
             byte[] data = persistentCache.getSessionData(sessionId);
             if (data != null) {
-                session = toSession(data, null, -1);
+                OpenSSLSessionImpl session = toSession(data, null, -1);
                 if (session != null && session.isValid()) {
                     super.putSession(session);
-                    return session;
+                    return Platform.wrapSSLSession(session);
                 }
             }
         }