From a5b890791839a23e6b7c345c2537d0381b89fe67 Mon Sep 17 00:00:00 2001
From: Davide De Rosa <keeshux@gmail.com>
Date: Wed, 20 Mar 2019 17:54:59 +0100
Subject: [PATCH] Postpone shutdown until notification is written

Otherwise socket might be force-closed while sending the packet.
---
 TunnelKit/Sources/Core/SessionProxy.swift | 34 ++++++++++++++++-------
 1 file changed, 24 insertions(+), 10 deletions(-)

diff --git a/TunnelKit/Sources/Core/SessionProxy.swift b/TunnelKit/Sources/Core/SessionProxy.swift
index ad4dc29..da820b6 100644
--- a/TunnelKit/Sources/Core/SessionProxy.swift
+++ b/TunnelKit/Sources/Core/SessionProxy.swift
@@ -1177,17 +1177,31 @@ public class SessionProxy {
     private func deferStop(_ method: StopMethod, _ error: Error?) {
         isStopping = true
 
-        // send exit notification if socket is unreliable (normally UDP)
-        if let link = link, !link.isReliable {
-            sendDataPackets([OCCPacket.exit.serialized()])
+        let completion = { [weak self] in
+            switch method {
+            case .shutdown:
+                self?.doShutdown(error: error)
+                
+            case .reconnect:
+                self?.doReconnect(error: error)
+            }
         }
-        
-        switch method {
-        case .shutdown:
-            doShutdown(error: error)
-        
-        case .reconnect:
-            doReconnect(error: error)
+
+        // shut down after sending exit notification if socket is unreliable (normally UDP)
+        if let link = link, !link.isReliable {
+            do {
+                guard let packets = try currentKey?.encrypt(packets: [OCCPacket.exit.serialized()]) else {
+                    completion()
+                    return
+                }
+                link.writePackets(packets) { (error) in
+                    completion()
+                }
+            } catch {
+                completion()
+            }
+        } else {
+            completion()
         }
     }