rebuilderd_common/api/
mod.rs

1use crate::auth;
2use crate::config::ConfigFile;
3use crate::errors::Error;
4use crate::utils::zstd_compress;
5use anyhow::{Context, anyhow};
6use async_trait::async_trait;
7use log::debug;
8use reqwest::header::CONTENT_ENCODING;
9use reqwest::{RequestBuilder, Response};
10use std::borrow::Cow;
11use std::env;
12use url::Url;
13
14pub mod v0;
15pub mod v1;
16
17pub const AUTH_COOKIE_HEADER: &str = "X-Auth-Cookie";
18pub const WORKER_KEY_HEADER: &str = "X-Worker-Key";
19pub const SIGNUP_SECRET_HEADER: &str = "X-Signup-Secret";
20
21pub struct Client {
22    endpoint: Url,
23    client: crate::http::Client,
24    is_default_endpoint: bool,
25    auth_cookie: Option<String>,
26    worker_key: Option<String>,
27    signup_secret: Option<String>,
28}
29
30impl Client {
31    pub fn new(config: ConfigFile, endpoint: Option<String>) -> anyhow::Result<Client> {
32        let (endpoint, auth_cookie, is_default_endpoint) = if let Some(endpoint) = endpoint {
33            let cookie = config
34                .endpoints
35                .get(&endpoint)
36                .map(|e| e.cookie.to_string());
37            (endpoint, cookie, false)
38        } else if let Some(endpoint) = config.http.endpoint {
39            (endpoint, None, true)
40        } else {
41            ("http://127.0.0.1:8484".to_string(), None, true)
42        };
43
44        let mut endpoint = endpoint
45            .parse::<Url>()
46            .with_context(|| anyhow!("Failed to parse endpoint as url: {:?}", endpoint))?;
47
48        // If the url ends with a slash, remove it
49        endpoint
50            .path_segments_mut()
51            .map_err(|_| anyhow!("Given endpoint url cannot be base"))?
52            .pop_if_empty();
53
54        debug!("Setting rebuilderd endpoint to {:?}", endpoint.as_str());
55        let client = crate::http::Client::builder().zstd(true).build()?;
56
57        Ok(Client {
58            endpoint,
59            client,
60            is_default_endpoint,
61            auth_cookie,
62            worker_key: None,
63            signup_secret: None,
64        })
65    }
66
67    pub fn with_auth_cookie(&mut self) -> anyhow::Result<&mut Self> {
68        if let Ok(cookie_path) = env::var("REBUILDERD_COOKIE_PATH") {
69            debug!("Found cookie path in environment: {:?}", cookie_path);
70            let auth_cookie =
71                auth::read_cookie_from_file(cookie_path).context("Failed to load auth cookie")?;
72            Ok(self.auth_cookie(auth_cookie))
73        } else if self.is_default_endpoint {
74            let auth_cookie = auth::find_auth_cookie().context("Failed to load auth cookie")?;
75            Ok(self.auth_cookie(auth_cookie))
76        } else {
77            Ok(self)
78        }
79    }
80
81    pub fn auth_cookie<I: Into<String>>(&mut self, cookie: I) -> &mut Self {
82        self.auth_cookie = Some(cookie.into());
83        self
84    }
85
86    pub fn worker_key<I: Into<String>>(&mut self, key: I) {
87        self.worker_key = Some(key.into());
88    }
89
90    pub fn signup_secret<I: Into<String>>(&mut self, secret: I) {
91        self.signup_secret = Some(secret.into());
92    }
93
94    fn url_join(&self, route: &str) -> Url {
95        let mut url = self.endpoint.clone();
96        {
97            // this unwrap is safe because we've called path_segments_mut in the constructor before
98            let mut path = url.path_segments_mut().expect("Url cannot be base");
99            for segment in route.split('/') {
100                path.push(segment);
101            }
102        }
103
104        url
105    }
106
107    fn authenticated(&self, mut req: RequestBuilder) -> RequestBuilder {
108        if let Some(auth_cookie) = &self.auth_cookie {
109            req = req.header(AUTH_COOKIE_HEADER, auth_cookie);
110        }
111
112        if let Some(worker_key) = &self.worker_key {
113            req = req.header(WORKER_KEY_HEADER, worker_key);
114        }
115
116        if let Some(signup_secret) = &self.signup_secret {
117            req = req.header(SIGNUP_SECRET_HEADER, signup_secret);
118        }
119
120        req
121    }
122
123    fn get(&self, path: Cow<'static, str>) -> crate::http::RequestBuilder {
124        let url = self.url_join(&path);
125        debug!("Sending GET request to {}", url.as_str());
126        let req = self.client.get(url);
127        self.authenticated(req)
128    }
129
130    fn post(&self, path: Cow<'static, str>) -> crate::http::RequestBuilder {
131        let url = self.url_join(&path);
132        debug!("Sending POST request to {}", url.as_str());
133        let req = self.client.post(url);
134        self.authenticated(req)
135    }
136
137    fn delete(&self, path: Cow<'static, str>) -> crate::http::RequestBuilder {
138        let url = self.url_join(&path);
139        debug!("Sending DELETE request to {}", url.as_str());
140        let req = self.client.delete(url);
141        self.authenticated(req)
142    }
143}
144
145#[async_trait]
146pub trait ZstdRequestBuilder {
147    async fn send_encoded(self) -> crate::errors::Result<Response>;
148}
149
150#[async_trait]
151impl ZstdRequestBuilder for RequestBuilder {
152    async fn send_encoded(self) -> crate::errors::Result<Response> {
153        if let Some(new_request) = self.try_clone() {
154            let mut request = self.build()?;
155
156            if let Some(body) = request.body_mut() {
157                if let Some(bytes) = body.as_bytes() {
158                    let encoded_body = zstd_compress(bytes).await?;
159
160                    new_request
161                        .body(encoded_body)
162                        .header(CONTENT_ENCODING, "zstd")
163                        .send()
164                        .await
165                        .map_err(Error::from)
166                } else {
167                    new_request.send().await.map_err(Error::from)
168                }
169            } else {
170                new_request.send().await.map_err(Error::from)
171            }
172        } else {
173            self.send().await.map_err(Error::from)
174        }
175    }
176}
rebuilderd_common/api/mod.rs

rebuilderd_common/api/
mod.rs
