diff --git a/services/tunnelbroker/src/constants.rs b/services/tunnelbroker/src/constants.rs --- a/services/tunnelbroker/src/constants.rs +++ b/services/tunnelbroker/src/constants.rs @@ -43,13 +43,15 @@ // This table holds a device token associated with a device. // // - (primary key) = (deviceID: Partition Key) - // - deviceID: The public key of a device's olm identity key + // - deviceID: The public key of a device's olm identity key. // - deviceToken: Token to push services uploaded by device. + // - tokenInvalid: Information is token is invalid. pub mod device_tokens { pub const TABLE_NAME: &str = "tunnelbroker-device-tokens"; pub const PARTITION_KEY: &str = "deviceID"; pub const DEVICE_ID: &str = "deviceID"; pub const DEVICE_TOKEN: &str = "deviceToken"; + pub const TOKEN_INVALID: &str = "tokenInvalid"; pub const DEVICE_TOKEN_INDEX_NAME: &str = "deviceToken-index"; } diff --git a/services/tunnelbroker/src/database/mod.rs b/services/tunnelbroker/src/database/mod.rs --- a/services/tunnelbroker/src/database/mod.rs +++ b/services/tunnelbroker/src/database/mod.rs @@ -36,6 +36,11 @@ } } +pub struct DeviceTokenEntry { + pub device_token: String, + pub token_invalid: Option, +} + impl DatabaseClient { pub fn new(aws_config: &AwsConfig) -> Self { let client = DynamoDBClient::new(aws_config); @@ -153,7 +158,7 @@ pub async fn get_device_token( &self, device_id: &str, - ) -> Result, Error> { + ) -> Result, Error> { let get_response = self .client .get_item() @@ -174,7 +179,17 @@ }; let device_token: String = item.take_attr(device_tokens::DEVICE_TOKEN)?; - Ok(Some(device_token)) + if let Ok(token_invalid) = item.take_attr(device_tokens::TOKEN_INVALID) { + Ok(Some(DeviceTokenEntry { + device_token, + token_invalid, + })) + } else { + Ok(Some(DeviceTokenEntry { + device_token, + token_invalid: None, + })) + } } pub async fn set_device_token( diff --git a/services/tunnelbroker/src/websockets/session.rs b/services/tunnelbroker/src/websockets/session.rs --- a/services/tunnelbroker/src/websockets/session.rs +++ b/services/tunnelbroker/src/websockets/session.rs @@ -69,6 +69,7 @@ MissingAPNsClient, MissingFCMClient, MissingDeviceToken, + InvalidDeviceToken, } // Parse a session request and retrieve the device information @@ -530,6 +531,19 @@ .get_device_token(&device_id) .await .map_err(SessionError::DatabaseError)?; - db_token.ok_or_else(|| SessionError::MissingDeviceToken) + + match db_token { + Some(token) => match token.token_invalid { + None => Ok(token.device_token), + Some(token_invalid) => { + if token_invalid { + Err(SessionError::InvalidDeviceToken) + } else { + Ok(token.device_token) + } + } + }, + None => Err(SessionError::MissingDeviceToken), + } } }