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 /// record paths into the inner response selecting targets for this payload
167 #[serde(default)]
168 targets: Option<Vec<String>>,
169}
170
171#[derive(Debug, Deserialize)]
172pub struct EnrichInput {
173 xrpc: String,
174 #[serde(default)]
175 params: Option<Map<String, Value>>,
176 enrich: Vec<LinkDescriptor>,
177 #[serde(default)]
178 viewer: Option<Did>,
179}
180
181pub async fn enrich(
182 State(state): State<AppState>,
183 Json(input): Json<EnrichInput>,
184) -> Result<Json<Value>, XrpcError> {
185 let mut seen_descriptors = HashSet::new();
186 let mut descriptors = Vec::new();
187 for (descriptor_index, descriptor) in input.enrich.iter().enumerate() {
188 if !KNOWN_TYPES.contains(&descriptor.ty.as_str()) {
189 return Err(descriptor_error(&descriptor.source, "unknown enrich type"));
190 }
191 validate_source(&descriptor.source, &descriptor.ty)?;
192 if descriptor.ty.as_str() == TYPE_VIEWER && input.viewer.is_none() {
193 return Err(descriptor_error(
194 &descriptor.source,
195 "viewer payloads require a viewer param",
196 ));
197 }
198 let targets = descriptor
199 .targets
200 .as_ref()
201 .map(|targets| {
202 targets
203 .iter()
204 .enumerate()
205 .map(|(target_index, path)| {
206 RecordPath::parse(path).map_err(|e| {
207 XrpcError::InvalidParams(format!(
208 "enrich[{descriptor_index}].targets[{target_index}] {path:?}: {e}"
209 ))
210 })
211 })
212 .collect::<Result<Vec<_>, _>>()
213 })
214 .transpose()?;
215 if seen_descriptors.insert((
216 &descriptor.source,
217 &descriptor.ty,
218 descriptor.targets.as_deref(),
219 )) {
220 descriptors.push((descriptor, targets));
221 }
222 }
223
224 let inner = run_inner(&state, &input.xrpc, input.params.unwrap_or_default()).await?;
225 if descriptors.is_empty() {
226 return Ok(Json(json!({ "output": inner, "data": {} })));
227 }
228
229 let mut data = Map::new();
230 let mut refs_by_targets: HashMap<Option<&[String]>, Vec<SubjectRef>> = HashMap::new();
231 let mut minidoc_targets: HashMap<Did<DefaultStr>, Vec<LinkSource>> = HashMap::new();
232 for (descriptor, targets) in descriptors {
233 let refs = refs_by_targets
234 .entry(descriptor.targets.as_deref())
235 .or_insert_with(|| {
236 let mut refs: Vec<SubjectRef> = Vec::new();
237 let mut seen: HashSet<SubjectRef> = HashSet::new();
238 match &targets {
239 Some(targets) => {
240 for path in targets {
241 for node in walk_path(path, [&inner]) {
242 collect_ref(node, &mut refs, &mut seen);
243 }
244 }
245 }
246 None => discover_refs(&inner, &mut refs, &mut seen),
247 }
248 // the authority of every at-uri is a reference too
249 // this lets each enrich address record authors
250 let authorities: Vec<SubjectRef> = refs
251 .iter()
252 .filter_map(|reference| repo_did(reference).map(SubjectRef::Did))
253 .collect();
254 for authority in authorities {
255 if seen.insert(authority.clone()) {
256 refs.push(authority);
257 }
258 }
259 refs
260 });
261 for reference in refs.iter() {
262 let Some(subject) = applicable_subject(descriptor, reference)? else {
263 continue;
264 };
265 match descriptor.ty.as_str() {
266 TYPE_COUNT | TYPE_DISTINCT_AUTHORS | TYPE_VIEWER => {
267 let kind = edge_kind(&descriptor.source)?;
268 let key = EdgeKey::new(kind, subject.clone());
269 let result = match descriptor.ty.as_str() {
270 TYPE_COUNT => Value::from(state.edges.count(&key)),
271 TYPE_DISTINCT_AUTHORS => {
272 Value::from(state.edges.count_distinct_authors(&key))
273 }
274 _ => {
275 let viewer = input.viewer.as_ref().expect("viewer param present");
276 match state.edges.viewer_source(&key, viewer.as_str()) {
277 Some(uri) => Value::String(uri.to_string()),
278 None => Value::Null,
279 }
280 }
281 };
282 put(
283 &mut data,
284 subject.as_str(),
285 descriptor.source.as_str(),
286 descriptor.ty.as_str(),
287 result,
288 );
289 }
290 TYPE_MINIDOC => {
291 if let Some(target) = repo_did(&subject) {
292 let sources = minidoc_targets.entry(target).or_default();
293 if !sources.contains(&descriptor.source) {
294 sources.push(descriptor.source.clone());
295 }
296 }
297 }
298 _ => unreachable!("validated above!"),
299 }
300 }
301 }
302
303 let docs = resolve_minidocs(&state, minidoc_targets).await;
304 for (target, sources, doc) in docs {
305 for source in sources {
306 put(
307 &mut data,
308 target.as_str(),
309 source.as_str(),
310 TYPE_MINIDOC,
311 doc.clone(),
312 );
313 }
314 }
315
316 Ok(Json(json!({ "output": inner, "data": data })))
317}
318
319fn put(data: &mut Map<String, Value>, reference: &str, source: &str, ty: &str, payload: Value) {
320 let per_ref = data
321 .entry(reference.to_owned())
322 .or_insert_with(|| Value::Object(Map::new()));
323 let Value::Object(per_ref) = per_ref else {
324 return;
325 };
326 let per_source = per_ref
327 .entry(source.to_owned())
328 .or_insert_with(|| Value::Object(Map::new()));
329 if let Value::Object(per_source) = per_source {
330 per_source.insert(ty.to_owned(), payload);
331 }
332}
333
334/// we drop failures, the client falls back to resolveMiniDoc for misses
335async fn resolve_minidocs(
336 state: &AppState,
337 targets: HashMap<Did<DefaultStr>, Vec<LinkSource>>,
338) -> Vec<(Did<DefaultStr>, Vec<LinkSource>, Value)> {
339 futures::stream::iter(targets)
340 .map(|(did, sources)| async move {
341 let doc = state
342 .slingshot
343 .resolve_mini_doc(&AtIdentifier::Did(did.clone()))
344 .await
345 .ok()
346 .and_then(|bytes| serde_json::from_slice::<Value>(&bytes).ok());
347 (did, sources, doc)
348 })
349 .buffer_unordered(MINIDOC_CONCURRENCY)
350 .filter_map(|(did, sources, doc)| async move { doc.map(|doc| (did, sources, doc)) })
351 .collect()
352 .await
353}
354
355fn descriptor_error(source: &LinkSource, msg: &str) -> XrpcError {
356 XrpcError::InvalidParams(format!("enrich {}: {msg}", source.as_str()))
357}
358
359fn validate_source(source: &LinkSource, ty: &PayloadType) -> Result<(), XrpcError> {
360 match source.path {
361 SourcePath::Subject => subject_shape(source.collection.as_ref())
362 .map(|_| ())
363 .ok_or_else(|| descriptor_error(source, "unknown collection")),
364 // minidocs read the uri's authority straight out of the response, only
365 // edge payloads need the author index
366 SourcePath::Repo if ty.as_str() == TYPE_MINIDOC => Ok(()),
367 SourcePath::Repo => mirror_kind(source.collection.as_ref())
368 .map(|_| ())
369 .ok_or_else(|| descriptor_error(source, "collection has no author index")),
370 }
371}
372
373fn repo_did(reference: &SubjectRef) -> Option<Did<DefaultStr>> {
374 match reference {
375 SubjectRef::Did(did) => Some(did.clone()),
376 SubjectRef::Uri(uri) => match uri.authority() {
377 AtIdentifier::Did(did) => Some(did.clone().into_static()),
378 AtIdentifier::Handle(_) => None,
379 },
380 SubjectRef::Global => None,
381 }
382}
383
384/// the subject an edge lookup or payload should be keyed by, or None if the
385/// descriptor's shape doesn't apply to this ref, eg. a did-only source asked
386/// about an at-uri. `.repo` sources normalize to the repo did
387fn applicable_subject(
388 descriptor: &LinkDescriptor,
389 reference: &SubjectRef,
390) -> Result<Option<SubjectRef>, XrpcError> {
391 match descriptor.source.path {
392 SourcePath::Subject => {
393 let (_, shape) = subject_shape(descriptor.source.collection.as_ref())
394 .ok_or_else(|| descriptor_error(&descriptor.source, "unknown collection"))?;
395 Ok(shape_accepts(shape, reference).then(|| reference.clone()))
396 }
397 SourcePath::Repo => Ok(repo_did(reference).map(SubjectRef::Did)),
398 }
399}
400
401/// the edge kind behind a source for stats payloads
402fn edge_kind(source: &LinkSource) -> Result<Nsid<DefaultStr>, XrpcError> {
403 let kind = match source.path {
404 SourcePath::Subject => subject_shape(source.collection.as_ref())
405 .map(|(nsid, _)| nsid)
406 .ok_or_else(|| descriptor_error(source, "unknown collection"))?,
407 SourcePath::Repo => mirror_kind(source.collection.as_ref())
408 .ok_or_else(|| descriptor_error(source, "collection has no author index"))?,
409 };
410 Ok(nsid_static(kind))
411}
412
413/// mismatch means skip, not reject
414fn shape_accepts(shape: SubjectShape, reference: &SubjectRef) -> bool {
415 match (shape, reference) {
416 (SubjectShape::BareDid, SubjectRef::Did(_)) => true,
417 (SubjectShape::Collection(expected), SubjectRef::Uri(uri)) => {
418 uri.collection().is_some_and(|c| c.as_ref() == expected)
419 }
420 (SubjectShape::OneOfCollections(allowed), SubjectRef::Uri(uri))
421 | (SubjectShape::BareDidOrOneOfCollections(allowed), SubjectRef::Uri(uri)) => uri
422 .collection()
423 .is_some_and(|c| allowed.contains(&c.as_ref())),
424 (SubjectShape::BareDidOrOneOfCollections(_), SubjectRef::Did(_)) => true,
425 (SubjectShape::AnyAtUri, SubjectRef::Uri(_)) => true,
426 _ => false,
427 }
428}
429
430/// a value counts as a reference if it is a did string or an at-uri string
431/// with collection and rkey, or a {uri, cid} strong ref
432fn collect_ref(value: &Value, refs: &mut Vec<SubjectRef>, seen: &mut HashSet<SubjectRef>) {
433 let candidate = match value {
434 Value::String(s) => Some(s.as_str()),
435 Value::Object(o) => o
436 .get("uri")
437 .and_then(Value::as_str)
438 .filter(|_| o.contains_key("cid")),
439 _ => None,
440 };
441 let Some(candidate) = candidate else { return };
442 let reference = if candidate.starts_with("did:") {
443 Did::<DefaultStr>::new_owned(candidate)
444 .ok()
445 .map(SubjectRef::Did)
446 } else if candidate.starts_with("at://") {
447 AtUri::<DefaultStr>::new_owned(candidate)
448 .ok()
449 .filter(|u| {
450 matches!(u.authority(), AtIdentifier::Did(_))
451 && u.collection().is_some()
452 && u.rkey().is_some()
453 })
454 .map(SubjectRef::Uri)
455 } else {
456 None
457 };
458 if let Some(reference) = reference
459 && seen.insert(reference.clone())
460 {
461 refs.push(reference);
462 }
463}
464
465fn discover_refs(value: &Value, refs: &mut Vec<SubjectRef>, seen: &mut HashSet<SubjectRef>) {
466 collect_ref(value, refs, seen);
467 match value {
468 Value::Array(items) => {
469 for item in items {
470 discover_refs(item, refs, seen);
471 }
472 }
473 Value::Object(map) => {
474 for v in map.values() {
475 discover_refs(v, refs, seen);
476 }
477 }
478 _ => {}
479 }
480}
481
482async fn run_inner(
483 state: &AppState,
484 nsid: &str,
485 params: Map<String, Value>,
486) -> Result<Value, XrpcError> {
487 let qs = encode_params(¶ms);
488 let uri = format!("/xrpc/{nsid}?{qs}");
489 let request = Request::builder()
490 .method("GET")
491 .uri(&uri)
492 .body(Body::empty())
493 .map_err(|e| XrpcError::Internal(format!("inner request: {e}")))?;
494 let response = state
495 .self_router()
496 .oneshot(request)
497 .await
498 .map_err(|e| XrpcError::Internal(format!("inner dispatch: {e}")))?;
499 if response.status() == StatusCode::NOT_FOUND {
500 let bytes = to_bytes(response.into_body(), usize::MAX)
501 .await
502 .map_err(|e| XrpcError::Internal(format!("inner response: {e}")))?;
503 return Err(bytes
504 .is_empty()
505 .then(|| XrpcError::InvalidParams(format!("unknown or unenrichable query: {nsid}")))
506 .unwrap_or(XrpcError::NotFound));
507 }
508 finish(response).await
509}
510
511fn encode_params(params: &Map<String, Value>) -> String {
512 let mut out = url::form_urlencoded::Serializer::new(String::new());
513 for (key, value) in params {
514 match value {
515 Value::Array(items) => {
516 for item in items {
517 if let Some(scalar) = scalar_str(item) {
518 out.append_pair(key, &scalar);
519 }
520 }
521 }
522 _ => {
523 if let Some(scalar) = scalar_str(value) {
524 out.append_pair(key, &scalar);
525 }
526 }
527 }
528 }
529 out.finish()
530}
531
532fn scalar_str(value: &Value) -> Option<String> {
533 match value {
534 Value::String(s) => Some(s.clone()),
535 Value::Number(n) => Some(n.to_string()),
536 Value::Bool(b) => Some(b.to_string()),
537 _ => None,
538 }
539}
540
541async fn finish(resp: Response) -> Result<Value, XrpcError> {
542 let status = resp.status();
543 let bytes = to_bytes(resp.into_body(), usize::MAX)
544 .await
545 .map_err(|e| XrpcError::Internal(format!("inner response: {e}")))?;
546 if !status.is_success() {
547 let msg = String::from_utf8_lossy(&bytes).into_owned();
548 return Err(match status {
549 StatusCode::BAD_REQUEST => XrpcError::InvalidParams(msg),
550 StatusCode::NOT_FOUND => XrpcError::NotFound,
551 StatusCode::SERVICE_UNAVAILABLE => XrpcError::Overloaded,
552 _ => XrpcError::UpstreamUnavailable(format!("inner query ({status}): {msg}")),
553 });
554 }
555 serde_json::from_slice(&bytes)
556 .map_err(|e| XrpcError::Internal(format!("inner response decode: {e}")))
557}
558
559#[cfg(test)]
560mod tests {
561 use super::*;
562 use serde_json::json;
563
564 #[test]
565 fn collects_dids_uris_and_strong_refs() {
566 let mut refs = Vec::new();
567 let mut seen = HashSet::new();
568 collect_ref(&json!("did:plc:abc"), &mut refs, &mut seen);
569 collect_ref(
570 &json!("at://did:plc:abc/sh.tangled.repo/x"),
571 &mut refs,
572 &mut seen,
573 );
574 collect_ref(
575 &json!({"uri": "at://did:plc:abc/sh.tangled.repo/y", "cid": "bafy"}),
576 &mut refs,
577 &mut seen,
578 );
579 collect_ref(&json!("at://did:plc:abc"), &mut refs, &mut seen);
580 collect_ref(&json!("oppi.li"), &mut refs, &mut seen);
581 collect_ref(&json!("did:plc:abc"), &mut refs, &mut seen);
582 assert_eq!(refs.len(), 3);
583 }
584
585 #[test]
586 fn repo_did_normalizes_uri_authorities() {
587 let uri = SubjectRef::Uri(AtUri::new_owned("at://did:plc:abc/sh.tangled.repo/x").unwrap());
588 let bare = SubjectRef::Did(Did::new_owned("did:plc:abc").unwrap());
589 assert_eq!(
590 repo_did(&uri).as_ref().map(|d| d.as_str()),
591 Some("did:plc:abc")
592 );
593 assert_eq!(
594 repo_did(&bare).as_ref().map(|d| d.as_str()),
595 Some("did:plc:abc")
596 );
597 assert_eq!(repo_did(&SubjectRef::Global), None);
598 }
599
600 #[test]
601 fn parses_link_sources_and_payload_types() {
602 let source = LinkSource::parse("sh.tangled.feed.star:subject").unwrap();
603 assert_eq!(source.collection.as_ref(), "sh.tangled.feed.star");
604 assert_eq!(source.path, SourcePath::Subject);
605 let source = LinkSource::parse("sh.tangled.graph.follow:.repo").unwrap();
606 assert_eq!(source.path, SourcePath::Repo);
607 assert!(LinkSource::parse("sh.tangled.feed.star").is_err());
608 assert!(LinkSource::parse("sh.tangled.feed.star:.rkey").is_err());
609 assert!(LinkSource::parse("not an nsid:subject").is_err());
610
611 let ty = PayloadType::parse(TYPE_MINIDOC).unwrap();
612 assert_eq!(ty.nsid.as_ref(), "com.bad-example.identity.miniDoc");
613 assert_eq!(ty.fragment.as_deref(), None);
614 let ty = PayloadType::parse(TYPE_COUNT).unwrap();
615 assert_eq!(ty.nsid.as_ref(), "sh.tangled.query.enrichResponse");
616 assert_eq!(ty.fragment.as_deref(), Some("count"));
617 assert!(PayloadType::parse("sh.tangled.query.enrichResponse#").is_err());
618 assert!(PayloadType::parse("nope#count").is_err());
619 }
620}