From 98c0379851838a91cf7f0ab4b9213088d935f7bf Mon Sep 17 00:00:00 2001
From: _ <53088156+636f7374@users.noreply.github.com>
Date: Fri, 18 Mar 2022 05:12:57 +0800
Subject: [PATCH] Improve OpenSSL::SSL::Socket.

---
 src/dns/address/https.cr            |  1 +
 src/dns/address/tls.cr              |  1 +
 src/dns/extra/openssl/ssl/socket.cr | 23 ++++++++++++++++++-----
 3 files changed, 20 insertions(+), 5 deletions(-)

diff --git a/src/dns/address/https.cr b/src/dns/address/https.cr
index 29f22c6..4d771d0 100644
--- a/src/dns/address/https.cr
+++ b/src/dns/address/https.cr
@@ -21,6 +21,7 @@ abstract struct DNS::Address
 
       begin
         tls_socket = OpenSSL::SSL::Socket::Client.new socket, context: tls_context, sync_close: true, hostname: tls.try &.hostname
+        tls_socket.close_after_finalize = true
         tls_socket.read_buffering = false
         tls_socket.ssl_context = tls_context
         tls_socket.sync = true
diff --git a/src/dns/address/tls.cr b/src/dns/address/tls.cr
index 3536e4b..f548ccb 100644
--- a/src/dns/address/tls.cr
+++ b/src/dns/address/tls.cr
@@ -18,6 +18,7 @@ abstract struct DNS::Address
 
       begin
         tls_socket = OpenSSL::SSL::Socket::Client.new socket, context: tls_context, sync_close: true, hostname: tls.try &.hostname
+        tls_socket.close_after_finalize = true
         tls_socket.read_buffering = false
         tls_socket.ssl_context = tls_context
         tls_socket.sync = true
diff --git a/src/dns/extra/openssl/ssl/socket.cr b/src/dns/extra/openssl/ssl/socket.cr
index b5dca3e..c97fbb9 100644
--- a/src/dns/extra/openssl/ssl/socket.cr
+++ b/src/dns/extra/openssl/ssl/socket.cr
@@ -1,8 +1,10 @@
 class OpenSSL::SSL::Socket
+  getter closeAfterFinalize : Bool
   getter freed : Bool
   getter freeMutex : Mutex
 
   protected def initialize(io, context : Context, @sync_close : Bool = false)
+    @closeAfterFinalize = false
     @freed = false
     @freeMutex = Mutex.new :unchecked
 
@@ -29,12 +31,20 @@ class OpenSSL::SSL::Socket
     @sslContext = value
   end
 
-  def ssl_context : Context?
+  def ssl_context? : Context?
     @sslContext
   end
 
+  def close_after_finalize=(value : Bool)
+    @closeAfterFinalize = value
+  end
+
+  def close_after_finalize? : Bool
+    @closeAfterFinalize
+  end
+
   private def free_ssl_context : Bool
-    ssl_context.try &.free
+    ssl_context?.try &.free
     @sslContext = nil
 
     true
@@ -97,10 +107,13 @@ class OpenSSL::SSL::Socket
         exception = ex
       end
 
-      LibSSL.ssl_free @ssl
-      free_ssl_context
+      if close_after_finalize?
+        LibSSL.ssl_free @ssl
+        free_ssl_context
+
+        @freed = true
+      end
 
-      @freed = true
       exception.try { |_exception| raise _exception }
     end
   end