fix(dgw): reuse credential injection across reconnects

devolutions-gateway/src/api/kdc_proxy.rs

@@ -5,9 +5,8 @@ use picky_krb::messages::KdcProxyMessage;
 use uuid::Uuid;
 
 use crate::DgwState;
-use crate::credential_injection_kdc::{
-    CredentialInjectionKdcInterception, CredentialInjectionKdcRequest, CredentialInjectionKdcResolveError,
-    kdc_proxy_message_realm,
+use crate::credential_injection::{
+    CredentialInjectionKdcInterception, CredentialInjectionKdcRequest, kdc_proxy_message_realm,
 };
 use crate::extract::KdcToken;
 use crate::http::HttpError;
@@ -22,7 +21,7 @@ pub fn make_router<S>(state: DgwState) -> Router<S> {
 async fn kdc_proxy(
     State(DgwState {
         conf_handle,
-        credentials,
+        synthetic_kdc_registry,
         agent_tunnel_handle,
         ..
     }): State<DgwState>,
@@ -47,7 +46,9 @@ async fn kdc_proxy(
         KdcDestination::Inject { jti } => {
             enforce_credential_injection_enabled(jti, conf.debug.enable_unstable)?;
 
-            let kdc = credentials.kdc_for(jti).map_err(credential_injection_resolve_error)?;
+            let kdc = synthetic_kdc_registry
+                .get(jti)
+                .ok_or_else(|| HttpError::bad_request().msg("no live synthetic KDC published for this session"))?;
 
             debug!(
                 jti = %kdc.jti(),
@@ -92,17 +93,6 @@ async fn kdc_proxy(
     }
 }
 
-fn credential_injection_resolve_error(error: CredentialInjectionKdcResolveError) -> HttpError {
-    match error {
-        CredentialInjectionKdcResolveError::BuildKdcConfig { .. } => HttpError::internal()
-            .with_msg("credential-injection KDC could not be initialized")
-            .build(error),
-        _ => HttpError::bad_request()
-            .with_msg("credential-injection state is not available")
-            .build(error),
-    }
-}
-
 // Forwards the request to the real KDC indicated by the token (or by the debug override) and
 // returns the response wrapped as a `KdcProxyMessage`.
 //
@@ -193,10 +183,9 @@ mod tests {
 
     #[test]
     fn enforce_realm_mismatch_passes_under_bypass() {
-        // `bypass=true` is the `__debug__.disable_token_validation` downgrade. CBenoit asked
-        // for explicit coverage of this branch because it is the only place the realm
-        // authorization is intentionally weakened, and slipping the gate (e.g. by inverting the
-        // condition) would only surface in production.
+        // `bypass=true` is the `__debug__.disable_token_validation` downgrade.
+        // This is the only branch where realm authorization is intentionally weakened, so pin it
+        // explicitly to catch an inverted gate.
         assert!(enforce_realm_token_match("ad.example", "evil.example", true).is_ok());
     }
 

devolutions-gateway/src/api/preflight.rs

@@ -11,10 +11,9 @@ use uuid::Uuid;
 
 use crate::DgwState;
 use crate::config::Conf;
-use crate::credential_injection_kdc::CredentialService;
 use crate::extract::PreflightScope;
 use crate::http::HttpError;
-use crate::provisioning::InsertError;
+use crate::provisioning::{InsertError, ProvisioningStore};
 use crate::session::SessionMessageSender;
 
 const OP_GET_VERSION: &str = "get-version";
@@ -204,7 +203,7 @@ pub(super) async fn post_preflight(
     State(DgwState {
         conf_handle,
         sessions,
-        credentials,
+        provisioning,
         ..
     }): State<DgwState>,
     _scope: PreflightScope,
@@ -231,13 +230,13 @@ pub(super) async fn post_preflight(
                 let outputs = outputs.clone();
                 let conf = conf_handle.get_conf();
                 let sessions = sessions.clone();
-                let credentials = credentials.clone();
+                let provisioning = provisioning.clone();
 
                 async move {
                     let operation_id = operation.id;
                     trace!(%operation.id, "Process preflight operation");
 
-                    if let Err(error) = handle_operation(operation, &outputs, &conf, &sessions, &credentials).await {
+                    if let Err(error) = handle_operation(operation, &outputs, &conf, &sessions, &provisioning).await {
                         outputs.push(PreflightOutput {
                             operation_id,
                             kind: PreflightOutputKind::Alert {
@@ -264,7 +263,7 @@ async fn handle_operation(
     outputs: &Outputs,
     conf: &Conf,
     sessions: &SessionMessageSender,
-    credentials: &CredentialService,
+    provisioning: &ProvisioningStore,
 ) -> Result<(), PreflightError> {
     match operation.kind.as_str() {
         OP_GET_VERSION => outputs.push(PreflightOutput {
@@ -355,7 +354,7 @@ async fn handle_operation(
                     })?;
             }
 
-            let replaced = credentials
+            let replaced = provisioning
                 .insert_credentials(token, mapping, time_to_live)
                 .inspect_err(|error| warn!(%operation.id, error = format!("{error:#}"), "Failed to insert credentials"))
                 .map_err(|error| match error {
@@ -397,7 +396,7 @@ async fn handle_operation(
                 PreflightError::new(PreflightAlertStatus::InvalidParams, format!("invalid token: {error:#}"))
             })?;
 
-            let replaced = credentials.insert_connection_options(jti, connection_options, time_to_live);
+            let replaced = provisioning.insert_connection_options(jti, connection_options, time_to_live);
 
             if replaced {
                 outputs.push(PreflightOutput {

devolutions-gateway/src/api/rdp.rs

@@ -25,7 +25,8 @@ pub async fn handler(
         subscriber_tx,
         recordings,
         shutdown_signal,
-        credentials,
+        provisioning,
+        synthetic_kdc_registry,
         agent_tunnel_handle,
         ..
     }): State<DgwState>,
@@ -46,7 +47,8 @@ pub async fn handler(
             subscriber_tx,
             recordings.active_recordings,
             source_addr,
-            credentials,
+            provisioning,
+            synthetic_kdc_registry,
             agent_tunnel_handle,
         )
         .instrument(span)
@@ -66,7 +68,8 @@ async fn handle_socket(
     subscriber_tx: SubscriberSender,
     active_recordings: Arc<ActiveRecordings>,
     source_addr: SocketAddr,
-    credentials: crate::credential_injection_kdc::CredentialService,
+    provisioning: crate::provisioning::ProvisioningStore,
+    synthetic_kdc_registry: crate::credential_injection::SyntheticKdcRegistry,
     agent_tunnel_handle: Option<Arc<agent_tunnel::AgentTunnelHandle>>,
 ) {
     let (stream, close_handle) = crate::ws::handle(
@@ -84,7 +87,8 @@ async fn handle_socket(
         sessions,
         subscriber_tx,
         &active_recordings,
-        &credentials,
+        &provisioning,
+        &synthetic_kdc_registry,
         agent_tunnel_handle,
     )
     .await;