fix: complete paste artifact storage contract

This commit is contained in:
2026-09-01 17:51:24 +09:00
parent 04e296a4ef
commit 2bb661f1cf
11 changed files with 188 additions and 15 deletions
+1
View File
@@ -8,6 +8,7 @@ license.workspace = true
[dependencies]
base64.workspace = true
agen = { workspace = true }
fs4.workspace = true
serde = { workspace = true, features = ["derive"] }
serde_json = { workspace = true }
sha2.workspace = true
+59
View File
@@ -450,6 +450,15 @@ mod tests {
.unwrap();
assert_eq!(reference.byte_len, content.len() as u64);
assert!(reference.created_at_ms > 0);
assert_eq!(
reference.media_type,
protocol::PasteArtifactMediaType::TextPlainUtf8
);
assert_eq!(
reference.availability,
protocol::PasteArtifactAvailability::Available
);
assert_eq!(reference.char_count, content.chars().count() as u64);
assert_eq!(reference.source_entry_id, "entry-1");
assert_eq!(
@@ -490,6 +499,55 @@ mod tests {
);
}
#[test]
fn concurrent_paste_writes_atomically_enforce_aggregate_caps() {
let tmp = tempfile::TempDir::new().unwrap();
let session_id = new_session_id();
let barrier = std::sync::Arc::new(std::sync::Barrier::new(3));
let limits = PasteArtifactLimits {
max_artifact_bytes: 4,
max_session_bytes: 8,
max_session_artifacts: 1,
};
let mut handles = Vec::new();
for entry_id in ["entry-1", "entry-2"] {
let root = tmp.path().to_path_buf();
let barrier = barrier.clone();
handles.push(std::thread::spawn(move || {
let store = FsStore::new(root).unwrap();
barrier.wait();
store.write_paste_artifact(session_id, entry_id, "1234", limits)
}));
}
barrier.wait();
let results = handles
.into_iter()
.map(|handle| handle.join().unwrap())
.collect::<Vec<_>>();
assert_eq!(results.iter().filter(|result| result.is_ok()).count(), 1);
assert_eq!(
results
.iter()
.filter(|result| matches!(result, Err(StoreError::PasteArtifactLimit(_))))
.count(),
1
);
assert_eq!(
std::fs::read_dir(
FsStore::new(tmp.path())
.unwrap()
.paste_artifact_dir(session_id)
)
.unwrap()
.filter_map(Result::ok)
.filter(
|entry| entry.path().extension().and_then(|value| value.to_str()) == Some("json")
)
.count(),
1
);
}
#[test]
fn paste_artifact_limits_and_corruption_fail_closed() {
let tmp = tempfile::TempDir::new().unwrap();
@@ -498,6 +556,7 @@ mod tests {
let limits = PasteArtifactLimits {
max_artifact_bytes: 5,
max_session_bytes: 8,
max_session_artifacts: 2,
};
let first = store
.write_paste_artifact(session_id, "entry-1", "1234", limits)
+31 -1
View File
@@ -3,8 +3,10 @@
use std::fs;
use std::io::Write as _;
use std::path::Path;
use std::time::{SystemTime, UNIX_EPOCH};
use protocol::PasteArtifactRef;
use fs4::fs_std::FileExt;
use protocol::{PasteArtifactAvailability, PasteArtifactMediaType, PasteArtifactRef};
use serde::{Deserialize, Serialize};
use sha2::{Digest, Sha256};
@@ -15,6 +17,7 @@ use crate::StoreError;
pub struct PasteArtifactLimits {
pub max_artifact_bytes: u64,
pub max_session_bytes: u64,
pub max_session_artifacts: u64,
}
impl Default for PasteArtifactLimits {
@@ -22,6 +25,7 @@ impl Default for PasteArtifactLimits {
Self {
max_artifact_bytes: 8 * 1024 * 1024,
max_session_bytes: 64 * 1024 * 1024,
max_session_artifacts: 1_024,
}
}
}
@@ -48,7 +52,14 @@ pub(crate) fn write_to_dir(
)));
}
fs::create_dir_all(artifact_dir)?;
let aggregate_lock = fs::OpenOptions::new()
.create(true)
.read(true)
.write(true)
.open(artifact_dir.join(".aggregate.lock"))?;
FileExt::lock_exclusive(&aggregate_lock)?;
let mut aggregate = 0_u64;
let mut artifact_count = 0_u64;
for entry in fs::read_dir(artifact_dir)? {
let path = entry?.path();
if path.extension().and_then(|value| value.to_str()) != Some("json") {
@@ -56,6 +67,9 @@ pub(crate) fn write_to_dir(
}
let stored: StoredPasteArtifact = serde_json::from_slice(&fs::read(&path)?)?;
verify(&stored, &stored.reference.artifact_id)?;
artifact_count = artifact_count.checked_add(1).ok_or_else(|| {
StoreError::PasteArtifactLimit("session artifact count overflow".to_string())
})?;
aggregate = aggregate
.checked_add(stored.reference.byte_len)
.ok_or_else(|| {
@@ -71,10 +85,23 @@ pub(crate) fn write_to_dir(
limits.max_session_bytes
)));
}
if artifact_count >= limits.max_session_artifacts {
return Err(StoreError::PasteArtifactLimit(format!(
"session already has {artifact_count} artifacts; maximum is {}",
limits.max_session_artifacts
)));
}
let artifact_id = uuid::Uuid::now_v7().to_string();
let created_at_ms = SystemTime::now()
.duration_since(UNIX_EPOCH)
.map_err(|error| StoreError::PasteArtifactIntegrity(error.to_string()))?
.as_millis() as u64;
let reference = PasteArtifactRef {
artifact_id: artifact_id.clone(),
created_at_ms,
media_type: PasteArtifactMediaType::TextPlainUtf8,
availability: PasteArtifactAvailability::Available,
byte_len,
char_count: content.chars().count() as u64,
line_count: line_count(content),
@@ -130,6 +157,9 @@ pub(crate) fn read_from_dir(
fn verify(stored: &StoredPasteArtifact, artifact_id: &str) -> Result<(), StoreError> {
let actual_digest = sha256_hex(&stored.content);
if stored.reference.artifact_id != artifact_id
|| stored.reference.created_at_ms == 0
|| stored.reference.media_type != PasteArtifactMediaType::TextPlainUtf8
|| stored.reference.availability != PasteArtifactAvailability::Available
|| stored.reference.byte_len != stored.content.len() as u64
|| stored.reference.char_count != stored.content.chars().count() as u64
|| stored.reference.line_count != line_count(&stored.content)