Merge pull request #903 from justcoding121/develop

beta
justcoding121 · Feb 1, 2022 · 178fa97 · 178fa97
2 parents 28ed4c9 + efecb69
commit 178fa97
Show file tree

Hide file tree

Showing 26 changed files with 222 additions and 154 deletions.
diff --git a/src/Titanium.Web.Proxy/Certificates/CertificateManager.cs b/src/Titanium.Web.Proxy/Certificates/CertificateManager.cs
@@ -120,7 +120,7 @@ private ICertificateMaker certEngine
         /// <param name="exceptionFunc"></param>
         internal CertificateManager(string? rootCertificateName, string? rootCertificateIssuerName,
             bool userTrustRootCertificate, bool machineTrustRootCertificate, bool trustRootCertificateAsAdmin,
-            ExceptionHandler exceptionFunc)
+            ExceptionHandler? exceptionFunc)
         {
             ExceptionFunc = exceptionFunc;
 
@@ -167,7 +167,7 @@ internal CertificateManager(string? rootCertificateName, string? rootCertificate
         /// <summary>
         /// Exception handler
         /// </summary>
-        internal ExceptionHandler ExceptionFunc { get; set; }
+        internal ExceptionHandler? ExceptionFunc { get; set; }
 
         /// <summary>
         ///     Select Certificate Engine.
@@ -339,7 +339,7 @@ private void installCertificate(StoreName storeName, StoreLocation storeLocation
             }
             catch (Exception e)
             {
-                ExceptionFunc(
+                onException(
                     new Exception("Failed to make system trust root certificate "
                                   + $" for {storeName}\\{storeLocation} store location. You may need admin rights.",
                         e));
@@ -360,7 +360,7 @@ private void uninstallCertificate(StoreName storeName, StoreLocation storeLocati
         {
             if (certificate == null)
             {
-                ExceptionFunc(new Exception("Could not remove certificate as it is null or empty."));
+                onException(new Exception("Could not remove certificate as it is null or empty."));
                 return;
             }
 
@@ -374,9 +374,8 @@ private void uninstallCertificate(StoreName storeName, StoreLocation storeLocati
             }
             catch (Exception e)
             {
-                ExceptionFunc(
-                    new Exception("Failed to remove root certificate trust "
-                                  + $" for {storeLocation} store location. You may need admin rights.", e));
+                onException(new Exception("Failed to remove root certificate trust "
+                                          + $" for {storeLocation} store location. You may need admin rights.", e));
             }
             finally
             {
@@ -408,6 +407,11 @@ private X509Certificate2 makeCertificate(string certificateName, bool isRootCert
             return certificate;
         }
 
+        private void onException(Exception exception)
+        {
+            ExceptionFunc?.Invoke(exception);
+        }
+
         private static ConcurrentDictionary<string, object> saveCertificateLocks
             = new ConcurrentDictionary<string, object>();
 
@@ -434,13 +438,13 @@ private static ConcurrentDictionary<string, object> saveCertificateLocks
 
                         if (certificate != null && certificate.NotAfter <= DateTime.Now)
                         {
-                            ExceptionFunc(new Exception($"Cached certificate for {subjectName} has expired."));
+                            onException(new Exception($"Cached certificate for {subjectName} has expired."));
                             certificate = null;
                         }
                     }
                     catch (Exception e)
                     {
-                        ExceptionFunc(new Exception("Failed to load fake certificate.", e));
+                        onException(new Exception("Failed to load fake certificate.", e));
                         certificate = null;
                     }
 
@@ -472,7 +476,7 @@ private static ConcurrentDictionary<string, object> saveCertificateLocks
                             }
                             catch (Exception e)
                             {
-                                ExceptionFunc(new Exception("Failed to save fake certificate.", e));
+                                onException(new Exception("Failed to save fake certificate.", e));
                             }
                         });
                     }
@@ -484,7 +488,7 @@ private static ConcurrentDictionary<string, object> saveCertificateLocks
             }
             catch (Exception e)
             {
-                ExceptionFunc(e);
+                onException(e);
                 certificate = null;
             }
 
@@ -628,7 +632,7 @@ public bool CreateRootCertificate(bool persistToFile = true)
 
                         if (rootCert != null && rootCert.NotAfter <= DateTime.Now)
                         {
-                            ExceptionFunc(new Exception("Loaded root certificate has expired."));
+                            onException(new Exception("Loaded root certificate has expired."));
                             return false;
                         }
 
@@ -641,7 +645,7 @@ public bool CreateRootCertificate(bool persistToFile = true)
                     catch (Exception e)
                     {
                         // root cert cannot be loaded
-                        ExceptionFunc(new Exception("Root cert cannot be loaded.", e));
+                        onException(new Exception("Root cert cannot be loaded.", e));
                     }
                 }
 
@@ -651,7 +655,7 @@ public bool CreateRootCertificate(bool persistToFile = true)
                 }
                 catch (Exception e)
                 {
-                    ExceptionFunc(e);
+                    onException(e);
                 }
 
                 if (persistToFile && RootCertificate != null)
@@ -664,14 +668,14 @@ public bool CreateRootCertificate(bool persistToFile = true)
                         }
                         catch (Exception e)
                         {
-                            ExceptionFunc(new Exception("An error happened when clearing certificate cache.", e));
+                            onException(new Exception("An error happened when clearing certificate cache.", e));
                         }
 
                         certificateCache.SaveRootCertificate(PfxFilePath, PfxPassword, RootCertificate);
                     }
                     catch (Exception e)
                     {
-                        ExceptionFunc(e);
+                        onException(e);
                     }
                 }
 
@@ -691,15 +695,15 @@ public bool CreateRootCertificate(bool persistToFile = true)
 
                 if (rootCert != null && rootCert.NotAfter <= DateTime.Now)
                 {
-                    ExceptionFunc(new ArgumentException("Loaded root certificate has expired."));
+                    onException(new ArgumentException("Loaded root certificate has expired."));
                     return null;
                 }
 
                 return rootCert;
             }
             catch (Exception e)
             {
-                ExceptionFunc(e);
+                onException(e);
                 return null;
             }
         }
@@ -808,7 +812,7 @@ public bool TrustRootCertificateAsAdmin(bool machineTrusted = false)
             }
             catch (Exception e)
             {
-                ExceptionFunc(e);
+                onException(e);
                 return false;
             }
 
@@ -1002,7 +1006,11 @@ void dispose(bool disposing)
                 return;
             }
 
-            clearCertificatesTokenSource.Dispose();
+            if (disposing)
+            {
+                clearCertificatesTokenSource.Dispose();
+            }
+
             disposed = true;
         }
 

diff --git a/src/Titanium.Web.Proxy/Certificates/Makers/BCCertificateMaker.cs b/src/Titanium.Web.Proxy/Certificates/Makers/BCCertificateMaker.cs
@@ -34,9 +34,9 @@ internal class BCCertificateMaker : ICertificateMaker
         // Set this flag to true when exception detected to avoid further exceptions
         private static bool doNotSetFriendlyName;
 
-        private readonly ExceptionHandler exceptionFunc;
+        private readonly ExceptionHandler? exceptionFunc;
 
-        internal BCCertificateMaker(ExceptionHandler exceptionFunc, int certificateValidDays)
+        internal BCCertificateMaker(ExceptionHandler? exceptionFunc, int certificateValidDays)
         {
             this.certificateValidDays = certificateValidDays;
             this.exceptionFunc = exceptionFunc;  

diff --git a/src/Titanium.Web.Proxy/Certificates/Makers/BCCertificateMakerFast.cs b/src/Titanium.Web.Proxy/Certificates/Makers/BCCertificateMakerFast.cs
@@ -34,11 +34,11 @@ internal class BCCertificateMakerFast : ICertificateMaker
         // Set this flag to true when exception detected to avoid further exceptions
         private static bool doNotSetFriendlyName;
 
-        private readonly ExceptionHandler exceptionFunc;
+        private readonly ExceptionHandler? exceptionFunc;
 
         public AsymmetricCipherKeyPair KeyPair { get; set; }
 
-        internal BCCertificateMakerFast(ExceptionHandler exceptionFunc, int certificateValidDays)
+        internal BCCertificateMakerFast(ExceptionHandler? exceptionFunc, int certificateValidDays)
         {
             this.certificateValidDays = certificateValidDays;
             this.exceptionFunc = exceptionFunc;

diff --git a/src/Titanium.Web.Proxy/Certificates/Makers/WinCertificateMaker.cs b/src/Titanium.Web.Proxy/Certificates/Makers/WinCertificateMaker.cs
@@ -17,7 +17,7 @@ internal class WinCertificateMaker : ICertificateMaker
         // Validity Days for Root Certificates Generated.
         private int certificateValidDays;
 
-        private readonly ExceptionHandler exceptionFunc;
+        private readonly ExceptionHandler? exceptionFunc;
 
         private readonly string sProviderName = "Microsoft Enhanced Cryptographic Provider v1.0";
 
@@ -53,7 +53,7 @@ internal class WinCertificateMaker : ICertificateMaker
         /// <summary>
         ///     Constructor.
         /// </summary>
-        internal WinCertificateMaker(ExceptionHandler exceptionFunc, int certificateValidDays)
+        internal WinCertificateMaker(ExceptionHandler? exceptionFunc, int certificateValidDays)
         {
             this.certificateValidDays = certificateValidDays;
             this.exceptionFunc = exceptionFunc;

diff --git a/src/Titanium.Web.Proxy/EventArguments/SessionEventArgs.cs b/src/Titanium.Web.Proxy/EventArguments/SessionEventArgs.cs
@@ -144,7 +144,7 @@ internal void OnMultipartRequestPartSent(ReadOnlySpan<char> boundary, HeaderColl
             }
             catch (Exception ex)
             {
-                ExceptionFunc(new Exception("Exception thrown in user event", ex));
+                OnException(new Exception("Exception thrown in user event", ex));
             }
         }
 
@@ -684,7 +684,8 @@ public void TerminateServerConnection()
             HttpClient.CloseServerConnection = true;
         }
 
-        private bool disposed = false;
+        private bool disposed;
+
         protected override void Dispose(bool disposing)
         {
             if (disposed)
@@ -700,6 +701,11 @@ protected override void Dispose(bool disposing)
 
         ~SessionEventArgs()
         {
+#if DEBUG
+            // Finalizer should not be called
+            System.Diagnostics.Debugger.Break();
+#endif
+
             Dispose(false);
         }
     }

diff --git a/src/Titanium.Web.Proxy/EventArguments/SessionEventArgsBase.cs b/src/Titanium.Web.Proxy/EventArguments/SessionEventArgsBase.cs
@@ -37,7 +37,7 @@ public abstract class SessionEventArgsBase : ProxyEventArgsBase, IDisposable
         public Guid ServerConnectionId => HttpClient.HasConnection ? ServerConnection.Id : Guid.Empty;
 
         protected readonly IBufferPool BufferPool;
-        protected readonly ExceptionHandler ExceptionFunc;
+        protected readonly ExceptionHandler? ExceptionFunc;
         private bool enableWinAuth;
 
         /// <summary>
@@ -150,6 +150,11 @@ public bool EnableWinAuth
         /// </summary>
         public Exception? Exception { get; internal set; }
 
+        protected void OnException(Exception exception)
+        {
+            ExceptionFunc?.Invoke(exception);
+        }
+
         private bool disposed = false;
 
         protected virtual void Dispose(bool disposing)
@@ -159,19 +164,18 @@ protected virtual void Dispose(bool disposing)
                 return;
             }
 
-            disposed = true;
-
             if (disposing)
             {
                 CustomUpStreamProxyUsed = null;
 
-                DataSent = null;
-                DataReceived = null;
-                Exception = null;
+                HttpClient.FinishSession();
             }
 
-            HttpClient.FinishSession();
+            DataSent = null;
+            DataReceived = null;
+            Exception = null;
 
+            disposed = true;
         }
 
         public void Dispose()
@@ -182,6 +186,11 @@ public void Dispose()
 
         ~SessionEventArgsBase()
         {
+#if DEBUG
+            // Finalizer should not be called
+            System.Diagnostics.Debugger.Break();
+#endif
+
             Dispose(false);
         }
 
@@ -203,7 +212,7 @@ internal void OnDataSent(byte[] buffer, int offset, int count)
             }
             catch (Exception ex)
             {
-                ExceptionFunc(new Exception("Exception thrown in user event", ex));
+                OnException(new Exception("Exception thrown in user event", ex));
             }
         }
 
@@ -215,7 +224,7 @@ internal void OnDataReceived(byte[] buffer, int offset, int count)
             }
             catch (Exception ex)
             {
-                ExceptionFunc(new Exception("Exception thrown in user event", ex));
+                OnException(new Exception("Exception thrown in user event", ex));
             }
         }
 

diff --git a/src/Titanium.Web.Proxy/EventArguments/TunnelConnectEventArgs.cs b/src/Titanium.Web.Proxy/EventArguments/TunnelConnectEventArgs.cs
@@ -60,7 +60,7 @@ internal void OnDecryptedDataSent(byte[] buffer, int offset, int count)
             }
             catch (Exception ex)
             {
-                ExceptionFunc(new Exception("Exception thrown in user event", ex));
+                OnException(new Exception("Exception thrown in user event", ex));
             }
         }
 
@@ -72,8 +72,18 @@ internal void OnDecryptedDataReceived(byte[] buffer, int offset, int count)
             }
             catch (Exception ex)
             {
-                ExceptionFunc(new Exception("Exception thrown in user event", ex));
+                OnException(new Exception("Exception thrown in user event", ex));
             }
         }
+
+        ~TunnelConnectSessionEventArgs()
+        {
+#if DEBUG
+            // Finalizer should not be called
+            System.Diagnostics.Debugger.Break();
+#endif
+
+            Dispose(false);
+        }
     }
 }
diff --git a/src/Titanium.Web.Proxy/ExplicitClientHandler.cs b/src/Titanium.Web.Proxy/ExplicitClientHandler.cs
@@ -38,10 +38,10 @@ private async Task handleClient(ExplicitProxyEndPoint endPoint, TcpClientConnect
             Task<TcpServerConnection?>? prefetchConnectionTask = null;
             bool closeServerConnection = false;
 
+            TunnelConnectSessionEventArgs? connectArgs = null;
+
             try
             {
-                TunnelConnectSessionEventArgs? connectArgs = null;
-
                 var method = await HttpHelper.GetMethod(clientStream, BufferPool, cancellationToken);
                 if (clientStream.IsClosed)
                 {
@@ -402,6 +402,7 @@ await Http2Helper.SendHttp2(clientStream, connection.Stream,
                 await tcpConnectionFactory.Release(prefetchConnectionTask, closeServerConnection);
 
                 clientStream.Dispose();
+                connectArgs?.Dispose();
             }
         }
     }