From 3c9b126993d406d510646b7e942389b38f31b8c4 Mon Sep 17 00:00:00 2001
From: Kirill Fomichev <fanatid@ya.ru>
Date: Thu, 7 May 2020 20:55:37 +0300
Subject: [PATCH] Drop context on MidHandshake success [tokio-native-tls] (#12)

---
 tokio-native-tls/src/lib.rs | 7 +++++--
 1 file changed, 5 insertions(+), 2 deletions(-)

diff --git a/tokio-native-tls/src/lib.rs b/tokio-native-tls/src/lib.rs
index 622a594..4e83c8e 100644
--- a/tokio-native-tls/src/lib.rs
+++ b/tokio-native-tls/src/lib.rs
@@ -348,13 +348,16 @@ impl<S: AsyncRead + AsyncWrite + Unpin> Future for MidHandshake<S> {
 
         s.get_mut().context = cx as *mut _ as *mut ();
         match s.handshake() {
-            Ok(stream) => Poll::Ready(Ok(TlsStream(stream))),
-            Err(HandshakeError::Failure(e)) => Poll::Ready(Err(e)),
+            Ok(mut s) => {
+                s.get_mut().context = null_mut();
+                Poll::Ready(Ok(TlsStream(s)))
+            }
             Err(HandshakeError::WouldBlock(mut s)) => {
                 s.get_mut().context = null_mut();
                 mut_self.0 = Some(s);
                 Poll::Pending
             }
+            Err(HandshakeError::Failure(e)) => Poll::Ready(Err(e)),
         }
     }
 }