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