fix(liveness): correct websocket retry logic

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)
-  [email protected] = webSocket
- } else {
-  // server will close this websocket, don't report that failure back
- offset = tempOffset
- start()
+ super.onOpen(webSocket, response)
+
+ [email protected] = 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