use super::depth_one::SHARED_AT_LEAST;
use super::depth_one::{map_path, prefix};
use super::hasher::Hasher;
use sha2::Digest as _;
use std::collections::HashSet;
use std::fs::File;
use std::num::NonZeroU32;
use std::os::unix::fs::FileExt as _;
use std::path::Path;
use anyhow::Context as _;
use gix::hash::ObjectId;
use knot_types::ObjectFormat;
use crate::reply_pack::{BLOCK, PACKET_DATA};
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub(crate) struct Packet {
pub(super) file: u64,
pub(super) pack: u64,
pub(super) len: u32,
}
impl Packet {
pub(crate) fn end(&self) -> u64 {
self.pack + u64::from(self.len)
}
pub(crate) fn file_end(&self) -> u64 {
self.file + 5 + u64::from(self.len)
}
pub(crate) fn header(&self) -> [u8; 5] {
const HEX: &[u8; 16] = b"0123456789abcdef";
let n = self.len + 5;
[
HEX[((n >> 12) & 0xf) as usize],
HEX[((n >> 8) & 0xf) as usize],
HEX[((n >> 4) & 0xf) as usize],
HEX[(n & 0xf) as usize],
1,
]
}
}
pub(crate) struct Share {
pub(super) packets: std::ops::Range<usize>,
pub(super) start: u64,
pub(super) end: u64,
pub(super) low: u64,
pub(super) high: u64,
}
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub(crate) struct MapEntry {
pub(super) oid: ObjectId,
pub(super) offset: u64,
pub(super) paths: NonZeroU32,
}
/// the only bytes in a reply are its prefix, band-1 packets, empty band-2 padding and the flush.
#[derive(Debug, PartialEq, Eq)]
pub(crate) struct Map {
pub(super) format: ObjectFormat,
pub(super) entries: Vec<MapEntry>,
pub(super) packets: Vec<Packet>,
pub(super) pack_len: u64,
pub(super) padding: u64,
pub(super) total_size: u64,
pub(super) trailer: ObjectId,
}
const MAGIC: &[u8; 8] = b"ksdepth1";
const DIGEST_LEN: usize = 32;
impl Map {
/// the map's digest covers only the map, and the reply is checked against it too.
pub(crate) fn of(path: &Path, format: ObjectFormat) -> Option<(File, Self)> {
let reply = File::open(path).ok()?;
let map = Self::parse(&std::fs::read(map_path(path)).ok()?, format)?;
let len = reply.metadata().ok()?.len();
let prefix_len = prefix(&ObjectId::null(format.kind())).len() as u64;
let mut at = prefix_len;
let mut padding = 0u64;
for packet in &map.packets {
let gap = packet.file.checked_sub(at)?;
if gap % 5 != 0 {
return None;
}
padding = padding.checked_add(gap)?;
at = packet.file.checked_add(5 + u64::from(packet.len))?;
}
let mut trailer = Vec::new();
let streamed = map.stream(&reply, map.body(), map.pack_len, &mut |bytes| {
trailer.extend_from_slice(bytes);
Ok(())
});
let mut head = Vec::new();
let headed = map.stream(&reply, 0, 12, &mut |bytes| {
head.extend_from_slice(bytes);
Ok(())
});
let mut flush = [0; 4];
reply.read_exact_at(&mut flush, at).ok()?;
let count = u32::try_from(map.entries.len()).ok()?;
let whole = streamed.is_ok()
&& headed.is_ok()
&& trailer == map.trailer.as_bytes()
&& head == crate::reply_pack::pack_header(count)
&& padding == map.padding
&& &flush == b"0000"
&& at.checked_add(4) == Some(len);
whole.then_some((reply, map))
}
pub(crate) fn body(&self) -> u64 {
self.pack_len - self.trailer.as_bytes().len() as u64
}
pub(crate) fn span(&self, (start, end): (usize, usize)) -> (u64, u64) {
let to = (self.entries.get(end)).map_or(self.body(), |entry| entry.offset);
(self.entries[start].offset, to)
}
pub(crate) fn share(&self, from: u64, to: u64) -> Option<Share> {
let first = self.packets.partition_point(|packet| packet.pack < from);
let last = self.packets.partition_point(|packet| packet.end() <= to);
if last <= first {
return None;
}
let (start, end) = (self.packets[first].file, self.packets[last - 1].file_end());
let (low, high) = (start.next_multiple_of(BLOCK), end / BLOCK * BLOCK);
(high.saturating_sub(low) >= SHARED_AT_LEAST).then_some(Share {
packets: first..last,
start,
end,
low,
high,
})
}
pub(crate) fn stream(
&self,
reply: &File,
start: u64,
end: u64,
each: &mut dyn FnMut(&[u8]) -> anyhow::Result<()>,
) -> anyhow::Result<()> {
let mut at = self.packets.partition_point(|packet| packet.end() <= start);
let mut buffer = vec![0; 5 + PACKET_DATA];
let mut pack_at = start;
while pack_at < end {
let packet = self.packets.get(at).context("a map short of its pack")?;
let skip = (pack_at - packet.pack) as usize;
let now = (packet.end().min(end) - pack_at) as usize;
let total = 5 + skip + now;
reply
.read_exact_at(&mut buffer[..total], packet.file)
.context("read a reply")?;
anyhow::ensure!(
buffer[..5] == packet.header(),
"a packet header in a reply doesn't match its map"
);
each(&buffer[5 + skip..total])?;
pack_at += now as u64;
at += 1;
}
Ok(())
}
pub(crate) fn bytes(&self) -> Vec<u8> {
let mut out = MAGIC.to_vec();
out.extend((self.entries.len() as u64).to_le_bytes());
out.extend((self.packets.len() as u64).to_le_bytes());
out.extend(self.pack_len.to_le_bytes());
out.extend(self.padding.to_le_bytes());
out.extend(self.total_size.to_le_bytes());
out.extend_from_slice(self.trailer.as_bytes());
for entry in &self.entries {
out.extend_from_slice(entry.oid.as_bytes());
out.extend(entry.offset.to_le_bytes());
out.extend(entry.paths.get().to_le_bytes());
}
for packet in &self.packets {
out.extend(packet.file.to_le_bytes());
out.extend(packet.pack.to_le_bytes());
out.extend(packet.len.to_le_bytes());
}
let digest = sha2::Sha256::digest(&out);
out.extend_from_slice(&digest);
out
}
pub(crate) fn parse(bytes: &[u8], format: ObjectFormat) -> Option<Self> {
let (bytes, digest) = bytes.split_at_checked(bytes.len().checked_sub(DIGEST_LEN)?)?;
if sha2::Sha256::digest(bytes)[..] != *digest {
return None;
}
let mut read = Reader(bytes.strip_prefix(MAGIC)?);
let hash_len = format.kind().len_in_bytes();
let (entries, packets) = (read.u64()?, read.u64()?);
let (pack_len, padding) = (read.u64()?, read.u64()?);
let total_size = read.u64()?;
let trailer = ObjectId::try_from(read.take(hash_len)?).ok()?;
let body = pack_len.checked_sub(12 + hash_len as u64)? + 12;
let mut map = Self {
format,
entries: Vec::new(),
packets: Vec::new(),
pack_len,
padding,
total_size,
trailer,
};
for _ in 0..entries {
let oid = ObjectId::try_from(read.take(hash_len)?).ok()?;
let offset = read.u64()?;
let paths = NonZeroU32::new(u32::from_le_bytes(read.take(4)?.try_into().ok()?))?;
map.entries.push(MapEntry { oid, offset, paths });
}
for _ in 0..packets {
let (file, pack) = (read.u64()?, read.u64()?);
let len = u32::from_le_bytes(read.take(4)?.try_into().ok()?);
if !(1..=PACKET_DATA as u32).contains(&len)
|| pack.checked_add(u64::from(len))? > pack_len
|| file.checked_add(5 + u64::from(len)).is_none()
{
return None;
}
map.packets.push(Packet { file, pack, len });
}
let ids: HashSet<&ObjectId> = map.entries.iter().map(|entry| &entry.oid).collect();
let entries_in_order = ids.len() == map.entries.len()
&& (map.entries.windows(2)).all(|pair| pair[0].offset < pair[1].offset)
&& map.entries.first().is_none_or(|entry| entry.offset == 12)
&& map.entries.last().is_none_or(|entry| entry.offset < body);
let packets_in_order = map.packets.first()?.pack == 0
&& map.packets.last()?.end() == pack_len
&& (map.packets.windows(2))
.all(|pair| pair[0].end() == pair[1].pack && pair[0].file_end() <= pair[1].file);
(read.0.is_empty() && entries_in_order && packets_in_order).then_some(map)
}
}
pub(crate) struct Reader<'a>(&'a [u8]);
impl<'a> Reader<'a> {
pub(crate) fn take(&mut self, n: usize) -> Option<&'a [u8]> {
let (now, rest) = self.0.split_at_checked(n)?;
self.0 = rest;
Some(now)
}
pub(crate) fn u64(&mut self) -> Option<u64> {
Some(u64::from_le_bytes(self.take(8)?.try_into().ok()?))
}
}
pub(crate) struct Writer<'a> {
pub(super) file: &'a File,
pub(super) format: ObjectFormat,
pub(super) at: u64,
pub(super) pack_at: u64,
pub(super) pending: Vec<u8>,
pub(super) hasher: Hasher,
pub(super) packets: Vec<Packet>,
pub(super) entries: Vec<MapEntry>,
pub(super) padding: u64,
pub(super) total_size: u64,
}
impl<'a> Writer<'a> {
pub(crate) fn new(file: &'a File, format: ObjectFormat, prefix: &[u8]) -> anyhow::Result<Self> {
file.write_all_at(prefix, 0).context("write a reply")?;
Ok(Self {
file,
format,
at: prefix.len() as u64,
pack_at: 0,
pending: Vec::with_capacity(2 * PACKET_DATA),
hasher: Hasher::new(format),
packets: Vec::new(),
entries: Vec::new(),
padding: 0,
total_size: 0,
})
}
pub(crate) fn position(&self) -> u64 {
self.pack_at + self.pending.len() as u64
}
pub(crate) fn pack(&mut self, bytes: &[u8]) -> anyhow::Result<()> {
self.hasher.update(bytes);
self.queue(bytes)
}
pub(crate) fn queue(&mut self, bytes: &[u8]) -> anyhow::Result<()> {
self.pending.extend_from_slice(bytes);
while self.pending.len() >= PACKET_DATA {
self.packet(PACKET_DATA)?;
}
Ok(())
}
pub(crate) fn packet(&mut self, len: usize) -> anyhow::Result<()> {
let mut packet = format!("{:04x}\u{1}", len + 5).into_bytes();
packet.extend_from_slice(&self.pending[..len]);
self.file
.write_all_at(&packet, self.at)
.context("write a reply")?;
self.packets.push(Packet {
file: self.at,
pack: self.pack_at,
len: len as u32,
});
self.at += packet.len() as u64;
self.pack_at += len as u64;
self.pending.drain(..len);
Ok(())
}
/// empty band-2 packets pad the run so the shared blocks land on a block boundary.
pub(crate) fn copy(
&mut self,
reply: &File,
map: &Map,
from: u64,
to: u64,
) -> anyhow::Result<()> {
let Some(Share {
packets,
start,
end,
low,
high,
}) = map.share(from, to)
else {
return map.stream(reply, from, to, &mut |bytes| self.pack(bytes));
};
let packets = &map.packets[packets];
let (a, b) = (packets[0], packets[packets.len() - 1]);
map.stream(reply, from, a.pack, &mut |bytes| self.pack(bytes))?;
if !self.pending.is_empty() {
self.packet(self.pending.len())?;
}
let lead = low - start;
let padding = crate::reply_pack::padding(self.at + lead);
self.file
.write_all_at(&padding, self.at)
.context("write a reply")?;
self.at += padding.len() as u64;
self.padding += padding.len() as u64;
let mut edge = vec![0; lead.max(end - high) as usize];
reply.read_exact_at(&mut edge[..lead as usize], start)?;
self.file.write_all_at(&edge[..lead as usize], self.at)?;
crate::reply_pack::clone_range(reply, low, self.file, self.at + lead, high - low)
.context("share a reply's bytes")?;
reply.read_exact_at(&mut edge[..(end - high) as usize], high)?;
self.file
.write_all_at(&edge[..(end - high) as usize], self.at + (high - start))?;
let carried: u64 = (packets.iter())
.map(|packet| 5 + u64::from(packet.len))
.sum();
self.padding += (end - start) - carried;
let first_shared = self.packets.len();
for packet in packets {
self.packets.push(Packet {
file: packet.file - start + self.at,
pack: packet.pack - a.pack + self.pack_at,
len: packet.len,
});
}
let mut buffer = vec![0; 5 + PACKET_DATA];
let mut pads = Vec::new();
for idx in first_shared..self.packets.len() {
let packet = self.packets[idx];
if idx > first_shared {
let prev_end = self.packets[idx - 1].file_end();
let gap = usize::try_from(packet.file - prev_end)?;
if gap > 0 {
pads.resize(gap, 0);
self.file
.read_exact_at(&mut pads, prev_end)
.context("read a reply")?;
anyhow::ensure!(
pads.chunks(5).all(|pad| pad == crate::reply_pack::PAD),
"padding in a reply doesn't match its map"
);
}
}
let total = 5 + packet.len as usize;
self.file
.read_exact_at(&mut buffer[..total], packet.file)
.context("read a reply")?;
anyhow::ensure!(
buffer[..5] == packet.header(),
"a packet header in a reply doesn't match its map"
);
self.hasher.update(&buffer[5..total]);
}
self.at += end - start;
self.pack_at += b.end() - a.pack;
map.stream(reply, b.end(), to, &mut |bytes| self.pack(bytes))
}
pub(crate) fn finish(mut self) -> anyhow::Result<Map> {
let trailer = std::mem::replace(&mut self.hasher, Hasher::new(self.format)).finish();
self.queue(&trailer)?;
let trailer = ObjectId::try_from(&trailer[..]).context("a pack's checksum")?;
if !self.pending.is_empty() {
self.packet(self.pending.len())?;
}
self.file
.write_all_at(b"0000", self.at)
.context("write a reply")?;
self.file.set_len(self.at + 4)?;
Ok(Map {
format: self.format,
entries: self.entries,
packets: self.packets,
pack_len: self.pack_at,
padding: self.padding,
total_size: self.total_size,
trailer,
})
}
}