Line data Source code
1 : use anyhow::Result;
2 : use axum::extract::{FromRequestParts, Path};
3 : use axum::response::{IntoResponse, Response};
4 : use axum::{RequestPartsExt, http::StatusCode, http::request::Parts};
5 : use axum_extra::TypedHeader;
6 : use axum_extra::headers::{Authorization, authorization::Bearer};
7 : use camino::Utf8PathBuf;
8 : use jsonwebtoken::{DecodingKey, Validation};
9 : use remote_storage::{GenericRemoteStorage, RemotePath};
10 : use serde::{Deserialize, Serialize};
11 : use std::fmt::Display;
12 : use std::result::Result as StdResult;
13 : use std::sync::Arc;
14 : use tokio_util::sync::CancellationToken;
15 : use tracing::{debug, error};
16 : use utils::id::{TenantId, TimelineId};
17 :
18 : // simplified version of utils::auth::JwtAuth
19 : pub struct JwtAuth {
20 : decoding_key: DecodingKey,
21 : validation: Validation,
22 : }
23 :
24 : pub const VALIDATION_ALGO: jsonwebtoken::Algorithm = jsonwebtoken::Algorithm::EdDSA;
25 : impl JwtAuth {
26 34 : pub fn new(key: &[u8]) -> Result<Self> {
27 34 : Ok(Self {
28 34 : decoding_key: DecodingKey::from_ed_pem(key)?,
29 34 : validation: Validation::new(VALIDATION_ALGO),
30 : })
31 34 : }
32 :
33 127 : pub fn decode<T: serde::de::DeserializeOwned>(&self, token: &str) -> Result<T> {
34 127 : Ok(jsonwebtoken::decode(token, &self.decoding_key, &self.validation).map(|t| t.claims)?)
35 127 : }
36 : }
37 :
38 49 : fn normalize_key(key: &str) -> StdResult<Utf8PathBuf, String> {
39 49 : let key = clean_utf8(&Utf8PathBuf::from(key));
40 49 : if key.starts_with("..") || key == "." || key == "/" {
41 9 : return Err(format!("invalid key {key}"));
42 40 : }
43 40 : match key.strip_prefix("/").map(Utf8PathBuf::from) {
44 1 : Ok(p) => Ok(p),
45 39 : _ => Ok(key),
46 : }
47 49 : }
48 :
49 : // Copied from path_clean crate with PathBuf->Utf8PathBuf
50 49 : fn clean_utf8(path: &camino::Utf8Path) -> Utf8PathBuf {
51 : use camino::Utf8Component as Comp;
52 49 : let mut out = Vec::new();
53 81 : for comp in path.components() {
54 81 : match comp {
55 1 : Comp::CurDir => (),
56 18 : Comp::ParentDir => match out.last() {
57 1 : Some(Comp::RootDir) => (),
58 12 : Some(Comp::Normal(_)) => {
59 12 : out.pop();
60 12 : }
61 : None | Some(Comp::CurDir) | Some(Comp::ParentDir) | Some(Comp::Prefix(_)) => {
62 5 : out.push(comp)
63 : }
64 : },
65 62 : comp => out.push(comp),
66 : }
67 : }
68 49 : if !out.is_empty() {
69 47 : out.iter().collect()
70 : } else {
71 2 : Utf8PathBuf::from(".")
72 : }
73 49 : }
74 :
75 : pub struct Storage {
76 : pub auth: JwtAuth,
77 : pub storage: GenericRemoteStorage,
78 : pub cancel: CancellationToken,
79 : pub max_upload_file_limit: usize,
80 : }
81 :
82 : pub type EndpointId = String; // If needed, reuse small string from proxy/src/types.rc
83 :
84 484 : #[derive(Deserialize, Serialize, PartialEq)]
85 : pub struct Claims {
86 : pub tenant_id: TenantId,
87 : pub timeline_id: TimelineId,
88 : pub endpoint_id: EndpointId,
89 : pub exp: u64,
90 : }
91 :
92 : impl Display for Claims {
93 0 : fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
94 0 : write!(
95 0 : f,
96 0 : "Claims(tenant_id {} timeline_id {} endpoint_id {} exp {})",
97 0 : self.tenant_id, self.timeline_id, self.endpoint_id, self.exp
98 0 : )
99 0 : }
100 : }
101 :
102 490 : #[derive(Deserialize, Serialize)]
103 : struct KeyRequest {
104 : tenant_id: TenantId,
105 : timeline_id: TimelineId,
106 : endpoint_id: EndpointId,
107 : path: String,
108 : }
109 :
110 : #[derive(Debug, PartialEq)]
111 : pub struct S3Path {
112 : pub path: RemotePath,
113 : }
114 :
115 : impl TryFrom<&KeyRequest> for S3Path {
116 : type Error = String;
117 40 : fn try_from(req: &KeyRequest) -> StdResult<Self, Self::Error> {
118 40 : let KeyRequest {
119 40 : tenant_id,
120 40 : timeline_id,
121 40 : endpoint_id,
122 40 : path,
123 40 : } = &req;
124 40 : let prefix = format!("{tenant_id}/{timeline_id}/{endpoint_id}",);
125 40 : let path = Utf8PathBuf::from(prefix).join(normalize_key(path)?);
126 37 : let path = RemotePath::new(&path).unwrap(); // unwrap() because the path is already relative
127 37 : Ok(S3Path { path })
128 40 : }
129 : }
130 :
131 84 : fn unauthorized(route: impl Display, claims: impl Display) -> Response {
132 84 : debug!(%route, %claims, "route doesn't match claims");
133 84 : StatusCode::UNAUTHORIZED.into_response()
134 84 : }
135 :
136 14 : pub fn bad_request(err: impl Display, desc: &'static str) -> Response {
137 14 : debug!(%err, desc);
138 14 : (StatusCode::BAD_REQUEST, err.to_string()).into_response()
139 14 : }
140 :
141 21 : pub fn ok() -> Response {
142 21 : StatusCode::OK.into_response()
143 21 : }
144 :
145 0 : pub fn internal_error(err: impl Display, path: impl Display, desc: &'static str) -> Response {
146 0 : error!(%err, %path, desc);
147 0 : StatusCode::INTERNAL_SERVER_ERROR.into_response()
148 0 : }
149 :
150 14 : pub fn not_found(key: impl ToString) -> Response {
151 14 : (StatusCode::NOT_FOUND, key.to_string()).into_response()
152 14 : }
153 :
154 : impl FromRequestParts<Arc<Storage>> for S3Path {
155 : type Rejection = Response;
156 127 : async fn from_request_parts(
157 127 : parts: &mut Parts,
158 127 : state: &Arc<Storage>,
159 127 : ) -> Result<Self, Self::Rejection> {
160 127 : let Path(path): Path<KeyRequest> = parts
161 127 : .extract()
162 127 : .await
163 127 : .map_err(|e| bad_request(e, "invalid route"))?;
164 121 : let TypedHeader(Authorization(bearer)) = parts
165 121 : .extract::<TypedHeader<Authorization<Bearer>>>()
166 121 : .await
167 121 : .map_err(|e| bad_request(e, "invalid token"))?;
168 121 : let claims: Claims = state
169 121 : .auth
170 121 : .decode(bearer.token())
171 121 : .map_err(|e| bad_request(e, "decoding token"))?;
172 121 : let route = Claims {
173 121 : tenant_id: path.tenant_id,
174 121 : timeline_id: path.timeline_id,
175 121 : endpoint_id: path.endpoint_id.clone(),
176 121 : exp: claims.exp,
177 121 : };
178 121 : if route != claims {
179 84 : return Err(unauthorized(route, claims));
180 37 : }
181 37 : (&path)
182 37 : .try_into()
183 37 : .map_err(|e| bad_request(e, "invalid route"))
184 127 : }
185 : }
186 :
187 44 : #[derive(Deserialize, Serialize, PartialEq)]
188 : pub struct PrefixKeyPath {
189 : pub tenant_id: TenantId,
190 : pub timeline_id: Option<TimelineId>,
191 : pub endpoint_id: Option<EndpointId>,
192 : }
193 :
194 : impl Display for PrefixKeyPath {
195 0 : fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
196 0 : write!(
197 0 : f,
198 0 : "PrefixKeyPath(tenant_id {} timeline_id {} endpoint_id {})",
199 0 : self.tenant_id,
200 0 : self.timeline_id
201 0 : .as_ref()
202 0 : .map(ToString::to_string)
203 0 : .unwrap_or("".to_string()),
204 0 : self.endpoint_id
205 0 : .as_ref()
206 0 : .map(ToString::to_string)
207 0 : .unwrap_or("".to_string())
208 0 : )
209 0 : }
210 : }
211 :
212 : #[derive(Debug, PartialEq)]
213 : pub struct PrefixS3Path {
214 : pub path: RemotePath,
215 : }
216 :
217 : impl From<&PrefixKeyPath> for PrefixS3Path {
218 9 : fn from(path: &PrefixKeyPath) -> Self {
219 9 : let timeline_id = path
220 9 : .timeline_id
221 9 : .as_ref()
222 9 : .map(ToString::to_string)
223 9 : .unwrap_or("".to_string());
224 9 : let endpoint_id = path
225 9 : .endpoint_id
226 9 : .as_ref()
227 9 : .map(ToString::to_string)
228 9 : .unwrap_or("".to_string());
229 9 : let path = Utf8PathBuf::from(path.tenant_id.to_string())
230 9 : .join(timeline_id)
231 9 : .join(endpoint_id);
232 9 : let path = RemotePath::new(&path).unwrap(); // unwrap() because the path is already relative
233 9 : PrefixS3Path { path }
234 9 : }
235 : }
236 :
237 : impl FromRequestParts<Arc<Storage>> for PrefixS3Path {
238 : type Rejection = Response;
239 12 : async fn from_request_parts(
240 12 : parts: &mut Parts,
241 12 : state: &Arc<Storage>,
242 12 : ) -> Result<Self, Self::Rejection> {
243 12 : let Path(path) = parts
244 12 : .extract::<Path<PrefixKeyPath>>()
245 12 : .await
246 12 : .map_err(|e| bad_request(e, "invalid route"))?;
247 6 : let TypedHeader(Authorization(bearer)) = parts
248 6 : .extract::<TypedHeader<Authorization<Bearer>>>()
249 6 : .await
250 6 : .map_err(|e| bad_request(e, "invalid token"))?;
251 6 : let claims: PrefixKeyPath = state
252 6 : .auth
253 6 : .decode(bearer.token())
254 6 : .map_err(|e| bad_request(e, "invalid token"))?;
255 6 : if path != claims {
256 0 : return Err(unauthorized(path, claims));
257 6 : }
258 6 : Ok((&path).into())
259 12 : }
260 : }
261 :
262 : #[cfg(test)]
263 : mod tests {
264 : use super::*;
265 :
266 : #[test]
267 1 : fn normalize_key() {
268 1 : let f = super::normalize_key;
269 1 : assert_eq!(f("hello/world/..").unwrap(), Utf8PathBuf::from("hello"));
270 1 : assert_eq!(
271 1 : f("ololo/1/../../not_ololo").unwrap(),
272 1 : Utf8PathBuf::from("not_ololo")
273 1 : );
274 1 : assert!(f("ololo/1/../../../").is_err());
275 1 : assert!(f(".").is_err());
276 1 : assert!(f("../").is_err());
277 1 : assert!(f("").is_err());
278 1 : assert_eq!(f("/1/2/3").unwrap(), Utf8PathBuf::from("1/2/3"));
279 1 : assert!(f("/1/2/3/../../../").is_err());
280 1 : assert!(f("/1/2/3/../../../../").is_err());
281 1 : }
282 :
283 : const TENANT_ID: TenantId =
284 : TenantId::from_array([1, 1, 2, 3, 4, 5, 6, 7, 8, 9, 1, 2, 3, 4, 5, 6]);
285 : const TIMELINE_ID: TimelineId =
286 : TimelineId::from_array([1, 1, 2, 3, 4, 5, 6, 7, 8, 9, 1, 2, 3, 4, 5, 7]);
287 : const ENDPOINT_ID: &str = "ep-winter-frost-a662z3vg";
288 :
289 : #[test]
290 1 : fn s3_path() {
291 1 : let auth = Claims {
292 1 : tenant_id: TENANT_ID,
293 1 : timeline_id: TIMELINE_ID,
294 1 : endpoint_id: ENDPOINT_ID.into(),
295 1 : exp: u64::MAX,
296 1 : };
297 2 : let s3_path = |key| {
298 2 : let path = &format!("{TENANT_ID}/{TIMELINE_ID}/{ENDPOINT_ID}/{key}");
299 2 : let path = RemotePath::from_string(path).unwrap();
300 2 : S3Path { path }
301 2 : };
302 :
303 1 : let path = "cache_key".to_string();
304 1 : let mut key_path = KeyRequest {
305 1 : path,
306 1 : tenant_id: auth.tenant_id,
307 1 : timeline_id: auth.timeline_id,
308 1 : endpoint_id: auth.endpoint_id,
309 1 : };
310 1 : assert_eq!(S3Path::try_from(&key_path).unwrap(), s3_path(key_path.path));
311 :
312 1 : key_path.path = "we/can/have/nested/paths".to_string();
313 1 : assert_eq!(S3Path::try_from(&key_path).unwrap(), s3_path(key_path.path));
314 :
315 1 : key_path.path = "../error/hello/../".to_string();
316 1 : assert!(S3Path::try_from(&key_path).is_err());
317 1 : }
318 :
319 : #[test]
320 1 : fn prefix_s3_path() {
321 1 : let mut path = PrefixKeyPath {
322 1 : tenant_id: TENANT_ID,
323 1 : timeline_id: None,
324 1 : endpoint_id: None,
325 1 : };
326 3 : let prefix_path = |s: String| RemotePath::from_string(&s).unwrap();
327 1 : assert_eq!(
328 1 : PrefixS3Path::from(&path).path,
329 1 : prefix_path(format!("{TENANT_ID}"))
330 1 : );
331 :
332 1 : path.timeline_id = Some(TIMELINE_ID);
333 1 : assert_eq!(
334 1 : PrefixS3Path::from(&path).path,
335 1 : prefix_path(format!("{TENANT_ID}/{TIMELINE_ID}"))
336 1 : );
337 :
338 1 : path.endpoint_id = Some(ENDPOINT_ID.into());
339 1 : assert_eq!(
340 1 : PrefixS3Path::from(&path).path,
341 1 : prefix_path(format!("{TENANT_ID}/{TIMELINE_ID}/{ENDPOINT_ID}"))
342 1 : );
343 1 : }
344 : }
|
