use super::{
capture::{CHECKPOINT_SCHEMA_VERSION, LedgerRead, SnapshotBlobRef, SnapshotLedgerRecord},
sanitized_diagnostic,
};
use crate::{
config::McPaths,
persistence::{CrossProcessFileLock, atomic_write_with_permissions, sync_parent_dir},
sessions::SESSION_FILE_MODE,
};
use sha2::{Digest, Sha256};
use std::{
error::Error,
fmt, fs,
io::{ErrorKind, Read, Seek, SeekFrom, Write},
path::{Path, PathBuf},
};
#[derive(Debug)]
pub(crate) enum CheckpointAppendFailure {
OutcomeUncertain {
path: PathBuf,
source: anyhow::Error,
},
CommittedButUndurable {
path: PathBuf,
source: anyhow::Error,
},
}
impl CheckpointAppendFailure {
fn uncertain(path: &Path, source: anyhow::Error) -> anyhow::Error {
anyhow::Error::new(Self::OutcomeUncertain {
path: path.to_path_buf(),
source,
})
}
fn committed_but_undurable(path: &Path, source: anyhow::Error) -> anyhow::Error {
anyhow::Error::new(Self::CommittedButUndurable {
path: path.to_path_buf(),
source,
})
}
}
impl fmt::Display for CheckpointAppendFailure {
fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
let (path, outcome) = match self {
Self::OutcomeUncertain { path, .. } => {
(path, "append outcome is uncertain; do not retry blindly")
}
Self::CommittedButUndurable { path, .. } => (
path,
"record is visible but final synchronization failed; do not retry blindly",
),
};
write!(
formatter,
"checkpoint ledger append failed for {}: {outcome}",
path.display()
)
}
}
impl Error for CheckpointAppendFailure {
fn source(&self) -> Option<&(dyn Error + 'static)> {
let source = match self {
Self::OutcomeUncertain { source, .. } | Self::CommittedButUndurable { source, .. } => {
source
}
};
Some(source.as_ref())
}
}
#[derive(Debug, Default)]
struct LedgerIoMetrics {
read_bytes: u64,
write_bytes: u64,
serialized_record_bytes: u64,
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub(crate) struct CheckpointStore {
pub(super) root: PathBuf,
}
impl CheckpointStore {
pub(crate) fn new(root: PathBuf) -> Self {
Self { root }
}
pub(crate) fn from_paths(paths: &McPaths) -> Self {
Self::new(paths.checkpoints.clone())
}
pub(crate) fn copy_before_turn(
&self,
source: &str,
target: &str,
turn: u64,
) -> anyhow::Result<()> {
crate::sessions::validate_session_id(source.to_owned())?;
crate::sessions::validate_session_id(target.to_owned())?;
let read = self.read_records(source);
anyhow::ensure!(
read.diagnostics.is_empty(),
"cannot copy unreadable file checkpoints"
);
for mut record in read
.records
.into_iter()
.filter(|record| record.event.user_turn < turn)
{
record.event.session_id = target.to_owned();
self.append_record(&record)?;
}
Ok(())
}
pub(super) fn blobs_dir(&self) -> PathBuf {
self.root.join("blobs")
}
pub(super) fn ledgers_dir(&self) -> PathBuf {
self.root.join("ledgers")
}
pub(super) fn ledger_path(&self, session_id: &str) -> PathBuf {
self.ledgers_dir().join(format!("{session_id}.jsonl"))
}
pub(super) fn blob_path(&self, sha256: &str) -> PathBuf {
self.blobs_dir().join(sha256)
}
fn mutation_lock_target(&self) -> PathBuf {
self.root.with_file_name("checkpoint-store")
}
pub(super) fn lock_mutations(&self) -> anyhow::Result<CrossProcessFileLock> {
CrossProcessFileLock::acquire(&self.mutation_lock_target())
}
pub(super) fn write_blob(&self, bytes: &[u8]) -> anyhow::Result<SnapshotBlobRef> {
let sha256 = sha256_hex(bytes);
let path = self.blob_path(&sha256);
if path.exists() {
return Ok(SnapshotBlobRef {
sha256,
bytes: bytes.len() as u64,
});
}
create_private_dir_all(&self.blobs_dir())?;
atomic_write_with_permissions(&path, bytes, Some(SESSION_FILE_MODE))?;
Ok(SnapshotBlobRef {
sha256,
bytes: bytes.len() as u64,
})
}
pub(super) fn read_blob(&self, blob: &SnapshotBlobRef) -> anyhow::Result<Vec<u8>> {
let bytes = fs::read(self.blob_path(&blob.sha256))?;
if sha256_hex(&bytes) != blob.sha256 {
anyhow::bail!("checkpoint blob hash mismatch");
}
Ok(bytes)
}
pub(super) fn append_record(&self, record: &SnapshotLedgerRecord) -> anyhow::Result<()> {
self.append_record_with(
record,
None,
|file| file.sync_all().map_err(anyhow::Error::from),
sync_parent_dir,
)
}
fn append_record_with(
&self,
record: &SnapshotLedgerRecord,
mut metrics: Option<&mut LedgerIoMetrics>,
sync_file: impl FnOnce(&fs::File) -> anyhow::Result<()>,
sync_parent: impl FnOnce(&Path) -> anyhow::Result<()>,
) -> anyhow::Result<()> {
crate::sessions::validate_session_id(record.event.session_id.clone())?;
let ledger = self.ledger_path(&record.event.session_id);
if let Some(parent) = ledger.parent() {
create_private_dir_all(parent)?;
}
let mut line = serde_json::to_vec(record)?;
line.push(b'\n');
if let Some(metrics) = metrics.as_deref_mut() {
metrics.serialized_record_bytes += line.len() as u64;
}
let _lock = CrossProcessFileLock::acquire(&ledger)?;
let (mut file, created) = open_ledger_for_append(&ledger)?;
let pre_append_len = file.metadata()?.len();
let needs_separator = if pre_append_len == 0 {
false
} else {
file.seek(SeekFrom::End(-1))?;
let mut tail = [0_u8; 1];
file.read_exact(&mut tail)?;
if let Some(metrics) = metrics.as_deref_mut() {
metrics.read_bytes += 1;
}
tail[0] != b'\n'
};
if needs_separator {
line.insert(0, b'\n');
}
if let Some(metrics) = metrics {
metrics.write_bytes += line.len() as u64;
}
if let Err(write_error) = file.write_all(&line) {
let rollback = file.set_len(pre_append_len).and_then(|()| file.sync_all());
return match rollback {
Ok(()) => Err(anyhow::Error::new(write_error).context(
"checkpoint ledger append failed; partial data was durably rolled back",
)),
Err(rollback_error) => Err(CheckpointAppendFailure::uncertain(
&ledger,
anyhow::anyhow!(
"append write failed ({write_error}); rollback failed ({rollback_error})"
),
)),
};
}
file.flush().map_err(|error| {
CheckpointAppendFailure::committed_but_undurable(&ledger, error.into())
})?;
sync_file(&file)
.map_err(|error| CheckpointAppendFailure::committed_but_undurable(&ledger, error))?;
if created {
let parent = ledger
.parent()
.ok_or_else(|| anyhow::anyhow!("checkpoint ledger has no parent"))?;
sync_parent(parent).map_err(|error| {
CheckpointAppendFailure::committed_but_undurable(&ledger, error)
})?;
}
Ok(())
}
pub(crate) fn read_records(&self, session_id: &str) -> LedgerRead {
let ledger = self.ledger_path(session_id);
let mut read = LedgerRead::default();
let bytes = match fs::read(&ledger) {
Ok(bytes) => bytes,
Err(error) if error.kind() == ErrorKind::NotFound => return read,
Err(error) => {
read.diagnostics.push(sanitized_diagnostic(format!(
"checkpoint ledger read failed: {error}"
)));
return read;
}
};
let final_segment = bytes.split(|byte| *byte == b'\n').count().saturating_sub(1);
let has_terminated_tail = bytes.ends_with(b"\n");
for (index, line) in bytes.split(|byte| *byte == b'\n').enumerate() {
if line.iter().all(u8::is_ascii_whitespace) {
continue;
}
match serde_json::from_slice::<SnapshotLedgerRecord>(line) {
Ok(record) if record.schema_version == CHECKPOINT_SCHEMA_VERSION => {
read.records.push(record);
}
Ok(_) => read.diagnostics.push(format!(
"ignored checkpoint ledger line {}: unsupported_schema_version",
index + 1
)),
Err(_) if index == final_segment && !has_terminated_tail => {
read.diagnostics.push(format!(
"ignored checkpoint ledger line {}: incomplete_tail",
index + 1
))
}
Err(_) => read.diagnostics.push(format!(
"ignored checkpoint ledger line {}: malformed_json",
index + 1
)),
}
}
read
}
pub(crate) fn prune_session(&self, session_id: &str) -> anyhow::Result<()> {
if !self.validate_deletion_directories()? {
return Ok(());
}
let _mutation_lock = self.lock_mutations()?;
if !self.validate_deletion_directories()? {
return Ok(());
}
let ledger = self.ledger_path(session_id);
if ledger.exists() {
fs::remove_file(&ledger)?;
}
self.prune_unreferenced_blobs_locked()
}
pub(crate) fn prune_unreferenced_blobs(&self) -> anyhow::Result<()> {
if !self.validate_deletion_directories()? {
return Ok(());
}
let _mutation_lock = self.lock_mutations()?;
if !self.validate_deletion_directories()? {
return Ok(());
}
self.prune_unreferenced_blobs_locked()
}
fn validate_deletion_directories(&self) -> anyhow::Result<bool> {
if !validate_owned_directory(&self.root)? {
return Ok(false);
}
for path in [self.ledgers_dir(), self.blobs_dir()] {
match fs::symlink_metadata(&path) {
Ok(_) => {
if !validate_owned_directory(&path)? {
anyhow::bail!("unsafe checkpoint directory: {}", path.display());
}
}
Err(error) if error.kind() == ErrorKind::NotFound => {}
Err(error) => return Err(error.into()),
}
}
Ok(true)
}
fn prune_unreferenced_blobs_locked(&self) -> anyhow::Result<()> {
let mut referenced = std::collections::BTreeSet::new();
let ledgers = self.ledgers_dir();
if ledgers.exists() {
for entry in fs::read_dir(&ledgers)? {
let entry = entry?;
let file_type = entry.file_type()?;
if file_type.is_symlink() {
anyhow::bail!("unsafe checkpoint ledger entry");
}
if !file_type.is_file() {
continue;
}
let file_name = entry
.file_name()
.into_string()
.map_err(|_| anyhow::anyhow!("checkpoint ledger name is not UTF-8"))?;
let Some(session_id) = file_name.strip_suffix(".jsonl") else {
continue;
};
crate::sessions::validate_session_id(session_id.to_string())?;
let read = self.read_records(session_id);
if !read.diagnostics.is_empty() {
anyhow::bail!("checkpoint ledger could not be read completely");
}
for record in read.records {
if let Some(blob) = record.event.pre {
referenced.insert(blob.sha256);
}
if let Some(blob) = record.event.post {
referenced.insert(blob.sha256);
}
}
}
}
let blobs = self.blobs_dir();
if !blobs.exists() {
return Ok(());
}
for entry in fs::read_dir(&blobs)? {
let entry = entry?;
if !entry.file_type()?.is_file() {
continue;
}
let name = entry.file_name().to_string_lossy().into_owned();
if !referenced.contains(&name) {
fs::remove_file(entry.path())?;
}
}
Ok(())
}
}
fn open_ledger_for_append(path: &Path) -> anyhow::Result<(fs::File, bool)> {
match fs::symlink_metadata(path) {
Ok(metadata) if metadata.file_type().is_symlink() || !metadata.is_file() => {
anyhow::bail!("unsafe checkpoint ledger entry")
}
Ok(_) => {
let file = fs::OpenOptions::new().read(true).append(true).open(path)?;
ensure_open_ledger_matches_path(&file, path)?;
Ok((file, false))
}
Err(error) if error.kind() == ErrorKind::NotFound => {
let file = fs::OpenOptions::new()
.read(true)
.append(true)
.create_new(true)
.open(path)?;
ensure_open_ledger_matches_path(&file, path)?;
Ok((file, true))
}
Err(error) => Err(error.into()),
}
}
fn ensure_open_ledger_matches_path(file: &fs::File, path: &Path) -> anyhow::Result<()> {
let path_metadata = fs::symlink_metadata(path)?;
let file_metadata = file.metadata()?;
if path_metadata.file_type().is_symlink()
|| !path_metadata.is_file()
|| !file_metadata.is_file()
{
anyhow::bail!("unsafe checkpoint ledger entry");
}
#[cfg(unix)]
{
use std::os::unix::fs::MetadataExt;
if path_metadata.dev() != file_metadata.dev()
|| path_metadata.ino() != file_metadata.ino()
|| file_metadata.uid() != unsafe { libc::geteuid() }
{
anyhow::bail!("unsafe checkpoint ledger entry");
}
}
Ok(())
}
fn create_private_dir_all(path: &Path) -> std::io::Result<()> {
#[cfg(unix)]
{
use std::os::unix::fs::DirBuilderExt;
fs::DirBuilder::new()
.recursive(true)
.mode(crate::sessions::SESSION_ROOT_MODE)
.create(path)
}
#[cfg(not(unix))]
{
fs::create_dir_all(path)
}
}
fn validate_owned_directory(path: &std::path::Path) -> anyhow::Result<bool> {
let metadata = match fs::symlink_metadata(path) {
Ok(metadata) => metadata,
Err(error) if error.kind() == ErrorKind::NotFound => return Ok(false),
Err(error) => return Err(error.into()),
};
if metadata.file_type().is_symlink() || !metadata.file_type().is_dir() {
anyhow::bail!("unsafe checkpoint directory: {}", path.display());
}
#[cfg(unix)]
{
use std::os::unix::fs::MetadataExt;
if metadata.uid() != unsafe { libc::geteuid() } {
anyhow::bail!("checkpoint directory is not owned by the current user");
}
}
Ok(true)
}
pub(crate) fn sha256_hex(bytes: &[u8]) -> String {
let digest = Sha256::digest(bytes);
crate::hex::lower_hex(digest)
}