This repository has no description
1use std::collections::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::types::did::Did;
13use jacquard_common::types::ident::AtIdentifier;
14use jacquard_common::types::string::AtUri;
15use serde::Deserialize;
16use serde_json::{Map, Value, json};
17use tower::ServiceExt;
18
19use crate::recordpath::{parse_record_path, walk_path};
20use crate::{AppState, SubjectShape, XrpcError, mirror_kind, subject_shape};
21
22#[derive(Debug, Clone, Copy, PartialEq, Eq, Deserialize)]
23#[serde(rename_all = "camelCase")]
24pub enum Aggregation {
25 Count,
26 DistinctAuthors,
27 Viewer,
28}
29
30impl Aggregation {
31 fn key(self) -> &'static str {
32 match self {
33 Aggregation::Count => "count",
34 Aggregation::DistinctAuthors => "distinctAuthors",
35 Aggregation::Viewer => "viewer",
36 }
37 }
38}
39
40#[derive(Debug, Deserialize)]
41pub struct LinkDescriptor {
42 /// constellation link source
43 source: String,
44 #[serde(rename = "type")]
45 aggregation: Aggregation,
46}
47
48impl LinkDescriptor {
49 fn parts(&self) -> Result<(&str, &str, Aggregation), XrpcError> {
50 match self.source.split_once(':') {
51 Some((collection, path)) if !collection.is_empty() && !path.is_empty() => {
52 Ok((collection, path, self.aggregation))
53 }
54 _ => Err(XrpcError::InvalidParams(format!(
55 "enrich {:?}: expected \"collection:path\"",
56 self.source
57 ))),
58 }
59 }
60}
61
62#[derive(Debug, Deserialize)]
63pub struct EnrichInput {
64 xrpc: String,
65 #[serde(default)]
66 params: Option<Map<String, Value>>,
67 enrich: Vec<LinkDescriptor>,
68 #[serde(default)]
69 sources: Option<Vec<String>>,
70 #[serde(default)]
71 viewer: Option<Did>,
72}
73
74pub async fn enrich(
75 State(state): State<AppState>,
76 Json(input): Json<EnrichInput>,
77) -> Result<Json<Value>, XrpcError> {
78 for descriptor in &input.enrich {
79 validate_descriptor(descriptor)?;
80 if descriptor.aggregation == Aggregation::Viewer && input.viewer.is_none() {
81 return Err(descriptor_error(
82 descriptor,
83 "viewer aggregation requires a viewer param",
84 ));
85 }
86 }
87
88 let inner = run_inner(&state, &input.xrpc, input.params.unwrap_or_default()).await?;
89
90 let mut refs: Vec<SubjectRef> = Vec::new();
91 let mut seen: HashSet<SubjectRef> = HashSet::new();
92 match &input.sources {
93 Some(sources) => {
94 for (i, path) in sources.iter().enumerate() {
95 let segs = parse_record_path(path)
96 .map_err(|e| XrpcError::InvalidParams(format!("sources[{i}] {path:?}: {e}")))?;
97 for node in walk_path(&segs, [&inner]) {
98 collect_ref(node, &mut refs, &mut seen);
99 }
100 }
101 }
102 None => discover_refs(&inner, &mut refs, &mut seen),
103 }
104
105 let mut stats = Map::new();
106 for reference in &refs {
107 let mut per_ref = Map::new();
108 for descriptor in &input.enrich {
109 let (.., aggregation) = descriptor.parts()?;
110 let Some(kind) = resolve_descriptor(descriptor, reference)? else {
111 continue;
112 };
113 let key = EdgeKey::new(kind, reference.clone());
114 let result = match aggregation {
115 Aggregation::Count => Value::from(state.edges.count(&key)),
116 Aggregation::DistinctAuthors => {
117 Value::from(state.edges.count_distinct_authors(&key))
118 }
119 Aggregation::Viewer => {
120 // viewer presence is validated up front
121 let viewer = input.viewer.as_ref().expect("viewer param present");
122 match state.edges.viewer_source(&key, viewer.as_str()) {
123 Some(uri) => Value::String(uri.to_string()),
124 None => Value::Null,
125 }
126 }
127 };
128 let entry = per_ref
129 .entry(descriptor.source.clone())
130 .or_insert_with(|| Value::Object(Map::new()));
131 if let Value::Object(counts) = entry {
132 counts.insert(aggregation.key().to_owned(), result);
133 }
134 }
135 if !per_ref.is_empty() {
136 stats.insert(reference.as_str().to_owned(), Value::Object(per_ref));
137 }
138 }
139
140 Ok(Json(json!({ "output": inner, "stats": stats })))
141}
142
143fn descriptor_error(descriptor: &LinkDescriptor, msg: &str) -> XrpcError {
144 XrpcError::InvalidParams(format!("enrich {}: {msg}", descriptor.source))
145}
146
147fn validate_descriptor(descriptor: &LinkDescriptor) -> Result<(), XrpcError> {
148 let (collection, path, _) = descriptor.parts()?;
149 match path {
150 "subject" => subject_shape(collection)
151 .map(|_| ())
152 .ok_or_else(|| descriptor_error(descriptor, "unknown collection")),
153 ".repo" => mirror_kind(collection)
154 .map(|_| ())
155 .ok_or_else(|| descriptor_error(descriptor, "collection has no author index")),
156 _ if path.starts_with('.') => Err(descriptor_error(
157 descriptor,
158 "envelope field not index-backed; try .repo",
159 )),
160 _ => Err(descriptor_error(
161 descriptor,
162 "path not index-backed; only `subject` and `.repo` are supported",
163 )),
164 }
165}
166
167/// the edge kind to look up for one reference, or None if the descriptor's shape
168/// doesn't apply to this ref kind, eg. a did-only descriptor asked about an at-uri
169fn resolve_descriptor(
170 descriptor: &LinkDescriptor,
171 reference: &SubjectRef,
172) -> Result<Option<jacquard_common::types::nsid::Nsid<DefaultStr>>, XrpcError> {
173 let (collection, path, _) = descriptor.parts()?;
174 match path {
175 "subject" => {
176 let Some((nsid, shape)) = subject_shape(collection) else {
177 return Err(descriptor_error(descriptor, "unknown collection"));
178 };
179 Ok(shape_accepts(shape, reference).then(|| nsid_static(nsid)))
180 }
181 ".repo" => {
182 let Some(kind) = mirror_kind(collection) else {
183 return Err(descriptor_error(
184 descriptor,
185 "collection has no author index",
186 ));
187 };
188 Ok(matches!(reference, SubjectRef::Did(_)).then(|| nsid_static(kind)))
189 }
190 _ => unreachable!("validated up front"),
191 }
192}
193
194/// mismatch means skip, not reject
195fn shape_accepts(shape: SubjectShape, reference: &SubjectRef) -> bool {
196 match (shape, reference) {
197 (SubjectShape::BareDid, SubjectRef::Did(_)) => true,
198 (SubjectShape::Collection(expected), SubjectRef::Uri(uri)) => {
199 uri.collection().is_some_and(|c| c.as_ref() == expected)
200 }
201 (SubjectShape::OneOfCollections(allowed), SubjectRef::Uri(uri))
202 | (SubjectShape::BareDidOrOneOfCollections(allowed), SubjectRef::Uri(uri)) => uri
203 .collection()
204 .is_some_and(|c| allowed.contains(&c.as_ref())),
205 (SubjectShape::BareDidOrOneOfCollections(_), SubjectRef::Did(_)) => true,
206 (SubjectShape::AnyAtUri, SubjectRef::Uri(_)) => true,
207 _ => false,
208 }
209}
210
211/// a value counts as a reference if it is a did string or an at-uri string
212/// with collection and rkey, or a {uri, cid} strong ref
213fn collect_ref(value: &Value, refs: &mut Vec<SubjectRef>, seen: &mut HashSet<SubjectRef>) {
214 let candidate = match value {
215 Value::String(s) => Some(s.as_str()),
216 Value::Object(o) => o
217 .get("uri")
218 .and_then(Value::as_str)
219 .filter(|_| o.contains_key("cid")),
220 _ => None,
221 };
222 let Some(candidate) = candidate else { return };
223 let reference = if candidate.starts_with("did:") {
224 Did::<DefaultStr>::new_owned(candidate)
225 .ok()
226 .map(SubjectRef::Did)
227 } else if candidate.starts_with("at://") {
228 AtUri::<DefaultStr>::new_owned(candidate)
229 .ok()
230 .filter(|u| {
231 matches!(u.authority(), AtIdentifier::Did(_))
232 && u.collection().is_some()
233 && u.rkey().is_some()
234 })
235 .map(SubjectRef::Uri)
236 } else {
237 None
238 };
239 if let Some(reference) = reference
240 && seen.insert(reference.clone())
241 {
242 refs.push(reference);
243 }
244}
245
246fn discover_refs(value: &Value, refs: &mut Vec<SubjectRef>, seen: &mut HashSet<SubjectRef>) {
247 collect_ref(value, refs, seen);
248 match value {
249 Value::Array(items) => {
250 for item in items {
251 discover_refs(item, refs, seen);
252 }
253 }
254 Value::Object(map) => {
255 for v in map.values() {
256 discover_refs(v, refs, seen);
257 }
258 }
259 _ => {}
260 }
261}
262
263async fn run_inner(
264 state: &AppState,
265 nsid: &str,
266 params: Map<String, Value>,
267) -> Result<Value, XrpcError> {
268 let qs = encode_params(¶ms);
269 let uri = format!("/xrpc/{nsid}?{qs}");
270 let request = Request::builder()
271 .method("GET")
272 .uri(&uri)
273 .body(Body::empty())
274 .map_err(|e| XrpcError::Internal(format!("inner request: {e}")))?;
275 let response = state
276 .self_router()
277 .oneshot(request)
278 .await
279 .map_err(|e| XrpcError::Internal(format!("inner dispatch: {e}")))?;
280 if response.status() == StatusCode::NOT_FOUND {
281 let bytes = to_bytes(response.into_body(), usize::MAX)
282 .await
283 .map_err(|e| XrpcError::Internal(format!("inner response: {e}")))?;
284 return Err(bytes
285 .is_empty()
286 .then(|| XrpcError::InvalidParams(format!("unknown or unenrichable query: {nsid}")))
287 .unwrap_or(XrpcError::NotFound));
288 }
289 finish(response).await
290}
291
292fn encode_params(params: &Map<String, Value>) -> String {
293 let mut out = url::form_urlencoded::Serializer::new(String::new());
294 for (key, value) in params {
295 match value {
296 Value::Array(items) => {
297 for item in items {
298 if let Some(scalar) = scalar_str(item) {
299 out.append_pair(key, &scalar);
300 }
301 }
302 }
303 _ => {
304 if let Some(scalar) = scalar_str(value) {
305 out.append_pair(key, &scalar);
306 }
307 }
308 }
309 }
310 out.finish()
311}
312
313fn scalar_str(value: &Value) -> Option<String> {
314 match value {
315 Value::String(s) => Some(s.clone()),
316 Value::Number(n) => Some(n.to_string()),
317 Value::Bool(b) => Some(b.to_string()),
318 _ => None,
319 }
320}
321
322async fn finish(resp: Response) -> Result<Value, XrpcError> {
323 let status = resp.status();
324 let bytes = to_bytes(resp.into_body(), usize::MAX)
325 .await
326 .map_err(|e| XrpcError::Internal(format!("inner response: {e}")))?;
327 if !status.is_success() {
328 let msg = String::from_utf8_lossy(&bytes).into_owned();
329 return Err(match status {
330 StatusCode::BAD_REQUEST => XrpcError::InvalidParams(msg),
331 StatusCode::NOT_FOUND => XrpcError::NotFound,
332 StatusCode::SERVICE_UNAVAILABLE => XrpcError::Overloaded,
333 _ => XrpcError::UpstreamUnavailable(format!("inner query ({status}): {msg}")),
334 });
335 }
336 serde_json::from_slice(&bytes)
337 .map_err(|e| XrpcError::Internal(format!("inner response decode: {e}")))
338}
339
340#[cfg(test)]
341mod tests {
342 use super::*;
343 use serde_json::json;
344
345 #[test]
346 fn collects_dids_uris_and_strong_refs() {
347 let mut refs = Vec::new();
348 let mut seen = HashSet::new();
349 collect_ref(&json!("did:plc:abc"), &mut refs, &mut seen);
350 collect_ref(
351 &json!("at://did:plc:abc/sh.tangled.repo/x"),
352 &mut refs,
353 &mut seen,
354 );
355 collect_ref(
356 &json!({"uri": "at://did:plc:abc/sh.tangled.repo/y", "cid": "bafy"}),
357 &mut refs,
358 &mut seen,
359 );
360 collect_ref(&json!("at://did:plc:abc"), &mut refs, &mut seen);
361 collect_ref(&json!("oppi.li"), &mut refs, &mut seen);
362 collect_ref(&json!("did:plc:abc"), &mut refs, &mut seen);
363 assert_eq!(refs.len(), 3);
364 }
365}