1use base64::Engine;
2use serde::{Deserialize, Deserializer, Serialize, Serializer};
3use std::fmt;
4
5pub const PROTOCOL_VERSION: u16 = 1;
6pub const MAX_FRAME_SIZE: usize = 16 * 1024 * 1024;
7pub const CONTENT_TYPE: &str = "application/vnd.fn0.doc-db+json";
8
9#[derive(Clone, Debug, Eq, PartialEq, Serialize, Deserialize)]
10#[serde(deny_unknown_fields)]
11pub struct DocDbRequest {
12 pub version: u16,
13 pub operation: DocDbOperation,
14}
15
16impl DocDbRequest {
17 pub fn new(operation: DocDbOperation) -> Self {
18 Self {
19 version: PROTOCOL_VERSION,
20 operation,
21 }
22 }
23}
24
25#[derive(Clone, Debug, Eq, PartialEq, Serialize, Deserialize)]
26#[serde(tag = "kind", content = "value", deny_unknown_fields)]
27pub enum DocDbOperation {
28 Get {
29 key: DocDbKey,
30 },
31 Put {
32 key: DocDbKey,
33 #[serde(with = "base64_bytes")]
34 data: Vec<u8>,
35 },
36 Delete {
37 key: DocDbKey,
38 },
39 Query {
40 pk: String,
41 after_sk: Option<String>,
42 limit: u64,
43 },
44 Scan {
45 after: Option<DocDbKey>,
46 limit: u64,
47 },
48 GetObserved {
49 key: DocDbKey,
50 },
51 Transact {
52 conditions: Vec<DocDbCondition>,
53 mutations: Vec<DocDbMutation>,
54 },
55 AdminPurgeProject {
56 project_id: String,
57 },
58}
59
60#[derive(Clone, Debug, Eq, PartialEq, Serialize, Deserialize)]
61#[serde(deny_unknown_fields)]
62pub struct DocDbKey {
63 pub pk: String,
64 pub sk: String,
65}
66
67impl DocDbKey {
68 pub fn new(pk: impl Into<String>, sk: impl Into<String>) -> Self {
69 Self {
70 pk: pk.into(),
71 sk: sk.into(),
72 }
73 }
74}
75
76#[derive(Clone, Copy, Debug, Eq, Hash, PartialEq, Serialize, Deserialize)]
77#[serde(transparent)]
78pub struct DocDbRevision(u64);
79
80impl DocDbRevision {
81 pub const fn new(value: u64) -> Self {
82 Self(value)
83 }
84
85 pub const fn value(self) -> u64 {
86 self.0
87 }
88}
89
90#[derive(Clone, Debug, Eq, PartialEq, Serialize, Deserialize)]
91#[serde(tag = "kind", content = "value", deny_unknown_fields)]
92pub enum DocDbCondition {
93 RevisionEquals {
94 key: DocDbKey,
95 expected_revision: DocDbRevision,
96 },
97 Exists {
98 key: DocDbKey,
99 },
100 NotExists {
101 key: DocDbKey,
102 },
103}
104
105#[derive(Clone, Debug, Eq, PartialEq, Serialize, Deserialize)]
106#[serde(tag = "kind", content = "value", deny_unknown_fields)]
107pub enum DocDbMutation {
108 Put {
109 key: DocDbKey,
110 #[serde(with = "base64_bytes")]
111 data: Vec<u8>,
112 },
113 Delete {
114 key: DocDbKey,
115 },
116}
117
118#[derive(Clone, Debug, Eq, PartialEq, Serialize, Deserialize)]
119#[serde(deny_unknown_fields)]
120pub struct DocDbResponse {
121 pub version: u16,
122 pub result: DocDbResult,
123}
124
125impl DocDbResponse {
126 pub fn new(result: DocDbResult) -> Self {
127 Self {
128 version: PROTOCOL_VERSION,
129 result,
130 }
131 }
132
133 pub fn error(error: DocDbError) -> Self {
134 Self::new(DocDbResult::Error { error })
135 }
136}
137
138#[derive(Clone, Debug, Eq, PartialEq, Serialize, Deserialize)]
139#[serde(tag = "kind", content = "value", deny_unknown_fields)]
140pub enum DocDbResult {
141 Get { data: Option<BinaryDocument> },
142 Put,
143 Delete,
144 Query { documents: Vec<DocDbDocument> },
145 Scan { documents: Vec<DocDbDocument> },
146 GetObserved { document: DocDbObservedDocument },
147 Transact { outcome: DocDbTransactOutcome },
148 AdminPurgeProject { deleted_rows: u64 },
149 Error { error: DocDbError },
150}
151
152#[derive(Clone, Debug, Eq, PartialEq, Serialize, Deserialize)]
153#[serde(deny_unknown_fields)]
154pub struct BinaryDocument {
155 #[serde(with = "base64_bytes")]
156 pub data: Vec<u8>,
157}
158
159#[derive(Clone, Debug, Eq, PartialEq, Serialize, Deserialize)]
160#[serde(deny_unknown_fields)]
161pub struct DocDbDocument {
162 pub key: DocDbKey,
163 #[serde(with = "base64_bytes")]
164 pub data: Vec<u8>,
165}
166
167#[derive(Clone, Debug, Eq, PartialEq, Serialize, Deserialize)]
168#[serde(tag = "kind", content = "value", deny_unknown_fields)]
169pub enum DocDbObservedDocument {
170 Present {
171 #[serde(with = "base64_bytes")]
172 data: Vec<u8>,
173 revision: DocDbRevision,
174 },
175 Missing {
176 revision: Option<DocDbRevision>,
177 },
178}
179
180#[derive(Clone, Debug, Eq, PartialEq, Serialize, Deserialize)]
181#[serde(tag = "kind", content = "value", deny_unknown_fields)]
182pub enum DocDbTransactOutcome {
183 Committed,
184 Conflict { condition_index: usize },
185}
186
187#[derive(Clone, Debug, Eq, PartialEq, Serialize, Deserialize)]
188#[serde(tag = "kind", content = "value", deny_unknown_fields)]
189pub enum DocDbError {
190 InvalidRequest { message: String },
191 Backend { message: String },
192 UnsupportedVersion { version: u16 },
193 Forbidden { message: String },
194}
195
196impl fmt::Display for DocDbError {
197 fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
198 match self {
199 Self::InvalidRequest { message } => write!(formatter, "invalid request: {message}"),
200 Self::Backend { message } => write!(formatter, "backend error: {message}"),
201 Self::UnsupportedVersion { version } => {
202 write!(formatter, "unsupported protocol version: {version}")
203 }
204 Self::Forbidden { message } => write!(formatter, "forbidden: {message}"),
205 }
206 }
207}
208
209#[derive(Debug)]
210pub enum CodecError {
211 Malformed(String),
212 UnsupportedVersion(u16),
213 Serialize(String),
214}
215
216impl fmt::Display for CodecError {
217 fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
218 match self {
219 Self::Malformed(message) => write!(formatter, "malformed doc-db payload: {message}"),
220 Self::UnsupportedVersion(version) => {
221 write!(formatter, "unsupported doc-db protocol version: {version}")
222 }
223 Self::Serialize(message) => {
224 write!(formatter, "doc-db payload serialization failed: {message}")
225 }
226 }
227 }
228}
229
230impl std::error::Error for CodecError {}
231
232pub fn encode_request(request: &DocDbRequest) -> Result<Vec<u8>, CodecError> {
233 serde_json::to_vec(request).map_err(|error| CodecError::Serialize(error.to_string()))
234}
235
236pub fn decode_request(bytes: &[u8]) -> Result<DocDbRequest, CodecError> {
237 let request: DocDbRequest =
238 serde_json::from_slice(bytes).map_err(|error| CodecError::Malformed(error.to_string()))?;
239 validate_version(request.version)?;
240 Ok(request)
241}
242
243pub fn encode_response(response: &DocDbResponse) -> Result<Vec<u8>, CodecError> {
244 serde_json::to_vec(response).map_err(|error| CodecError::Serialize(error.to_string()))
245}
246
247pub fn decode_response(bytes: &[u8]) -> Result<DocDbResponse, CodecError> {
248 let response: DocDbResponse =
249 serde_json::from_slice(bytes).map_err(|error| CodecError::Malformed(error.to_string()))?;
250 validate_version(response.version)?;
251 Ok(response)
252}
253
254fn validate_version(version: u16) -> Result<(), CodecError> {
255 if version == PROTOCOL_VERSION {
256 Ok(())
257 } else {
258 Err(CodecError::UnsupportedVersion(version))
259 }
260}
261
262mod base64_bytes {
263 use super::*;
264
265 pub fn serialize<S>(bytes: &[u8], serializer: S) -> Result<S::Ok, S::Error>
266 where
267 S: Serializer,
268 {
269 serializer.serialize_str(&base64::engine::general_purpose::STANDARD.encode(bytes))
270 }
271
272 pub fn deserialize<'de, D>(deserializer: D) -> Result<Vec<u8>, D::Error>
273 where
274 D: Deserializer<'de>,
275 {
276 let value = String::deserialize(deserializer)?;
277 base64::engine::general_purpose::STANDARD
278 .decode(value.as_bytes())
279 .map_err(serde::de::Error::custom)
280 }
281}
282
283#[cfg(test)]
284mod tests {
285 use super::*;
286
287 fn binary_data() -> Vec<u8> {
288 let mut data = vec![0, 1, 2, 127, 128, 255];
289 data.extend_from_slice(&[0; 4096]);
290 data.extend_from_slice("invalid utf-8: \u{00e9}".as_bytes());
291 data
292 }
293
294 #[test]
295 fn round_trips_all_operations_and_binary_data() {
296 let data = binary_data();
297 let requests = vec![
298 DocDbOperation::Get {
299 key: DocDbKey::new("pk", "sk"),
300 },
301 DocDbOperation::Put {
302 key: DocDbKey::new("pk", "put"),
303 data: data.clone(),
304 },
305 DocDbOperation::Delete {
306 key: DocDbKey::new("pk", "delete"),
307 },
308 DocDbOperation::Query {
309 pk: "pk".to_string(),
310 after_sk: Some("after".to_string()),
311 limit: 17,
312 },
313 DocDbOperation::Scan {
314 after: Some(DocDbKey::new("after-pk", "after-sk")),
315 limit: 19,
316 },
317 DocDbOperation::GetObserved {
318 key: DocDbKey::new("pk", "observed"),
319 },
320 DocDbOperation::Transact {
321 conditions: vec![
322 DocDbCondition::RevisionEquals {
323 key: DocDbKey::new("pk", "version"),
324 expected_revision: DocDbRevision::new(4),
325 },
326 DocDbCondition::Exists {
327 key: DocDbKey::new("pk", "missing"),
328 },
329 DocDbCondition::NotExists {
330 key: DocDbKey::new("pk", "insert"),
331 },
332 ],
333 mutations: vec![
334 DocDbMutation::Put {
335 key: DocDbKey::new("pk", "put"),
336 data: data.clone(),
337 },
338 DocDbMutation::Delete {
339 key: DocDbKey::new("pk", "delete"),
340 },
341 ],
342 },
343 DocDbOperation::AdminPurgeProject {
344 project_id: "abcdefgh".to_string(),
345 },
346 ];
347
348 for operation in requests {
349 let request = DocDbRequest::new(operation);
350 let encoded = encode_request(&request).unwrap();
351 let decoded = decode_request(&encoded).unwrap();
352 assert_eq!(decoded, request);
353 }
354
355 let responses = vec![
356 DocDbResponse::new(DocDbResult::GetObserved {
357 document: DocDbObservedDocument::Present {
358 data: data.clone(),
359 revision: DocDbRevision::new(9),
360 },
361 }),
362 DocDbResponse::new(DocDbResult::Transact {
363 outcome: DocDbTransactOutcome::Conflict { condition_index: 3 },
364 }),
365 ];
366
367 for response in responses {
368 let encoded = encode_response(&response).unwrap();
369 let decoded = decode_response(&encoded).unwrap();
370 assert_eq!(decoded, response);
371 }
372 }
373
374 #[test]
375 fn encodes_binary_data_as_base64_string() {
376 let request = DocDbRequest::new(DocDbOperation::Put {
377 key: DocDbKey::new("pk", "sk"),
378 data: binary_data(),
379 });
380 let encoded = String::from_utf8(encode_request(&request).unwrap()).unwrap();
381 assert!(encoded.contains("\"data\":\""));
382 assert!(!encoded.contains("[0,1,2"));
383 }
384
385 #[test]
386 fn rejects_malformed_payload() {
387 assert!(matches!(
388 decode_request(b"{not-json"),
389 Err(CodecError::Malformed(_))
390 ));
391 }
392
393 #[test]
394 fn rejects_unsupported_protocol_version() {
395 let request = DocDbRequest {
396 version: PROTOCOL_VERSION + 1,
397 operation: DocDbOperation::Get {
398 key: DocDbKey::new("pk", "sk"),
399 },
400 };
401 let encoded = encode_request(&request).unwrap();
402 assert!(matches!(
403 decode_request(&encoded),
404 Err(CodecError::UnsupportedVersion(version)) if version == PROTOCOL_VERSION + 1
405 ));
406 }
407
408 #[test]
409 fn round_trips_revisions_above_i64_max() {
410 let revision = DocDbRevision::new(i64::MAX as u64 + 1);
411 let response = DocDbResponse::new(DocDbResult::GetObserved {
412 document: DocDbObservedDocument::Missing {
413 revision: Some(revision),
414 },
415 });
416 let encoded = encode_response(&response).unwrap();
417 assert_eq!(decode_response(&encoded).unwrap(), response);
418 }
419}