From adba7368c0be5e8ab5840d893d05ab3504f9a796 Mon Sep 17 00:00:00 2001
From: Kyle Browning <kylebrowning@me.com>
Date: Tue, 26 Nov 2019 08:50:45 -0800
Subject: [PATCH] Add Channel Handler Removed. (#63)

* Fail if we dont get a proper response

* add channel handlers removed.

* Typo in test.

* Addressing comments.

* Update to use a new Response Struct

* Update to use correct errors.

* Update to use new decoding API from swift nio

* Consistent internal naming of errors.

* Update swift nio package version.

* Adding comment.
---
 Package.swift                                 |  2 +-
 Sources/APNSwift/APNSwiftErrors.swift         |  6 ++++
 Sources/APNSwift/APNSwiftStreamHandler.swift  | 15 +++++++--
 .../APNSwiftTests/APNSwiftRequestTests.swift  | 33 ++++++++++++++++++-
 4 files changed, 52 insertions(+), 4 deletions(-)

diff --git a/Package.swift b/Package.swift
index c9314d8e..186cf88b 100644
--- a/Package.swift
+++ b/Package.swift
@@ -8,7 +8,7 @@ let package = Package(
         .library(name: "APNSwift", targets: ["APNSwift"]),
     ],
     dependencies: [
-        .package(url: "https://github.com/apple/swift-nio.git", .upToNextMajor(from: "2.8.0")),
+        .package(url: "https://github.com/apple/swift-nio.git", .upToNextMajor(from: "2.10.1")),
         .package(url: "https://github.com/apple/swift-nio-ssl.git", .upToNextMajor(from: "2.4.0")),
         .package(url: "https://github.com/apple/swift-nio-http2.git", .upToNextMajor(from: "1.6.0")),
     ],
diff --git a/Sources/APNSwift/APNSwiftErrors.swift b/Sources/APNSwift/APNSwiftErrors.swift
index 11ce5988..d66bc40e 100644
--- a/Sources/APNSwift/APNSwiftErrors.swift
+++ b/Sources/APNSwift/APNSwiftErrors.swift
@@ -14,6 +14,11 @@
 
 import Foundation
 
+/// An error for when the connections ends, but we still have promises in our queue.
+public struct NoResponseReceivedBeforeConnectionEnded: Error, Equatable {}
+/// An error where a request was made to Apple, but the response body buffer was nil
+public struct NoResponseBodyFromApple: Error, Equatable {}
+
 /// This is an enum that provides the possible responses from Apple
 public struct APNSwiftError: Equatable {
     public enum ResponseError: Error, Equatable {
@@ -59,6 +64,7 @@ public struct APNSwiftError: Equatable {
             return rawValue
         }
     }
+    
     public enum SigningError: Error {
         case invalidAuthKey
         case invalidASN1
diff --git a/Sources/APNSwift/APNSwiftStreamHandler.swift b/Sources/APNSwift/APNSwiftStreamHandler.swift
index 440e72f7..8e8805dc 100644
--- a/Sources/APNSwift/APNSwiftStreamHandler.swift
+++ b/Sources/APNSwift/APNSwiftStreamHandler.swift
@@ -33,10 +33,15 @@ final class APNSwiftStreamHandler: ChannelDuplexHandler {
         let res = unwrapInboundIn(data)
         guard let current = self.queue.popLast() else { return }
         guard res.header.status == .ok else {
-            if var buffer = res.byteBuffer, let data = buffer.readData(length: buffer.readableBytes), let error = try? JSONDecoder().decode(APNSwiftError.ResponseStruct.self, from: data) {
+            guard let buffer = res.byteBuffer else {
+                return current.responsePromise.fail(NoResponseBodyFromApple())
+            }
+            do {
+                let error = try JSONDecoder().decode(APNSwiftError.ResponseStruct.self, from: buffer)
                 return current.responsePromise.fail(APNSwiftError.ResponseError.badRequest(error.reason))
+            } catch {
+                return current.responsePromise.fail(error)
             }
-            return
         }
         current.responsePromise.succeed(Void())
     }
@@ -46,4 +51,10 @@ final class APNSwiftStreamHandler: ChannelDuplexHandler {
         queue.insert(input, at: 0)
         context.write(wrapOutboundOut(input.request), promise: promise)
     }
+    
+    func handlerRemoved(context: ChannelHandlerContext) {
+        while let context = queue.popLast() {
+            context.responsePromise.fail(NoResponseReceivedBeforeConnectionEnded())
+        }
+    }
 }
diff --git a/Tests/APNSwiftTests/APNSwiftRequestTests.swift b/Tests/APNSwiftTests/APNSwiftRequestTests.swift
index 00e1db99..1af41735 100644
--- a/Tests/APNSwiftTests/APNSwiftRequestTests.swift
+++ b/Tests/APNSwiftTests/APNSwiftRequestTests.swift
@@ -219,9 +219,40 @@ final class APNSwiftRequestTests: XCTestCase {
                     XCTFail("response is: \(error), should be: \(expected)")
                 }
             default:
-                XCTFail("response should not success")
+                XCTFail("response should not succeed")
             }
         }
+        XCTAssertNoThrow(XCTAssertNotNil(try channel.readOutbound()))
+        XCTAssertNoThrow(XCTAssertTrue(try channel.finish().isClean))
+
+    }
+    
+    func testErrorsFromAPNSLeak() throws {
+        let encodedData = try JSONEncoder().encode(["test string"])
+        let allocator = ByteBufferAllocator()
+        var errorBuffer = allocator.buffer(capacity: encodedData.count)
+        errorBuffer.writeBytes(encodedData)
+              
+        let channel = EmbeddedChannel(handler: APNSwiftStreamHandler())
+        let responsePromise = channel.eventLoop.makePromise(of: Void.self)
+        let context = APNSwiftRequestContext(
+           request: errorBuffer,
+           responsePromise: responsePromise
+        )
+        try channel.writeOutbound(context)
+        responsePromise.futureResult.whenComplete { temp in
+            switch temp {
+            case .failure(let error):
+                let error = error as! NoResponseReceivedBeforeConnectionEnded
+                let expected = NoResponseReceivedBeforeConnectionEnded()
+                if error != expected {
+                    XCTFail("response is: \(error), should be: \(expected)")
+                }
+            default:
+                XCTFail("response should not succeed")
+            }
+        }
+        XCTAssertNoThrow(XCTAssertNotNil(try channel.readOutbound()))
         XCTAssertNoThrow(XCTAssertTrue(try channel.finish().isClean))
     }