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