feat: Tokenserver: Add validations and user updating for generation, keys_changed_at, and client_state (#1145)

Closes #866

.circleci/config.yml

@@ -38,7 +38,7 @@ commands:
           command: |
             cargo fmt -- --check
             # https://github.com/bodil/sized-chunks/issues/11
-            cargo audit --ignore RUSTSEC-2020-0041 --ignore RUSTSEC-2021-0078 --ignore RUSTSEC-2021-0079
+            cargo audit --ignore RUSTSEC-2020-0041 --ignore RUSTSEC-2021-0078 --ignore RUSTSEC-2021-0079 --ignore RUSTSEC-2020-0159 --ignore RUSTSEC-2020-0071
   python-check:
     steps:
       - run:

src/tokenserver/db/mock.rs

@@ -57,6 +57,10 @@ impl Db for MockDb {
         Box::pin(future::ok(true))
     }
 
+    fn get_node_id(&self, _params: params::GetNodeId) -> DbFuture<'_, results::GetNodeId> {
+        Box::pin(future::ok(results::GetNodeId::default()))
+    }
+
     #[cfg(test)]
     fn set_user_created_at(
         &self,

src/tokenserver/db/models.rs

@@ -10,11 +10,7 @@ use diesel_logger::LoggingConnection;
 use futures::future::LocalBoxFuture;
 use futures::TryFutureExt;
 
-use std::{
-    result,
-    sync::Arc,
-    time::{SystemTime, UNIX_EPOCH},
-};
+use std::{result, sync::Arc};
 
 use super::{params, results};
 use crate::db::error::{DbError, DbErrorKind};
@@ -94,6 +90,7 @@ impl TokenserverDb {
         }
 
         raw_users.sort_by_key(|raw_user| (raw_user.generation, raw_user.created_at));
+        raw_users.reverse();
 
         // The user with the greatest `generation` and `created_at` is the current user
         let raw_user = raw_users[0].clone();
@@ -132,6 +129,21 @@ impl TokenserverDb {
         Ok(user)
     }
 
+    fn get_node_id_sync(&self, params: params::GetNodeId) -> DbResult<results::GetNodeId> {
+        const QUERY: &str = r#"
+            SELECT id
+              FROM nodes
+             WHERE service = ?
+               AND node = ?
+        "#;
+
+        diesel::sql_query(QUERY)
+            .bind::<Integer, _>(params.service_id)
+            .bind::<Text, _>(&params.node)
+            .get_result(&self.inner.conn)
+            .map_err(Into::into)
+    }
+
     /// Mark users matching the given email and service ID as replaced.
     fn replace_users_sync(&self, params: params::ReplaceUsers) -> DbResult<results::ReplaceUsers> {
         const QUERY: &str = r#"
@@ -142,13 +154,12 @@ impl TokenserverDb {
                AND replaced_at IS NULL
                AND created_at < ?
         "#;
-        let timestamp = Self::get_timestamp_in_milliseconds();
 
         diesel::sql_query(QUERY)
-            .bind::<Bigint, _>(timestamp)
+            .bind::<Bigint, _>(params.replaced_at)
             .bind::<Integer, _>(&params.service_id)
             .bind::<Text, _>(&params.email)
-            .bind::<Bigint, _>(timestamp)
+            .bind::<Bigint, _>(params.replaced_at)
             .execute(&self.inner.conn)
             .map(|_| ())
             .map_err(Into::into)
@@ -214,7 +225,7 @@ impl TokenserverDb {
             .bind::<Text, _>(&user.email)
             .bind::<Bigint, _>(user.generation)
             .bind::<Text, _>(&user.client_state)
-            .bind::<Bigint, _>(Self::get_timestamp_in_milliseconds())
+            .bind::<Bigint, _>(user.created_at)
             .bind::<Bigint, _>(user.node_id)
             .bind::<Nullable<Bigint>, _>(user.keys_changed_at)
             .execute(&self.inner.conn)?;
@@ -231,13 +242,6 @@ impl TokenserverDb {
         Ok(result as u64 > 0)
     }
 
-    fn get_timestamp_in_milliseconds() -> i64 {
-        SystemTime::now()
-            .duration_since(UNIX_EPOCH)
-            .unwrap()
-            .as_millis() as i64
-    }
-
     #[cfg(test)]
     fn set_user_created_at_sync(
         &self,
@@ -315,6 +319,7 @@ impl Db for TokenserverDb {
     sync_db_method!(replace_users, replace_users_sync, ReplaceUsers);
     sync_db_method!(post_user, post_user_sync, PostUser);
     sync_db_method!(put_user, put_user_sync, PutUser);
+    sync_db_method!(get_node_id, get_node_id_sync, GetNodeId);
 
     fn check(&self) -> DbFuture<'_, results::Check> {
         let db = self.clone();
@@ -349,6 +354,8 @@ pub trait Db {
 
     fn check(&self) -> DbFuture<'_, results::Check>;
 
+    fn get_node_id(&self, params: params::GetNodeId) -> DbFuture<'_, results::GetNodeId>;
+
     #[cfg(test)]
     fn set_user_created_at(
         &self,
@@ -369,6 +376,8 @@ pub trait Db {
 mod tests {
     use super::*;
 
+    use std::time::{SystemTime, UNIX_EPOCH};
+
     use crate::settings::test_settings;
     use crate::tokenserver::db;
     use crate::tokenserver::db::pool::{DbPool, TokenserverPool};
@@ -704,6 +713,7 @@ mod tests {
         db.replace_users(params::ReplaceUsers {
             service_id: db::SYNC_1_5_SERVICE_ID,
             email: email1.to_owned(),
+            replaced_at: now,
         })
         .await?;
 
@@ -797,7 +807,45 @@ mod tests {
         Ok(())
     }
 
-    pub async fn db_pool() -> DbResult<TokenserverPool> {
+    #[tokio::test]
+    async fn get_node_id() -> Result<()> {
+        let pool = db_pool().await?;
+        let db = pool.get()?;
+
+        // Add a node
+        let node_id1 = db
+            .post_node(params::PostNode {
+                service_id: db::SYNC_1_5_SERVICE_ID,
+                node: "https://node1".to_owned(),
+                ..Default::default()
+            })
+            .await?
+            .id;
+
+        // Add another node
+        db.post_node(params::PostNode {
+            service_id: db::SYNC_1_5_SERVICE_ID,
+            node: "https://node2".to_owned(),
+            ..Default::default()
+        })
+        .await?;
+
+        // Get the ID of the first node
+        let id = db
+            .get_node_id(params::GetNodeId {
+                service_id: db::SYNC_1_5_SERVICE_ID,
+                node: "https://node1".to_owned(),
+            })
+            .await?
+            .id;
+
+        // The ID should match that of the first node
+        assert_eq!(node_id1, id);
+
+        Ok(())
+    }
+
+    async fn db_pool() -> DbResult<TokenserverPool> {
         let _ = env_logger::try_init();
 
         let tokenserver_settings = test_settings().tokenserver;

src/tokenserver/db/params.rs

@@ -29,6 +29,7 @@ pub struct PostUser {
     pub email: String,
     pub generation: i64,
     pub client_state: String,
+    pub created_at: i64,
     pub replaced_at: Option<i64>,
     pub node_id: i64,
     pub keys_changed_at: Option<i64>,
@@ -48,6 +49,7 @@ pub struct PutUser {
 pub struct ReplaceUsers {
     pub email: String,
     pub service_id: i32,
+    pub replaced_at: i64,
 }
 
 #[derive(Default)]
@@ -57,6 +59,12 @@ pub struct ReplaceUser {
     pub replaced_at: i64,
 }
 
+#[derive(Debug, Default)]
+pub struct GetNodeId {
+    pub service_id: i32,
+    pub node: String,
+}
+
 #[cfg(test)]
 pub type GetRawUsers = String;
 

src/tokenserver/db/pool.rs

@@ -66,7 +66,7 @@ impl DbPool for TokenserverPool {
     fn get(&self) -> Result<Box<dyn Db>, DbError> {
         self.inner
             .get()
-            .map(|db_pool| Box::new(TokenserverDb::new(db_pool)) as Box<dyn Db>)
+            .map(|conn| Box::new(TokenserverDb::new(conn)) as Box<dyn Db>)
             .map_err(DbError::from)
     }
 

src/tokenserver/db/results.rs

@@ -67,6 +67,12 @@ pub type ReplaceUsers = ();
 pub type ReplaceUser = ();
 pub type PutUser = ();
 
+#[derive(Default, QueryableByName)]
+pub struct GetNodeId {
+    #[sql_type = "Bigint"]
+    pub id: i64,
+}
+
 #[cfg(test)]
 pub type SetUserCreatedAt = ();
 

src/tokenserver/error.rs

@@ -6,11 +6,11 @@ use serde::{
     Serialize,
 };
 
-#[derive(Debug)]
+#[derive(Debug, PartialEq)]
 pub struct TokenserverError {
     status: &'static str,
     location: ErrorLocation,
-    name: &'static str,
+    name: String,
     description: &'static str,
     http_status: StatusCode,
 }
@@ -20,7 +20,7 @@ impl Default for TokenserverError {
         Self {
             status: "",
             location: ErrorLocation::default(),
-            name: "",
+            name: "".to_owned(),
             description: "Unauthorized",
             http_status: StatusCode::UNAUTHORIZED,
         }
@@ -63,7 +63,9 @@ impl TokenserverError {
     pub fn invalid_client_state(description: &'static str) -> Self {
         Self {
             status: "invalid-client-state",
+            location: ErrorLocation::Body,
             description,
+            name: "X-Client-State".to_owned(),
             ..Self::default()
         }
     }
@@ -78,13 +80,13 @@ impl TokenserverError {
         }
     }
 
-    pub fn unsupported(description: &'static str) -> Self {
+    pub fn unsupported(description: &'static str, name: String) -> Self {
         Self {
             status: "error",
             location: ErrorLocation::Url,
             description,
+            name,
             http_status: StatusCode::NOT_FOUND,
-            ..Self::default()
         }
     }
 
@@ -97,7 +99,7 @@ impl TokenserverError {
     }
 }
 
-#[derive(Clone, Copy, Debug)]
+#[derive(Clone, Copy, Debug, PartialEq)]
 pub enum ErrorLocation {
     Header,
     Url,
@@ -149,7 +151,7 @@ struct ErrorResponse {
 
 struct ErrorInstance {
     location: ErrorLocation,
-    name: &'static str,
+    name: String,
     description: &'static str,
 }
 
@@ -159,7 +161,7 @@ impl From<&TokenserverError> for ErrorResponse {
             status: error.status,
             errors: [ErrorInstance {
                 location: error.location,
-                name: error.name,
+                name: error.name.clone(),
                 description: error.description,
             }],
         }

src/tokenserver/extractors.rs

@@ -25,11 +25,12 @@ use crate::settings::Secrets;
 const DEFAULT_TOKEN_DURATION: u64 = 5 * 60;
 
 /// Information from the request needed to process a Tokenserver request.
-#[derive(Debug, PartialEq)]
+#[derive(Debug, Default, PartialEq)]
 pub struct TokenserverRequest {
     pub user: results::GetUser,
     pub fxa_uid: String,
-    pub generation: i64,
+    pub email: String,
+    pub generation: Option<i64>,
     pub keys_changed_at: i64,
     pub client_state: String,
     pub shared_secret: String,
@@ -39,6 +40,90 @@ pub struct TokenserverRequest {
     pub duration: u64,
 }
 
+impl TokenserverRequest {
+    /// Performs an elaborate set of consistency checks on the
+    /// provided claims, which we expect to behave as follows:
+    ///
+    ///   * `generation` is a monotonic timestamp, and increases every time
+    ///     there is an authentication-related change on the user's account.
+    ///
+    ///   * `keys_changed_at` is a monotonic timestamp, and increases every time
+    ///     the user's keys change. This is a type of auth-related change, so
+    ///     `keys_changed_at` <= `generation` at all times.
+    ///
+    ///   * `client_state` is a key fingerprint and should never change back
+    ///      to a previously-seen value.
+    ///
+    /// Callers who provide identity claims that violate any of these rules
+    /// either have stale credetials (in which case they should re-authenticate)
+    /// or are buggy (in which case we deny them access to the user's data).
+    ///
+    /// The logic here is slightly complicated by the fact that older versions
+    /// of the FxA server may not have been sending all the expected fields, and
+    /// that some clients do not report the `generation` timestamp.
+    fn validate(&self) -> Result<(), TokenserverError> {
+        // If the caller reports a generation number, then a change
+        // in keys should correspond to a change in generation number.
+        // Unfortunately a previous version of the server that didn't
+        // have `keys_changed_at` support may have already seen and
+        // written the new value of `generation`. The best we can do
+        // here is enforce that `keys_changed_at` <= `generation`.
+        if let (Some(generation), Some(user_keys_changed_at)) =
+            (self.generation, self.user.keys_changed_at)
+        {
+            if self.keys_changed_at > user_keys_changed_at && generation < self.keys_changed_at {
+                return Err(TokenserverError::invalid_keys_changed_at());
+            }
+        }
+
+        // The client state on the request must not have been used in the past.
+        if self.user.old_client_states.contains(&self.client_state) {
+            let error_message = "Unacceptable client-state value stale value";
+            return Err(TokenserverError::invalid_client_state(error_message));
+        }
+
+        // If the client state on the request differs from the most recently-used client state, it must
+        // be accompanied by a valid change in generation (if the client reports a generation).
+        if let Some(generation) = self.generation {
+            if self.client_state != self.user.client_state && generation <= self.user.generation {
+                let error_message =
+                    "Unacceptable client-state value new value with no generation change";
+                return Err(TokenserverError::invalid_client_state(error_message));
+            }
+        }
+
+        // If the client state on the request differs from the most recently-used client state, it must
+        // be accompanied by a valid change in keys_changed_at
+        if let Some(user_keys_changed_at) = self.user.keys_changed_at {
+            if self.client_state != self.user.client_state
+                && self.keys_changed_at <= user_keys_changed_at
+            {
+                let error_message =
+                    "Unacceptable client-state value new value with no keys_changed_at change";
+                return Err(TokenserverError::invalid_client_state(error_message));
+            }
+        }
+
+        // The generation on the request cannot be earlier than the generation stored on the user
+        // record.
+        if let Some(generation) = self.generation {
+            if self.user.generation > generation {
+                return Err(TokenserverError::invalid_generation());
+            }
+        }
+
+        // The keys_changed_at on the request cannot be earlier than the keys_changed_at stored on
+        // the user record.
+        if let Some(user_keys_changed_at) = self.user.keys_changed_at {
+            if user_keys_changed_at > self.keys_changed_at {
+                return Err(TokenserverError::invalid_keys_changed_at());
+            }
+        }
+
+        Ok(())
+    }
+}
+
 impl FromRequest for TokenserverRequest {
     type Config = ();
     type Error = Error;
@@ -79,22 +164,36 @@ impl FromRequest for TokenserverRequest {
                     } else {
                         return Err(TokenserverError::unsupported(
                             "Unsupported application version",
+                            version.to_owned(),
                         )
                         .into());
                     }
                 } else {
-                    return Err(TokenserverError::unsupported("Unsupported application").into());
+                    // NOTE: It would probably be better to include the name of the unsupported
+                    // application in the error message, but the old Tokenserver only includes
+                    // "application" in the error message. To keep the APIs between the old and
+                    // new Tokenservers as close as possible, we defer to the error message from
+                    // the old Tokenserver.
+                    return Err(TokenserverError::unsupported(
+                        "Unsupported application",
+                        "application".to_owned(),
+                    )
+                    .into());
                 }
             };
+            let email = format!("{}@{}", fxa_uid, state.fxa_email_domain);
             let user = {
                 let db = state.db_pool.get().map_err(|_| {
                     error!("⚠️ Could not acquire database connection");
 
                     TokenserverError::internal_error()
                 })?;
-                let email = format!("{}@{}", fxa_uid, state.fxa_email_domain);
 
-                db.get_user(params::GetUser { email, service_id }).await?
+                db.get_user(params::GetUser {
+                    email: email.clone(),
+                    service_id,
+                })
+                .await?
             };
             let duration = {
                 let params = Query::<QueryParams>::extract(&req).await?;
@@ -114,7 +213,8 @@ impl FromRequest for TokenserverRequest {
             let tokenserver_request = TokenserverRequest {
                 user,
                 fxa_uid,
-                generation: token_data.generation.unwrap_or(0),
+                email,
+                generation: token_data.generation,
                 keys_changed_at: key_id.keys_changed_at,
                 client_state: key_id.client_state,
                 shared_secret,
@@ -124,6 +224,8 @@ impl FromRequest for TokenserverRequest {
                 duration: duration.unwrap_or(DEFAULT_TOKEN_DURATION),
             };
 
+            tokenserver_request.validate()?;
+
             Ok(tokenserver_request)
         })
     }
@@ -355,7 +457,8 @@ mod tests {
         let expected_tokenserver_request = TokenserverRequest {
             user: results::GetUser::default(),
             fxa_uid: fxa_uid.to_owned(),
-            generation: 1234,
+            email: "test123@test.com".to_owned(),
+            generation: Some(1234),
             keys_changed_at: 1234,
             client_state: "616161".to_owned(),
             shared_secret: "Ted Koppel is a robot".to_owned(),
@@ -372,14 +475,12 @@ mod tests {
     async fn test_invalid_auth_token() {
         let fxa_uid = "test123";
         let verifier = {
-            let start = SystemTime::now();
-            let current_time = start.duration_since(UNIX_EPOCH).unwrap();
             let token_data = TokenData {
                 user: fxa_uid.to_owned(),
                 client_id: "client id".to_owned(),
                 scope: vec!["scope".to_owned()],
-                generation: Some(current_time.as_secs() as i64),
-                profile_changed_at: Some(current_time.as_secs() as i64),
+                generation: Some(1234),
+                profile_changed_at: None,
             };
             let valid = false;
 
@@ -415,14 +516,12 @@ mod tests {
         fn build_request() -> TestRequest {
             let fxa_uid = "test123";
             let verifier = {
-                let start = SystemTime::now();
-                let current_time = start.duration_since(UNIX_EPOCH).unwrap();
                 let token_data = TokenData {
                     user: fxa_uid.to_owned(),
                     client_id: "client id".to_owned(),
                     scope: vec!["scope".to_owned()],
-                    generation: Some(current_time.as_secs() as i64),
-                    profile_changed_at: Some(current_time.as_secs() as i64),
+                    generation: Some(1234),
+                    profile_changed_at: None,
                 };
                 let valid = true;
 
@@ -452,7 +551,8 @@ mod tests {
 
             assert_eq!(response.status(), StatusCode::NOT_FOUND);
 
-            let expected_error = TokenserverError::unsupported("Unsupported application version");
+            let expected_error =
+                TokenserverError::unsupported("Unsupported application version", "1.0".to_owned());
             let body = extract_body_as_str(ServiceResponse::new(request, response));
             assert_eq!(body, serde_json::to_string(&expected_error).unwrap());
         }
@@ -471,7 +571,8 @@ mod tests {
 
             assert_eq!(response.status(), StatusCode::NOT_FOUND);
 
-            let expected_error = TokenserverError::unsupported("Unsupported application");
+            let expected_error =
+                TokenserverError::unsupported("Unsupported application", "application".to_owned());
             let body = extract_body_as_str(ServiceResponse::new(request, response));
             assert_eq!(body, serde_json::to_string(&expected_error).unwrap());
         }
@@ -490,7 +591,8 @@ mod tests {
 
             assert_eq!(response.status(), StatusCode::NOT_FOUND);
 
-            let expected_error = TokenserverError::unsupported("Unsupported application");
+            let expected_error =
+                TokenserverError::unsupported("Unsupported application", "application".to_owned());
             let body = extract_body_as_str(ServiceResponse::new(request, response));
             assert_eq!(body, serde_json::to_string(&expected_error).unwrap());
         }
@@ -668,6 +770,187 @@ mod tests {
         }
     }
 
+    #[actix_rt::test]
+    async fn test_old_generation() {
+        // The request includes a generation that is less than the generation currently stored on
+        // the user record
+        let tokenserver_request = TokenserverRequest {
+            user: results::GetUser {
+                uid: 1,
+                client_state: "616161".to_owned(),
+                generation: 1234,
+                node: "node".to_owned(),
+                keys_changed_at: Some(1234),
+                created_at: 1234,
+                old_client_states: vec![],
+            },
+            fxa_uid: "test".to_owned(),
+            email: "test@test.com".to_owned(),
+            generation: Some(1233),
+            keys_changed_at: 1234,
+            client_state: "616161".to_owned(),
+            shared_secret: "secret".to_owned(),
+            hashed_fxa_uid: "abcdef".to_owned(),
+            hashed_device_id: "abcdef".to_owned(),
+            service_id: 1,
+            duration: DEFAULT_TOKEN_DURATION,
+        };
+
+        let error = tokenserver_request.validate().unwrap_err();
+        assert_eq!(error, TokenserverError::invalid_generation());
+    }
+
+    #[actix_rt::test]
+    async fn test_old_keys_changed_at() {
+        // The request includes a keys_changed_at that is less than the keys_changed_at currently
+        // stored on the user record
+        let tokenserver_request = TokenserverRequest {
+            user: results::GetUser {
+                uid: 1,
+                client_state: "616161".to_owned(),
+                generation: 1234,
+                node: "node".to_owned(),
+                keys_changed_at: Some(1234),
+                created_at: 1234,
+                old_client_states: vec![],
+            },
+            fxa_uid: "test".to_owned(),
+            email: "test@test.com".to_owned(),
+            generation: Some(1234),
+            keys_changed_at: 1233,
+            client_state: "616161".to_owned(),
+            shared_secret: "secret".to_owned(),
+            hashed_fxa_uid: "abcdef".to_owned(),
+            hashed_device_id: "abcdef".to_owned(),
+            service_id: 1,
+            duration: DEFAULT_TOKEN_DURATION,
+        };
+
+        let error = tokenserver_request.validate().unwrap_err();
+        assert_eq!(error, TokenserverError::invalid_keys_changed_at());
+    }
+
+    #[actix_rt::test]
+    async fn test_keys_changed_without_generation_change() {
+        // The request includes a new value for keys_changed_at without a new value for generation
+        let tokenserver_request = TokenserverRequest {
+            user: results::GetUser {
+                uid: 1,
+                client_state: "616161".to_owned(),
+                generation: 1234,
+                node: "node".to_owned(),
+                keys_changed_at: Some(1234),
+                created_at: 1234,
+                old_client_states: vec![],
+            },
+            fxa_uid: "test".to_owned(),
+            email: "test@test.com".to_owned(),
+            generation: Some(1234),
+            keys_changed_at: 1235,
+            client_state: "616161".to_owned(),
+            shared_secret: "secret".to_owned(),
+            hashed_fxa_uid: "abcdef".to_owned(),
+            hashed_device_id: "abcdef".to_owned(),
+            service_id: 1,
+            duration: DEFAULT_TOKEN_DURATION,
+        };
+
+        let error = tokenserver_request.validate().unwrap_err();
+        assert_eq!(error, TokenserverError::invalid_keys_changed_at());
+    }
+
+    #[actix_rt::test]
+    async fn test_old_client_state() {
+        // The request includes a previously-used client state that is not the user's current
+        // client state
+        let tokenserver_request = TokenserverRequest {
+            user: results::GetUser {
+                uid: 1,
+                client_state: "616161".to_owned(),
+                generation: 1234,
+                node: "node".to_owned(),
+                keys_changed_at: Some(1234),
+                created_at: 1234,
+                old_client_states: vec!["626262".to_owned()],
+            },
+            fxa_uid: "test".to_owned(),
+            email: "test@test.com".to_owned(),
+            generation: Some(1234),
+            keys_changed_at: 1234,
+            client_state: "626262".to_owned(),
+            shared_secret: "secret".to_owned(),
+            hashed_fxa_uid: "abcdef".to_owned(),
+            hashed_device_id: "abcdef".to_owned(),
+            service_id: 1,
+            duration: DEFAULT_TOKEN_DURATION,
+        };
+
+        let error = tokenserver_request.validate().unwrap_err();
+        let error_message = "Unacceptable client-state value stale value";
+        assert_eq!(error, TokenserverError::invalid_client_state(error_message));
+    }
+
+    #[actix_rt::test]
+    async fn test_new_client_state_without_generation_change() {
+        // The request includes a new client state without a new generation value
+        let tokenserver_request = TokenserverRequest {
+            user: results::GetUser {
+                uid: 1,
+                client_state: "616161".to_owned(),
+                generation: 1234,
+                node: "node".to_owned(),
+                keys_changed_at: Some(1234),
+                created_at: 1234,
+                old_client_states: vec![],
+            },
+            fxa_uid: "test".to_owned(),
+            email: "test@test.com".to_owned(),
+            generation: Some(1234),
+            keys_changed_at: 1234,
+            client_state: "626262".to_owned(),
+            shared_secret: "secret".to_owned(),
+            hashed_fxa_uid: "abcdef".to_owned(),
+            hashed_device_id: "abcdef".to_owned(),
+            service_id: 1,
+            duration: DEFAULT_TOKEN_DURATION,
+        };
+
+        let error = tokenserver_request.validate().unwrap_err();
+        let error_message = "Unacceptable client-state value new value with no generation change";
+        assert_eq!(error, TokenserverError::invalid_client_state(error_message));
+    }
+
+    #[actix_rt::test]
+    async fn test_new_client_state_without_key_change() {
+        // The request includes a new client state without a new keys_changed_at value
+        let tokenserver_request = TokenserverRequest {
+            user: results::GetUser {
+                uid: 1,
+                client_state: "616161".to_owned(),
+                generation: 1234,
+                node: "node".to_owned(),
+                keys_changed_at: Some(1234),
+                created_at: 1234,
+                old_client_states: vec![],
+            },
+            fxa_uid: "test".to_owned(),
+            email: "test@test.com".to_owned(),
+            generation: Some(1235),
+            keys_changed_at: 1234,
+            client_state: "626262".to_owned(),
+            shared_secret: "secret".to_owned(),
+            hashed_fxa_uid: "abcdef".to_owned(),
+            hashed_device_id: "abcdef".to_owned(),
+            service_id: 1,
+            duration: DEFAULT_TOKEN_DURATION,
+        };
+
+        let error = tokenserver_request.validate().unwrap_err();
+        let error_message =
+            "Unacceptable client-state value new value with no keys_changed_at change";
+        assert_eq!(error, TokenserverError::invalid_client_state(error_message));
+    }
+
     fn extract_body_as_str(sresponse: ServiceResponse) -> String {
         String::from_utf8(block_on(test::read_body(sresponse)).to_vec()).unwrap()
     }

src/tokenserver/handlers.rs

@@ -1,27 +1,18 @@
 use std::{
-    sync::Arc,
+    cmp,
+    collections::HashMap,
     time::{Duration, SystemTime, UNIX_EPOCH},
 };
 
-use actix_web::http::StatusCode;
-use actix_web::web::Data;
-use actix_web::Error;
-use actix_web::{HttpRequest, HttpResponse};
-use hmac::{Hmac, Mac, NewMac};
+use actix_web::{http::StatusCode, Error, HttpResponse};
 use serde::Serialize;
 use serde_json::Value;
-use sha2::Sha256;
-use std::collections::HashMap;
 
-use super::db::{self, models::Db, params::GetUser};
+use super::db::models::Db;
+use super::db::params::{GetNodeId, PostUser, PutUser, ReplaceUsers};
 use super::extractors::TokenserverRequest;
 use super::support::Tokenlib;
-use super::ServerState;
-use crate::{
-    error::{ApiError, ApiErrorKind},
-    settings::Secrets,
-    tokenserver::support::MakeTokenPlaintext,
-};
+use crate::tokenserver::support::MakeTokenPlaintext;
 
 #[derive(Debug, Serialize)]
 pub struct TokenserverResult {
@@ -34,121 +25,137 @@ pub struct TokenserverResult {
 }
 
 pub async fn get_tokenserver_result(
-    tokenserver_request: TokenserverRequest,
-    request: HttpRequest,
+    req: TokenserverRequest,
+    db: Box<dyn Db>,
 ) -> Result<HttpResponse, Error> {
-    let state = request
-        .app_data::<Data<Option<ServerState>>>()
-        .ok_or_else(|| internal_error("Could not load the app state"))?
-        .as_ref()
-        .as_ref()
-        .unwrap();
-    let db = {
-        let db_pool = state.db_pool.clone();
-        db_pool.get().map_err(ApiError::from)?
-    };
-
-    let user_email = format!("{}@{}", tokenserver_request.fxa_uid, state.fxa_email_domain);
-    let tokenserver_user = {
-        let params = GetUser {
-            email: user_email.clone(),
-            service_id: db::SYNC_1_5_SERVICE_ID,
-        };
-
-        db.get_user(params).await?
-    };
-
-    let fxa_metrics_hash_secret = state.fxa_metrics_hash_secret.clone().into_bytes();
-
-    let hashed_fxa_uid_full =
-        fxa_metrics_hash(&tokenserver_request.fxa_uid, &fxa_metrics_hash_secret)?;
-    let hashed_fxa_uid = &hashed_fxa_uid_full[0..32];
-    let hashed_device_id = {
-        let device_id = "none".to_string();
-        hash_device_id(hashed_fxa_uid, &device_id, &fxa_metrics_hash_secret)?
-    };
-
-    let fxa_kid = {
-        let client_state_b64 =
-            base64::encode_config(&tokenserver_user.client_state, base64::URL_SAFE_NO_PAD);
-
-        format!(
-            "{:013}-{:}",
-            tokenserver_user
-                .keys_changed_at
-                .unwrap_or(tokenserver_request.generation),
-            client_state_b64
-        )
-    };
+    let updates = update_user(&req, db).await?;
 
     let (token, derived_secret) = {
-        let shared_secret = String::from_utf8(
-            request
-                .app_data::<Data<Arc<Secrets>>>()
-                .ok_or_else(|| internal_error("Could not load the app secrets"))?
-                .master_secret
-                .clone(),
-        )
-        .map_err(|_| internal_error("Failed to read master secret"))?;
-
-        let make_token_plaintext = {
-            let expires = {
-                let start = SystemTime::now();
-                let current_time = start.duration_since(UNIX_EPOCH).unwrap();
-                let expires = current_time + Duration::new(tokenserver_request.duration, 0);
-
-                expires.as_secs()
-            };
-
-            MakeTokenPlaintext {
-                node: tokenserver_user.node.clone(),
-                fxa_kid,
-                fxa_uid: tokenserver_request.fxa_uid.clone(),
-                hashed_device_id,
-                hashed_fxa_uid: hashed_fxa_uid.to_owned(),
-                expires,
-                uid: tokenserver_user.uid,
-            }
-        };
-
-        Tokenlib::get_token_and_derived_secret(make_token_plaintext, &shared_secret)?
+        let token_plaintext = get_token_plaintext(&req, &updates);
+        Tokenlib::get_token_and_derived_secret(token_plaintext, &req.shared_secret)?
     };
 
-    let api_endpoint = format!("{:}/1.5/{:}", tokenserver_user.node, tokenserver_user.uid);
-
     let result = TokenserverResult {
         id: token,
         key: derived_secret,
-        uid: tokenserver_user.uid,
-        api_endpoint,
-        duration: tokenserver_request.duration,
-        hashed_fxa_uid: hashed_fxa_uid.to_owned(),
+        uid: updates.uid,
+        api_endpoint: format!("{:}/1.5/{:}", req.user.node, req.user.uid),
+        duration: req.duration,
+        hashed_fxa_uid: req.hashed_fxa_uid,
     };
 
     Ok(HttpResponse::build(StatusCode::OK).json(result))
 }
 
-fn fxa_metrics_hash(fxa_uid: &str, hmac_key: &[u8]) -> Result<String, Error> {
-    let mut mac = Hmac::<Sha256>::new_from_slice(hmac_key)
-        .map_err::<ApiError, _>(|err| ApiErrorKind::Internal(err.to_string()).into())?;
-    mac.update(fxa_uid.as_bytes());
+fn get_token_plaintext(req: &TokenserverRequest, updates: &UserUpdates) -> MakeTokenPlaintext {
+    let fxa_kid = {
+        let client_state_b64 = base64::encode_config(&req.client_state, base64::URL_SAFE_NO_PAD);
 
-    let result = mac.finalize().into_bytes();
-    Ok(hex::encode(result))
+        format!("{:013}-{:}", updates.keys_changed_at, client_state_b64)
+    };
+
+    let expires = {
+        let start = SystemTime::now();
+        let current_time = start.duration_since(UNIX_EPOCH).unwrap();
+        let expires = current_time + Duration::from_secs(req.duration);
+
+        expires.as_secs()
+    };
+
+    MakeTokenPlaintext {
+        node: req.user.node.to_owned(),
+        fxa_kid,
+        fxa_uid: req.fxa_uid.clone(),
+        hashed_device_id: req.hashed_device_id.clone(),
+        hashed_fxa_uid: req.hashed_fxa_uid.clone(),
+        expires,
+        uid: updates.uid.to_owned(),
+    }
 }
 
-fn hash_device_id(fxa_uid: &str, device: &str, hmac_key: &[u8]) -> Result<String, Error> {
-    let mut to_hash = String::from(fxa_uid);
-    to_hash.push_str(device);
-    let fxa_metrics_hash = fxa_metrics_hash(&to_hash, hmac_key)?;
-
-    Ok(String::from(&fxa_metrics_hash[0..32]))
+struct UserUpdates {
+    keys_changed_at: i64,
+    uid: i64,
 }
 
-fn internal_error(message: &str) -> HttpResponse {
-    error!("⚠️ {}", message);
+async fn update_user(req: &TokenserverRequest, db: Box<dyn Db>) -> Result<UserUpdates, Error> {
+    // If the keys_changed_at in the request is larger than that stored on the user record,
+    // update to the value in the request.
+    let keys_changed_at = if let Some(user_keys_changed_at) = req.user.keys_changed_at {
+        cmp::max(req.keys_changed_at, user_keys_changed_at)
+    } else {
+        req.keys_changed_at
+    };
 
-    HttpResponse::InternalServerError().body("")
+    let generation = if let Some(generation) = req.generation {
+        // If there's a generation on the request, choose the larger of that and the generation
+        // already stored on the user record.
+        cmp::max(generation, req.user.generation)
+    } else if req.keys_changed_at > req.user.generation {
+        // If there's not a generation on the request and the keys_changed_at on the request is
+        // larger than the generation stored on the user record, set the user's generation to be
+        // the keys_changed_at on the request.
+        req.keys_changed_at
+    } else {
+        // As a fallback, set the user's generation to be 0.
+        0
+    };
+
+    // If the client state changed, we need to mark the current user as "replaced" and create a
+    // new user record. Otherwise, we can update the user in place.
+    if req.client_state != req.user.client_state {
+        let timestamp = SystemTime::now()
+            .duration_since(UNIX_EPOCH)
+            .unwrap()
+            .as_millis() as i64;
+
+        // Create new user record with updated generation/keys_changed_at
+        let post_user_params = PostUser {
+            service_id: req.service_id,
+            email: req.email.clone(),
+            generation,
+            client_state: req.client_state.clone(),
+            replaced_at: None,
+            node_id: db
+                .get_node_id(GetNodeId {
+                    service_id: req.service_id,
+                    node: req.user.node.clone(),
+                })
+                .await?
+                .id,
+            keys_changed_at: Some(keys_changed_at),
+            created_at: timestamp,
+        };
+        let uid = db.post_user(post_user_params).await?.id;
+
+        // Make sure each old row is marked as replaced (they might not be, due to races in row
+        // creation)
+        db.replace_users(ReplaceUsers {
+            email: req.email.clone(),
+            service_id: req.service_id,
+            replaced_at: timestamp,
+        })
+        .await?;
+
+        Ok(UserUpdates {
+            keys_changed_at,
+            uid,
+        })
+    } else {
+        let params = PutUser {
+            email: req.email.clone(),
+            service_id: req.service_id,
+            generation,
+            keys_changed_at: Some(keys_changed_at),
+        };
+
+        db.put_user(params).await?;
+
+        Ok(UserUpdates {
+            keys_changed_at,
+            uid: req.user.uid,
+        })
+    }
 }
 
 pub async fn heartbeat(db: Box<dyn Db>) -> Result<HttpResponse, Error> {

tools/integration_tests/tokenserver/test_authorization.py

@@ -517,3 +517,21 @@ class TestAuthorization(TestCase, unittest.TestCase):
         self.assertEqual(res.json, expected_error_response)
         headers['X-Client-State'] = '616161'
         res = self.app.get('/1.0/sync/1.5', headers=headers)
+
+    def test_x_key_id_header_required(self):
+        headers = {
+            'Authorization': 'Bearer %s' % self._forge_oauth_token()
+        }
+        # A request without an X-KeyID header should fail
+        res = self.app.get('/1.0/sync/1.5', headers=headers, status=401)
+        expected_error_response = {
+            'errors': [
+                {
+                    'description': 'Missing X-KeyID header',
+                    'location': 'header',
+                    'name': ''
+                }
+            ],
+            'status': 'invalid-key-id'
+        }
+        self.assertEqual(res.json, expected_error_response)