fix(liveness): correct websocket retry logic (#2634)

Co-authored-by: Thomas Leing <[email protected]>
Co-authored-by: tjroach <[email protected]>

Author: 3 people

aws-predictions/src/main/java/com/amplifyframework/predictions/aws/http/LivenessWebSocket.kt

@@ -37,6 +37,7 @@ import com.amplifyframework.predictions.aws.models.liveness.ColorDisplayed
 import com.amplifyframework.predictions.aws.models.liveness.FaceMovementAndLightClientChallenge
 import com.amplifyframework.predictions.aws.models.liveness.FreshnessColor
 import com.amplifyframework.predictions.aws.models.liveness.InitialFace
+import com.amplifyframework.predictions.aws.models.liveness.InvalidSignatureException
 import com.amplifyframework.predictions.aws.models.liveness.LivenessResponseStream
 import com.amplifyframework.predictions.aws.models.liveness.SessionInformation
 import com.amplifyframework.predictions.aws.models.liveness.TargetFace
@@ -78,23 +79,13 @@ internal class LivenessWebSocket(
     private val signer = AWSV4Signer()
     private var credentials: Credentials? = null
 
-    internal var offset = 0L
-    internal enum class ReconnectState {
-        INITIAL,
-        RECONNECTING,
-        RECONNECTING_AGAIN;
-
-        companion object {
-            fun next(state: ReconnectState): ReconnectState {
-                return when (state) {
-                    INITIAL -> RECONNECTING
-                    RECONNECTING -> RECONNECTING_AGAIN
-                    RECONNECTING_AGAIN -> RECONNECTING_AGAIN
-                }
-            }
-        }
+    // The reported time difference between the server and client. Only set if diff is higher than 4 minutes
+    internal var timeDiffOffsetInMillis = 0L
+    internal enum class ConnectionState {
+        NORMAL,
+        ATTEMPT_RECONNECT,
     }
-    internal var reconnectState = ReconnectState.INITIAL
+    internal var reconnectState = ConnectionState.NORMAL
 
     @VisibleForTesting
     internal var webSocket: WebSocket? = null
@@ -103,7 +94,7 @@ internal class LivenessWebSocket(
     private var faceDetectedStart = 0L
     private var videoStartTimestamp = 0L
     private var videoEndTimestamp = 0L
-    private var webSocketError: PredictionsException? = null
+    @VisibleForTesting internal var webSocketError: PredictionsException? = null
     internal var clientStoppedSession = false
     val json = Json { ignoreUnknownKeys = true }
 
@@ -119,15 +110,15 @@ internal class LivenessWebSocket(
                 date.time - adjustedDate()
             } else 0
 
-            reconnectState = ReconnectState.next(reconnectState)
-            // if offset is > 5 minutes, server will reject the request
-            if (kotlin.math.abs(tempOffset) < FIVE_MINUTES) {
-                super.onOpen(webSocket, response)
-                this@LivenessWebSocket.webSocket = webSocket
-            } else {
-                // server will close this websocket, don't report that failure back
-                offset = tempOffset
-                start()
+            super.onOpen(webSocket, response)
+
+            this@LivenessWebSocket.webSocket = webSocket
+
+            // If offset is > 4 minutes, server may reject the request
+            // The real allowed diff from serer is < 5 but we check for 4 to add a buffer
+            if (!isTimeDiffSafe(tempOffset)) {
+                LOG.info("Server reported a time difference between client and server of > 4 minutes")
+                timeDiffOffsetInMillis = tempOffset
             }
         }
 
@@ -169,14 +160,29 @@ internal class LivenessWebSocket(
         override fun onClosed(webSocket: WebSocket, code: Int, reason: String) {
             LOG.debug("WebSocket onClosed")
             super.onClosed(webSocket, code, reason)
-            if (reconnectState == ReconnectState.RECONNECTING) {
-                // do nothing; we expected the server to close the connection
-            } else if (code != NORMAL_SOCKET_CLOSURE_STATUS_CODE && !clientStoppedSession) {
-                val faceLivenessException = webSocketError ?: PredictionsException(
-                    "An error occurred during the face liveness check.",
-                    reason
-                )
-                onErrorReceived.accept(faceLivenessException)
+            if (code != NORMAL_SOCKET_CLOSURE_STATUS_CODE && !clientStoppedSession) {
+                val recordedError = webSocketError
+
+                /*
+                If the server reports an invalid signature due to a time difference between the local clock and the
+                server clock, AND we haven't already tried to reconnect, then we should try to reconnect with an offset
+                */
+                if (reconnectState == ConnectionState.NORMAL &&
+                    !isTimeDiffSafe(timeDiffOffsetInMillis) &&
+                    recordedError is PredictionsException &&
+                    recordedError.cause is InvalidSignatureException
+                ) {
+                    LOG.info("The server rejected the connection due to a likely time difference. Attempting reconnect")
+                    reconnectState = ConnectionState.ATTEMPT_RECONNECT
+                    webSocketError = null
+                    start()
+                } else {
+                    val faceLivenessException = recordedError ?: PredictionsException(
+                        "An error occurred during the face liveness check.",
+                        reason
+                    )
+                    onErrorReceived.accept(faceLivenessException)
+                }
             } else {
                 onComplete.call()
             }
@@ -197,14 +203,6 @@ internal class LivenessWebSocket(
     }
 
     fun start() {
-        if (reconnectState == ReconnectState.RECONNECTING_AGAIN) {
-            onErrorReceived.accept(
-                PredictionsException(
-                    "Invalid device time",
-                    "Too many attempts were made to correct device time"
-                )
-            )
-        }
         val userAgent = getUserAgent()
 
         val okHttpClient = OkHttpClient.Builder()
@@ -312,6 +310,18 @@ internal class LivenessWebSocket(
             AccessDeniedException(
                 cause = livenessResponse.accessDeniedException
             )
+        } else if (livenessResponse.unrecognizedClientException != null) {
+            PredictionsException(
+                "Unrecognized client",
+                livenessResponse.unrecognizedClientException,
+                "Please check your credentials"
+            )
+        } else if (livenessResponse.invalidSignatureException != null) {
+            PredictionsException(
+                "Invalid signature",
+                livenessResponse.invalidSignatureException,
+                "Please check your credentials"
+            )
         } else {
             PredictionsException(
                 "An unknown error occurred during the Liveness flow.",
@@ -477,12 +487,14 @@ internal class LivenessWebSocket(
     }
 
     fun adjustedDate(date: Long = Date().time): Long {
-        return date + offset
+        return date + timeDiffOffsetInMillis
     }
 
+    private fun isTimeDiffSafe(diffInMillis: Long) = kotlin.math.abs(diffInMillis) < FOUR_MINUTES
+
     companion object {
         private const val NORMAL_SOCKET_CLOSURE_STATUS_CODE = 1000
-        private val FIVE_MINUTES = 1000 * 60 * 5
+        private const val FOUR_MINUTES = 1000 * 60 * 4
         @VisibleForTesting val datePattern = "EEE, d MMM yyyy HH:mm:ss z"
         private val LOG = Amplify.Logging.logger(CategoryType.PREDICTIONS, "amplify:aws-predictions")
     }

aws-predictions/src/main/java/com/amplifyframework/predictions/aws/models/liveness/LivenessResponseStream.kt

@@ -30,5 +30,6 @@ internal data class LivenessResponseStream(
     @SerialName("ServiceUnavailableException") val serviceUnavailableException: ServiceUnavailableException? = null,
     @SerialName("SessionNotFoundException") val sessionNotFoundException: SessionNotFoundException? = null,
     @SerialName("AccessDeniedException") val accessDeniedException: AccessDeniedException? = null,
-    @SerialName("InvalidSignatureException") val invalidSignatureException: InvalidSignatureException? = null
+    @SerialName("InvalidSignatureException") val invalidSignatureException: InvalidSignatureException? = null,
+    @SerialName("UnrecognizedClientException") val unrecognizedClientException: UnrecognizedClientException? = null
 )

aws-predictions/src/main/java/com/amplifyframework/predictions/aws/models/liveness/UnrecognizedClientException.kt

@@ -0,0 +1,28 @@
+/*
+ * Copyright 2023 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License").
+ * You may not use this file except in compliance with the License.
+ * A copy of the License is located at
+ *
+ *  http://aws.amazon.com/apache2.0
+ *
+ * or in the "license" file accompanying this file. This file is distributed
+ * on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either
+ * express or implied. See the License for the specific language governing
+ * permissions and limitations under the License.
+ */
+package com.amplifyframework.predictions.aws.models.liveness
+
+import kotlinx.serialization.SerialName
+import kotlinx.serialization.Serializable
+
+/**
+ * Constructs a new UnrecognizedClientException with the specified error message.
+ *
+ * @param message Describes the error encountered.
+ */
+@Serializable
+internal data class UnrecognizedClientException(
+    @SerialName("Message") override val message: String
+) : Exception(message)

aws-predictions/src/test/java/com/amplifyframework/predictions/aws/http/LivenessWebSocketTest.kt

@@ -28,6 +28,7 @@ import com.amplifyframework.predictions.aws.models.liveness.ColorSequence
 import com.amplifyframework.predictions.aws.models.liveness.DisconnectionEvent
 import com.amplifyframework.predictions.aws.models.liveness.FaceMovementAndLightServerChallenge
 import com.amplifyframework.predictions.aws.models.liveness.FreshnessColor
+import com.amplifyframework.predictions.aws.models.liveness.InvalidSignatureException
 import com.amplifyframework.predictions.aws.models.liveness.LightChallengeType
 import com.amplifyframework.predictions.aws.models.liveness.OvalParameters
 import com.amplifyframework.predictions.aws.models.liveness.ServerChallenge
@@ -326,7 +327,9 @@ internal class LivenessWebSocketTest {
     fun `web socket detects clock skew from server response`() {
         val livenessWebSocket = createLivenessWebSocket()
         mockkConstructor(WebSocket::class)
-        val socket: WebSocket = mockk()
+        val socket: WebSocket = mockk {
+            every { close(any(), any()) } returns true
+        }
         livenessWebSocket.webSocket = socket
         val sdf = SimpleDateFormat(LivenessWebSocket.datePattern, Locale.US)
 
@@ -339,7 +342,7 @@ internal class LivenessWebSocketTest {
         livenessWebSocket.webSocketListener.onOpen(socket, response)
 
         // now we should restart the websocket with an adjusted time
-        val openLatch = CountDownLatch(1)
+        val openLatch = CountDownLatch(2)
         val latchingListener = LatchingWebSocketResponseListener(
             livenessWebSocket.webSocketListener,
             openLatch = openLatch
@@ -349,25 +352,32 @@ internal class LivenessWebSocketTest {
         server.enqueue(MockResponse().withWebSocketUpgrade(ServerWebSocketListener()))
         server.start()
 
+        livenessWebSocket.webSocketError = PredictionsException(
+            "invalid signature",
+            InvalidSignatureException("invalid signature"),
+            "invalid signature"
+        )
+        livenessWebSocket.webSocketListener.onClosed(mockk(), 1011, "closing")
+
         openLatch.await(3, TimeUnit.SECONDS)
 
         assertTrue(livenessWebSocket.webSocket != null)
-        val originalRequest = livenessWebSocket.webSocket!!.request()
+        val reconnectRequest = livenessWebSocket.webSocket!!.request()
 
         // make sure that followup request sends offset date
         val sdfGMT = SimpleDateFormat("yyyyMMdd'T'HHmmss'Z'", Locale.US)
         sdfGMT.timeZone = TimeZone.getTimeZone("GMT")
-        val sentDate = originalRequest.url.queryParameter("X-Amz-Date") ?.let { sdfGMT.parse(it) }
+        val sentDate = reconnectRequest.url.queryParameter("X-Amz-Date") ?.let { sdfGMT.parse(it) }
         val diff = abs(Date().time - sentDate?.time!!)
         assert(oneHour - 10000 < diff && diff < oneHour + 10000)
 
         // also make sure that followup request is valid
         assertTrue(
-            originalRequest.url.queryParameter("X-Amz-Credential")!!.endsWith("//rekognition/aws4_request")
+            reconnectRequest.url.queryParameter("X-Amz-Credential")!!.endsWith("//rekognition/aws4_request")
         )
-        assertEquals("299", originalRequest.url.queryParameter("X-Amz-Expires"))
-        assertEquals("host", originalRequest.url.queryParameter("X-Amz-SignedHeaders"))
-        assertEquals("AWS4-HMAC-SHA256", originalRequest.url.queryParameter("X-Amz-Algorithm"))
+        assertEquals("299", reconnectRequest.url.queryParameter("X-Amz-Expires"))
+        assertEquals("host", reconnectRequest.url.queryParameter("X-Amz-SignedHeaders"))
+        assertEquals("AWS4-HMAC-SHA256", reconnectRequest.url.queryParameter("X-Amz-Algorithm"))
     }
 
     @Test