Skip to main content

little_durable_objects/control_plane/
process.rs

1use std::{collections::HashMap, env, future::Future, net::SocketAddr, sync::Arc, time::Duration};
2
3use anyhow::{Context, Result, ensure};
4use tracing::info;
5
6use crate::{
7    host_leases::PostgresHostLeaseStore,
8    placement::PostgresObjectPlacementStore,
9    postgres::PostgresDatabase,
10    sandbox::{CommandSandboxProvider, HostSandboxRuntimeConfig},
11    storage_urls::{GcsStorageUrlSigner, validate_buckets},
12};
13
14use super::{ActorJwtVerifier, ControlPlaneService};
15
16const DEFAULT_JWT_ISSUER: &str = "durable-object-control-plane";
17const DEFAULT_AUTHORITY_AUDIENCE: &str = "durable-object-authority";
18const DEFAULT_INVOCATION_AUDIENCE: &str = "durable-object-invoke";
19const DEFAULT_JWT_TTL_SECONDS: u64 = 1_800;
20const DEFAULT_ACTOR_IDLE_TIMEOUT_MS: u64 = 60_000;
21const DEFAULT_HOST_IDLE_TIMEOUT_MS: u64 = 300_000;
22const MAX_IDLE_TIMEOUT_MS: u64 = 86_400_000;
23
24pub struct ControlPlaneProcessConfig {
25    pub bind: SocketAddr,
26    pub jwt_signing_key: String,
27    pub jwt_key_id: String,
28    pub jwt_issuer: String,
29    pub authority_audience: String,
30    pub invocation_audience: String,
31    pub jwt_max_lifetime: Duration,
32    pub admin_token: String,
33    pub storage: ControlPlaneStorageConfig,
34    pub sandbox_provider: Option<SandboxProviderConfig>,
35    pub socket_event_sink: Option<SocketEventSinkConfig>,
36    pub socket_authenticator: Option<SocketAuthenticatorConfig>,
37}
38
39pub struct ControlPlaneStorageConfig {
40    pub postgres_url: String,
41    pub standard_buckets: HashMap<String, String>,
42}
43
44pub struct SandboxProviderConfig {
45    pub provider_name: String,
46    pub command: String,
47    pub environment: HashMap<String, String>,
48    pub runtime: HostSandboxRuntimeConfig,
49}
50
51pub struct SocketEventSinkConfig {
52    pub url: String,
53    pub token: String,
54}
55
56pub struct SocketAuthenticatorConfig {
57    pub url: String,
58    pub token: String,
59}
60
61impl ControlPlaneProcessConfig {
62    pub fn from_env() -> Result<Self> {
63        Self::from_lookup(|name| env::var(name).ok())
64    }
65}
66
67pub async fn serve_control_plane(
68    config: ControlPlaneProcessConfig,
69    shutdown: impl Future<Output = ()> + Send + 'static,
70) -> Result<()> {
71    let bind = config.bind;
72    let routes = control_plane_routes(config).await?;
73    info!(bind = %bind, "durable-object control plane is ready");
74    let listener = tokio::net::TcpListener::bind(bind)
75        .await
76        .context("bind durable-object control plane")?;
77    serve_routes(listener, routes, shutdown).await
78}
79
80async fn serve_routes(
81    listener: tokio::net::TcpListener,
82    routes: tonic::service::Routes,
83    shutdown: impl Future<Output = ()> + Send + 'static,
84) -> Result<()> {
85    axum::serve(listener, routes.into_axum_router())
86        .with_graceful_shutdown(shutdown)
87        .await
88        .context("serve durable-object control plane")
89}
90
91async fn control_plane_routes(config: ControlPlaneProcessConfig) -> Result<tonic::service::Routes> {
92    let issuer = super::ActorJwtIssuer::from_base64_pkcs8(
93        &config.jwt_signing_key,
94        config.jwt_key_id,
95        config.jwt_issuer.clone(),
96        config.authority_audience.clone(),
97        config.invocation_audience,
98        config.jwt_max_lifetime,
99    )?;
100    let auth = ActorJwtVerifier::for_scope(
101        issuer.verifier_keys_json()?,
102        config.jwt_issuer,
103        config.authority_audience,
104        super::ActorTokenPurpose::ControlPlane,
105        config.jwt_max_lifetime,
106    )?;
107    let database = PostgresDatabase::connect(&config.storage.postgres_url).await?;
108    let leases = Arc::new(PostgresHostLeaseStore::from_database(database.clone()));
109    let placements = Arc::new(PostgresObjectPlacementStore::from_database(
110        database.clone(),
111    ));
112    let registry = Arc::new(super::PostgresAdminRegistry::from_database(database));
113    let storage_urls = Arc::new(GcsStorageUrlSigner::from_adc(
114        config.storage.standard_buckets,
115    )?);
116    let provisioner = sandbox_provisioner(config.sandbox_provider, &issuer, &leases)?;
117    let socket_events = config
118        .socket_event_sink
119        .map(|sink| super::event_sink::HttpSocketMessageEventSink::new(sink.url, sink.token))
120        .transpose()?
121        .map(|sink| Arc::new(sink) as Arc<dyn super::event_sink::SocketMessageEventSink>);
122    let socket_authenticator = config
123        .socket_authenticator
124        .map(|auth| super::socket_auth::HttpSocketAuthenticator::new(auth.url, auth.token))
125        .transpose()?
126        .map(|auth| Arc::new(auth) as Arc<dyn super::socket_auth::SocketAuthenticator>);
127    let service = ControlPlaneService::new(
128        leases,
129        placements,
130        storage_urls,
131        auth,
132        registry.clone(),
133        issuer.clone(),
134        provisioner,
135    )
136    .with_socket_event_sink(socket_events)
137    .with_socket_authenticator(socket_authenticator);
138    let admin = super::admin::AdminService::new(config.admin_token, registry, issuer)?;
139    let public_api = super::public_api::router(service.clone(), admin);
140    let internal_api = service.into_internal_service();
141    Ok(tonic::service::Routes::from(public_api).add_service(internal_api))
142}
143
144fn sandbox_provisioner(
145    config: Option<SandboxProviderConfig>,
146    issuer: &super::ActorJwtIssuer,
147    leases: &Arc<PostgresHostLeaseStore>,
148) -> Result<Option<Arc<dyn super::service::HostProvisioner>>> {
149    config
150        .map(
151            |config| -> Result<Arc<dyn super::service::HostProvisioner>> {
152                let provider = Arc::new(CommandSandboxProvider::new(
153                    config.provider_name,
154                    config.command,
155                    config.environment,
156                )?);
157                Ok(Arc::new(super::service::SandboxHostProvisioner::new(
158                    provider,
159                    config.runtime,
160                    issuer.clone(),
161                    leases.clone(),
162                )))
163            },
164        )
165        .transpose()
166}
167
168impl ControlPlaneProcessConfig {
169    fn from_lookup(mut get: impl FnMut(&str) -> Option<String>) -> Result<Self> {
170        let bind = get("DURABLE_OBJECT_CONTROL_PLANE_BIND")
171            .unwrap_or_else(|| "127.0.0.1:7100".into())
172            .parse()
173            .context("DURABLE_OBJECT_CONTROL_PLANE_BIND must be a socket address")?;
174        let jwt_signing_key = required(&mut get, "DURABLE_OBJECT_JWT_SIGNING_KEY")?;
175        let jwt_key_id = get("DURABLE_OBJECT_JWT_KEY_ID").unwrap_or_else(|| "primary".into());
176        let jwt_issuer =
177            get("DURABLE_OBJECT_JWT_ISSUER").unwrap_or_else(|| DEFAULT_JWT_ISSUER.into());
178        let authority_audience = get("DURABLE_OBJECT_AUTHORITY_JWT_AUDIENCE")
179            .unwrap_or_else(|| DEFAULT_AUTHORITY_AUDIENCE.into());
180        let invocation_audience = get("DURABLE_OBJECT_INVOKE_JWT_AUDIENCE")
181            .unwrap_or_else(|| DEFAULT_INVOCATION_AUDIENCE.into());
182        let jwt_max_lifetime = Duration::from_secs(
183            get("DURABLE_OBJECT_JWT_MAX_TTL_SECONDS")
184                .map(|value| value.parse())
185                .transpose()
186                .context("DURABLE_OBJECT_JWT_MAX_TTL_SECONDS must be an integer")?
187                .unwrap_or(DEFAULT_JWT_TTL_SECONDS),
188        );
189        ensure!(
190            !jwt_max_lifetime.is_zero(),
191            "DURABLE_OBJECT_JWT_MAX_TTL_SECONDS must be positive"
192        );
193        let admin_token = required(&mut get, "DURABLE_OBJECT_ADMIN_TOKEN")?;
194        ensure!(
195            admin_token.trim() == admin_token,
196            "DURABLE_OBJECT_ADMIN_TOKEN has surrounding whitespace"
197        );
198        let standard_buckets: HashMap<String, String> =
199            serde_json::from_str(&required(&mut get, "DURABLE_OBJECT_STANDARD_BUCKETS")?)
200                .context("DURABLE_OBJECT_STANDARD_BUCKETS must be a JSON region-to-bucket map")?;
201        validate_buckets(&standard_buckets)?;
202        let storage = ControlPlaneStorageConfig {
203            postgres_url: required(&mut get, "DURABLE_OBJECT_POSTGRES_URL")?,
204            standard_buckets,
205        };
206        let sandbox_provider =
207            sandbox_provider_config(&mut get, &jwt_issuer, &invocation_audience)?;
208        let socket_event_sink = socket_event_sink_config(&mut get)?;
209        let socket_authenticator = socket_authenticator_config(&mut get)?;
210        Ok(Self {
211            bind,
212            jwt_signing_key,
213            jwt_key_id,
214            jwt_issuer,
215            authority_audience,
216            invocation_audience,
217            jwt_max_lifetime,
218            admin_token,
219            storage,
220            sandbox_provider,
221            socket_event_sink,
222            socket_authenticator,
223        })
224    }
225}
226
227fn socket_authenticator_config(
228    get: &mut impl FnMut(&str) -> Option<String>,
229) -> Result<Option<SocketAuthenticatorConfig>> {
230    let url = get("DURABLE_OBJECT_SOCKET_AUTH_URL");
231    let token = get("DURABLE_OBJECT_SOCKET_AUTH_TOKEN");
232    match (url, token) {
233        (None, None) => Ok(None),
234        (Some(url), Some(token)) => Ok(Some(SocketAuthenticatorConfig {
235            url: validated_http_url(&url, "DURABLE_OBJECT_SOCKET_AUTH_URL")?,
236            token,
237        })),
238        _ => anyhow::bail!(
239            "DURABLE_OBJECT_SOCKET_AUTH_URL and DURABLE_OBJECT_SOCKET_AUTH_TOKEN must be configured together"
240        ),
241    }
242}
243
244fn socket_event_sink_config(
245    get: &mut impl FnMut(&str) -> Option<String>,
246) -> Result<Option<SocketEventSinkConfig>> {
247    let url = get("DURABLE_OBJECT_SOCKET_EVENT_URL");
248    let token = get("DURABLE_OBJECT_SOCKET_EVENT_TOKEN");
249    match (url, token) {
250        (None, None) => Ok(None),
251        (Some(url), Some(token)) => Ok(Some(SocketEventSinkConfig {
252            url: validated_http_url(&url, "DURABLE_OBJECT_SOCKET_EVENT_URL")?,
253            token,
254        })),
255        _ => anyhow::bail!(
256            "DURABLE_OBJECT_SOCKET_EVENT_URL and DURABLE_OBJECT_SOCKET_EVENT_TOKEN must be configured together"
257        ),
258    }
259}
260
261fn sandbox_provider_config(
262    get: &mut impl FnMut(&str) -> Option<String>,
263    jwt_issuer: &str,
264    invocation_audience: &str,
265) -> Result<Option<SandboxProviderConfig>> {
266    let Some(provider_name) = get("DURABLE_OBJECT_SANDBOX_PROVIDER") else {
267        ensure!(
268            get("DURABLE_OBJECT_SANDBOX_COMMAND").is_none(),
269            "sandbox command requires a provider"
270        );
271        return Ok(None);
272    };
273    ensure!(
274        provider_name == "modal",
275        "unsupported sandbox provider {provider_name:?}"
276    );
277    let environment = HashMap::from([
278        (
279            "MODAL_TOKEN_ID".into(),
280            provider_credential(get, "MODAL_TOKEN_ID")?,
281        ),
282        (
283            "MODAL_TOKEN_SECRET".into(),
284            provider_credential(get, "MODAL_TOKEN_SECRET")?,
285        ),
286    ]);
287    let control_plane_url = validated_http_url(
288        &required(get, "DURABLE_OBJECT_CONTROL_PLANE_URL")?,
289        "DURABLE_OBJECT_CONTROL_PLANE_URL",
290    )?;
291    Ok(Some(SandboxProviderConfig {
292        provider_name,
293        command: get("DURABLE_OBJECT_SANDBOX_COMMAND")
294            .unwrap_or_else(|| "little-durable-objects-modal".into()),
295        environment,
296        runtime: HostSandboxRuntimeConfig {
297            control_plane_url,
298            jwt_issuer: jwt_issuer.into(),
299            invocation_jwt_audience: invocation_audience.into(),
300            actor_idle_timeout_ms: idle_timeout(
301                get,
302                "DURABLE_OBJECT_ACTOR_IDLE_TIMEOUT_MS",
303                DEFAULT_ACTOR_IDLE_TIMEOUT_MS,
304            )?,
305            host_idle_timeout_ms: idle_timeout(
306                get,
307                "DURABLE_OBJECT_HOST_IDLE_TIMEOUT_MS",
308                DEFAULT_HOST_IDLE_TIMEOUT_MS,
309            )?,
310        },
311    }))
312}
313
314fn provider_credential(get: &mut impl FnMut(&str) -> Option<String>, name: &str) -> Result<String> {
315    let value = required(get, name)?;
316    ensure!(value.trim() == value, "{name} has surrounding whitespace");
317    Ok(value)
318}
319
320fn required(get: &mut impl FnMut(&str) -> Option<String>, name: &str) -> Result<String> {
321    let value = get(name).with_context(|| format!("{name} is required"))?;
322    ensure!(!value.is_empty(), "{name} must not be empty");
323    Ok(value)
324}
325
326fn validated_http_url(value: &str, name: &str) -> Result<String> {
327    let url = reqwest::Url::parse(value).with_context(|| format!("{name} must be a URL"))?;
328    ensure!(
329        matches!(url.scheme(), "http" | "https") && url.host_str().is_some(),
330        "{name} must be HTTP or HTTPS"
331    );
332    Ok(url.to_string())
333}
334
335fn idle_timeout(
336    get: &mut impl FnMut(&str) -> Option<String>,
337    name: &str,
338    default: u64,
339) -> Result<u64> {
340    let value = get(name)
341        .map(|value| value.parse())
342        .transpose()
343        .with_context(|| format!("{name} must be an integer"))?
344        .unwrap_or(default);
345    ensure!(
346        (1..=MAX_IDLE_TIMEOUT_MS).contains(&value),
347        "{name} is outside the supported range"
348    );
349    Ok(value)
350}
351
352#[cfg(test)]
353mod tests {
354    use axum::{Router, extract::WebSocketUpgrade, response::Response, routing::get};
355    use futures_util::{SinkExt, StreamExt};
356    use tokio::sync::oneshot;
357    use tokio_tungstenite::{connect_async, tungstenite::Message};
358
359    use super::*;
360
361    #[tokio::test]
362    async fn server_carries_websocket_upgrades() -> Result<()> {
363        let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await?;
364        let address = listener.local_addr()?;
365        let routes =
366            tonic::service::Routes::from(Router::new().route("/socket", get(echo_websocket)));
367        let (shutdown_tx, shutdown_rx) = oneshot::channel();
368        let server = tokio::spawn(serve_routes(listener, routes, async {
369            let _ = shutdown_rx.await;
370        }));
371
372        let (mut socket, _) = connect_async(format!("ws://{address}/socket")).await?;
373        socket.send(Message::Text("hello".into())).await?;
374        assert_eq!(
375            socket.next().await.transpose()?,
376            Some(Message::Text("hello".into()))
377        );
378        socket.close(None).await?;
379        let _ = shutdown_tx.send(());
380        server.await??;
381        Ok(())
382    }
383
384    #[test]
385    fn parses_the_minimal_storage_configuration() -> Result<()> {
386        let values = HashMap::from([
387            ("DURABLE_OBJECT_JWT_SIGNING_KEY", "c2lnbmluZw=="),
388            ("DURABLE_OBJECT_ADMIN_TOKEN", "admin-token"),
389            (
390                "DURABLE_OBJECT_POSTGRES_URL",
391                "postgresql://localhost/actors",
392            ),
393            (
394                "DURABLE_OBJECT_STANDARD_BUCKETS",
395                "{\"us-east\":\"actor-state-test\"}",
396            ),
397        ]);
398        let config = ControlPlaneProcessConfig::from_lookup(|name| {
399            values.get(name).map(|value| (*value).into())
400        })?;
401        assert_eq!(
402            config.storage.standard_buckets["us-east"],
403            "actor-state-test"
404        );
405        Ok(())
406    }
407
408    #[test]
409    fn parses_socket_event_sink_only_when_url_and_token_are_present() -> Result<()> {
410        let mut complete = HashMap::from([
411            (
412                "DURABLE_OBJECT_SOCKET_EVENT_URL",
413                "https://api.example.com/events",
414            ),
415            ("DURABLE_OBJECT_SOCKET_EVENT_TOKEN", "event-token"),
416        ]);
417        let sink =
418            socket_event_sink_config(&mut |name| complete.get(name).map(|value| (*value).into()))?
419                .context("socket event sink was not configured")?;
420        assert_eq!(sink.url, "https://api.example.com/events");
421        assert_eq!(sink.token, "event-token");
422
423        complete.remove("DURABLE_OBJECT_SOCKET_EVENT_TOKEN");
424        assert!(
425            socket_event_sink_config(&mut |name| complete.get(name).map(|value| (*value).into()))
426                .is_err()
427        );
428        Ok(())
429    }
430
431    #[test]
432    fn parses_socket_authenticator_only_when_url_and_token_are_present() -> Result<()> {
433        let mut complete = HashMap::from([
434            (
435                "DURABLE_OBJECT_SOCKET_AUTH_URL",
436                "https://api.example.com/authorize",
437            ),
438            ("DURABLE_OBJECT_SOCKET_AUTH_TOKEN", "auth-token"),
439        ]);
440        let auth = socket_authenticator_config(&mut |name| {
441            complete.get(name).map(|value| (*value).into())
442        })?
443        .context("socket authenticator was not configured")?;
444        assert_eq!(auth.url, "https://api.example.com/authorize");
445        assert_eq!(auth.token, "auth-token");
446
447        complete.remove("DURABLE_OBJECT_SOCKET_AUTH_TOKEN");
448        assert!(
449            socket_authenticator_config(&mut |name| complete
450                .get(name)
451                .map(|value| (*value).into()))
452            .is_err()
453        );
454        Ok(())
455    }
456
457    async fn echo_websocket(upgrade: WebSocketUpgrade) -> Response {
458        upgrade.on_upgrade(async |mut socket| {
459            if let Some(Ok(message)) = socket.recv().await {
460                let _ = socket.send(message).await;
461            }
462        })
463    }
464}