refactor(dgw): give the gateway state the provisioner task store

devolutions-gateway/Cargo.toml

@@ -37,6 +37,7 @@ sysevent.path = "../crates/sysevent"
 sysevent-codes.path = "../crates/sysevent-codes"
 agent-tunnel = { path = "../crates/agent-tunnel" }
 gateway-db.path = "../crates/gateway-db"
+provisioner-task.path = "../crates/provisioner-task"
 ironrdp-pdu = { version = "0.9", features = ["std"] }
 ironrdp-core = { version = "0.2", features = ["std"] }
 ironrdp-rdcleanpath = "0.2"

devolutions-gateway/src/job_queue.rs

@@ -11,6 +11,9 @@ use job_queue_libsql::libsql;
 use time::OffsetDateTime;
 use tokio::sync::{Notify, mpsc};
 
+/// Runs of a job before the job queue gives up on it.
+pub const MAX_ATTEMPTS: u32 = 5;
+
 pub struct JobQueueCtx {
     notify_runner: Arc<Notify>,
     runner_waker: RunnerWaker,
@@ -59,6 +62,7 @@ impl JobQueueCtx {
         let queue = job_queue_libsql::LibSqlJobQueue::builder()
             .runner_waker(runner_waker.clone())
             .conn(conn)
+            .max_attempts(MAX_ATTEMPTS)
             .build();
 
         let queue = Arc::new(queue);

devolutions-gateway/src/lib.rs

@@ -68,6 +68,7 @@ pub struct DgwState {
     pub monitoring_state: Arc<network_monitor::State>,
     pub traffic_audit_handle: traffic_audit::TrafficAuditHandle,
     pub agent_tunnel_handle: Option<Arc<agent_tunnel::AgentTunnelHandle>>,
+    pub provisioner_tasks: provisioner_task::DynProvisionerTaskStore,
 }
 
 #[doc(hidden)]
@@ -82,7 +83,7 @@ pub struct MockHandles {
 
 impl DgwState {
     #[doc(hidden)]
-    pub fn mock(json_config: &str) -> anyhow::Result<(Self, MockHandles)> {
+    pub async fn mock(json_config: &str) -> anyhow::Result<(Self, MockHandles)> {
         let conf_handle = config::ConfHandle::mock(json_config)?;
         let token_cache = Arc::new(token::new_token_cache());
         let jrl = Arc::new(parking_lot::Mutex::new(token::JrlTokenClaims::default()));
@@ -95,6 +96,10 @@ impl DgwState {
         let provisioning = provisioning::ProvisioningStore::new();
         let synthetic_kdc_registry = credential_injection::SyntheticKdcRegistry::new();
         let monitoring_state = Arc::new(network_monitor::State::new(Arc::new(MockMonitorsCache))?);
+        let provisioner_tasks = gateway_db::GatewayDb::open_path(":memory:").await?.connect().await?;
+        let provisioner_tasks = Arc::new(gateway_db::provisioner_task::LibSqlProvisionerTaskStore::new(
+            provisioner_tasks,
+        ));
 
         let state = Self {
             conf_handle,
@@ -110,6 +115,7 @@ impl DgwState {
             synthetic_kdc_registry,
             monitoring_state,
             agent_tunnel_handle: None,
+            provisioner_tasks,
         };
 
         let handles = MockHandles {

devolutions-gateway/src/service.rs

@@ -337,6 +337,14 @@ async fn spawn_tasks(conf_handle: ConfHandle) -> anyhow::Result<Tasks> {
         None
     };
 
+    let provisioner_tasks_conn = gateway_db
+        .connect()
+        .await
+        .context("failed to connect the provisioner tasks to the gateway database")?;
+    let provisioner_tasks = Arc::new(gateway_db::provisioner_task::LibSqlProvisionerTaskStore::new(
+        provisioner_tasks_conn,
+    ));
+
     let state = DgwState {
         conf_handle: conf_handle.clone(),
         token_cache: Arc::clone(&token_cache),
@@ -351,6 +359,7 @@ async fn spawn_tasks(conf_handle: ConfHandle) -> anyhow::Result<Tasks> {
         monitoring_state,
         traffic_audit_handle: traffic_audit_task.handle(),
         agent_tunnel_handle,
+        provisioner_tasks,
     };
 
     for listener in &conf.listeners {

devolutions-gateway/tests/preflight.rs

@@ -58,15 +58,15 @@ fn unsigned_jws(payload: serde_json::Value) -> anyhow::Result<String> {
     Ok(format!("{header}.{payload}.{signature}"))
 }
 
-fn make_router() -> anyhow::Result<(Router, DgwState, MockHandles)> {
-    let (state, handles) = DgwState::mock(CONFIG)?;
+async fn make_router() -> anyhow::Result<(Router, DgwState, MockHandles)> {
+    let (state, handles) = DgwState::mock(CONFIG).await?;
     let app = devolutions_gateway::make_http_service(state.clone())
         .layer(MockConnectInfo(SocketAddr::from(([0, 0, 0, 0], 3000))));
     Ok((app, state, handles))
 }
 
-fn make_router_with_config(config: &str) -> anyhow::Result<(Router, DgwState, MockHandles)> {
-    let (state, handles) = DgwState::mock(config)?;
+async fn make_router_with_config(config: &str) -> anyhow::Result<(Router, DgwState, MockHandles)> {
+    let (state, handles) = DgwState::mock(config).await?;
     let app = devolutions_gateway::make_http_service(state.clone())
         .layer(MockConnectInfo(SocketAddr::from(([0, 0, 0, 0], 3000))));
     Ok((app, state, handles))
@@ -83,7 +83,7 @@ fn init_logger() -> tracing::subscriber::DefaultGuard {
 async fn test_provision_credentials_success() -> anyhow::Result<()> {
     let _guard = init_logger();
 
-    let (app, _state, _handles) = make_router()?;
+    let (app, _state, _handles) = make_router().await?;
 
     let jti = Uuid::new_v4();
     let token = unsigned_jws(json!({
@@ -124,7 +124,7 @@ async fn test_provision_credentials_success_when_unstable_disabled() -> anyhow::
     // `provision-credentials` is protocol-neutral: NTLM credential injection relies on this
     // path even when the unstable feature flag is off.
     let config = CONFIG.replace("\"enable_unstable\": true", "\"enable_unstable\": false");
-    let (app, _state, _handles) = make_router_with_config(&config)?;
+    let (app, _state, _handles) = make_router_with_config(&config).await?;
 
     let jti = Uuid::new_v4();
     let token = unsigned_jws(json!({
@@ -161,7 +161,7 @@ async fn test_provision_credentials_success_when_unstable_disabled() -> anyhow::
 async fn test_provision_credentials_rejects_alternate_targets() -> anyhow::Result<()> {
     let _guard = init_logger();
 
-    let (app, _state, _handles) = make_router()?;
+    let (app, _state, _handles) = make_router().await?;
 
     let token = unsigned_jws(json!({
         "jti": "5e3e833f-84c7-4541-b676-acc3299e39b8",
@@ -205,7 +205,7 @@ async fn test_provision_credentials_rejects_alternate_targets() -> anyhow::Resul
 async fn test_provision_credentials_rejects_missing_target_hostname() -> anyhow::Result<()> {
     let _guard = init_logger();
 
-    let (app, _state, _handles) = make_router()?;
+    let (app, _state, _handles) = make_router().await?;
 
     let token = unsigned_jws(json!({
         "jti": "5e3e833f-84c7-4541-b676-acc3299e39b8"
@@ -247,7 +247,7 @@ async fn test_provision_credentials_rejects_missing_target_hostname() -> anyhow:
 async fn test_provision_token_overwrite_alert() -> anyhow::Result<()> {
     let _guard = init_logger();
 
-    let (app, _state, _handles) = make_router()?;
+    let (app, _state, _handles) = make_router().await?;
 
     // Same JTI twice: second provision-token replaces the stored token-only entry (master behavior).
     let token = unsigned_jws(json!({
@@ -288,7 +288,7 @@ async fn test_provision_token_overwrite_alert() -> anyhow::Result<()> {
 async fn test_provision_connection_options_success() -> anyhow::Result<()> {
     let _guard = init_logger();
 
-    let (app, _state, _handles) = make_router()?;
+    let (app, _state, _handles) = make_router().await?;
 
     let jti = Uuid::new_v4();
     let token = unsigned_jws(json!({
@@ -319,7 +319,7 @@ async fn test_provision_connection_options_success() -> anyhow::Result<()> {
 async fn test_provision_credentials_and_connection_options_fold() -> anyhow::Result<()> {
     let _guard = init_logger();
 
-    let (app, _state, _handles) = make_router()?;
+    let (app, _state, _handles) = make_router().await?;
 
     let jti = Uuid::new_v4();
     let token = unsigned_jws(json!({
@@ -361,7 +361,7 @@ async fn test_provision_credentials_and_connection_options_fold() -> anyhow::Res
 async fn test_provision_invalid_params() -> anyhow::Result<()> {
     let _guard = init_logger();
 
-    let (app, _state, _handles) = make_router()?;
+    let (app, _state, _handles) = make_router().await?;
 
     let token = "eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9.eyJqdGkiOiI1ZTNlODMzZi04NGM3LTQ1NDEtYjY3Ni1hY2MzMjk5ZTM5YjgifQ.1qECGlrW7y9HWFArc6GPHLGTOY7PhAvzKJ5XMRBg4k4";
 

devolutions-gateway/tests/traffic.rs

@@ -71,7 +71,7 @@ impl Drop for HandlesGuard {
 /// Build a Router with a live TrafficAuditManagerTask and real handle,
 /// similar to tests/preflight.rs harness.
 async fn make_router() -> anyhow::Result<(Router, DgwState, HandlesGuard)> {
-    let (mut state, handles) = DgwState::mock(CONFIG)?;
+    let (mut state, handles) = DgwState::mock(CONFIG).await?;
 
     // Start the manager with an in-memory DB.
     let manager = TrafficAuditManagerTask::init(":memory:").await?;

devolutions-gateway/tests/webapp.rs

@@ -60,7 +60,7 @@ fn initialize_conf() {
 async fn custom_authentication_flow() -> anyhow::Result<()> {
     let (cov, _guard) = init_cov_mark();
     initialize_conf();
-    let (state, _handle) = devolutions_gateway::DgwState::mock(CONFIG)?;
+    let (state, _handle) = devolutions_gateway::DgwState::mock(CONFIG).await?;
 
     let mut app =
         devolutions_gateway::make_http_service(state).layer(MockConnectInfo(SocketAddr::from(([0, 0, 0, 0], 3000))));
@@ -191,7 +191,7 @@ async fn static_assets_are_served(#[case] subdir: &str, #[case] uri_prefix: &str
     config["WebApp"]["StaticRootPath"] = json!(static_root);
     let config = config.to_string();
 
-    let (state, _handle) = devolutions_gateway::DgwState::mock(&config)?;
+    let (state, _handle) = devolutions_gateway::DgwState::mock(&config).await?;
 
     let mut app =
         devolutions_gateway::make_http_service(state).layer(MockConnectInfo(SocketAddr::from(([0, 0, 0, 0], 3000))));
@@ -229,7 +229,7 @@ async fn static_assets_are_served(#[case] subdir: &str, #[case] uri_prefix: &str
 async fn sign_app_token_bad_password() -> anyhow::Result<()> {
     let (cov, _guard) = init_cov_mark();
     initialize_conf();
-    let (state, _handle) = devolutions_gateway::DgwState::mock(CONFIG)?;
+    let (state, _handle) = devolutions_gateway::DgwState::mock(CONFIG).await?;
 
     let app =
         devolutions_gateway::make_http_service(state).layer(MockConnectInfo(SocketAddr::from(([0, 0, 0, 0], 3000))));
@@ -267,7 +267,7 @@ async fn sign_app_token_bad_password() -> anyhow::Result<()> {
 async fn sign_app_token_username_mismatch() -> anyhow::Result<()> {
     let (cov, _guard) = init_cov_mark();
     initialize_conf();
-    let (state, _handles) = devolutions_gateway::DgwState::mock(CONFIG)?;
+    let (state, _handles) = devolutions_gateway::DgwState::mock(CONFIG).await?;
 
     let app =
         devolutions_gateway::make_http_service(state).layer(MockConnectInfo(SocketAddr::from(([0, 0, 0, 0], 3000))));