Add check for SSL tunnel before closing proxy server

sswguo · Aug 1, 2024 · c7b9033 · c7b9033
1 parent 337ca53
commit c7b9033
Show file tree

Hide file tree

Showing 3 changed files with 48 additions and 1 deletion.
diff --git a/src/main/java/org/commonjava/indy/service/httprox/handler/ProxyMITMSSLServer.java b/src/main/java/org/commonjava/indy/service/httprox/handler/ProxyMITMSSLServer.java
@@ -75,6 +75,10 @@ public class ProxyMITMSSLServer implements Runnable
 
     private final HttpConduitWrapper httpConduitWrapper;
 
+    private ProxySSLTunnel sslTunnel;
+
+    private final static long MAX_WAIT_TIME_IN_MILLIS = 60 * 1000;
+
     public ProxyMITMSSLServer( String host, int port, String trackingId, UserPass proxyUserPass,
                                ProxyResponseHelper proxyResponseHelper, ProxyConfiguration config, ProxyMeter meter, HttpConduitWrapper httpConduitWrapper)
     {
@@ -101,7 +105,39 @@ public void run()
         }
         finally
         {
+            if ( sslTunnel != null )
+            {
+                long startTime = System.currentTimeMillis();
+                while ( !sslTunnel.isClosed() )
+                {
+                    if (System.currentTimeMillis() - startTime > MAX_WAIT_TIME_IN_MILLIS)
+                    {
+                        logger.warn("Maximum wait time exceeded, stopping wait for SSL tunnel to close.");
+                        break;
+                    }
+                    try
+                    {
+                        logger.info("Waiting ssl tunnel to finish...");
+                        TimeUnit.MILLISECONDS.sleep( GET_SOCKET_CHANNEL_WAIT_TIME_IN_MILLISECONDS );
+                    }
+                    catch (InterruptedException e)
+                    {
+                        e.printStackTrace();
+                    }
+                }
+
+                if (sslTunnel.isClosed())
+                {
+                    logger.info("SSL tunnel is closed.");
+                }
+                else
+                {
+                    logger.warn("SSL tunnel is still not closed after maximum wait time.");
+                }
+            }
+
             closeProperly();
+            logger.debug( "MITM server closed" );
         }
     }
 
@@ -256,7 +292,6 @@ else if ( line.isEmpty() )
                     }
                 }
             }
-            logger.debug( "MITM server closed" );
         }
         finally
         {
@@ -361,4 +396,9 @@ SSLServerSocketFactory getSslSocketFactory() {
         }
 
     }
+
+    public void setProxySSLTunnel( ProxySSLTunnel sslTunnel )
+    {
+        this.sslTunnel = sslTunnel;
+    }
 }
diff --git a/src/main/java/org/commonjava/indy/service/httprox/handler/ProxyResponseWriter.java b/src/main/java/org/commonjava/indy/service/httprox/handler/ProxyResponseWriter.java
@@ -283,6 +283,7 @@ private void doHandleEvent(final ConduitStreamSinkChannel sinkChannel)
                                 sslTunnel = new ProxySSLTunnel( sinkChannel, socketChannel, config );
                                 tunnelAndMITMExecutor.submit( sslTunnel );
                                 proxyRequestReader.setProxySSLTunnel( sslTunnel ); // client input will be directed to target socket
+                                svr.setProxySSLTunnel( sslTunnel );
 
                                 // When all is ready, send the 200 to client. Client send the SSL handshake to reader,
                                 // reader direct it to tunnel to MITM. MITM finish the handshake and read the request data,

diff --git a/src/main/java/org/commonjava/indy/service/httprox/handler/ProxySSLTunnel.java b/src/main/java/org/commonjava/indy/service/httprox/handler/ProxySSLTunnel.java
@@ -156,4 +156,10 @@ public void close()
             logger.error( "Close tunnel selector failed", e );
         }
     }
+
+    public boolean isClosed()
+    {
+        return closed;
+    }
+
 }