spacetimedb_client_api/
util.rs1mod 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 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 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}