This repository has no description
1use std::collections::{HashMap, HashSet};
2
3use axum::{
4 Json,
5 body::{Body, to_bytes},
6 extract::State,
7 http::{Request, StatusCode},
8 response::Response,
9};
10use bobbin_types::ids::{EdgeKey, SubjectRef, nsid_static};
11use futures::StreamExt;
12use jacquard_common::DefaultStr;
13use jacquard_common::IntoStatic;
14use jacquard_common::types::did::Did;
15use jacquard_common::types::ident::AtIdentifier;
16use jacquard_common::types::nsid::Nsid;
17use jacquard_common::types::string::AtUri;
18use serde::{Deserialize, Serialize};
19use serde_json::{Map, Value, json};
20use tower::ServiceExt;
21
22use crate::recordpath::{RecordPath, walk_path};
23use crate::{AppState, SubjectShape, XrpcError, mirror_kind, subject_shape};
24
25pub const TYPE_COUNT: &str = "sh.tangled.query.enrichResponse#count";
26pub const TYPE_DISTINCT_AUTHORS: &str = "sh.tangled.query.enrichResponse#distinctAuthors";
27pub const TYPE_VIEWER: &str = "sh.tangled.query.enrichResponse#viewer";
28/// vendored from upstream com.bad-example.identity.resolveMiniDoc's output
29pub const TYPE_MINIDOC: &str = "com.bad-example.identity.miniDoc";
30
31const KNOWN_TYPES: [&str; 4] = [TYPE_COUNT, TYPE_DISTINCT_AUTHORS, TYPE_VIEWER, TYPE_MINIDOC];
32const MINIDOC_CONCURRENCY: usize = 32;
33
34/// a payload type nsid, with an optional #fragment for lexicon defs. the raw
35/// string is kept because it echoes into the data map as the payload key
36#[derive(Debug, Clone, PartialEq, Eq, Hash)]
37pub struct PayloadType {
38 raw: String,
39 nsid: Nsid<DefaultStr>,
40 #[allow(dead_code)]
41 fragment: Option<String>,
42}
43
44impl PayloadType {
45 pub fn parse(raw: &str) -> Result<Self, String> {
46 let (nsid, fragment) = match raw.split_once('#') {
47 Some((nsid, fragment)) => {
48 let valid = !fragment.is_empty()
49 && fragment.starts_with(|c: char| c.is_ascii_alphabetic())
50 && fragment.chars().all(|c| c.is_ascii_alphanumeric());
51 if !valid {
52 return Err(format!("invalid fragment #{fragment}"));
53 }
54 (nsid, Some(fragment.to_owned()))
55 }
56 None => (raw, None),
57 };
58 let nsid = Nsid::new_owned(nsid).map_err(|e| format!("invalid nsid: {e}"))?;
59 Ok(Self {
60 raw: raw.to_owned(),
61 nsid,
62 fragment,
63 })
64 }
65
66 fn as_str(&self) -> &str {
67 &self.raw
68 }
69}
70
71impl<'de> Deserialize<'de> for PayloadType {
72 fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
73 where
74 D: serde::Deserializer<'de>,
75 {
76 let raw = String::deserialize(deserializer)?;
77 Self::parse(&raw).map_err(serde::de::Error::custom)
78 }
79}
80
81impl Serialize for PayloadType {
82 fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
83 where
84 S: serde::Serializer,
85 {
86 serializer.serialize_str(&self.raw)
87 }
88}
89
90/// the paths the edge index can back: a record's reference field, or its
91/// authoring repo
92#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
93pub enum SourcePath {
94 Subject,
95 Repo,
96}
97
98/// a "collection:path" link source. raw echoes into the data map verbatim
99#[derive(Debug, Clone, PartialEq, Eq, Hash)]
100pub struct LinkSource {
101 raw: String,
102 collection: Nsid<DefaultStr>,
103 path: SourcePath,
104}
105
106impl LinkSource {
107 pub fn parse(raw: &str) -> Result<Self, String> {
108 let (collection, path) = match raw.split_once(':') {
109 Some((collection, path)) if !collection.is_empty() && !path.is_empty() => {
110 (collection, path)
111 }
112 _ => return Err("expected \"collection:path\"".to_owned()),
113 };
114 let collection =
115 Nsid::new_owned(collection).map_err(|e| format!("invalid collection: {e}"))?;
116 let path = match path {
117 "subject" => SourcePath::Subject,
118 ".repo" => SourcePath::Repo,
119 _ if path.starts_with('.') => {
120 return Err("envelope field not index-backed; try .repo".to_owned());
121 }
122 _ => {
123 return Err(
124 "path not index-backed; only `subject` and `.repo` are supported".to_owned(),
125 );
126 }
127 };
128 Ok(Self {
129 raw: raw.to_owned(),
130 collection,
131 path,
132 })
133 }
134
135 fn as_str(&self) -> &str {
136 &self.raw
137 }
138}
139
140impl<'de> Deserialize<'de> for LinkSource {
141 fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
142 where
143 D: serde::Deserializer<'de>,
144 {
145 let raw = String::deserialize(deserializer)?;
146 Self::parse(&raw).map_err(serde::de::Error::custom)
147 }
148}
149
150impl Serialize for LinkSource {
151 fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
152 where
153 S: serde::Serializer,
154 {
155 serializer.serialize_str(&self.raw)
156 }
157}
158
159#[derive(Debug, Deserialize)]
160pub struct LinkDescriptor {
161 /// constellation link source
162 source: LinkSource,
163 /// nsid of the sidecar payload to produce for this source
164 #[serde(rename = "type")]
165 ty: PayloadType,
166}
167
168#[derive(Debug, Deserialize)]
169pub struct EnrichInput {
170 xrpc: String,
171 #[serde(default)]
172 params: Option<Map<String, Value>>,
173 enrich: Vec<LinkDescriptor>,
174 #[serde(default)]
175 sources: Option<Vec<String>>,
176 #[serde(default)]
177 viewer: Option<Did>,
178}
179
180pub async fn enrich(
181 State(state): State<AppState>,
182 Json(input): Json<EnrichInput>,
183) -> Result<Json<Value>, XrpcError> {
184 let mut seen_descriptors = HashSet::new();
185 let mut descriptors = Vec::new();
186 for descriptor in &input.enrich {
187 if !KNOWN_TYPES.contains(&descriptor.ty.as_str()) {
188 return Err(descriptor_error(&descriptor.source, "unknown enrich type"));
189 }
190 validate_source(&descriptor.source, &descriptor.ty)?;
191 if descriptor.ty.as_str() == TYPE_VIEWER && input.viewer.is_none() {
192 return Err(descriptor_error(
193 &descriptor.source,
194 "viewer payloads require a viewer param",
195 ));
196 }
197 if seen_descriptors.insert((&descriptor.source, &descriptor.ty)) {
198 descriptors.push(descriptor);
199 }
200 }
201
202 let sources = input
203 .sources
204 .as_ref()
205 .map(|sources| {
206 sources
207 .iter()
208 .enumerate()
209 .map(|(i, path)| {
210 RecordPath::parse(path).map_err(|e| {
211 XrpcError::InvalidParams(format!("sources[{i}] {path:?}: {e}"))
212 })
213 })
214 .collect::<Result<Vec<_>, _>>()
215 })
216 .transpose()?;
217
218 let inner = run_inner(&state, &input.xrpc, input.params.unwrap_or_default()).await?;
219 if descriptors.is_empty() {
220 return Ok(Json(json!({ "output": inner, "data": {} })));
221 }
222
223 let mut refs: Vec<SubjectRef> = Vec::new();
224 let mut seen: HashSet<SubjectRef> = HashSet::new();
225 match &sources {
226 Some(sources) => {
227 for path in sources {
228 for node in walk_path(path, [&inner]) {
229 collect_ref(node, &mut refs, &mut seen);
230 }
231 }
232 }
233 None => discover_refs(&inner, &mut refs, &mut seen),
234 }
235 // the authority of every at-uri is a reference too: it's embedded in the
236 // response, and it's the only way record authors get payloads
237 let authorities: Vec<SubjectRef> = refs
238 .iter()
239 .filter_map(|reference| repo_did(reference).map(SubjectRef::Did))
240 .collect();
241 for authority in authorities {
242 if seen.insert(authority.clone()) {
243 refs.push(authority);
244 }
245 }
246
247 let mut data = Map::new();
248 let mut minidoc_targets: HashMap<Did<DefaultStr>, Vec<LinkSource>> = HashMap::new();
249 for descriptor in descriptors {
250 for reference in &refs {
251 let Some(subject) = applicable_subject(descriptor, reference)? else {
252 continue;
253 };
254 match descriptor.ty.as_str() {
255 TYPE_COUNT | TYPE_DISTINCT_AUTHORS | TYPE_VIEWER => {
256 let kind = edge_kind(&descriptor.source)?;
257 let key = EdgeKey::new(kind, subject.clone());
258 let result = match descriptor.ty.as_str() {
259 TYPE_COUNT => Value::from(state.edges.count(&key)),
260 TYPE_DISTINCT_AUTHORS => {
261 Value::from(state.edges.count_distinct_authors(&key))
262 }
263 _ => {
264 // viewer presence is validated up front
265 let viewer = input.viewer.as_ref().expect("viewer param present");
266 match state.edges.viewer_source(&key, viewer.as_str()) {
267 Some(uri) => Value::String(uri.to_string()),
268 None => Value::Null,
269 }
270 }
271 };
272 put(
273 &mut data,
274 subject.as_str(),
275 descriptor.source.as_str(),
276 descriptor.ty.as_str(),
277 result,
278 );
279 }
280 TYPE_MINIDOC => {
281 if let Some(target) = repo_did(&subject) {
282 let sources = minidoc_targets.entry(target).or_default();
283 if !sources.contains(&descriptor.source) {
284 sources.push(descriptor.source.clone());
285 }
286 }
287 }
288 _ => unreachable!("validated up front"),
289 }
290 }
291 }
292
293 let docs = resolve_minidocs(&state, minidoc_targets).await;
294 for (target, sources, doc) in docs {
295 for source in sources {
296 put(
297 &mut data,
298 target.as_str(),
299 source.as_str(),
300 TYPE_MINIDOC,
301 doc.clone(),
302 );
303 }
304 }
305
306 Ok(Json(json!({ "output": inner, "data": data })))
307}
308
309fn put(data: &mut Map<String, Value>, reference: &str, source: &str, ty: &str, payload: Value) {
310 let per_ref = data
311 .entry(reference.to_owned())
312 .or_insert_with(|| Value::Object(Map::new()));
313 let Value::Object(per_ref) = per_ref else {
314 return;
315 };
316 let per_source = per_ref
317 .entry(source.to_owned())
318 .or_insert_with(|| Value::Object(Map::new()));
319 if let Value::Object(per_source) = per_source {
320 per_source.insert(ty.to_owned(), payload);
321 }
322}
323
324/// we drop failures, the client falls back to resolveMiniDoc for misses
325async fn resolve_minidocs(
326 state: &AppState,
327 targets: HashMap<Did<DefaultStr>, Vec<LinkSource>>,
328) -> Vec<(Did<DefaultStr>, Vec<LinkSource>, Value)> {
329 futures::stream::iter(targets)
330 .map(|(did, sources)| async move {
331 let doc = state
332 .slingshot
333 .resolve_mini_doc(&AtIdentifier::Did(did.clone()))
334 .await
335 .ok()
336 .and_then(|bytes| serde_json::from_slice::<Value>(&bytes).ok());
337 (did, sources, doc)
338 })
339 .buffer_unordered(MINIDOC_CONCURRENCY)
340 .filter_map(|(did, sources, doc)| async move { doc.map(|doc| (did, sources, doc)) })
341 .collect()
342 .await
343}
344
345fn descriptor_error(source: &LinkSource, msg: &str) -> XrpcError {
346 XrpcError::InvalidParams(format!("enrich {}: {msg}", source.as_str()))
347}
348
349fn validate_source(source: &LinkSource, ty: &PayloadType) -> Result<(), XrpcError> {
350 match source.path {
351 SourcePath::Subject => subject_shape(source.collection.as_ref())
352 .map(|_| ())
353 .ok_or_else(|| descriptor_error(source, "unknown collection")),
354 // minidocs read the uri's authority straight out of the response, only
355 // edge payloads need the author index
356 SourcePath::Repo if ty.as_str() == TYPE_MINIDOC => Ok(()),
357 SourcePath::Repo => mirror_kind(source.collection.as_ref())
358 .map(|_| ())
359 .ok_or_else(|| descriptor_error(source, "collection has no author index")),
360 }
361}
362
363fn repo_did(reference: &SubjectRef) -> Option<Did<DefaultStr>> {
364 match reference {
365 SubjectRef::Did(did) => Some(did.clone()),
366 SubjectRef::Uri(uri) => match uri.authority() {
367 AtIdentifier::Did(did) => Some(did.clone().into_static()),
368 AtIdentifier::Handle(_) => None,
369 },
370 SubjectRef::Global => None,
371 }
372}
373
374/// the subject an edge lookup or payload should be keyed by, or None if the
375/// descriptor's shape doesn't apply to this ref, eg. a did-only source asked
376/// about an at-uri. `.repo` sources normalize to the repo did
377fn applicable_subject(
378 descriptor: &LinkDescriptor,
379 reference: &SubjectRef,
380) -> Result<Option<SubjectRef>, XrpcError> {
381 match descriptor.source.path {
382 SourcePath::Subject => {
383 let (_, shape) = subject_shape(descriptor.source.collection.as_ref())
384 .ok_or_else(|| descriptor_error(&descriptor.source, "unknown collection"))?;
385 Ok(shape_accepts(shape, reference).then(|| reference.clone()))
386 }
387 SourcePath::Repo => Ok(repo_did(reference).map(SubjectRef::Did)),
388 }
389}
390
391/// the edge kind behind a source for stats payloads
392fn edge_kind(source: &LinkSource) -> Result<Nsid<DefaultStr>, XrpcError> {
393 let kind = match source.path {
394 SourcePath::Subject => subject_shape(source.collection.as_ref())
395 .map(|(nsid, _)| nsid)
396 .ok_or_else(|| descriptor_error(source, "unknown collection"))?,
397 SourcePath::Repo => mirror_kind(source.collection.as_ref())
398 .ok_or_else(|| descriptor_error(source, "collection has no author index"))?,
399 };
400 Ok(nsid_static(kind))
401}
402
403/// mismatch means skip, not reject
404fn shape_accepts(shape: SubjectShape, reference: &SubjectRef) -> bool {
405 match (shape, reference) {
406 (SubjectShape::BareDid, SubjectRef::Did(_)) => true,
407 (SubjectShape::Collection(expected), SubjectRef::Uri(uri)) => {
408 uri.collection().is_some_and(|c| c.as_ref() == expected)
409 }
410 (SubjectShape::OneOfCollections(allowed), SubjectRef::Uri(uri))
411 | (SubjectShape::BareDidOrOneOfCollections(allowed), SubjectRef::Uri(uri)) => uri
412 .collection()
413 .is_some_and(|c| allowed.contains(&c.as_ref())),
414 (SubjectShape::BareDidOrOneOfCollections(_), SubjectRef::Did(_)) => true,
415 (SubjectShape::AnyAtUri, SubjectRef::Uri(_)) => true,
416 _ => false,
417 }
418}
419
420/// a value counts as a reference if it is a did string or an at-uri string
421/// with collection and rkey, or a {uri, cid} strong ref
422fn collect_ref(value: &Value, refs: &mut Vec<SubjectRef>, seen: &mut HashSet<SubjectRef>) {
423 let candidate = match value {
424 Value::String(s) => Some(s.as_str()),
425 Value::Object(o) => o
426 .get("uri")
427 .and_then(Value::as_str)
428 .filter(|_| o.contains_key("cid")),
429 _ => None,
430 };
431 let Some(candidate) = candidate else { return };
432 let reference = if candidate.starts_with("did:") {
433 Did::<DefaultStr>::new_owned(candidate)
434 .ok()
435 .map(SubjectRef::Did)
436 } else if candidate.starts_with("at://") {
437 AtUri::<DefaultStr>::new_owned(candidate)
438 .ok()
439 .filter(|u| {
440 matches!(u.authority(), AtIdentifier::Did(_))
441 && u.collection().is_some()
442 && u.rkey().is_some()
443 })
444 .map(SubjectRef::Uri)
445 } else {
446 None
447 };
448 if let Some(reference) = reference
449 && seen.insert(reference.clone())
450 {
451 refs.push(reference);
452 }
453}
454
455fn discover_refs(value: &Value, refs: &mut Vec<SubjectRef>, seen: &mut HashSet<SubjectRef>) {
456 collect_ref(value, refs, seen);
457 match value {
458 Value::Array(items) => {
459 for item in items {
460 discover_refs(item, refs, seen);
461 }
462 }
463 Value::Object(map) => {
464 for v in map.values() {
465 discover_refs(v, refs, seen);
466 }
467 }
468 _ => {}
469 }
470}
471
472async fn run_inner(
473 state: &AppState,
474 nsid: &str,
475 params: Map<String, Value>,
476) -> Result<Value, XrpcError> {
477 let qs = encode_params(¶ms);
478 let uri = format!("/xrpc/{nsid}?{qs}");
479 let request = Request::builder()
480 .method("GET")
481 .uri(&uri)
482 .body(Body::empty())
483 .map_err(|e| XrpcError::Internal(format!("inner request: {e}")))?;
484 let response = state
485 .self_router()
486 .oneshot(request)
487 .await
488 .map_err(|e| XrpcError::Internal(format!("inner dispatch: {e}")))?;
489 if response.status() == StatusCode::NOT_FOUND {
490 let bytes = to_bytes(response.into_body(), usize::MAX)
491 .await
492 .map_err(|e| XrpcError::Internal(format!("inner response: {e}")))?;
493 return Err(bytes
494 .is_empty()
495 .then(|| XrpcError::InvalidParams(format!("unknown or unenrichable query: {nsid}")))
496 .unwrap_or(XrpcError::NotFound));
497 }
498 finish(response).await
499}
500
501fn encode_params(params: &Map<String, Value>) -> String {
502 let mut out = url::form_urlencoded::Serializer::new(String::new());
503 for (key, value) in params {
504 match value {
505 Value::Array(items) => {
506 for item in items {
507 if let Some(scalar) = scalar_str(item) {
508 out.append_pair(key, &scalar);
509 }
510 }
511 }
512 _ => {
513 if let Some(scalar) = scalar_str(value) {
514 out.append_pair(key, &scalar);
515 }
516 }
517 }
518 }
519 out.finish()
520}
521
522fn scalar_str(value: &Value) -> Option<String> {
523 match value {
524 Value::String(s) => Some(s.clone()),
525 Value::Number(n) => Some(n.to_string()),
526 Value::Bool(b) => Some(b.to_string()),
527 _ => None,
528 }
529}
530
531async fn finish(resp: Response) -> Result<Value, XrpcError> {
532 let status = resp.status();
533 let bytes = to_bytes(resp.into_body(), usize::MAX)
534 .await
535 .map_err(|e| XrpcError::Internal(format!("inner response: {e}")))?;
536 if !status.is_success() {
537 let msg = String::from_utf8_lossy(&bytes).into_owned();
538 return Err(match status {
539 StatusCode::BAD_REQUEST => XrpcError::InvalidParams(msg),
540 StatusCode::NOT_FOUND => XrpcError::NotFound,
541 StatusCode::SERVICE_UNAVAILABLE => XrpcError::Overloaded,
542 _ => XrpcError::UpstreamUnavailable(format!("inner query ({status}): {msg}")),
543 });
544 }
545 serde_json::from_slice(&bytes)
546 .map_err(|e| XrpcError::Internal(format!("inner response decode: {e}")))
547}
548
549#[cfg(test)]
550mod tests {
551 use super::*;
552 use serde_json::json;
553
554 #[test]
555 fn collects_dids_uris_and_strong_refs() {
556 let mut refs = Vec::new();
557 let mut seen = HashSet::new();
558 collect_ref(&json!("did:plc:abc"), &mut refs, &mut seen);
559 collect_ref(
560 &json!("at://did:plc:abc/sh.tangled.repo/x"),
561 &mut refs,
562 &mut seen,
563 );
564 collect_ref(
565 &json!({"uri": "at://did:plc:abc/sh.tangled.repo/y", "cid": "bafy"}),
566 &mut refs,
567 &mut seen,
568 );
569 collect_ref(&json!("at://did:plc:abc"), &mut refs, &mut seen);
570 collect_ref(&json!("oppi.li"), &mut refs, &mut seen);
571 collect_ref(&json!("did:plc:abc"), &mut refs, &mut seen);
572 assert_eq!(refs.len(), 3);
573 }
574
575 #[test]
576 fn repo_did_normalizes_uri_authorities() {
577 let uri = SubjectRef::Uri(AtUri::new_owned("at://did:plc:abc/sh.tangled.repo/x").unwrap());
578 let bare = SubjectRef::Did(Did::new_owned("did:plc:abc").unwrap());
579 assert_eq!(
580 repo_did(&uri).as_ref().map(|d| d.as_str()),
581 Some("did:plc:abc")
582 );
583 assert_eq!(
584 repo_did(&bare).as_ref().map(|d| d.as_str()),
585 Some("did:plc:abc")
586 );
587 assert_eq!(repo_did(&SubjectRef::Global), None);
588 }
589
590 #[test]
591 fn parses_link_sources_and_payload_types() {
592 let source = LinkSource::parse("sh.tangled.feed.star:subject").unwrap();
593 assert_eq!(source.collection.as_ref(), "sh.tangled.feed.star");
594 assert_eq!(source.path, SourcePath::Subject);
595 let source = LinkSource::parse("sh.tangled.graph.follow:.repo").unwrap();
596 assert_eq!(source.path, SourcePath::Repo);
597 assert!(LinkSource::parse("sh.tangled.feed.star").is_err());
598 assert!(LinkSource::parse("sh.tangled.feed.star:.rkey").is_err());
599 assert!(LinkSource::parse("not an nsid:subject").is_err());
600
601 let ty = PayloadType::parse(TYPE_MINIDOC).unwrap();
602 assert_eq!(ty.nsid.as_ref(), "com.bad-example.identity.miniDoc");
603 assert_eq!(ty.fragment.as_deref(), None);
604 let ty = PayloadType::parse(TYPE_COUNT).unwrap();
605 assert_eq!(ty.nsid.as_ref(), "sh.tangled.query.enrichResponse");
606 assert_eq!(ty.fragment.as_deref(), Some("count"));
607 assert!(PayloadType::parse("sh.tangled.query.enrichResponse#").is_err());
608 assert!(PayloadType::parse("nope#count").is_err());
609 }
610}