A vibe coded tangled fork which supports pijul.
1use std::collections::HashMap;
2
3use serde::Deserialize;
4use worker::*;
5
6/// The JSON value stored in Workers KV, keyed by domain.
7///
8/// Example KV entry:
9/// key: "foo.example.com"
10/// value: {"did": "did:plc:...", "repos": {"my_repo": true, "other_repo": false}}
11///
12/// The boolean on each repo indicates whether it is the index site for the
13/// domain (true) or a sub-path site (false). At most one repo may be true.
14#[derive(Deserialize)]
15struct DomainMapping {
16 did: String,
17 /// repo name → is_index
18 repos: HashMap<String, bool>,
19}
20
21impl DomainMapping {
22 /// Returns the repo that is marked as the index site, if any.
23 fn index_repo(&self) -> Option<&str> {
24 self.repos
25 .iter()
26 .find_map(|(name, &is_index)| if is_index { Some(name.as_str()) } else { None })
27 }
28}
29
30/// Build the R2 object key for a given did/repo and intra-site path.
31/// `site_path` should start with a `/` or be empty.
32fn r2_key(did: &str, repo: &str, site_path: &str) -> String {
33 let base = format!("{}/{}/", did, repo);
34 if site_path.is_empty() || site_path == "/" {
35 format!("{}index.html", base)
36 } else {
37 let trimmed = site_path.trim_start_matches('/');
38 if trimmed.is_empty() || trimmed.ends_with('/') {
39 format!("{}{}index.html", base, trimmed)
40 } else {
41 format!("{}{}", base, trimmed)
42 }
43 }
44}
45
46/// Fetch an object from R2, falling back to appending /index.html if the
47/// key looks like a directory (no file extension in the last segment).
48async fn fetch_from_r2(bucket: &Bucket, key: &str) -> Result<Option<Object>> {
49 if let Some(obj) = bucket.get(key).execute().await? {
50 return Ok(Some(obj));
51 }
52
53 let last_segment = key.rsplit('/').next().unwrap_or(key);
54 if !last_segment.contains('.') {
55 let index_key = format!("{}/index.html", key.trim_end_matches('/'));
56 if let Some(obj) = bucket.get(&index_key).execute().await? {
57 return Ok(Some(obj));
58 }
59 }
60
61 Ok(None)
62}
63
64/// Build a Response from an R2 Object, forwarding the content-type header.
65fn response_from_object(obj: Object) -> Result<Response> {
66 let content_type = obj
67 .http_metadata()
68 .content_type
69 .unwrap_or_else(|| "application/octet-stream".to_string());
70
71 let body = obj.body().ok_or_else(|| Error::RustError("empty R2 body".into()))?;
72 let mut resp = Response::from_body(body.response_body()?)?;
73 resp.headers_mut().set("Content-Type", &content_type)?;
74 resp.headers_mut().set("Cache-Control", "public, max-age=60")?;
75 Ok(resp)
76}
77
78fn is_excluded(path: &str) -> bool {
79 let excluded = ["/.well-known/atproto-did"];
80 excluded.iter().any(|&prefix| path.starts_with(prefix))
81}
82
83#[event(fetch)]
84async fn fetch(req: Request, env: Env, _ctx: Context) -> Result<Response> {
85 let kv = env.kv("SITES")?;
86 let bucket = env.bucket("SITES_BUCKET")?;
87
88 // Extract host, stripping any port.
89 let host = req.headers().get("host")?.unwrap_or_default();
90 let host = host.split(':').next().unwrap_or("").to_string();
91
92 if host.is_empty() {
93 return Response::error("Bad Request: missing host", 400);
94 }
95
96 let url = req.url()?;
97 let path = url.path();
98
99 if is_excluded(path) {
100 return Fetch::Request(req).send().await;
101 }
102
103 // Single KV lookup for the whole domain.
104 let mapping = match kv.get(&host).text().await? {
105 Some(raw) => match serde_json::from_str::<DomainMapping>(&raw) {
106 Ok(m) => m,
107 Err(_) => return Response::error("Internal Error: bad mapping", 500),
108 },
109 None => return Response::error("site not found!", 404),
110 };
111
112 let path = url.path(); // always starts with "/"
113
114 // First path segment, e.g. "my_repo" from "/my_repo/page.html"
115 let first_segment = path
116 .trim_start_matches('/')
117 .split('/')
118 .next()
119 .unwrap_or("")
120 .to_string();
121
122 // 1. sub-path site
123 // If the first path segment matches a non-index repo, serve from it.
124 if !first_segment.is_empty() {
125 if let Some(&is_index) = mapping.repos.get(&first_segment) {
126 if !is_index {
127 // Strip the leading "/{first_segment}" to get the intra-site path.
128 let site_path = path
129 .trim_start_matches('/')
130 .trim_start_matches(&first_segment)
131 .to_string();
132
133 let key = r2_key(&mapping.did, &first_segment, &site_path);
134 return match fetch_from_r2(&bucket, &key).await? {
135 Some(obj) => response_from_object(obj),
136 None => Response::error("Not Found", 404),
137 };
138 }
139 }
140 }
141
142 // 2. index site
143 // Fall back to the repo marked as the index site, serving the full path.
144 if let Some(index_repo) = mapping.index_repo() {
145 let key = r2_key(&mapping.did, index_repo, path);
146 return match fetch_from_r2(&bucket, &key).await? {
147 Some(obj) => response_from_object(obj),
148 None => Response::error("Not Found", 404),
149 };
150 }
151
152 Response::error("Not Found", 404)
153}