From 03e83e2332d13a3c06d5cd70948166e5724c43c7 Mon Sep 17 00:00:00 2001
From: Mads Odgaard <mads@madsodgaard.com>
Date: Mon, 1 Feb 2021 16:25:09 +0100
Subject: [PATCH] Add timeout configuration if the connection does not receive
 a response (#101)

---
 Sources/APNSwift/APNSwiftConfiguration.swift  |  6 ++-
 Sources/APNSwift/APNSwiftConnection.swift     | 45 ++++++++++++++-----
 Sources/APNSwift/APNSwiftErrors.swift         |  2 +
 .../APNSwiftConfigurationTests.swift          |  4 +-
 .../APNSwiftTests/APNSwiftRequestTests.swift  |  1 -
 5 files changed, 44 insertions(+), 14 deletions(-)

diff --git a/Sources/APNSwift/APNSwiftConfiguration.swift b/Sources/APNSwift/APNSwiftConfiguration.swift
index 51c605cf..d24a818d 100644
--- a/Sources/APNSwift/APNSwiftConfiguration.swift
+++ b/Sources/APNSwift/APNSwiftConfiguration.swift
@@ -168,6 +168,8 @@ public struct APNSwiftConfiguration {
     public var topic: String
     public var environment: Environment
     internal var logger: Logger?
+    /// Optional timeout time if the connection does not receive a response.
+    public var timeout: TimeAmount? = nil
 
     public var url: URL {
         switch environment {
@@ -197,7 +199,8 @@ public struct APNSwiftConfiguration {
         authenticationMethod: AuthenticationMethod,
         topic: String,
         environment: APNSwiftConfiguration.Environment,
-        logger: Logger? = nil
+        logger: Logger? = nil,
+        timeout: TimeAmount? = nil
     ) {
         self.topic = topic
         self.authenticationMethod = authenticationMethod
@@ -206,6 +209,7 @@ public struct APNSwiftConfiguration {
             logger[metadataKey: "origin"] = "APNSwift"
             self.logger = logger
         }
+        self.timeout = timeout
     }
 }
 
diff --git a/Sources/APNSwift/APNSwiftConnection.swift b/Sources/APNSwift/APNSwiftConnection.swift
index a81a8773..06a52c33 100644
--- a/Sources/APNSwift/APNSwiftConnection.swift
+++ b/Sources/APNSwift/APNSwiftConnection.swift
@@ -200,20 +200,43 @@ public final class APNSwiftConnection: APNSwiftClient {
             request: payload,
             responsePromise: responsePromise
         )
-
         streamPromise.futureResult.cascadeFailure(to: responsePromise)
-        return streamPromise.futureResult.flatMap { stream in
-            logger?.info("Send - sending")
-            return stream.writeAndFlush(context).flatMapErrorThrowing { error in
-                logger?.info("Send - sending - failed - \(error)")
-                responsePromise.fail(error)
-                throw error
+        
+        let timeoutPromise = self.channel.eventLoop.makePromise(of: Void.self)
+        responsePromise.futureResult.cascade(to: timeoutPromise)
+        timeoutPromise.futureResult.cascadeFailure(to: responsePromise)
+        var timeoutTask: Scheduled<Any>? = nil
+        let timeoutTime = configuration.timeout
+        
+        return streamPromise.futureResult
+            .flatMap { stream in
+                logger?.info("Send - sending")
+                if let timeoutTime = timeoutTime {
+                    timeoutTask = stream.eventLoop.scheduleTask(in: timeoutTime) {
+                        logger?.warning("Send - sending - failed - No response was received within the timeout.")
+                        return timeoutPromise.fail(NoResponseWithinTimeoutError())
+                    }
+                } else {
+                    timeoutPromise.succeed(())
+                }
+                
+                return stream.writeAndFlush(context).flatMapErrorThrowing { error in
+                    logger?.info("Send - sending - failed - \(error)")
+                    responsePromise.fail(error)
+                    throw error
+                }
+            }
+            .flatMap {
+                responsePromise
+                    .futureResult
+                    .and(timeoutPromise.futureResult)
+                    .map { _ in () }
+            }
+            .always { _ in
+                timeoutTask?.cancel()
             }
-        }.flatMap {
-            responsePromise.futureResult
-        }
     }
-
+    
     var onClose: EventLoopFuture<Void> {
         logger?.debug("Connection - closed")
         return self.channel.closeFuture
diff --git a/Sources/APNSwift/APNSwiftErrors.swift b/Sources/APNSwift/APNSwiftErrors.swift
index 01603fce..5999fddd 100644
--- a/Sources/APNSwift/APNSwiftErrors.swift
+++ b/Sources/APNSwift/APNSwiftErrors.swift
@@ -18,6 +18,8 @@ import Foundation
 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 {}
+/// An error where no the connection received no response within the timeout period
+public struct NoResponseWithinTimeoutError: Error, Equatable {}
 
 public struct APNSwiftError: Equatable {
     public enum ResponseError: Error, Equatable {
diff --git a/Tests/APNSwiftTests/APNSwiftConfigurationTests.swift b/Tests/APNSwiftTests/APNSwiftConfigurationTests.swift
index 376638c3..677b7358 100644
--- a/Tests/APNSwiftTests/APNSwiftConfigurationTests.swift
+++ b/Tests/APNSwiftTests/APNSwiftConfigurationTests.swift
@@ -27,7 +27,8 @@ class APNSwiftConfigurationTests: XCTestCase {
                 teamIdentifier: "MY_TEAM_ID"
             ),
             topic: "MY_TOPIC",
-            environment: environment
+            environment: environment,
+            timeout: .seconds(5)
         )
 
         switch environment {
@@ -46,6 +47,7 @@ class APNSwiftConfigurationTests: XCTestCase {
             XCTFail("expected JWT auth method")
         }
         XCTAssertEqual(apnsConfiguration.topic, "MY_TOPIC")
+        XCTAssertEqual(apnsConfiguration.timeout, .seconds(5))
 
     }
 
diff --git a/Tests/APNSwiftTests/APNSwiftRequestTests.swift b/Tests/APNSwiftTests/APNSwiftRequestTests.swift
index 7ea06562..a94f20aa 100644
--- a/Tests/APNSwiftTests/APNSwiftRequestTests.swift
+++ b/Tests/APNSwiftTests/APNSwiftRequestTests.swift
@@ -304,7 +304,6 @@ final class APNSwiftRequestTests: XCTestCase {
         // Should have changed
         XCTAssertFalse(newCachedToken == bearerToken.currentBearerToken)
         bearerToken.cancel()
-        
     }
 
     let validAuthKey = """