#![cfg_attr(not(test), no_std)]
use byteorder::{BigEndian, ByteOrder as _};
use data_encoding::BASE64;
use heapless::Vec as HVec;
use sha1::{Digest as _, Sha1};
pub const MAX_REQ_BYTES: usize = 4096;
const WS_GUID: &[u8; 36] = b"258EAFA5-E914-47DA-95CA-C5AB0DC85B11";
const REQUEST_PATH: &str = "/xrpc/com.atproto.sync.subscribeRepos";
const RESPONSE_PREFIX: &[u8; 97] = b"HTTP/1.1 101 Switching Protocols\r\n\
Upgrade: websocket\r\n\
Connection: Upgrade\r\n\
Sec-WebSocket-Accept: ";
const RESPONSE_SUFFIX: &[u8; 4] = b"\r\n\r\n";
pub const HANDSHAKE_RESPONSE_LEN: usize = 129;
pub const PING_FRAME: [u8; 2] = [0x89, 0x00];
pub const FUTURE_CURSOR_AND_CLOSE: [u8; 31] = [
0x82, 0x19, 0xa1, 0x62, b'o', b'p', 0x20, 0xa1, 0x65, b'e', b'r', b'r', b'o', b'r', 0x6c, b'F',
b'u', b't', b'u', b'r', b'e', b'C', b'u', b'r', b's', b'o', b'r', 0x88, 0x02, 0x03, 0xf0,
];
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct ParsedRequest {
pub key: [u8; 24],
pub cursor: Option<u64>,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum ParseStatus {
Incomplete,
Invalid,
Complete(ParsedRequest),
}
pub fn accept_key(key: &[u8; 24]) -> [u8; 28] {
let mut hasher = Sha1::new();
hasher.update(key);
hasher.update(WS_GUID);
let digest = hasher.finalize();
let mut out = [0u8; 28];
BASE64.encode_mut(&digest, &mut out);
out
}
pub fn format_handshake_response(key: &[u8; 24]) -> [u8; HANDSHAKE_RESPONSE_LEN] {
let accept = accept_key(key);
let mut out = [0u8; HANDSHAKE_RESPONSE_LEN];
out[..RESPONSE_PREFIX.len()].copy_from_slice(RESPONSE_PREFIX);
out[RESPONSE_PREFIX.len()..RESPONSE_PREFIX.len() + 28].copy_from_slice(&accept);
out[RESPONSE_PREFIX.len() + 28..].copy_from_slice(RESPONSE_SUFFIX);
out
}
fn is_valid_key(val: &[u8]) -> Option<[u8; 24]> {
let key: [u8; 24] = val.try_into().ok()?;
let mut decoded = [0u8; 18];
(BASE64.decode_mut(&key, &mut decoded).ok()? == 16).then_some(key)
}
fn parse_cursor(val: &[u8]) -> Option<Option<u64>> {
if val.is_empty() {
return Some(None);
}
let text = core::str::from_utf8(val).ok()?;
if !text.bytes().all(|b| b.is_ascii_digit()) {
return None;
}
text.parse::<u64>().ok().map(Some)
}
pub fn parse_request(buf: &[u8]) -> ParseStatus {
if buf.len() > MAX_REQ_BYTES {
return ParseStatus::Invalid;
}
let mut headers = [httparse::EMPTY_HEADER; 32];
let mut req = httparse::Request::new(&mut headers);
let end = match req.parse(buf) {
Ok(httparse::Status::Complete(end)) => end,
Ok(httparse::Status::Partial) => {
return if buf.len() >= MAX_REQ_BYTES {
ParseStatus::Invalid
} else {
ParseStatus::Incomplete
};
}
Err(_) => return ParseStatus::Invalid,
};
if end != buf.len()
|| req.method != Some("GET")
|| req.path != Some(REQUEST_PATH)
|| req.version != Some(1)
{
return ParseStatus::Invalid;
}
let mut ks_key: Option<[u8; 24]> = None;
let mut ks_cursor: Option<Option<u64>> = None;
for header in req.headers.iter() {
let val = header.value.trim_ascii();
if header.name.eq_ignore_ascii_case("ks-key") {
if ks_key.is_some() {
return ParseStatus::Invalid;
}
let Some(k) = is_valid_key(val) else {
return ParseStatus::Invalid;
};
ks_key = Some(k);
} else if header.name.eq_ignore_ascii_case("ks-cursor") {
if ks_cursor.is_some() {
return ParseStatus::Invalid;
}
let Some(c) = parse_cursor(val) else {
return ParseStatus::Invalid;
};
ks_cursor = Some(c);
}
}
let Some(key) = ks_key else {
return ParseStatus::Invalid;
};
ParseStatus::Complete(ParsedRequest {
key,
cursor: ks_cursor.unwrap_or(None),
})
}
pub const CHECKPOINT_STRIDE: u64 = 4096;
pub const MAX_CHECKPOINTS: usize = 65536;
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum FrameHeader {
Incomplete,
Invalid,
Complete { header_len: usize, total_len: u64 },
}
/// only unmasked binary frames with FIN set are valid, byte 0 is `0x82`.
pub fn parse_ws_frame_header(bytes: &[u8]) -> FrameHeader {
if bytes.len() < 2 {
return FrameHeader::Incomplete;
}
if bytes[0] != 0x82 || (bytes[1] & 0x80) != 0 {
return FrameHeader::Invalid;
}
match bytes[1] & 0x7f {
len7 @ 0..=125 => FrameHeader::Complete {
header_len: 2,
total_len: 2 + u64::from(len7),
},
126 => {
let Some(len_bytes) = bytes.get(2..4) else {
return FrameHeader::Incomplete;
};
FrameHeader::Complete {
header_len: 4,
total_len: 4 + u64::from(BigEndian::read_u16(len_bytes)),
}
}
_ => {
let Some(len_bytes) = bytes.get(2..10) else {
return FrameHeader::Incomplete;
};
let Some(total_len) = BigEndian::read_u64(len_bytes).checked_add(10) else {
return FrameHeader::Invalid;
};
FrameHeader::Complete {
header_len: 10,
total_len,
}
}
}
}
pub struct FrameIndex {
pub checkpoints: HVec<u64, MAX_CHECKPOINTS>,
pub tip_index: u64,
pub tip_offset: u64,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum CursorTarget {
StreamFrom { offset: u64 },
FutureCursor,
InvalidStream,
}
impl FrameIndex {
pub const fn new() -> Self {
Self {
checkpoints: HVec::new(),
tip_index: 0,
tip_offset: 0,
}
}
pub fn reset(&mut self) {
self.checkpoints.clear();
self.tip_index = 0;
self.tip_offset = 0;
}
/// the tip moves past complete frames only, and a half written one waits for its payload.
pub fn advance<F>(&mut self, file_size: u64, mut read_at: F)
where
F: FnMut(u64, &mut [u8]) -> usize,
{
let mut hdr = [0u8; 10];
while self.tip_offset < file_size {
if self.tip_index.is_multiple_of(CHECKPOINT_STRIDE)
&& (self.tip_index / CHECKPOINT_STRIDE) as usize == self.checkpoints.len()
{
let _ = self.checkpoints.push(self.tip_offset);
}
let avail = file_size - self.tip_offset;
if avail < 2 {
break;
}
let to_read = (avail as usize).min(10);
let n = read_at(self.tip_offset, &mut hdr[..to_read]);
if n < 2 {
break;
}
match parse_ws_frame_header(&hdr[..n]) {
FrameHeader::Complete { total_len, .. } => {
let Some(end) = self.tip_offset.checked_add(total_len) else {
break;
};
if end <= file_size {
self.tip_offset = end;
self.tip_index += 1;
} else {
break;
}
}
FrameHeader::Incomplete | FrameHeader::Invalid => break,
}
}
}
pub fn resolve_cursor<F>(&self, cursor: Option<u64>, mut read_at: F) -> CursorTarget
where
F: FnMut(u64, &mut [u8]) -> usize,
{
let Some(c) = cursor else {
return CursorTarget::StreamFrom {
offset: self.tip_offset,
};
};
if c > self.tip_index {
return CursorTarget::FutureCursor;
}
if c == self.tip_index {
return CursorTarget::StreamFrom {
offset: self.tip_offset,
};
}
if c == 0 {
return CursorTarget::StreamFrom { offset: 0 };
}
let slot = ((c / CHECKPOINT_STRIDE) as usize).min(self.checkpoints.len().saturating_sub(1));
let mut idx = (slot as u64) * CHECKPOINT_STRIDE;
let mut off = self.checkpoints.get(slot).copied().unwrap_or(0);
let mut hdr = [0u8; 10];
while idx < c {
let n = read_at(off, &mut hdr);
match parse_ws_frame_header(&hdr[..n]) {
FrameHeader::Complete { total_len, .. } => {
let Some(next_off) = off.checked_add(total_len) else {
return CursorTarget::InvalidStream;
};
if next_off > self.tip_offset {
return CursorTarget::InvalidStream;
}
off = next_off;
idx += 1;
}
FrameHeader::Incomplete | FrameHeader::Invalid => {
return CursorTarget::InvalidStream;
}
}
}
CursorTarget::StreamFrom { offset: off }
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn rfc6455_accept_key_matches_spec_example() {
let key = b"dGhlIHNhbXBsZSBub25jZQ==";
assert_eq!(&accept_key(key), b"s3pPLMBiTxaQ9kYGzzhZRbK+xOo=");
let resp = format_handshake_response(key);
assert_eq!(
&resp[..],
b"HTTP/1.1 101 Switching Protocols\r\nUpgrade: websocket\r\nConnection: Upgrade\r\nSec-WebSocket-Accept: s3pPLMBiTxaQ9kYGzzhZRbK+xOo=\r\n\r\n"
);
}
#[test]
fn parse_valid_requests_with_and_without_cursor() {
let req_empty = b"GET /xrpc/com.atproto.sync.subscribeRepos HTTP/1.1\r\n\
Host: git.klbr.net\r\n\
Upgrade: websocket\r\n\
Connection: Upgrade\r\n\
Ks-Key: dGhlIHNhbXBsZSBub25jZQ==\r\n\
Ks-Cursor: \r\n\r\n";
assert_eq!(
parse_request(req_empty),
ParseStatus::Complete(ParsedRequest {
key: *b"dGhlIHNhbXBsZSBub25jZQ==",
cursor: None,
})
);
let req_absent = b"GET /xrpc/com.atproto.sync.subscribeRepos HTTP/1.1\r\n\
Host: git.klbr.net\r\n\
Upgrade: websocket\r\n\
Connection: Upgrade\r\n\
Ks-Key: dGhlIHNhbXBsZSBub25jZQ==\r\n\r\n";
assert_eq!(
parse_request(req_absent),
ParseStatus::Complete(ParsedRequest {
key: *b"dGhlIHNhbXBsZSBub25jZQ==",
cursor: None,
})
);
let req_zero = b"GET /xrpc/com.atproto.sync.subscribeRepos HTTP/1.1\r\n\
ks-key: dGhlIHNhbXBsZSBub25jZQ==\r\n\
ks-cursor: 0\r\n\r\n";
assert_eq!(
parse_request(req_zero),
ParseStatus::Complete(ParsedRequest {
key: *b"dGhlIHNhbXBsZSBub25jZQ==",
cursor: Some(0),
})
);
let req_num = b"GET /xrpc/com.atproto.sync.subscribeRepos HTTP/1.1\r\n\
Ks-Key: dGhlIHNhbXBsZSBub25jZQ==\r\n\
Ks-Cursor: 18446744073709551615\r\n\r\n";
assert_eq!(
parse_request(req_num),
ParseStatus::Complete(ParsedRequest {
key: *b"dGhlIHNhbXBsZSBub25jZQ==",
cursor: Some(u64::MAX),
})
);
}
#[test]
fn parse_rejects_incomplete_or_malformed_requests() {
let partial = b"GET /xrpc/com.atproto.sync.subscribeRepos HTTP/1.1\r\nKs-Key: dGhlIHNhbXBsZSBub25jZQ==\r\n";
assert_eq!(parse_request(partial), ParseStatus::Incomplete);
let wrong_path = b"GET /xrpc/com.atproto.sync.subscribeRepos?cursor=1 HTTP/1.1\r\n\
Ks-Key: dGhlIHNhbXBsZSBub25jZQ==\r\n\
Ks-Cursor: 1\r\n\r\n";
assert_eq!(parse_request(wrong_path), ParseStatus::Invalid);
let no_key_hdr = b"GET /xrpc/com.atproto.sync.subscribeRepos HTTP/1.1\r\n\
Ks-Cursor: 0\r\n\r\n";
assert_eq!(parse_request(no_key_hdr), ParseStatus::Invalid);
let bad_key = b"GET /xrpc/com.atproto.sync.subscribeRepos HTTP/1.1\r\n\
Ks-Key: dGhlIHNhbXBsZSBub25jZQo=\r\n\
Ks-Cursor: 0\r\n\r\n";
assert_eq!(parse_request(bad_key), ParseStatus::Invalid);
let bad_cursor = b"GET /xrpc/com.atproto.sync.subscribeRepos HTTP/1.1\r\n\
Ks-Key: dGhlIHNhbXBsZSBub25jZQ==\r\n\
Ks-Cursor: -1\r\n\r\n";
assert_eq!(parse_request(bad_cursor), ParseStatus::Invalid);
let overflow_cursor = b"GET /xrpc/com.atproto.sync.subscribeRepos HTTP/1.1\r\n\
Ks-Key: dGhlIHNhbXBsZSBub25jZQ==\r\n\
Ks-Cursor: 18446744073709551616\r\n\r\n";
assert_eq!(parse_request(overflow_cursor), ParseStatus::Invalid);
let trailing = b"GET /xrpc/com.atproto.sync.subscribeRepos HTTP/1.1\r\n\
Ks-Key: dGhlIHNhbXBsZSBub25jZQ==\r\n\
Ks-Cursor: 0\r\n\r\njunk";
assert_eq!(parse_request(trailing), ParseStatus::Invalid);
let huge = [b'A'; MAX_REQ_BYTES];
assert_eq!(parse_request(&huge), ParseStatus::Invalid);
}
fn build_ws_frame(payload_len: usize, fill: u8) -> std::vec::Vec<u8> {
let mut out = std::vec::Vec::new();
out.push(0x82);
if payload_len < 126 {
out.push(payload_len as u8);
} else if payload_len <= 65535 {
out.push(126);
out.extend_from_slice(&(payload_len as u16).to_be_bytes());
} else {
out.push(127);
out.extend_from_slice(&(payload_len as u64).to_be_bytes());
}
out.resize(out.len() + payload_len, fill);
out
}
fn slice_reader<'a>(data: &'a [u8]) -> impl FnMut(u64, &mut [u8]) -> usize + 'a {
move |off, buf| {
let off = off as usize;
if off >= data.len() {
return 0;
}
let avail = (data.len() - off).min(buf.len());
buf[..avail].copy_from_slice(&data[off..off + avail]);
avail
}
}
#[test]
fn parse_ws_frame_headers_across_7_16_64_bit_lengths() {
let f7 = build_ws_frame(125, 0xaa);
assert_eq!(
parse_ws_frame_header(&f7[..10.min(f7.len())]),
FrameHeader::Complete {
header_len: 2,
total_len: 127,
}
);
let f16 = build_ws_frame(1000, 0xbb);
assert_eq!(parse_ws_frame_header(&f16[..3]), FrameHeader::Incomplete);
assert_eq!(
parse_ws_frame_header(&f16[..10]),
FrameHeader::Complete {
header_len: 4,
total_len: 1004,
}
);
let f64 = build_ws_frame(70_000, 0xcc);
assert_eq!(parse_ws_frame_header(&f64[..9]), FrameHeader::Incomplete);
assert_eq!(
parse_ws_frame_header(&f64[..10]),
FrameHeader::Complete {
header_len: 10,
total_len: 70_010,
}
);
assert_eq!(parse_ws_frame_header(&[0x81, 0x05]), FrameHeader::Invalid);
assert_eq!(parse_ws_frame_header(&[0x82, 0x85]), FrameHeader::Invalid);
}
#[test]
fn index_stops_before_incomplete_frame_and_resolves_checkpoints() {
let mut stream = std::vec::Vec::new();
let mut expected_offsets = std::vec::Vec::new();
for i in 0..4105usize {
expected_offsets.push(stream.len() as u64);
let payload_len = match i % 5 {
0 => 10,
1 => 125,
2 => 126,
3 => 300,
_ => 42,
};
let f = build_ws_frame(payload_len, (i & 0xff) as u8);
stream.extend_from_slice(&f);
}
let complete_end = stream.len() as u64;
let partial = build_ws_frame(500, 0xff);
stream.extend_from_slice(&partial[..104]);
let mut idx = Box::new(FrameIndex::new());
idx.advance(stream.len() as u64, slice_reader(&stream));
assert_eq!(idx.tip_index, 4105);
assert_eq!(idx.tip_offset, complete_end);
assert_eq!(idx.checkpoints[0], 0);
assert_eq!(idx.checkpoints[1], expected_offsets[4096]);
assert_eq!(
idx.resolve_cursor(None, slice_reader(&stream)),
CursorTarget::StreamFrom {
offset: complete_end
}
);
assert_eq!(
idx.resolve_cursor(Some(0), slice_reader(&stream)),
CursorTarget::StreamFrom { offset: 0 }
);
assert_eq!(
idx.resolve_cursor(Some(2048), slice_reader(&stream)),
CursorTarget::StreamFrom {
offset: expected_offsets[2048]
}
);
assert_eq!(
idx.resolve_cursor(Some(4096), slice_reader(&stream)),
CursorTarget::StreamFrom {
offset: expected_offsets[4096]
}
);
assert_eq!(
idx.resolve_cursor(Some(4104), slice_reader(&stream)),
CursorTarget::StreamFrom {
offset: expected_offsets[4104]
}
);
assert_eq!(
idx.resolve_cursor(Some(4105), slice_reader(&stream)),
CursorTarget::StreamFrom {
offset: complete_end
}
);
assert_eq!(
idx.resolve_cursor(Some(4106), slice_reader(&stream)),
CursorTarget::FutureCursor
);
stream.extend_from_slice(&partial[104..]);
idx.advance(stream.len() as u64, slice_reader(&stream));
assert_eq!(idx.tip_index, 4106);
assert_eq!(idx.tip_offset, stream.len() as u64);
assert_eq!(
idx.resolve_cursor(Some(4105), slice_reader(&stream)),
CursorTarget::StreamFrom {
offset: complete_end
}
);
}
}