Skip to main content

little_durable_objects/
state_transport.rs

1use anyhow::{Context, Result, ensure};
2use async_trait::async_trait;
3
4use crate::{state_log::StateLog, storage_urls::STATE_CONTENT_TYPE};
5
6const GENERATION_HEADER: &str = "x-goog-generation";
7
8#[derive(Debug)]
9pub struct LoadedState {
10    pub log: StateLog,
11    pub generation: String,
12}
13
14#[derive(Clone, Debug, PartialEq, Eq)]
15pub enum StateWrite {
16    Written,
17    GenerationMismatch,
18}
19
20#[async_trait]
21pub trait StateTransport: Send + Sync {
22    async fn read(&self, signed_url: &str) -> Result<LoadedState>;
23    async fn write(&self, signed_url: &str, bytes: Vec<u8>) -> Result<StateWrite>;
24}
25
26#[derive(Clone, Default)]
27pub struct HttpStateTransport {
28    client: reqwest::Client,
29}
30
31impl HttpStateTransport {
32    pub fn new() -> Self {
33        Self {
34            client: reqwest::Client::new(),
35        }
36    }
37}
38
39#[async_trait]
40impl StateTransport for HttpStateTransport {
41    async fn read(&self, signed_url: &str) -> Result<LoadedState> {
42        validate_url(signed_url)?;
43        let response = self
44            .client
45            .get(signed_url)
46            .send()
47            .await
48            .context("read actor state through signed URL")?;
49        if response.status() == reqwest::StatusCode::NOT_FOUND {
50            return Ok(LoadedState {
51                log: StateLog::default(),
52                generation: "0".into(),
53            });
54        }
55        ensure!(
56            response.status().is_success(),
57            "actor-state read failed with HTTP {}",
58            response.status()
59        );
60        let generation = generation_header(&response)?;
61        let bytes = response
62            .bytes()
63            .await
64            .context("read actor-state response body")?;
65        Ok(LoadedState {
66            log: StateLog::decode(&bytes)?,
67            generation,
68        })
69    }
70
71    async fn write(&self, signed_url: &str, bytes: Vec<u8>) -> Result<StateWrite> {
72        validate_url(signed_url)?;
73        let response = self
74            .client
75            .put(signed_url)
76            .header(reqwest::header::CONTENT_TYPE, STATE_CONTENT_TYPE)
77            .body(bytes)
78            .send()
79            .await
80            .context("write actor state through signed URL")?;
81        if response.status() == reqwest::StatusCode::PRECONDITION_FAILED {
82            return Ok(StateWrite::GenerationMismatch);
83        }
84        ensure!(
85            response.status().is_success(),
86            "actor-state write failed with HTTP {}",
87            response.status()
88        );
89        Ok(StateWrite::Written)
90    }
91}
92
93fn generation_header(response: &reqwest::Response) -> Result<String> {
94    let generation = response
95        .headers()
96        .get(GENERATION_HEADER)
97        .context("GCS response omitted its object generation")?
98        .to_str()
99        .context("GCS object generation is not ASCII")?;
100    ensure!(
101        !generation.is_empty() && generation.bytes().all(|byte| byte.is_ascii_digit()),
102        "GCS object generation is invalid"
103    );
104    Ok(generation.to_owned())
105}
106
107fn validate_url(url: &str) -> Result<()> {
108    let url = reqwest::Url::parse(url).context("parse signed actor-state URL")?;
109    ensure!(
110        matches!(url.scheme(), "http" | "https") && url.host_str().is_some(),
111        "signed actor-state URL must be HTTP or HTTPS"
112    );
113    Ok(())
114}
115
116#[cfg(test)]
117mod tests {
118    use super::*;
119
120    #[test]
121    fn rejects_non_http_capabilities() {
122        assert!(validate_url("file:///tmp/state").is_err());
123        assert!(validate_url("https://storage.googleapis.com/bucket/object").is_ok());
124    }
125}