Page MenuHomePhabricator
Files F3333712

D10121.id34080.diff
No OneTemporary
Actions

D10121.id34080.diff
View Options

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 @@
&registration_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<T>(
+ response: &Response<T>,
+) -> Result<MetadataValue<Ascii>, 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",
+ ))
+ })
+}

File Metadata

Mime Type
text/plain
Expires
Fri, Nov 22, 5:01 AM (18 h, 51 m)
Storage Engine
blob
Storage Format
Raw Data
Storage Handle
2560027
Default Alt Text
D10121.id34080.diff (10 KB)

Event Timeline