fix(predictions): remove liveness event listeners on session completion

AmplifyPlugins/Predictions/AWSPredictionsPlugin/Liveness/Service/FaceLivenessSession.swift

@@ -45,18 +45,27 @@ public final class FaceLivenessSession: LivenessService {
         self.websocket = websocket
 
         websocket.onMessageReceived { [weak self] result in
-            self?.receive(result: result) ?? .stopAndInvalidateSession
+            guard let receiveResult = self?.receive(result: result) else {
+                self?.removeLivenessEventListeners()
+                return .stopAndInvalidateSession
+            }
+            return receiveResult
         }
 
         websocket.onSocketClosed { [weak self] closeCode in
             self?.onComplete(.unexpectedClosure(closeCode))
+            self?.removeLivenessEventListeners()
         }
 
         websocket.onServerDateReceived { [weak self] serverDate in
             self?.serverDate = serverDate
         }
     }
 
+    deinit {
+        Amplify.log.verbose("\(#fileID)-\(#function)")
+    }
+
     public var onServiceException: (FaceLivenessSessionError) -> Void = { _ in }
 
     public func register(
@@ -77,18 +86,19 @@ public final class FaceLivenessSession: LivenessService {
     }
 
     public func closeSocket(with code: URLSessionWebSocketTask.CloseCode) {
-        livenessServiceDispatchQueue.async {
-            self.websocket.close(with: code)
+        Amplify.log.verbose("\(#fileID)-\(#function): closeSocket with code: \(code)")
+        livenessServiceDispatchQueue.async { [weak self] in
+            self?.websocket.close(with: code)
         }
     }
 
     public func initializeLivenessStream(
         withSessionID sessionID: String,
-
-                                         userAgent: String = "",
+        userAgent: String = "",
         challenges: [Challenge] = FaceLivenessSession.supportedChallenges,
         options: FaceLivenessSession.Options
     ) throws {
+        Amplify.log.verbose("\(#fileID)-\(#function): Initialize liveness stream")
         var components = URLComponents(url: baseURL, resolvingAgainstBaseURL: false)
         components?.queryItems = [
             URLQueryItem(name: "session-id", value: sessionID),
@@ -108,17 +118,19 @@ public final class FaceLivenessSession: LivenessService {
 
         savedURLForReconnect = url
         let signedConnectionURL = signer.sign(url: url)
-        livenessServiceDispatchQueue.async {
-            self.websocket.open(url: signedConnectionURL)
+        livenessServiceDispatchQueue.async { [weak self] in
+            self?.websocket.open(url: signedConnectionURL)
         }
     }
 
     public func send(
         _ event: LivenessEvent<some Any>,
         eventDate: @escaping () -> Date = Date.init
     ) {
-        livenessServiceDispatchQueue.async {
-            let encodedPayload = self.eventStreamEncoder.encode(
+        Amplify.log.verbose("\(#fileID)-\(#function): Sending websocket event: \(event)")
+        livenessServiceDispatchQueue.async { [weak self] in
+            guard let self else { return }
+            let encodedPayload = eventStreamEncoder.encode(
                 payload: event.payload,
                 headers: [
                     ":content-type": .string("application/json"),
@@ -128,44 +140,49 @@ public final class FaceLivenessSession: LivenessService {
             )
 
             let dateForSigning: Date
-            if let serverDate = self.serverDate {
+            if let serverDate {
                 dateForSigning = serverDate
             } else {
                 dateForSigning = eventDate()
             }
 
-            let signedPayload = self.signer.signWithPreviousSignature(
+            let signedPayload = signer.signWithPreviousSignature(
                 payload: encodedPayload,
                 dateHeader: (key: ":date", value: dateForSigning)
             )
 
-            let encodedEvent = self.eventStreamEncoder.encode(
+            let encodedEvent = eventStreamEncoder.encode(
                 payload: encodedPayload,
                 headers: [
                     ":date": .timestamp(dateForSigning),
                     ":chunk-signature": .data(signedPayload)
                 ]
             )
 
-            self.websocket.send(
+            websocket.send(
                 message: .data(encodedEvent),
-                onError: { _ in }
+                onError: { error in
+                    Amplify.log.verbose("\(#fileID)-\(#function): Error sending web socket message: \(error)")
+                }
             )
         }
     }
 
     private func fallbackDecoding(_ message: EventStream.Message) -> WebSocketSession.WebSocketMessageResult {
-        // We only care about two events above.
         // Just in case the header value changes (it shouldn't)
         // We'll try to decode each of these events
         if let payload = try? JSONDecoder().decode(ServerSessionInformationEvent.self, from: message.payload) {
+            Amplify.log.verbose("\(#fileID)-\(#function): Fallback decoding server session information: \(payload)")
             let sessionConfiguration = sessionConfiguration(from: payload)
             serverEventListeners[.challenge]?(sessionConfiguration)
         } else if let payload = try? JSONDecoder().decode(ChallengeEvent.self, from: message.payload) {
+            Amplify.log.verbose("\(#fileID)-\(#function): Fallback decoding challenge: \(payload)")
             let challenge = challenge(from: payload)
             challengeTypeListeners[.challenge]?(challenge)
-        } else if (try? JSONDecoder().decode(DisconnectEvent.self, from: message.payload)) != nil {
+        } else if let payload = try? JSONDecoder().decode(DisconnectEvent.self, from: message.payload) {
+            Amplify.log.verbose("\(#fileID)-\(#function): Fallback decoding disconnect: \(payload)")
             onComplete(.disconnectionEvent)
+            removeLivenessEventListeners()
             return .stopAndInvalidateSession
         }
         return .continueToReceive
@@ -179,6 +196,7 @@ public final class FaceLivenessSession: LivenessService {
 
                 if let eventType = message.headers.first(where: { $0.name == ":event-type" }) {
                     let serverEvent = LivenessEventKind.Server(rawValue: eventType.value)
+                    Amplify.log.verbose("\(#fileID)-\(#function): Received server event: \(serverEvent)")
                     switch serverEvent {
                     case .challenge:
                         // :event-type ChallengeEvent
@@ -199,22 +217,24 @@ public final class FaceLivenessSession: LivenessService {
                     case .disconnect:
                         // :event-type DisconnectionEvent
                         onComplete(.disconnectionEvent)
+                        removeLivenessEventListeners()
                         return .stopAndInvalidateSession
                     default:
                         return .continueToReceive
                     }
                 } else if let exceptionType = message.headers.first(where: { $0.name == ":exception-type" }) {
                     let exceptionEvent = LivenessEventKind.Exception(rawValue: exceptionType.value)
-                    Amplify.log.verbose("\(#function): Received exception: \(exceptionEvent)")
+                    Amplify.log.verbose("\(#fileID)-\(#function): Received exception: \(exceptionEvent)")
                     guard exceptionEvent == .invalidSignature,
                           connectingState == .normal,
                           let savedURLForReconnect,
                           let serverDate else {
                         if let runtimeError = URLSessionWebSocketTask.CloseCode(rawValue: 4_005) {
-                            Amplify.log.verbose("\(#function): Closing websocket with runtime error")
+                            Amplify.log.verbose("\(#fileID)-\(#function): Closing websocket with runtime error")
                             closeSocket(with: runtimeError)
                         }
                         onServiceException(.init(event: exceptionEvent))
+                        removeLivenessEventListeners()
                         return .stopAndInvalidateSession
                     }
 
@@ -228,12 +248,21 @@ public final class FaceLivenessSession: LivenessService {
                     return fallbackDecoding(message)
                 }
             } catch {
+                Amplify.log.verbose("\(#fileID)-\(#function): Error decoding web socket message: \(error)")
+                removeLivenessEventListeners()
                 return .stopAndInvalidateSession
             }
         case .success:
             return .continueToReceive
-        case .failure:
+        case .failure(let error):
+            Amplify.log.verbose("\(#fileID)-\(#function): Failure result in web socket message: \(error)")
+            removeLivenessEventListeners()
             return .stopAndInvalidateSession
         }
     }
+
+    private func removeLivenessEventListeners() {
+        serverEventListeners.removeAll()
+        challengeTypeListeners.removeAll()
+    }
 }

AmplifyPlugins/Predictions/AWSPredictionsPlugin/Liveness/Service/FaceLivenessSessionRepresentable.swift

@@ -21,8 +21,7 @@ public protocol LivenessService {
 
     func initializeLivenessStream(
         withSessionID sessionID: String,
-
-                                  userAgent: String,
+        userAgent: String,
         challenges: [Challenge],
         options: FaceLivenessSession.Options
     ) throws

AmplifyPlugins/Predictions/AWSPredictionsPlugin/Liveness/Service/WebSocketSession.swift

@@ -31,6 +31,11 @@ final class WebSocketSession {
         )
     }
 
+    deinit {
+        Amplify.log.verbose("\(#fileID)-\(#function)")
+        task?.cancel(with: .normalClosure, reason: nil)
+    }
+
     func onMessageReceived(_ receive: @escaping (Result<URLSessionWebSocketTask.Message, Error>) -> WebSocketMessageResult) {
         receiveMessage = receive
     }
@@ -92,17 +97,22 @@ final class WebSocketSession {
         )
     }
 
-    final class Delegate: NSObject, URLSessionWebSocketDelegate, URLSessionTaskDelegate {
+    final class Delegate: NSObject, URLSessionWebSocketDelegate {
         var onClose: (URLSessionWebSocketTask.CloseCode) -> Void = { _ in }
         var onOpen: () -> Void = {}
         var onServerDateReceived: (Date?) -> Void = { _ in }
 
+        deinit {
+            Amplify.log.verbose("\(#fileID).Delegate-\(#function)")
+        }
+
         // MARK: - URLSessionWebSocketDelegate methods
         func urlSession(
             _ session: URLSession,
             webSocketTask: URLSessionWebSocketTask,
             didOpenWithProtocol protocol: String?
         ) {
+            Amplify.log.verbose("\(#fileID)-\(#function): Web socket task didOpen")
             onOpen()
         }
 
@@ -112,6 +122,7 @@ final class WebSocketSession {
             didCloseWith closeCode: URLSessionWebSocketTask.CloseCode,
             reason: Data?
         ) {
+            Amplify.log.verbose("\(#fileID)-\(#function): Web socket task didCloseWith: \(closeCode)")
             onClose(closeCode)
         }
 
@@ -123,7 +134,7 @@ final class WebSocketSession {
         ) {
             guard let httpResponse = metrics.transactionMetrics.first?.response as? HTTPURLResponse,
                   let dateString = httpResponse.value(forHTTPHeaderField: "Date") else {
-                Amplify.log.verbose("\(#function): Couldn't find Date header in URLSession metrics")
+                Amplify.log.verbose("\(#fileID)-\(#function): Couldn't find Date header in URLSession metrics")
                 onServerDateReceived(nil)
                 return
             }
@@ -133,13 +144,29 @@ final class WebSocketSession {
             dateFormatter.dateFormat = "EEE, d MMM yyyy HH:mm:ss z"
             dateFormatter.timeZone = TimeZone(secondsFromGMT: 0)
             guard let serverDate = dateFormatter.date(from: dateString) else {
-                Amplify.log.verbose("\(#function): Error parsing Date header in expected format")
+                Amplify.log.verbose("\(#fileID)-\(#function): Error parsing Date header in expected format")
                 onServerDateReceived(nil)
                 return
             }
 
             onServerDateReceived(serverDate)
         }
+
+        func urlSession(
+            _ session: URLSession,
+            task: URLSessionTask,
+            didCompleteWithError error: Error?
+        ) {
+            Amplify.log.verbose("\(#fileID)-\(#function): Session task didCompleteWithError : \(error)")
+        }
+
+        // MARK: - URLSessionDelegate methods
+        func urlSession(
+            _ session: URLSession,
+            didBecomeInvalidWithError error: Error?
+        ) {
+            Amplify.log.verbose("\(#fileID)-\(#function): Session task didBecomeInvalidWithError : \(error)")
+        }
     }
 
     enum WebSocketMessageResult {