Bug 1146649 - Retry token server token fetches to avoid 401/invalid-t…

…imestamp. r=rnewman

Account/SyncAuthState.swift

@@ -26,6 +26,34 @@ public class SyncAuthState {
         self.cache = nil
     }
 
+    // Generate an assertion and try to fetch a token server token, retrying at most a fixed number
+    // of times.
+    //
+    // It's tricky to get Swift to recurse into a closure that captures from the environment without
+    // segfaulting the compiler, so we pass everything around, like barbarians.
+    private func generateAssertionAndFetchTokenAt(audience: String, client: TokenServerClient, clientState: String?, married: MarriedState,
+            now: Timestamp, retryCount: Int) -> Deferred<Result<TokenServerToken>> {
+        let assertion = married.generateAssertionForAudience(audience, now: now)
+        return client.token(assertion, clientState: clientState).bind { result in
+            if retryCount > 0 {
+                if let tokenServerError = result.failureValue as? TokenServerError {
+                    switch tokenServerError {
+                    case let .Remote(code, status, remoteTimestamp) where code == 401 && status == "invalid-timestamp":
+                        if let remoteTimestamp = remoteTimestamp {
+                            let skew = Int64(remoteTimestamp) - Int64(now) // Without casts, runtime crash due to overflow.
+                            log.info("Token server responded with 401/invalid-timestamp: retrying with remote timestamp \(remoteTimestamp), which is local timestamp + skew = \(now) + \(skew).")
+                            return self.generateAssertionAndFetchTokenAt(audience, client: client, clientState: clientState, married: married, now: remoteTimestamp, retryCount: retryCount - 1)
+                        }
+                    default:
+                        break
+                    }
+                }
+            }
+            // Fall-through.
+            return Deferred(value: result)
+        }
+    }
+
     public func token(now: Timestamp, canBeExpired: Bool) -> Deferred<Result<(token: TokenServerToken, forKey: NSData)>> {
         if let (token, forKey, expiresAt) = cache {
             // Give ourselves some room to do work.
@@ -49,11 +77,11 @@ public class SyncAuthState {
         return account.marriedState().bind { result in
             if let married = result.successValue {
                 log.info("Account is in Married state; generating assertion.")
-                let assertion = married.generateAssertionForAudience(TokenServerClient.getAudienceForURL(self.tokenServerURL), now: now)
+                let audience = TokenServerClient.getAudienceForURL(self.tokenServerURL)
                 let client = TokenServerClient(URL: self.tokenServerURL)
                 let clientState = FxAClient10.computeClientState(married.kB)
-                let deferred = client.token(assertion, clientState: clientState)
                 log.debug("Fetching token server token.")
+                let deferred = self.generateAssertionAndFetchTokenAt(audience, client: client, clientState: clientState, married: married, now: now, retryCount: 1)
                 deferred.upon { result in
                     // This could race to update the cache with multiple token results.
                     // One racer will win -- that's fine, presumably she has the freshest token.
@@ -65,14 +93,9 @@ public class SyncAuthState {
                         self.cache = newCache
                     }
                 }
-                return deferred.map { result in
-                    return result.map { token in
-                        (token: token, forKey: married.kB)
-                    }
-                }
-            } else {
-                return Deferred(value: Result(failure: result.failureValue!))
+                return chain(deferred, { (token: $0, forKey: married.kB) })
             }
+            return deferResult(result.failureValue!)
         }
     }
 }

Account/TokenServerClient.swift

@@ -17,14 +17,8 @@ public struct TokenServerToken {
     public let api_endpoint: String
     public let uid: UInt64
     public let durationInSeconds: UInt64
-
-    init(id: String, key: String, api_endpoint: String, uid: UInt64, durationInSeconds: UInt64) {
-        self.id = id
-        self.key = key
-        self.api_endpoint = api_endpoint
-        self.uid = uid
-        self.durationInSeconds = durationInSeconds
-    }
+    // A healthy token server reports its timestamp.
+    public let remoteTimestamp: Timestamp
 
     /**
      * Return true if this token points to the same place as the other token.
@@ -36,14 +30,16 @@ public struct TokenServerToken {
 }
 
 enum TokenServerError {
-    case Remote(code: Int32, status: String?)
+    // A Remote error definitely has a status code, but we may not have a well-formed JSON response
+    // with a status; and we could have an unhealthy server that is not reporting its timestamp.
+    case Remote(code: Int32, status: String?, remoteTimestamp: Timestamp?)
     case Local(NSError)
 }
 
 extension TokenServerError: Printable, ErrorType {
     var description: String {
         switch self {
-        case let Remote(code: code, status: status):
+        case let Remote(code: code, status: status, remoteTimestamp: remoteTimestamp):
             if let status = status {
                 return "<TokenServerError.Remote \(code): \(status)>"
             } else {
@@ -70,33 +66,39 @@ public class TokenServerClient {
         }
     }
 
-    private class func remoteErrorFromJSON(json: JSON, statusCode: Int) -> TokenServerError? {
+    private class func parseTimestampHeader(header: String?) -> Timestamp? {
+        if let timestampString = header {
+            return decimalSecondsStringToTimestamp(timestampString)
+        } else {
+            return nil
+        }
+    }
+
+    private class func remoteErrorFromJSON(json: JSON, statusCode: Int, remoteTimestampHeader: String?) -> TokenServerError? {
         if json.isError {
             return nil
         }
         if 200 <= statusCode && statusCode <= 299 {
             return nil
         }
-        return TokenServerError.Remote(code: Int32(statusCode), status: json["status"].asString)
+        return TokenServerError.Remote(code: Int32(statusCode), status: json["status"].asString,
+            remoteTimestamp: parseTimestampHeader(remoteTimestampHeader))
     }
 
-    private class func tokenFromJSON(json: JSON) -> TokenServerToken? {
+    private class func tokenFromJSON(json: JSON, remoteTimestampHeader: String?) -> TokenServerToken? {
         if json.isError {
             return nil
         }
-        if let id = json["id"].asString {
-            if let key = json["key"].asString {
-                if let api_endpoint = json["api_endpoint"].asString {
-                    if let uid = json["uid"].asInt {
-                        if let durationInSeconds = json["duration"].asInt64 {
-                            if durationInSeconds > 0 {
-                                return TokenServerToken(id: id, key: key, api_endpoint: api_endpoint, uid: UInt64(uid),
-                                        durationInSeconds: UInt64(durationInSeconds))
-                            }
-                        }
-                    }
-                }
-            }
+        if let
+            remoteTimestamp = parseTimestampHeader(remoteTimestampHeader), // A token server that is not providing its timestamp is not healthy.
+            id = json["id"].asString,
+            key = json["key"].asString,
+            api_endpoint = json["api_endpoint"].asString,
+            uid = json["uid"].asInt,
+            durationInSeconds = json["duration"].asInt64
+            where durationInSeconds > 0 {
+            return TokenServerToken(id: id, key: key, api_endpoint: api_endpoint, uid: UInt64(uid),
+                durationInSeconds: UInt64(durationInSeconds), remoteTimestamp: remoteTimestamp)
         }
         return nil
     }
@@ -120,12 +122,15 @@ public class TokenServerClient {
 
                 if let data: AnyObject = data { // Declaring the type quiets a Swift warning about inferring AnyObject.
                     let json = JSON(data)
-                    if let remoteError = TokenServerClient.remoteErrorFromJSON(json, statusCode: response!.statusCode) {
+                    let remoteTimestampHeader = response?.allHeaderFields["x-timestamp"] as? String
+
+                    if let remoteError = TokenServerClient.remoteErrorFromJSON(json, statusCode: response!.statusCode,
+                        remoteTimestampHeader: remoteTimestampHeader) {
                         deferred.fill(Result(failure: remoteError))
                         return
                     }
 
-                    if let token = TokenServerClient.tokenFromJSON(json) {
+                    if let token = TokenServerClient.tokenFromJSON(json, remoteTimestampHeader: remoteTimestampHeader) {
                         deferred.fill(Result(success: token))
                         return
                     }

AccountTests/TokenServerClientTests.swift

@@ -62,6 +62,7 @@ class TokenServerClientTests: LiveAccountTest {
                     XCTAssertNotNil(token.api_endpoint)
                     XCTAssertTrue(token.uid >= 0)
                     XCTAssertTrue(token.api_endpoint.hasSuffix(String(token.uid)))
+                    XCTAssertTrue(token.remoteTimestamp >= 1429121686000) // Not a special timestamp; just a sanity check.
                 } else {
                     XCTAssertEqual(result.failureValue!.description, "")
                 }
@@ -85,9 +86,11 @@ class TokenServerClientTests: LiveAccountTest {
                 } else {
                     if let error = result.failureValue as? TokenServerError {
                         switch error {
-                        case let .Remote(code: code, status: status):
+                        case let .Remote(code, status, remoteTimestamp):
                             XCTAssertEqual(code, Int32(401)) // Bad auth.
                             XCTAssertEqual(status!, "error")
+                            XCTAssertFalse(remoteTimestamp == nil)
+                            XCTAssertTrue(remoteTimestamp >= 1429121686000) // Not a special timestamp; just a sanity check.
                         case let .Local(error):
                             XCTAssertNil(error)
                         }