1use sha2::{Digest, Sha256};
24use thiserror::Error;
25
26use crate::filter::Filter;
27use crate::message::SubscriptionId;
28use crate::util::hex;
29
30const FINGERPRINT_BYTES: usize = 16;
31const ID_BYTES: usize = 32;
32const INFINITY_TIMESTAMP: u64 = u64::MAX;
33const RESERVED_TIMESTAMP_INFINITY_OFFSET: u64 = 0;
34
35#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
37pub struct NegProtocolVersion(pub u8);
38
39impl NegProtocolVersion {
40 pub const V1: Self = Self(0x61);
42
43 #[must_use]
47 pub const fn from_byte(byte: u8) -> Option<Self> {
48 if byte >= 0x60 && byte < 0x70 {
49 Some(Self(byte))
50 } else {
51 None
52 }
53 }
54
55 #[must_use]
57 pub const fn as_byte(self) -> u8 {
58 self.0
59 }
60
61 #[must_use]
63 pub const fn version(self) -> u8 {
64 self.0.wrapping_sub(0x60)
65 }
66}
67
68#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
70pub struct NegItem {
71 pub timestamp: u64,
74 pub id: [u8; ID_BYTES],
76}
77
78impl NegItem {
79 #[must_use]
81 pub const fn new(timestamp: u64, id: [u8; ID_BYTES]) -> Self {
82 Self { timestamp, id }
83 }
84
85 #[must_use]
87 pub const fn infinity() -> Self {
88 Self {
89 timestamp: INFINITY_TIMESTAMP,
90 id: [0u8; ID_BYTES],
91 }
92 }
93}
94
95#[derive(Debug, Clone, PartialEq, Eq)]
98pub struct NegBound {
99 pub timestamp: u64,
102 pub id_prefix: Vec<u8>,
105}
106
107impl NegBound {
108 #[must_use]
111 pub const fn infinity() -> Self {
112 Self {
113 timestamp: INFINITY_TIMESTAMP,
114 id_prefix: Vec::new(),
115 }
116 }
117}
118
119#[derive(Debug, Clone, PartialEq, Eq)]
121pub enum NegRangeMode {
122 Skip,
124 Fingerprint([u8; FINGERPRINT_BYTES]),
127 IdList(Vec<[u8; ID_BYTES]>),
129}
130
131#[derive(Debug, Clone, PartialEq, Eq)]
133pub struct NegRange {
134 pub upper_bound: NegBound,
137 pub mode: NegRangeMode,
139}
140
141#[derive(Debug, Clone, PartialEq, Eq)]
143pub struct NegPayload {
144 pub version: NegProtocolVersion,
146 pub ranges: Vec<NegRange>,
148}
149
150#[derive(Debug, Clone, PartialEq, Eq)]
152pub struct NegOpen {
153 pub subscription_id: SubscriptionId,
155 pub filter: Filter,
157 pub payload: NegPayload,
159}
160
161#[derive(Debug, Clone, PartialEq, Eq)]
163pub enum NegMessage {
164 Open(Box<NegOpen>),
166 Msg {
168 subscription_id: SubscriptionId,
170 payload: NegPayload,
172 },
173 Close {
175 subscription_id: SubscriptionId,
177 },
178 Err {
180 subscription_id: SubscriptionId,
182 reason: String,
185 },
186}
187
188#[derive(Debug, Error)]
190#[non_exhaustive]
191pub enum NegentropyError {
192 #[error("unexpected end of input while decoding varint")]
194 UnexpectedEof,
195 #[error("varint exceeds 10 bytes")]
197 VarintOverflow,
198 #[error("buffer ended {expected} bytes before payload completed")]
200 PayloadTruncated {
201 expected: usize,
203 },
204 #[error("unknown range mode {0}")]
206 UnknownRangeMode(u64),
207 #[error("unsupported Negentropy protocol version 0x{0:02x}")]
209 UnsupportedVersion(u8),
210 #[error("hex decode failure: {0}")]
212 Hex(String),
213}
214
215fn write_varint(value: u64, out: &mut Vec<u8>) {
216 let mut value = value;
217 let mut tmp: Vec<u8> = Vec::with_capacity(10);
218 loop {
219 let byte = u8::try_from(value & 0x7f).unwrap_or(0);
220 tmp.push(byte);
221 value >>= 7;
222 if value == 0 {
223 break;
224 }
225 }
226 for (i, byte) in tmp.iter().rev().enumerate() {
229 let with_continuation = if i + 1 < tmp.len() {
230 *byte | 0x80
231 } else {
232 *byte
233 };
234 out.push(with_continuation);
235 }
236}
237
238fn read_varint(buf: &[u8], cursor: &mut usize) -> Result<u64, NegentropyError> {
239 let mut value: u64 = 0;
240 for _ in 0..10 {
241 let byte = *buf.get(*cursor).ok_or(NegentropyError::UnexpectedEof)?;
242 *cursor += 1;
243 value = value
244 .checked_shl(7)
245 .ok_or(NegentropyError::VarintOverflow)?
246 | u64::from(byte & 0x7f);
247 if byte & 0x80 == 0 {
248 return Ok(value);
249 }
250 }
251 Err(NegentropyError::VarintOverflow)
252}
253
254fn encode_bound(bound: &NegBound, prev_timestamp: &mut u64, out: &mut Vec<u8>) {
255 let encoded_ts = if bound.timestamp == INFINITY_TIMESTAMP {
256 RESERVED_TIMESTAMP_INFINITY_OFFSET
257 } else {
258 bound
261 .timestamp
262 .saturating_sub(*prev_timestamp)
263 .saturating_add(1)
264 };
265 write_varint(encoded_ts, out);
266 if bound.timestamp != INFINITY_TIMESTAMP {
267 *prev_timestamp = bound.timestamp;
268 }
269 let len = bound.id_prefix.len().min(ID_BYTES);
270 write_varint(len as u64, out);
271 out.extend_from_slice(bound.id_prefix.get(..len).unwrap_or(&[]));
272}
273
274fn decode_bound(
275 buf: &[u8],
276 cursor: &mut usize,
277 prev_timestamp: &mut u64,
278) -> Result<NegBound, NegentropyError> {
279 let ts_field = read_varint(buf, cursor)?;
280 let timestamp = if ts_field == RESERVED_TIMESTAMP_INFINITY_OFFSET {
281 INFINITY_TIMESTAMP
282 } else {
283 let value = prev_timestamp.saturating_add(ts_field.saturating_sub(1));
284 *prev_timestamp = value;
285 value
286 };
287 let len_u64 = read_varint(buf, cursor)?;
288 let len = usize::try_from(len_u64).map_err(|_| NegentropyError::VarintOverflow)?;
289 if len > ID_BYTES {
290 return Err(NegentropyError::PayloadTruncated { expected: len });
291 }
292 let chunk = read_chunk(buf, cursor, len)?;
293 Ok(NegBound {
294 timestamp,
295 id_prefix: chunk.to_vec(),
296 })
297}
298
299fn encode_range(range: &NegRange, prev_timestamp: &mut u64, out: &mut Vec<u8>) {
300 encode_bound(&range.upper_bound, prev_timestamp, out);
301 match &range.mode {
302 NegRangeMode::Skip => write_varint(0, out),
303 NegRangeMode::Fingerprint(fp) => {
304 write_varint(1, out);
305 out.extend_from_slice(fp);
306 }
307 NegRangeMode::IdList(ids) => {
308 write_varint(2, out);
309 write_varint(ids.len() as u64, out);
310 for id in ids {
311 out.extend_from_slice(id);
312 }
313 }
314 }
315}
316
317fn read_chunk<'a>(
318 buf: &'a [u8],
319 cursor: &mut usize,
320 len: usize,
321) -> Result<&'a [u8], NegentropyError> {
322 let end = cursor
323 .checked_add(len)
324 .ok_or(NegentropyError::VarintOverflow)?;
325 let chunk = buf.get(*cursor..end);
326 chunk.map_or_else(
327 || {
328 Err(NegentropyError::PayloadTruncated {
329 expected: end.saturating_sub(buf.len()),
330 })
331 },
332 |chunk| {
333 *cursor = end;
334 Ok(chunk)
335 },
336 )
337}
338
339fn decode_range(
340 buf: &[u8],
341 cursor: &mut usize,
342 prev_timestamp: &mut u64,
343) -> Result<NegRange, NegentropyError> {
344 let upper_bound = decode_bound(buf, cursor, prev_timestamp)?;
345 let mode = read_varint(buf, cursor)?;
346 let mode = match mode {
347 0 => NegRangeMode::Skip,
348 1 => {
349 let chunk = read_chunk(buf, cursor, FINGERPRINT_BYTES)?;
350 let mut fp = [0u8; FINGERPRINT_BYTES];
351 fp.copy_from_slice(chunk);
352 NegRangeMode::Fingerprint(fp)
353 }
354 2 => {
355 let count_u64 = read_varint(buf, cursor)?;
356 let count = usize::try_from(count_u64).map_err(|_| NegentropyError::VarintOverflow)?;
357 let mut ids = Vec::with_capacity(count);
358 for _ in 0..count {
359 let chunk = read_chunk(buf, cursor, ID_BYTES)?;
360 let mut id = [0u8; ID_BYTES];
361 id.copy_from_slice(chunk);
362 ids.push(id);
363 }
364 NegRangeMode::IdList(ids)
365 }
366 other => return Err(NegentropyError::UnknownRangeMode(other)),
367 };
368 Ok(NegRange { upper_bound, mode })
369}
370
371#[must_use]
374pub fn encode_payload(payload: &NegPayload) -> Vec<u8> {
375 let mut out = Vec::with_capacity(1 + payload.ranges.len() * 32);
376 out.push(payload.version.as_byte());
377 let mut prev_timestamp: u64 = 0;
378 for range in &payload.ranges {
379 encode_range(range, &mut prev_timestamp, &mut out);
380 }
381 out
382}
383
384pub fn decode_payload(buf: &[u8]) -> Result<NegPayload, NegentropyError> {
390 let (version_byte, rest) = buf.split_first().ok_or(NegentropyError::UnexpectedEof)?;
391 let version = NegProtocolVersion::from_byte(*version_byte)
392 .ok_or(NegentropyError::UnsupportedVersion(*version_byte))?;
393 let mut cursor = 0;
394 let mut prev_timestamp: u64 = 0;
395 let mut ranges = Vec::new();
396 while cursor < rest.len() {
397 ranges.push(decode_range(rest, &mut cursor, &mut prev_timestamp)?);
398 }
399 Ok(NegPayload { version, ranges })
400}
401
402#[must_use]
405pub fn encode_payload_hex(payload: &NegPayload) -> String {
406 hex::encode(encode_payload(payload))
407}
408
409pub fn decode_payload_hex(hex_str: &str) -> Result<NegPayload, NegentropyError> {
416 let bytes = hex::decode(hex_str).map_err(|e| NegentropyError::Hex(e.to_string()))?;
417 decode_payload(&bytes)
418}
419
420#[must_use]
428pub fn fingerprint(ids: &[[u8; ID_BYTES]]) -> [u8; FINGERPRINT_BYTES] {
429 let mut sum = [0u8; ID_BYTES];
430 for id in ids {
431 let mut carry: u16 = 0;
432 for (sum_byte, id_byte) in sum.iter_mut().zip(id.iter()) {
433 let total = u16::from(*sum_byte) + u16::from(*id_byte) + carry;
434 *sum_byte = u8::try_from(total & 0xff).unwrap_or(0);
435 carry = total >> 8;
436 }
437 }
438 let mut hasher = Sha256::new();
439 hasher.update(sum);
440 let mut count_buf = Vec::new();
441 write_varint(ids.len() as u64, &mut count_buf);
442 hasher.update(&count_buf);
443 let digest = hasher.finalize();
444 let mut out = [0u8; FINGERPRINT_BYTES];
445 for (slot, byte) in out.iter_mut().zip(digest.iter()) {
446 *slot = *byte;
447 }
448 out
449}
450
451#[cfg(test)]
452mod tests {
453 use super::*;
454
455 #[test]
456 fn varint_round_trip() {
457 for &value in &[0u64, 1, 127, 128, 300, 1_000_000, u64::MAX / 2] {
458 let mut buf = Vec::new();
459 write_varint(value, &mut buf);
460 let mut cursor = 0;
461 assert_eq!(read_varint(&buf, &mut cursor).unwrap(), value);
462 assert_eq!(cursor, buf.len());
463 }
464 }
465
466 #[test]
467 fn payload_round_trip() {
468 let payload = NegPayload {
469 version: NegProtocolVersion::V1,
470 ranges: vec![
471 NegRange {
472 upper_bound: NegBound {
473 timestamp: 1_700_000_000,
474 id_prefix: vec![0xab, 0xcd],
475 },
476 mode: NegRangeMode::Fingerprint([0xff; FINGERPRINT_BYTES]),
477 },
478 NegRange {
479 upper_bound: NegBound::infinity(),
480 mode: NegRangeMode::Skip,
481 },
482 ],
483 };
484 let bytes = encode_payload(&payload);
485 let parsed = decode_payload(&bytes).unwrap();
486 assert_eq!(parsed, payload);
487 }
488
489 #[test]
490 fn id_list_payload_round_trip() {
491 let ids = vec![[0x11; 32], [0x22; 32]];
492 let payload = NegPayload {
493 version: NegProtocolVersion::V1,
494 ranges: vec![NegRange {
495 upper_bound: NegBound::infinity(),
496 mode: NegRangeMode::IdList(ids.clone()),
497 }],
498 };
499 let hex_str = encode_payload_hex(&payload);
500 let parsed = decode_payload_hex(&hex_str).unwrap();
501 match &parsed.ranges[0].mode {
502 NegRangeMode::IdList(decoded) => assert_eq!(decoded, &ids),
503 other => panic!("unexpected mode {other:?}"),
504 }
505 }
506
507 #[test]
508 fn unsupported_version_is_rejected() {
509 let bytes = vec![0x10];
510 assert!(matches!(
511 decode_payload(&bytes),
512 Err(NegentropyError::UnsupportedVersion(0x10))
513 ));
514 }
515
516 #[test]
517 fn fingerprint_is_stable() {
518 let ids = vec![[1u8; 32], [2u8; 32]];
519 let fp = fingerprint(&ids);
520 let fp_again = fingerprint(&ids);
521 assert_eq!(fp, fp_again);
522 let fp_other = fingerprint(&[[3u8; 32]]);
523 assert_ne!(fp, fp_other);
524 }
525}