spacetimedb_client_api/
util.rs

1mod flat_csv;
2pub(crate) mod serde;
3pub mod websocket;
4
5use core::fmt;
6use std::net::IpAddr;
7
8use axum::body::Bytes;
9use axum::extract::{FromRequest, Request};
10use axum::response::IntoResponse;
11use bytestring::ByteString;
12use futures::TryStreamExt;
13use http::{HeaderName, HeaderValue, StatusCode};
14
15use hyper::body::Body;
16use spacetimedb::Identity;
17use spacetimedb_client_api_messages::name::DatabaseName;
18
19use crate::routes::identity::IdentityForUrl;
20use crate::{log_and_500, ControlStateReadAccess};
21
22pub struct ByteStringBody(pub ByteString);
23
24#[async_trait::async_trait]
25impl<S: Send + Sync> FromRequest<S> for ByteStringBody {
26    type Rejection = axum::response::Response;
27
28    async fn from_request(req: Request, state: &S) -> Result<Self, Self::Rejection> {
29        let bytes = Bytes::from_request(req, state)
30            .await
31            .map_err(IntoResponse::into_response)?;
32
33        let string = bytes
34            .try_into()
35            .map_err(|_| (StatusCode::BAD_REQUEST, "Request body didn't contain valid UTF-8").into_response())?;
36
37        Ok(ByteStringBody(string))
38    }
39}
40
41pub struct XForwardedFor(pub IpAddr);
42
43impl headers::Header for XForwardedFor {
44    fn name() -> &'static HeaderName {
45        static NAME: HeaderName = HeaderName::from_static("x-forwarded-for");
46        &NAME
47    }
48
49    fn decode<'i, I: Iterator<Item = &'i HeaderValue>>(values: &mut I) -> Result<Self, headers::Error> {
50        let val = values.next().ok_or_else(headers::Error::invalid)?;
51        let val = val.to_str().map_err(|_| headers::Error::invalid())?;
52        let (first, _) = val.split_once(',').ok_or_else(headers::Error::invalid)?;
53        let ip = first.trim().parse().map_err(|_| headers::Error::invalid())?;
54        Ok(XForwardedFor(ip))
55    }
56
57    fn encode<E: Extend<HeaderValue>>(&self, values: &mut E) {
58        values.extend([self.0.to_string().try_into().unwrap()])
59    }
60}
61
62#[derive(Clone, Debug)]
63pub enum NameOrIdentity {
64    Identity(IdentityForUrl),
65    Name(DatabaseName),
66}
67
68impl NameOrIdentity {
69    pub fn into_string(self) -> String {
70        match self {
71            NameOrIdentity::Identity(addr) => Identity::from(addr).to_hex().to_string(),
72            NameOrIdentity::Name(name) => name.into(),
73        }
74    }
75
76    pub fn name(&self) -> Option<&DatabaseName> {
77        if let Self::Name(name) = self {
78            Some(name)
79        } else {
80            None
81        }
82    }
83
84    /// Resolve this [`NameOrIdentity`].
85    ///
86    /// If `self` is a [`NameOrIdentity::Identity`], the inner [`Identity`] is
87    /// returned directly.
88    ///
89    /// Otherwise, if `self` is a [`NameOrIdentity::Name`], the [`Identity`] is
90    /// looked up by that name in the SpacetimeDB DNS and returned.
91    ///
92    /// Errors are returned if [`NameOrIdentity::Name`] the DNS lookup fails.
93    ///
94    /// An `Ok` result is itself a [`Result`], which is `Err(DatabaseName)` if the
95    /// given [`NameOrIdentity::Name`] is not registered in the SpacetimeDB DNS,
96    /// i.e. no corresponding [`Identity`] exists.
97    pub async fn try_resolve(
98        &self,
99        ctx: &(impl ControlStateReadAccess + ?Sized),
100    ) -> axum::response::Result<Result<Identity, &DatabaseName>> {
101        Ok(match self {
102            Self::Identity(identity) => Ok(Identity::from(*identity)),
103            Self::Name(name) => ctx.lookup_identity(name.as_ref()).map_err(log_and_500)?.ok_or(name),
104        })
105    }
106
107    /// A variant of [`Self::try_resolve()`] which maps to a 404 (Not Found)
108    /// response if `self` is a [`NameOrIdentity::Name`] for which no
109    /// corresponding [`Identity`] is found in the SpacetimeDB DNS.
110    pub async fn resolve(&self, ctx: &(impl ControlStateReadAccess + ?Sized)) -> axum::response::Result<Identity> {
111        self.try_resolve(ctx).await?.map_err(|_| StatusCode::NOT_FOUND.into())
112    }
113}
114
115impl<'de> ::serde::Deserialize<'de> for NameOrIdentity {
116    fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
117    where
118        D: ::serde::Deserializer<'de>,
119    {
120        let s = String::deserialize(deserializer)?;
121        if let Ok(addr) = Identity::from_hex(&s) {
122            Ok(NameOrIdentity::Identity(IdentityForUrl::from(addr)))
123        } else {
124            let name: DatabaseName = s.try_into().map_err(::serde::de::Error::custom)?;
125            Ok(NameOrIdentity::Name(name))
126        }
127    }
128}
129
130impl fmt::Display for NameOrIdentity {
131    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
132        match self {
133            Self::Identity(addr) => f.write_str(addr.into_inner().to_hex().as_str()),
134            Self::Name(name) => f.write_str(name.as_ref()),
135        }
136    }
137}
138
139pub struct EmptyBody;
140
141#[async_trait::async_trait]
142impl<S> FromRequest<S> for EmptyBody {
143    type Rejection = axum::response::Response;
144    async fn from_request(req: Request, _state: &S) -> Result<Self, Self::Rejection> {
145        let body = req.into_body();
146        if body.is_end_stream() {
147            return Ok(Self);
148        }
149
150        if body
151            .into_data_stream()
152            .try_any(|data| futures::future::ready(!data.is_empty()))
153            .await
154            .map_err(|_| (StatusCode::BAD_REQUEST, "Failed to buffer the request body").into_response())?
155        {
156            return Err((StatusCode::BAD_REQUEST, "body must be empty").into_response());
157        }
158        Ok(Self)
159    }
160}