diff --git a/keyserver/addons/rust-node-addon/src/identity_client/login.rs b/keyserver/addons/rust-node-addon/src/identity_client/login.rs --- a/keyserver/addons/rust-node-addon/src/identity_client/login.rs +++ b/keyserver/addons/rust-node-addon/src/identity_client/login.rs @@ -54,20 +54,29 @@ }; debug!("Starting login to identity service"); - let login_start_response = identity_client + let response = identity_client .login_password_user_start(login_start_request) .await - .map_err(handle_grpc_error)? - .into_inner(); - + .map_err(handle_grpc_error)?; debug!("Received login response from identity service"); + + // We need to get the load balancer cookie from from the response and send it + // in the subsequent request to ensure it is routed to the same identity + // service instance as the first request + let cookie = get_lb_cookie(&response) + .map_err(|_| Error::from_reason("Failed to get lb cookie from response"))?; + + let login_start_response = response.into_inner(); + let opaque_login_upload = client_login .finish(&login_start_response.opaque_login_response) .map_err(|_| Error::from_reason("Failed to finish opaque login request"))?; - let login_finish_request = OpaqueLoginFinishRequest { + + let mut login_finish_request = Request::new(OpaqueLoginFinishRequest { session_id: login_start_response.session_id, opaque_login_upload, - }; + }); + login_finish_request.metadata_mut().insert("cookie", cookie); debug!("Attempting to finalize opaque login exchange with identity service"); let login_finish_response = identity_client diff --git a/keyserver/addons/rust-node-addon/src/identity_client/mod.rs b/keyserver/addons/rust-node-addon/src/identity_client/mod.rs --- a/keyserver/addons/rust-node-addon/src/identity_client/mod.rs +++ b/keyserver/addons/rust-node-addon/src/identity_client/mod.rs @@ -14,7 +14,7 @@ }; use grpc_clients::identity::authenticated::ChainedInterceptedAuthClient; use grpc_clients::identity::protos::unauthenticated as client_proto; -use grpc_clients::identity::shared::CodeVersionLayer; +use grpc_clients::identity::shared::{get_lb_cookie, CodeVersionLayer}; use lazy_static::lazy_static; use napi::bindgen_prelude::*; use serde::{Deserialize, Serialize}; diff --git a/keyserver/addons/rust-node-addon/src/identity_client/register_user.rs b/keyserver/addons/rust-node-addon/src/identity_client/register_user.rs --- a/keyserver/addons/rust-node-addon/src/identity_client/register_user.rs +++ b/keyserver/addons/rust-node-addon/src/identity_client/register_user.rs @@ -50,13 +50,20 @@ }); // Finish OPAQUE registration and send final registration request - let registration_start_response = identity_client + let response = identity_client .register_password_user_start(registration_start_request) .await - .map_err(handle_grpc_error)? - .into_inner(); + .map_err(handle_grpc_error)?; debug!("Received registration start response"); + // We need to get the load balancer cookie from from the response and send it + // in the subsequent request to ensure it is routed to the same identity + // service instance as the first request + let cookie = get_lb_cookie(&response) + .map_err(|_| Error::from_reason("Failed to get lb cookie from response"))?; + + let registration_start_response = response.into_inner(); + let opaque_registration_upload = opaque_registration .finish( &password, @@ -64,10 +71,14 @@ ) .map_err(|_| Error::from_status(Status::GenericFailure))?; - let registration_finish_request = Request::new(RegistrationFinishRequest { - session_id: registration_start_response.session_id, - opaque_registration_upload, - }); + let mut registration_finish_request = + Request::new(RegistrationFinishRequest { + session_id: registration_start_response.session_id, + opaque_registration_upload, + }); + registration_finish_request + .metadata_mut() + .insert("cookie", cookie); let registration_response = identity_client .register_password_user_finish(registration_finish_request) diff --git a/native/native_rust_library/src/lib.rs b/native/native_rust_library/src/lib.rs --- a/native/native_rust_library/src/lib.rs +++ b/native/native_rust_library/src/lib.rs @@ -1,7 +1,6 @@ use crate::ffi::{bool_callback, string_callback, void_callback}; use comm_opaque2::client::{Login, Registration}; use comm_opaque2::grpc::opaque_error_to_grpc_status as handle_error; -use grpc_clients::identity::get_unauthenticated_client; use grpc_clients::identity::protos::client::{ outbound_keys_for_user_request::Identifier, DeleteUserRequest, DeviceKeyUpload, DeviceType, Empty, IdentityKeyInfo, @@ -10,11 +9,12 @@ RegistrationStartRequest, UpdateUserPasswordFinishRequest, UpdateUserPasswordStartRequest, WalletLoginRequest, }; +use grpc_clients::identity::{get_lb_cookie, get_unauthenticated_client}; use lazy_static::lazy_static; use serde::Serialize; use std::sync::Arc; use tokio::runtime::{Builder, Runtime}; -use tonic::Status; +use tonic::{Request, Status}; use tracing::instrument; mod argon2_tools; @@ -321,10 +321,16 @@ DEVICE_TYPE.as_str_name().to_lowercase(), ) .await?; - let registration_start_response = identity_client + let response = identity_client .register_password_user_start(registration_start_request) - .await? - .into_inner(); + .await?; + + // We need to get the load balancer cookie from from the response and send it + // in the subsequent request to ensure it is routed to the same identity + // service instance as the first request + let cookie = get_lb_cookie(&response)?; + + let registration_start_response = response.into_inner(); let opaque_registration_upload = client_registration .finish( @@ -332,13 +338,17 @@ ®istration_start_response.opaque_registration_response, ) .map_err(handle_error)?; + let registration_finish_request = RegistrationFinishRequest { session_id: registration_start_response.session_id, opaque_registration_upload, }; + let mut finish_request = Request::new(registration_finish_request); + finish_request.metadata_mut().insert("cookie", cookie); + let registration_finish_response = identity_client - .register_password_user_finish(registration_finish_request) + .register_password_user_finish(finish_request) .await? .into_inner(); let user_id_and_access_token = UserIDAndDeviceAccessToken { @@ -417,21 +427,31 @@ ) .await?; - let login_start_response = identity_client + let response = identity_client .login_password_user_start(login_start_request) - .await? - .into_inner(); + .await?; + + // We need to get the load balancer cookie from from the response and send it + // in the subsequent request to ensure it is routed to the same identity + // service instance as the first request + let cookie = get_lb_cookie(&response)?; + + let login_start_response = response.into_inner(); let opaque_login_upload = client_login .finish(&login_start_response.opaque_login_response) .map_err(handle_error)?; + let login_finish_request = OpaqueLoginFinishRequest { session_id: login_start_response.session_id, opaque_login_upload, }; + let mut finish_request = Request::new(login_finish_request); + finish_request.metadata_mut().insert("cookie", cookie); + let login_finish_response = identity_client - .login_password_user_finish(login_finish_request) + .login_password_user_finish(finish_request) .await? .into_inner(); let user_id_and_access_token = UserIDAndDeviceAccessToken { @@ -580,10 +600,16 @@ ) .await?; - let update_password_start_response = identity_client + let response = identity_client .update_user_password_start(update_password_start_request) - .await? - .into_inner(); + .await?; + + // We need to get the load balancer cookie from from the response and send it + // in the subsequent request to ensure it is routed to the same identity + // service instance as the first request + let cookie = get_lb_cookie(&response)?; + + let update_password_start_response = response.into_inner(); let opaque_registration_upload = client_registration .finish( @@ -591,13 +617,17 @@ &update_password_start_response.opaque_registration_response, ) .map_err(handle_error)?; + let update_password_finish_request = UpdateUserPasswordFinishRequest { session_id: update_password_start_response.session_id, opaque_registration_upload, }; + let mut finish_request = Request::new(update_password_finish_request); + finish_request.metadata_mut().insert("cookie", cookie); + identity_client - .update_user_password_finish(update_password_finish_request) + .update_user_password_finish(finish_request) .await?; Ok(()) diff --git a/shared/grpc_clients/src/error.rs b/shared/grpc_clients/src/error.rs --- a/shared/grpc_clients/src/error.rs +++ b/shared/grpc_clients/src/error.rs @@ -12,6 +12,8 @@ GrpcStatus(Status), #[display(fmt = "Invalid Device Type")] InvalidDeviceType, + #[display(fmt = "Cookie Error: {}", _0)] + CookieError(CookieError), } pub fn unsupported_version() -> Status { @@ -22,3 +24,16 @@ status.code() == Code::Unimplemented && status.message() == "Unsupported version" } + +#[derive(Debug, derive_more::Display, derive_more::Error)] +pub struct CookieError { + message: String, +} + +impl CookieError { + pub fn new(message: &str) -> CookieError { + CookieError { + message: message.to_string(), + } + } +} diff --git a/shared/grpc_clients/src/identity/mod.rs b/shared/grpc_clients/src/identity/mod.rs --- a/shared/grpc_clients/src/identity/mod.rs +++ b/shared/grpc_clients/src/identity/mod.rs @@ -17,4 +17,5 @@ pub use authenticated::get_auth_client; pub use device::DeviceType; +pub use shared::get_lb_cookie; pub use unauthenticated::get_unauthenticated_client; diff --git a/shared/grpc_clients/src/identity/shared.rs b/shared/grpc_clients/src/identity/shared.rs --- a/shared/grpc_clients/src/identity/shared.rs +++ b/shared/grpc_clients/src/identity/shared.rs @@ -1,7 +1,8 @@ +use crate::error::{CookieError, Error}; use tonic::{ metadata::{errors::InvalidMetadataValue, Ascii, MetadataValue}, service::Interceptor, - Request, Status, + Request, Response, Status, }; pub struct CodeVersionLayer { @@ -55,3 +56,19 @@ self.second.call(request) } } + +pub fn get_lb_cookie( + response: &Response, +) -> Result, Error> { + let cookie_meta = response.metadata().get("set-cookie").ok_or_else(|| { + Error::from(CookieError::new("No cookie found in response")) + })?; + let cookie_str = cookie_meta.to_str().map_err(|_| { + Error::from(CookieError::new("Failed to convert cookie to string")) + })?; + MetadataValue::try_from(cookie_str).map_err(|_| { + Error::from(CookieError::new( + "Failed to create MetadataValue from cookie string", + )) + }) +}