| 1 | use crate::*; |
| 2 | use rusqlite::OptionalExtension; |
| 3 | use std::collections::HashSet; |
| 4 | use axum::body::Body; |
| 5 | use lol_html::{RewriteStrSettings, element, html_content::ContentType, text}; |
| 6 | |
| 7 | struct Profile { |
| 8 | name: String, |
| 9 | url: Option<String>, |
| 10 | icon: &'static str, |
| 11 | picture: Option<String>, |
| 12 | owner: Option<(String, Option<String>)>, |
| 13 | } |
| 14 | |
| 15 | fn profiles(auth: &auth::Store) -> Result<HashMap<String, Profile>> { |
| 16 | let db = auth.db.lock().unwrap(); |
| 17 | let mut query = db.prepare("SELECT json_extract(profile,'$.username'),coalesce(json_extract(profile,'$.firstName'),''),provider,subject,json_extract(profile,'$.attributes.picture[0]') FROM users JOIN external_identities ON user_id=users.id WHERE json_extract(profile,'$.kind')='guest'")?; |
| 18 | let rows = query.query_map([], |row| { |
| 19 | Ok(( |
| 20 | row.get::<_, String>(0)?, |
| 21 | row.get::<_, String>(1)?, |
| 22 | row.get::<_, String>(2)?, |
| 23 | row.get::<_, String>(3)?, |
| 24 | row.get::<_, Option<String>>(4)?, |
| 25 | )) |
| 26 | })?; |
| 27 | let mut profiles = HashMap::new(); |
| 28 | for row in rows { |
| 29 | let (username, name, provider, subject, picture) = row?; |
| 30 | if name.is_empty() { continue; } |
| 31 | let (url, icon) = match provider.as_str() { |
| 32 | "github" => { |
| 33 | if matches!(name.as_str(), "." | "..") { continue; } |
| 34 | let mut url = url::Url::parse("https://github.com/")?; |
| 35 | url.path_segments_mut().unwrap().pop_if_empty().push(&name); |
| 36 | (Some(String::from(url)), include_str!("../web/sso/github.svg")) |
| 37 | } |
| 38 | "astheno" => (None, include_str!("astheno.svg")), |
| 39 | _ => continue, |
| 40 | }; |
| 41 | profiles.insert( |
| 42 | username, |
| 43 | Profile { |
| 44 | name, |
| 45 | url, |
| 46 | icon, |
| 47 | owner: None, |
| 48 | picture: if provider == "github" { |
| 49 | Some(format!( |
| 50 | "https://avatars.githubusercontent.com/u/{subject}?s=64" |
| 51 | )) |
| 52 | } else { |
| 53 | picture |
| 54 | }, |
| 55 | }, |
| 56 | ); |
| 57 | } |
| 58 | let config: Option<String> = db.query_row("SELECT config FROM oidc_clients WHERE id='shale'", [], |r| r.get(0)).optional()?; |
| 59 | let config: Value = config.map(|config| serde_json::from_str(&config)).transpose()?.unwrap_or(Value::Null); |
| 60 | let mut query = db.prepare("SELECT u.username,json_extract(u.profile,'$.attributes.picture[0]'),i.identity FROM users u JOIN ( |
| 61 | SELECT id AS user_id,'email:'||lower(json_extract(profile,'$.email')) AS identity FROM users WHERE json_extract(profile,'$.emailVerified')=1 AND json_extract(profile,'$.email') IS NOT NULL |
| 62 | UNION SELECT user_id,'email:'||email FROM commit_emails WHERE verified=1 |
| 63 | UNION SELECT user_id,'key:'||fingerprint FROM signing_keys WHERE verified=1 |
| 64 | ) i ON i.user_id=u.id WHERE json_extract(u.profile,'$.enabled')=1 AND coalesce(json_extract(u.profile,'$.kind'),'user') NOT IN ('ai','guest')")?; |
| 65 | for row in query.query_map([], |r| Ok((r.get::<_,String>(0)?,r.get::<_,Option<String>>(1)?,r.get::<_,String>(2)?)))? { |
| 66 | let (username,picture,identity) = row?; |
| 67 | let name = config["usernameAliases"][&username].as_str().unwrap_or(&username); |
| 68 | let mut url = url::Url::parse("https://shale.invalid/")?; |
| 69 | url.path_segments_mut().unwrap().pop_if_empty().push(&format!("~{name}")); |
| 70 | profiles.insert(identity, Profile { name: format!("~{name}"),url: Some(url.path().into()),icon:"",picture,owner:None }); |
| 71 | } |
| 72 | Ok(profiles) |
| 73 | } |
| 74 | |
| 75 | pub(crate) fn assisted_model(body: &str) -> Option<&str> { |
| 76 | let line = body.trim_end().lines().last()?; |
| 77 | let model = line.strip_prefix("Assisted-by: ")?; |
| 78 | if model.is_empty() || model.len() > 128 || !model.bytes().all(|b| b.is_ascii_alphanumeric() || matches!(b, b'-' | b'.' | b'_')) { return None; } |
| 79 | let before = &body[..body.trim_end().len() - line.len()]; |
| 80 | let mut fence = None; |
| 81 | for line in before.lines() { |
| 82 | let line = line.trim_start(); |
| 83 | if line.starts_with("```") || line.starts_with("~~~") { |
| 84 | let marker = line.as_bytes()[0]; |
| 85 | let count = line.bytes().take_while(|&b| b == marker).count(); |
| 86 | match fence { |
| 87 | None => fence = Some((marker, count)), |
| 88 | Some((open, length)) if marker == open && count >= length && line[count..].trim().is_empty() => fence = None, |
| 89 | _ => (), |
| 90 | } |
| 91 | } |
| 92 | } |
| 93 | fence.is_none().then_some(model) |
| 94 | } |
| 95 | |
| 96 | fn ai_profile(model: Option<&str>) -> Profile { |
| 97 | let (name, icon) = match model { |
| 98 | None => ("ai".into(), include_str!("bot.svg")), |
| 99 | Some(id) => { |
| 100 | let normalize_version = id.starts_with("claude-") || id.starts_with("gpt-"); |
| 101 | let (id, icon, gpt) = if let Some(rest) = id.strip_prefix("claude-") { |
| 102 | let rest = rest.rsplit_once('-').filter(|(_, tail)| tail.len() == 8 && tail.bytes().all(|b| b.is_ascii_digit())).map(|(rest, _)| rest).unwrap_or(rest); |
| 103 | (rest, include_str!("claude.svg"), false) |
| 104 | } else if let Some(rest) = id.strip_prefix("gpt-") { |
| 105 | (rest, include_str!("openai.svg"), true) |
| 106 | } else if id.starts_with("qwen") { |
| 107 | (id, include_str!("qwen.svg"), false) |
| 108 | } else { |
| 109 | (id, include_str!("bot.svg"), false) |
| 110 | }; |
| 111 | let mut words: Vec<String> = Vec::new(); |
| 112 | for part in id.split('-').filter(|part| !part.is_empty()) { |
| 113 | let number = part.find(|c: char| c.is_ascii_digit()).unwrap_or(part.len()); |
| 114 | let (word, version) = part.split_at(number); |
| 115 | if !word.is_empty() { |
| 116 | let mut letters = word.chars(); |
| 117 | words.push(letters.next().unwrap().to_uppercase().collect::<String>() + letters.as_str()); |
| 118 | } |
| 119 | if !version.is_empty() { |
| 120 | let numeric = |value: &str| value.bytes().all(|b| b.is_ascii_digit() || b == b'.'); |
| 121 | if normalize_version && numeric(version) && words.last().is_some_and(|word| numeric(word)) { |
| 122 | let previous = words.last_mut().unwrap(); |
| 123 | previous.push('.'); |
| 124 | previous.push_str(version); |
| 125 | } else { |
| 126 | words.push(version.into()); |
| 127 | } |
| 128 | } |
| 129 | } |
| 130 | let name = words.join(" "); |
| 131 | (if gpt { format!("GPT-{name}") } else { name }, icon) |
| 132 | } |
| 133 | }; |
| 134 | Profile { name, url: None, icon, picture: None, owner: None } |
| 135 | } |
| 136 | |
| 137 | fn account_path(path: &str) -> Option<&str> { |
| 138 | let name = path.strip_prefix("/~")?; |
| 139 | let name = name.strip_suffix('/').unwrap_or(name); |
| 140 | (!name.is_empty() && !name.contains('/')).then_some(name) |
| 141 | } |
| 142 | |
| 143 | fn avatar(picture: &str) -> String { |
| 144 | let picture = picture.replace('&', "&").replace('"', """).replace('<', "<"); |
| 145 | format!("<img src=\"{picture}\" alt=\"\" height=\"16\" referrerpolicy=\"no-referrer\">") |
| 146 | } |
| 147 | |
| 148 | fn rewrite(html: &str, origin: &url::Url, profiles: &HashMap<String, Profile>, avatars: &HashMap<String, String>, open_count: Option<u64>) -> Result<String> { |
| 149 | use std::{cell::Cell, rc::Rc}; |
| 150 | let markdown = Rc::new(Cell::new(0)); |
| 151 | let bodies = markdown.clone(); |
| 152 | let active = Rc::new(Cell::new(None::<bool>)); |
| 153 | let images = active.clone(); |
| 154 | let labels = active.clone(); |
| 155 | let compact = Rc::new(Cell::new(false)); |
| 156 | let author_column = compact.clone(); |
| 157 | let commit_author = Rc::new(std::cell::RefCell::new(None::<(String,Option<String>)>)); |
| 158 | let commit_images = commit_author.clone(); |
| 159 | let commit_names = commit_author.clone(); |
| 160 | let signature = Rc::new(std::cell::RefCell::new(String::new())); |
| 161 | Ok(lol_html::rewrite_str(html, RewriteStrSettings::new() |
| 162 | .append_element_content_handler(text!("header .usa-nav__primary a[href]", |text| { |
| 163 | if markdown.get() == 0 && let Some(count) = open_count && let Some(start) = text.as_str().find("issues (") |
| 164 | && let Some(end) = text.as_str()[start..].find(')') |
| 165 | { |
| 166 | text.replace(&format!("{}issues ({count}){}", &text.as_str()[..start], &text.as_str()[start + end + 1..]), ContentType::Text); |
| 167 | } |
| 168 | Ok(()) |
| 169 | })) |
| 170 | .append_element_content_handler(element!("#m-metainfo td", |element| { |
| 171 | *commit_author.borrow_mut() = None; |
| 172 | let closing = commit_author.clone(); |
| 173 | element.on_end_tag(lol_html::end_tag!(move |_| { *closing.borrow_mut() = None; Ok(()) }))?; |
| 174 | Ok(()) |
| 175 | })) |
| 176 | .append_element_content_handler(element!("#page-log table tbody td[title]", |element| { |
| 177 | if markdown.get() != 0 { return Ok(()); } |
| 178 | let Some(profile) = element.get_attribute("title").and_then(|email| profiles.get(&format!("email:{}", email.trim().to_lowercase()))) else { return Ok(()); }; |
| 179 | let Some(url) = &profile.url else { return Ok(()); }; |
| 180 | *commit_author.borrow_mut() = Some((url.clone(),profile.picture.clone())); |
| 181 | let closing = commit_author.clone(); |
| 182 | element.on_end_tag(lol_html::end_tag!(move |_| { *closing.borrow_mut() = None; Ok(()) }))?; |
| 183 | element.prepend(&format!("<a class=\"usa-link\" href=\"{url}\">"),ContentType::Html); |
| 184 | element.append("</a>",ContentType::Html); |
| 185 | Ok(()) |
| 186 | })) |
| 187 | .append_element_content_handler(element!("#page-log td[title] img, #m-metainfo td img[alt]", |element| { |
| 188 | if markdown.get() != 0 { return Ok(()); } |
| 189 | let profile = element.get_attribute("alt").and_then(|alt| alt.strip_prefix("gravatar for ").and_then(|email| profiles.get(&format!("email:{}",email.trim().to_lowercase())))); |
| 190 | if element.get_attribute("alt").is_some() && commit_images.borrow().is_none() { |
| 191 | *commit_images.borrow_mut() = profile.and_then(|profile| profile.url.as_ref().map(|url| (url.clone(),profile.picture.clone()))); |
| 192 | } |
| 193 | if let Some((_,Some(picture))) = &*commit_images.borrow() { |
| 194 | element.set_attribute("src",picture)?; |
| 195 | element.set_attribute("width", "14")?; |
| 196 | element.set_attribute("height", "14")?; |
| 197 | element.set_attribute("alt", "")?; |
| 198 | element.set_attribute("referrerpolicy", "no-referrer")?; |
| 199 | } |
| 200 | Ok(()) |
| 201 | })) |
| 202 | .append_element_content_handler(element!("#m-metainfo td span", |element| { |
| 203 | if markdown.get() != 0 { return Ok(()); } |
| 204 | if let Some((url,_)) = commit_names.borrow_mut().take() { |
| 205 | element.prepend(&format!("<a class=\"usa-link\" href=\"{url}\">"),ContentType::Html); |
| 206 | element.append("</a>",ContentType::Html); |
| 207 | } |
| 208 | Ok(()) |
| 209 | })) |
| 210 | .append_element_content_handler(text!("#m-metainfo td span", |chunk| { |
| 211 | if markdown.get() != 0 { return Ok(()); } |
| 212 | signature.borrow_mut().push_str(chunk.as_str()); |
| 213 | if !chunk.last_in_text_node() { chunk.remove(); return Ok(()); } |
| 214 | let value = std::mem::take(&mut *signature.borrow_mut()); |
| 215 | let profile = value.trim().strip_prefix("Signed by SSH key ").and_then(|fingerprint| profiles.get(&format!("key:{fingerprint}"))); |
| 216 | if let Some(profile) = profile { |
| 217 | let name = profile.name.replace('&',"&").replace('<',"<").replace('>',">"); |
| 218 | let original = &value; |
| 219 | chunk.replace(&format!("{original} (<a class=\"usa-link\" href=\"{}\">{name}</a>)",profile.url.as_deref().unwrap_or("/")),ContentType::Html); |
| 220 | } else { chunk.replace(&value,ContentType::Html); } |
| 221 | Ok(()) |
| 222 | })) |
| 223 | .append_element_content_handler(element!(".markdown", move |element| { |
| 224 | bodies.set(bodies.get() + 1); |
| 225 | let bodies = bodies.clone(); |
| 226 | element.on_end_tag(lol_html::end_tag!(move |_| { bodies.set(bodies.get() - 1); Ok(()) }))?; |
| 227 | Ok(()) |
| 228 | })) |
| 229 | .append_element_content_handler(element!(".timeline > li.spanspacer, .timeline .n-card__header > .spanspacer", |element| { |
| 230 | let style = element.get_attribute("style").unwrap_or_default(); |
| 231 | element.set_attribute("style", &format!("{style};flex-wrap:wrap"))?; |
| 232 | Ok(()) |
| 233 | })) |
| 234 | .append_element_content_handler(element!("#page-issues table td:nth-child(4)", move |element| { |
| 235 | author_column.set(true); |
| 236 | let author_column = author_column.clone(); |
| 237 | element.on_end_tag(lol_html::end_tag!(move |_| { author_column.set(false); Ok(()) }))?; |
| 238 | Ok(()) |
| 239 | })) |
| 240 | .append_element_content_handler(element!("img[src]", |element| { |
| 241 | if markdown.get() == 0 |
| 242 | && let Some(url) = element.get_attribute("src").and_then(|src| url::Url::parse(&src).ok()) |
| 243 | && url.host_str() == Some("api.dicebear.com") |
| 244 | && let Some(picture) = url.query_pairs().find(|(key, _)| key == "seed").and_then(|(_, seed)| avatars.get(seed.as_ref())) |
| 245 | { |
| 246 | element.set_attribute("src", picture)?; |
| 247 | element.set_attribute("alt", "")?; |
| 248 | element.set_attribute("referrerpolicy", "no-referrer")?; |
| 249 | } |
| 250 | Ok(()) |
| 251 | })) |
| 252 | .append_element_content_handler(element!("a[href]", |element| { |
| 253 | let target = element.get_attribute("href").and_then(|href| origin.join(&href).ok()).filter(|target| target.origin() == origin.origin()); |
| 254 | let ai_profile = target.as_ref().filter(|target| markdown.get() == 0 && target.path().trim_end_matches('/') == "/~ai") |
| 255 | .map(|_| ai_profile(None)); |
| 256 | let profile = ai_profile.as_ref().or_else(|| target.and_then(|target| account_path(target.path()).and_then(|name| profiles.get(name))).filter(|profile| profile.owner.is_none() || markdown.get() == 0)); |
| 257 | active.set(profile.map(|profile| profile.picture.is_some() || ai_profile.is_some() || profile.owner.is_some())); |
| 258 | let Some(profile) = profile else { return Ok(()); }; |
| 259 | let closing = active.clone(); |
| 260 | element.on_end_tag(lol_html::end_tag!(move |_| { closing.set(None); Ok(()) }))?; |
| 261 | if let Some(url) = &profile.url { |
| 262 | element.set_attribute("href", url)?; |
| 263 | element.set_attribute("target", "_blank")?; |
| 264 | element.set_attribute("rel", "noreferrer")?; |
| 265 | } else { |
| 266 | for attribute in ["href", "target", "rel", "tabindex"] { element.remove_attribute(attribute); } |
| 267 | } |
| 268 | if profile.owner.is_some() || ai_profile.is_some() { element.set_tag_name("span")?; } |
| 269 | let style = element.get_attribute("style").unwrap_or_default(); |
| 270 | let cursor = if profile.url.is_some() { "" } else { ";cursor:default" }; |
| 271 | let spacing = if profile.owner.is_some() { format!(";display:inline-flex;align-items:baseline;flex-wrap:{};gap:.25em;margin-right:.25em;max-width:100%", if compact.get() { "nowrap" } else { "wrap" }) } else { String::new() }; |
| 272 | element.set_attribute("style", &format!("{style};text-decoration:none{cursor}{spacing}"))?; |
| 273 | if let Some(picture) = &profile.picture { |
| 274 | element.prepend(&avatar(picture), ContentType::Html); |
| 275 | } |
| 276 | if profile.owner.is_some() { element.append("<span class=\"va-middle-childs\" style=\"white-space:nowrap\">", ContentType::Html); } |
| 277 | let dimensions = if ai_profile.is_some() || profile.owner.is_some() { "width:1em;height:1em;margin-right:.25em" } else { "width:.85em;height:.85em;vertical-align:-.1em;margin-right:.05em" }; |
| 278 | let icon = profile.icon.trim().replace("<svg ", &format!("<svg aria-hidden=\"true\" focusable=\"false\" style=\"color:inherit;stroke:currentColor;{dimensions}\" ")); |
| 279 | let icon = if profile.url.is_some() { icon.replace("<path ", "<path style=\"fill:currentColor;stroke:none\" ") } else { icon }; |
| 280 | element.append(&icon, ContentType::Html); |
| 281 | if profile.url.is_some() { |
| 282 | element.append("<span style=\"text-decoration:underline;text-decoration-color:currentColor\">", ContentType::Html); |
| 283 | } |
| 284 | let name = if compact.get() && profile.owner.is_some() { |
| 285 | let number = profile.name.find(|c: char| c.is_ascii_digit()).unwrap_or(profile.name.len()); |
| 286 | profile.name[..number].trim_end_matches([' ', '-']) |
| 287 | } else { &profile.name }; |
| 288 | element.append(name, ContentType::Text); |
| 289 | if profile.url.is_some() { element.append("</span>", ContentType::Html); } |
| 290 | if let Some((owner, picture)) = &profile.owner { |
| 291 | element.append("</span> <small style=\"font-size:.8em;white-space:nowrap\">(", ContentType::Html); |
| 292 | if !compact.get() { element.append("on behalf of ", ContentType::Text); } |
| 293 | let mut owner_url = origin.clone(); |
| 294 | owner_url.set_path(&format!("/~{owner}")); owner_url.set_query(None); owner_url.set_fragment(None); |
| 295 | let href = owner_url.path().replace('&', "&").replace('"', """).replace('<', "<"); |
| 296 | element.append(&format!("<a class=\"usa-link va-middle-childs\" href=\"{href}\">"), ContentType::Html); |
| 297 | if let Some(picture) = picture { element.append(&avatar(picture), ContentType::Html); } |
| 298 | element.append("<span>", ContentType::Html); |
| 299 | element.append(&format!("~{owner}"), ContentType::Text); |
| 300 | element.append("</span></a>)</small> ", ContentType::Html); |
| 301 | } |
| 302 | Ok(()) |
| 303 | })) |
| 304 | .append_element_content_handler(element!("a[href] img, a[href] span", move |child| { |
| 305 | if let Some(replace_picture) = images.get() { |
| 306 | if child.tag_name() == "span" { child.remove_and_keep_content(); } |
| 307 | else if replace_picture { child.remove(); } |
| 308 | else { |
| 309 | child.set_attribute("alt", "")?; |
| 310 | child.set_attribute("referrerpolicy", "no-referrer")?; |
| 311 | } |
| 312 | } |
| 313 | Ok(()) |
| 314 | })) |
| 315 | .append_element_content_handler(text!("a[href]", move |text| { |
| 316 | if labels.get().is_some() { text.remove(); } |
| 317 | Ok(()) |
| 318 | })))?) |
| 319 | } |
| 320 | |
| 321 | fn strip_hop_headers(headers: &mut HeaderMap) { |
| 322 | let named: Vec<_> = headers |
| 323 | .get_all("connection") |
| 324 | .iter() |
| 325 | .filter_map(|v| v.to_str().ok()) |
| 326 | .flat_map(|v| v.split(',')) |
| 327 | .map(|v| v.trim().to_owned()) |
| 328 | .collect(); |
| 329 | for name in named { |
| 330 | headers.remove(name); |
| 331 | } |
| 332 | for name in [ |
| 333 | "connection", |
| 334 | "keep-alive", |
| 335 | "proxy-authenticate", |
| 336 | "proxy-authorization", |
| 337 | "te", |
| 338 | "trailer", |
| 339 | "transfer-encoding", |
| 340 | "upgrade", |
| 341 | ] { |
| 342 | headers.remove(name); |
| 343 | } |
| 344 | } |
| 345 | |
| 346 | pub async fn proxy(app: Arc<App>, request: Request) -> Result<Response> { |
| 347 | let (mut parts, body) = request.into_parts(); |
| 348 | let host = parts |
| 349 | .headers |
| 350 | .get("host") |
| 351 | .and_then(|v| v.to_str().ok()) |
| 352 | .unwrap_or_default(); |
| 353 | let domain = env("STUDIO_DOMAIN", "studio.test"); |
| 354 | let service = host.strip_suffix(&format!(".{domain}")).unwrap_or_default().to_owned(); |
| 355 | if service != "shale" |
| 356 | && !(service |
| 357 | .strip_prefix("shale-preview-") |
| 358 | .is_some_and(|id| id.len() == 8 && id.bytes().all(|b| b.is_ascii_hexdigit()))) |
| 359 | { |
| 360 | return Err(Error::new(403, "Open Shale to continue.")); |
| 361 | } |
| 362 | let upstream: std::net::SocketAddr = parts |
| 363 | .headers |
| 364 | .get("studio-shale-upstream") |
| 365 | .and_then(|v| v.to_str().ok()) |
| 366 | .ok_or_else(|| Error::new(403, "Open Shale to continue."))? |
| 367 | .parse()?; |
| 368 | let uri: axum::http::Uri = parts |
| 369 | .headers |
| 370 | .get("studio-shale-uri") |
| 371 | .and_then(|v| v.to_str().ok()) |
| 372 | .ok_or_else(|| Error::new(403, "Open Shale to continue."))? |
| 373 | .parse()?; |
| 374 | if uri.scheme().is_some() || uri.authority().is_some() || !uri.path().starts_with('/') { |
| 375 | return Err(Error::new(403, "Open Shale to continue.")); |
| 376 | } |
| 377 | let origin = url::Url::parse(&format!("https://{host}{uri}"))?; |
| 378 | let profiles = profiles(&app.auth)?; |
| 379 | if matches!(parts.method, Method::GET | Method::HEAD) |
| 380 | && let Some(url) = account_path(uri.path()).and_then(|name| profiles.get(name)).and_then(|profile| profile.url.as_deref()) |
| 381 | { |
| 382 | return Ok(( |
| 383 | StatusCode::FOUND, |
| 384 | [ |
| 385 | ("location", url), |
| 386 | ("cache-control", "no-store"), |
| 387 | ("referrer-policy", "no-referrer"), |
| 388 | ], |
| 389 | ) |
| 390 | .into_response()); |
| 391 | } |
| 392 | let target = match &app.internal { |
| 393 | Some((base, _)) => format!("{base}services/{service}{uri}"), |
| 394 | None => format!("http://{upstream}{uri}"), |
| 395 | }; |
| 396 | if app.internal.is_some() { |
| 397 | parts.headers.remove("host"); |
| 398 | } |
| 399 | strip_hop_headers(&mut parts.headers); |
| 400 | for name in [ |
| 401 | "studio-shale-upstream", |
| 402 | "studio-shale-uri", |
| 403 | "studio-proxy-token", |
| 404 | "user-name", |
| 405 | "user-groups", |
| 406 | "user-id", |
| 407 | "if-none-match", |
| 408 | "if-modified-since", |
| 409 | "accept-encoding", |
| 410 | ] { |
| 411 | parts.headers.remove(name); |
| 412 | } |
| 413 | parts |
| 414 | .headers |
| 415 | .insert("accept-encoding", "identity".parse().unwrap()); |
| 416 | let mut upstream_response = app |
| 417 | .shale_http |
| 418 | .request(parts.method.clone(), target) |
| 419 | .headers(parts.headers) |
| 420 | .body(reqwest::Body::wrap_stream(body.into_data_stream())) |
| 421 | .send() |
| 422 | .await?; |
| 423 | let status = upstream_response.status(); |
| 424 | let mut headers = upstream_response.headers().clone(); |
| 425 | strip_hop_headers(&mut headers); |
| 426 | let html = parts.method != Method::HEAD |
| 427 | && headers |
| 428 | .get("content-type") |
| 429 | .and_then(|v| v.to_str().ok()) |
| 430 | .is_some_and(|v| { |
| 431 | v.split(';') |
| 432 | .next() |
| 433 | .unwrap_or_default() |
| 434 | .trim() |
| 435 | .eq_ignore_ascii_case("text/html") |
| 436 | }) |
| 437 | && headers |
| 438 | .get("content-encoding") |
| 439 | .is_none_or(|v| v == "identity"); |
| 440 | let mut prefix = Vec::new(); |
| 441 | if html { |
| 442 | while let Some(chunk) = upstream_response.chunk().await? { |
| 443 | prefix.extend_from_slice(&chunk); |
| 444 | if prefix.len() > 8 * 1024 * 1024 { |
| 445 | break; |
| 446 | } |
| 447 | } |
| 448 | if prefix.len() <= 8 * 1024 * 1024 |
| 449 | && let Ok(text) = std::str::from_utf8(&prefix) |
| 450 | { |
| 451 | let metadata = if service == "shale" && status.is_success() { |
| 452 | let (repository, accounts) = { |
| 453 | let document = scraper::Html::parse_document(text); |
| 454 | let repository = document.select(&scraper::Selector::parse("meta[name='astheno.shale.repo.name']").unwrap()).next().and_then(|e| e.attr("content")).map(str::to_owned); |
| 455 | let mut accounts: HashSet<String> = document.select(&scraper::Selector::parse("a[href]").unwrap()).filter_map(|element| { |
| 456 | let url = origin.join(element.attr("href")?).ok()?; |
| 457 | if url.origin() != origin.origin() { return None; } |
| 458 | account_path(url.path()).map(str::to_owned) |
| 459 | }).take(1024).collect(); |
| 460 | if accounts.len() < 1024 && let Some(name) = account_path(origin.path()) { accounts.insert(name.to_owned()); } |
| 461 | (repository, accounts) |
| 462 | }; |
| 463 | match tokio::time::timeout(Duration::from_secs(3), host::call(json!({"operation":"shale.metadata","repository":repository,"accounts":accounts}))).await { |
| 464 | Ok(Ok(value)) => value, |
| 465 | _ => Value::Null, |
| 466 | } |
| 467 | } else { Value::Null }; |
| 468 | let mut profiles = profiles; |
| 469 | let mut avatars = HashMap::new(); |
| 470 | if let Some(accounts) = metadata["accounts"].as_object() { |
| 471 | for (name, account) in accounts { |
| 472 | if let (Some(seed), Some(picture)) = (account["avatar"].as_str(), account["picture"].as_str()) { |
| 473 | avatars.insert(seed.to_owned(), picture.to_owned()); |
| 474 | } |
| 475 | let Some(owner) = account["owner"].as_str() else { continue; }; |
| 476 | let mut profile = ai_profile(account["model"].as_str()); |
| 477 | profile.owner = Some((owner.into(), account["picture"].as_str().map(str::to_owned))); |
| 478 | profiles.insert(name.clone(), profile); |
| 479 | } |
| 480 | } |
| 481 | let rewritten = rewrite(text, &origin, &profiles, &avatars, metadata["openIssues"].as_u64())?; |
| 482 | for name in [ |
| 483 | "content-length", |
| 484 | "etag", |
| 485 | "last-modified", |
| 486 | "content-md5", |
| 487 | "digest", |
| 488 | "content-digest", |
| 489 | "accept-ranges", |
| 490 | ] { |
| 491 | headers.remove(name); |
| 492 | } |
| 493 | headers.insert("cache-control", "no-store".parse().unwrap()); |
| 494 | return Ok((status, headers, rewritten).into_response()); |
| 495 | } |
| 496 | } |
| 497 | let stream = futures::stream::once(async move { Ok::<_, reqwest::Error>(Bytes::from(prefix)) }) |
| 498 | .chain(upstream_response.bytes_stream()); |
| 499 | Ok((status, headers, Body::from_stream(stream)).into_response()) |
| 500 | } |
| 501 | |
| 502 | #[cfg(test)] |
| 503 | mod tests { |
| 504 | use super::*; |
| 505 | use scraper::{Html, Selector}; |
| 506 | |
| 507 | #[test] |
| 508 | fn commit_identity_keeps_entities_and_cannot_escape_its_metadata_cell() { |
| 509 | let origin = url::Url::parse("https://shale.test/repo/commit/abc").unwrap(); |
| 510 | let profile = || Profile { name: "~clover".into(), url: Some("/~clover".into()), icon:"",picture:Some("https://snow.test/picture".into()),owner:None }; |
| 511 | let profiles = HashMap::from([("email:known@example.test".into(),profile()),("key:SHA256:known".into(),profile())]); |
| 512 | let html = "<body id='page-commit'><table id=m-metainfo><tr><td><img src=old alt='gravatar for known@example.test'></td></tr><tr><td><img src=unknown alt='gravatar for other@example.test'><span>John & Jane <other@example.test></span></td></tr><tr><td><span> Signed by SSH key SHA256:known</span></td></tr></table><div class=markdown><table id=m-metainfo><tr><td><img src=example alt='gravatar for known@example.test'><span> Signed by SSH key SHA256:known</span></td></tr></table></div></body>"; |
| 513 | let output = rewrite(html,&origin,&profiles,&HashMap::new(),None).unwrap(); |
| 514 | let document = Html::parse_document(&output); |
| 515 | assert!(!output.contains("&amp;")); |
| 516 | assert!(output.contains("John & Jane <other@example.test>")); |
| 517 | assert_eq!(document.select(&Selector::parse("a[href='/~clover']").unwrap()).count(),1); |
| 518 | assert!(output.contains("src=unknown")); |
| 519 | assert!(output.contains("src=example")); |
| 520 | assert_eq!(document.select(&Selector::parse(".markdown a").unwrap()).count(),0); |
| 521 | let html = "<body id=page-log><table><tbody><tr><td title='known@example.test'><img src=old><span>Author & name</span></td></tr><tr><td title='other@example.test'><img src=unknown><span>Other</span></td></tr></tbody></table></body>"; |
| 522 | let output = rewrite(html,&origin,&profiles,&HashMap::new(),None).unwrap(); |
| 523 | assert!(output.contains("Author & name")); |
| 524 | let document = Html::parse_document(&output); |
| 525 | assert_eq!(document.select(&Selector::parse("a[href='/~clover']").unwrap()).count(),1); |
| 526 | assert_eq!(document.select(&Selector::parse("img[src='https://snow.test/picture']").unwrap()).count(),1); |
| 527 | assert_eq!(document.select(&Selector::parse("img[src='https://snow.test/picture'][width='14'][height='14']").unwrap()).count(),1); |
| 528 | } |
| 529 | |
| 530 | #[test] |
| 531 | fn navigation_uses_the_current_repository_count() { |
| 532 | let html = "<header><ul class='usa-nav__primary'><li><a href='/chat/issues/'>issues (9)</a></li></ul></header><div class='markdown'><a href='/chat/issues/'>issues (9)</a></div>"; |
| 533 | let origin = url::Url::parse("https://shale.paperclover.net/chat/issues/25").unwrap(); |
| 534 | let output = rewrite(html, &origin, &HashMap::new(), &HashMap::new(), Some(8)).unwrap(); |
| 535 | assert!(output.contains(">issues (8)</a>")); |
| 536 | assert!(output.contains("<div class='markdown'><a href='/chat/issues/'>issues (9)</a></div>")); |
| 537 | } |
| 538 | |
| 539 | #[test] |
| 540 | fn reconstructed_markdown_anchors_cannot_relabel_humans() { |
| 541 | let html = "<div class='markdown'><b><a href='https://example.com'>external</b>again</a></div><a href='/~ai'>~ai</a><a href='/~clo'>~clo</a>"; |
| 542 | let origin = url::Url::parse("https://shale.paperclover.net/chat/issues/25").unwrap(); |
| 543 | let output = rewrite(html, &origin, &HashMap::new(), &HashMap::new(), None).unwrap(); |
| 544 | assert!(output.contains("<a href='/~clo'>~clo</a>")); |
| 545 | assert!(output.contains("</svg>ai")); |
| 546 | } |
| 547 | |
| 548 | #[test] |
| 549 | fn terminal_trailers_ignore_examples_and_unclosed_fences() { |
| 550 | assert_eq!(assisted_model("Body\n\nAssisted-by: gpt-6-astra\n"), Some("gpt-6-astra")); |
| 551 | for body in ["Assisted-by: <model-id>", "> Assisted-by: claude-opus-5.5", " Assisted-by: gpt-6-astra", "```\nAssisted-by: gpt-6-astra", "~~~~\n```\nAssisted-by: gpt-6-astra", "Assisted-by: gpt-6-astra\nMore text", "Assisted-by: gpt-6-astra extra"] { |
| 552 | assert_eq!(assisted_model(body), None, "{body}"); |
| 553 | } |
| 554 | assert_eq!(assisted_model("```\nexample\n```\n\nAssisted-by: claude-opus-5.5"), Some("claude-opus-5.5")); |
| 555 | } |
| 556 | |
| 557 | #[test] |
| 558 | fn ai_identity_is_durable_and_never_reads_comment_trailers() { |
| 559 | let origin = url::Url::parse("https://shale.paperclover.net/chat/issues/25").unwrap(); |
| 560 | let mut opus = ai_profile(Some("claude-opus-5.5")); |
| 561 | opus.owner = Some(("clover".into(), Some("https://example.com/picture?a=1&b=2".into()))); |
| 562 | let mut sol = ai_profile(Some("gpt-6.1-sol")); |
| 563 | sol.owner = Some(("clover".into(), None)); |
| 564 | let profiles = HashMap::from([("ai-clover-claude-opus-5.5".into(), opus), ("ai-clover-gpt-6.1-sol".into(), sol)]); |
| 565 | let html = "<a href='/~ai-clover-claude-opus-5.5'><img src='old'><span>ai account</span></a><a href='/~ai-clover-gpt-6.1-sol'>ai account</a><div class='markdown'><a href='/~ai-clover-claude-opus-5.5'>Mention</a><p>Assisted-by: unknown-model</p></div><a href='/~clover'>~clover</a>"; |
| 566 | let output = rewrite(html, &origin, &profiles, &HashMap::new(), None).unwrap(); |
| 567 | let document = Html::parse_document(&output); |
| 568 | assert!(output.contains("Opus 5.5</span> <small")); |
| 569 | assert!(output.contains(">(on behalf of <a class=")); |
| 570 | assert!(output.contains("GPT-6.1 Sol</span> <small")); |
| 571 | assert!(!output.contains("src=\"old\"")); |
| 572 | assert!(output.contains("<a href='/~clover'>~clover</a>")); |
| 573 | assert!(output.contains(">Mention</a>")); |
| 574 | assert_eq!(document.select(&Selector::parse("a[href='/~clover']").unwrap()).count(), 3); |
| 575 | assert!(document.select(&Selector::parse("svg").unwrap()).all(|e| e.attr("aria-hidden") == Some("true") && e.attr("focusable") == Some("false"))); |
| 576 | assert_eq!(ai_profile(Some("gpt-6.1-sol")).name, "GPT-6.1 Sol"); |
| 577 | assert_eq!(ai_profile(Some("claude-opus-5.5")).name, "Opus 5.5"); |
| 578 | assert_eq!(ai_profile(Some("claude-haiku-5-5")).name, "Haiku 5.5"); |
| 579 | assert_eq!(ai_profile(Some("claude-haiku-5.5")).name, "Haiku 5.5"); |
| 580 | assert_eq!(ai_profile(Some("claude-sonnet-4-5-20250929")).name, "Sonnet 4.5"); |
| 581 | for (id, name) in [ |
| 582 | ("claude-new-variant-5-5", "New Variant 5.5"), |
| 583 | ("gpt-6-1-new-variant", "GPT-6.1 New Variant"), |
| 584 | ("qwen-3.8-27b", "Qwen 3.8 27b"), |
| 585 | ("qwen3.8-27b", "Qwen 3.8 27b"), |
| 586 | ("big-pickle-2", "Big Pickle 2"), |
| 587 | ("llama-3-70", "Llama 3 70"), |
| 588 | ] { assert_eq!(ai_profile(Some(id)).name, name, "{id}"); } |
| 589 | } |
| 590 | |
| 591 | #[test] |
| 592 | fn compact_authors_keep_native_owner_links_and_avatars() { |
| 593 | let origin = url::Url::parse("https://shale.paperclover.net/chat/issues/").unwrap(); |
| 594 | let mut profile = ai_profile(Some("gpt-6.1-sol")); |
| 595 | profile.owner = Some(("clover".into(), Some("https://snowglobe.test/picture".into()))); |
| 596 | let profiles = HashMap::from([("ai-clover-gpt-6.1-sol".into(), profile)]); |
| 597 | let avatars = HashMap::from([("native-clover".into(), "https://snowglobe.test/picture".into())]); |
| 598 | let html = "<body id='page-issues'><table><tr><td>1</td><td>Task</td><td>Todo</td><td><a class='usa-link va-middle-childs' href='/~ai-clover-gpt-6.1-sol'>ai</a></td></tr></table><a class='usa-link va-middle-childs' href='/~clover'><img height='16' src='https://api.dicebear.com/10.x/line-face/svg?seed=native-clover'><span>~clover</span></a><div class='markdown'><img src='https://api.dicebear.com/10.x/line-face/svg?seed=native-clover'></div>"; |
| 599 | let output = rewrite(html, &origin, &profiles, &avatars, None).unwrap(); |
| 600 | let document = Html::parse_document(&output); |
| 601 | let author = document.select(&Selector::parse("td:nth-child(4)").unwrap()).next().unwrap(); |
| 602 | assert_eq!(author.text().collect::<String>().trim(), "GPT (~clover)"); |
| 603 | assert!(!output.contains("on behalf of")); |
| 604 | let owners: Vec<_> = document.select(&Selector::parse("a[href='/~clover']").unwrap()).collect(); |
| 605 | assert_eq!(owners.len(), 2); |
| 606 | for owner in owners { |
| 607 | assert_eq!(owner.attr("class"), Some("usa-link va-middle-childs")); |
| 608 | let image = owner.select(&Selector::parse("img").unwrap()).next().unwrap(); |
| 609 | assert_eq!(image.attr("src"), Some("https://snowglobe.test/picture")); |
| 610 | assert_eq!(image.attr("height"), Some("16")); |
| 611 | assert_eq!(image.attr("style"), None); |
| 612 | assert_eq!(image.attr("alt"), Some("")); |
| 613 | } |
| 614 | assert!(output.contains("<div class='markdown'><img src='https://api.dicebear.com/10.x/line-face/svg?seed=native-clover'>")); |
| 615 | for (model, short) in [("qwen3.8-27b", "Qwen"), ("big-pickle-2", "Big Pickle"), ("claude-new-variant-5-5", "New Variant")] { |
| 616 | let mut profile = ai_profile(Some(model)); |
| 617 | profile.owner = Some(("clover".into(), None)); |
| 618 | let profiles = HashMap::from([("ai-model".into(), profile)]); |
| 619 | let html = "<body id='page-issues'><div class='grid-container'><table><tr><td>1</td><td>Task</td><td>Todo</td><td><a href='/~ai-model'>ai</a></td></tr></table></div></body>"; |
| 620 | let output = rewrite(html, &origin, &profiles, &HashMap::new(), None).unwrap(); |
| 621 | let document = Html::parse_document(&output); |
| 622 | assert_eq!(document.select(&Selector::parse("td:nth-child(4)").unwrap()).next().unwrap().text().collect::<String>().trim(), format!("{short} (~clover)")); |
| 623 | } |
| 624 | } |
| 625 | |
| 626 | #[test] |
| 627 | fn rewrites_only_known_local_account_links_and_escapes_provider_names() { |
| 628 | let origin = url::Url::parse("https://shale.paperclover.net/").unwrap(); |
| 629 | let profiles = HashMap::from([ |
| 630 | ( |
| 631 | "guest-github-123".into(), |
| 632 | Profile { |
| 633 | name: "<&\"clover".into(), |
| 634 | url: Some("https://github.com/clover".into()), |
| 635 | icon: include_str!("../web/sso/github.svg"), |
| 636 | owner: None, |
| 637 | picture: Some("https://avatars.githubusercontent.com/u/123?s=64".into()), |
| 638 | }, |
| 639 | ), |
| 640 | ( |
| 641 | "guest-astheno-abc".into(), |
| 642 | Profile { |
| 643 | name: "Astheno user".into(), |
| 644 | url: None, |
| 645 | icon: include_str!("astheno.svg"), |
| 646 | owner: None, |
| 647 | picture: None, |
| 648 | }, |
| 649 | ), |
| 650 | ]); |
| 651 | let html = r#"<a class="usa-nav-link" href="/~guest-github-123"><img src="avatar"><span>~guest-github-123</span></a><a href="https://shale.paperclover.net/~guest-astheno-abc/"><img src="astheno-avatar" alt="avatar">guest</a><a href="/~guest-github-unknown">unknown</a><a href="https://other.example/~guest-github-123">external</a><a href="/~guest-github-123/repo">repo</a><pre>~guest-github-123</pre>"#; |
| 652 | let rewritten = rewrite(html, &origin, &profiles, &HashMap::new(), None).unwrap(); |
| 653 | let page = Html::parse_document(&rewritten); |
| 654 | let links: Vec<_> = page.select(&Selector::parse("body > a").unwrap()).collect(); |
| 655 | assert_eq!(links[0].attr("target"), Some("_blank")); |
| 656 | assert_eq!(links[0].attr("rel"), Some("noreferrer")); |
| 657 | assert!(links[0].attr("style").unwrap().contains("text-decoration:none")); |
| 658 | assert!(links[0].select(&Selector::parse("span[style]").unwrap()).next().unwrap().attr("style").unwrap().contains("text-decoration:underline")); |
| 659 | assert_eq!(links[1].value().name(), "a"); |
| 660 | assert_eq!(links[1].attr("target"), None); |
| 661 | assert_eq!(links[1].attr("rel"), None); |
| 662 | assert_eq!(links[1].text().collect::<String>().trim(), "Astheno user"); |
| 663 | for link in &links[..2] { |
| 664 | assert_eq!( |
| 665 | link.select(&Selector::parse("svg[aria-hidden=true]").unwrap()) |
| 666 | .count(), |
| 667 | 1 |
| 668 | ); |
| 669 | } |
| 670 | assert_eq!( |
| 671 | links[1] |
| 672 | .select(&Selector::parse("img").unwrap()) |
| 673 | .next() |
| 674 | .unwrap() |
| 675 | .attr("src"), |
| 676 | Some("astheno-avatar") |
| 677 | ); |
| 678 | assert_eq!(links[0].select(&Selector::parse("img").unwrap()).count(), 1); |
| 679 | assert_eq!( |
| 680 | links[0] |
| 681 | .select(&Selector::parse("img").unwrap()) |
| 682 | .next() |
| 683 | .unwrap() |
| 684 | .attr("src"), |
| 685 | Some("https://avatars.githubusercontent.com/u/123?s=64") |
| 686 | ); |
| 687 | assert_eq!(links[0].text().collect::<String>().trim(), "<&\"clover"); |
| 688 | assert_eq!(links[0].attr("class"), Some("usa-nav-link")); |
| 689 | assert_eq!(links[1].attr("href"), None); |
| 690 | for link in &links[2..] { |
| 691 | assert_eq!(link.attr("target"), None); |
| 692 | } |
| 693 | assert!(rewritten.contains("<pre>~guest-github-123</pre>")); |
| 694 | assert_eq!(account_path("/~guest-github-123//"), None); |
| 695 | } |
| 696 | } |