From c7b9033a2cacde56bc77027cdc233811943ba21b Mon Sep 17 00:00:00 2001
From: Wenjie Guo <wguo@redhat.com>
Date: Thu, 1 Aug 2024 13:47:14 +0800
Subject: [PATCH] Add check for SSL tunnel before closing proxy server

---
 .../httprox/handler/ProxyMITMSSLServer.java   | 42 ++++++++++++++++++-
 .../httprox/handler/ProxyResponseWriter.java  |  1 +
 .../httprox/handler/ProxySSLTunnel.java       |  6 +++
 3 files changed, 48 insertions(+), 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
index 7923416..ea73165 100644
--- 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
index 528816d..6b4132e 100644
--- 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
index 342ca79..87c2c7d 100644
--- 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;
+    }
+
 }