mdebray/tvl-depot
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

refactor(nix-compat/nar/reader): always enable poisoning *API*

Browse source 
The poisoning API is now always available, whether debug_assertions is
enabled or not. When debug assertions are not enabled, it is equivalent
to a unit struct, and is always considered ready and unpoisoned.

Change-Id: I950237f4816d480330d9acab32661ed4f1663049
Reviewed-on: https://cl.tvl.fyi/c/depot/+/11525
Tested-by: BuildkiteCI
Reviewed-by: flokli <flokli@flokli.de>
This commit is contained in:
edef 2024-04-26 09:52:26 +00:00
parent 98b85b4580
commit 5070f9eff7
1 changed files with 66 additions and 64 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

110
tvix/nix-compat/src/nar/reader/mod.rs
View file

@ -10,6 +10,9 @@ use std::io::{
Read, Write,
};
#[cfg(not(debug_assertions))]
use std::marker::PhantomData;
// Required reading for understanding this module.
use crate::nar::wire;
@ -27,25 +30,15 @@ struct ArchiveReader<'a, 'r> {
/// * An error is encountered at any point
/// * A file or directory reader is dropped before being read entirely.
/// All of these checks vanish in release mode.
#[cfg(debug_assertions)]
status: ArchiveReaderStatus<'a>,
}
macro_rules! poison {
($it:expr) => {
#[cfg(debug_assertions)]
{
$it.status.poison();
}
};
}
macro_rules! try_or_poison {
($it:expr, $ex:expr) => {
match $ex {
Ok(x) => x,
Err(e) => {
poison!($it);
$it.status.poison();
return Err(e.into());
}
}
@ -56,11 +49,7 @@ pub fn open<'a, 'r>(reader: &'a mut Reader<'r>) -> io::Result<Node<'a, 'r>> {
read::token(reader, &wire::TOK_NAR)?;
Node::new(ArchiveReader {
inner: reader,
#[cfg(debug_assertions)]
status: ArchiveReaderStatus::StackTop {
poisoned: false,
ready: true,
},
status: ArchiveReaderStatus::top(),
})
}
@ -80,7 +69,6 @@ impl<'a, 'r> Node<'a, 'r> {
///
/// Reading the terminating [wire::TOK_PAR] is done immediately for [Node::Symlink],
/// but is otherwise left to [DirReader] or [FileReader].
#[allow(unused_mut)] // due to debug_assertions code
fn new(mut reader: ArchiveReader<'a, 'r>) -> io::Result<Self> {
Ok(match read::tag(reader.inner)? {
wire::Node::Sym => {
@ -88,15 +76,12 @@ impl<'a, 'r> Node<'a, 'r> {
try_or_poison!(reader, read::bytes(reader.inner, wire::MAX_TARGET_LEN));
if target.is_empty() || target.contains(&0) {
poison!(reader);
reader.status.poison();
return Err(InvalidData.into());
}
try_or_poison!(reader, read::token(reader.inner, &wire::TOK_PAR));
#[cfg(debug_assertions)]
{
reader.status.ready_parent(); // Immediately allow reading from parent again
}
Node::Symlink { target }
}
@ -131,18 +116,14 @@ pub struct FileReader<'a, 'r> {
impl<'a, 'r> FileReader<'a, 'r> {
/// Instantiate a new reader, starting after [wire::TOK_REG] or [wire::TOK_EXE].
/// We handle the terminating [wire::TOK_PAR] on semantic EOF.
#[allow(unused_mut)] // due to debug_assertions code
fn new(mut reader: ArchiveReader<'a, 'r>, len: u64) -> io::Result<Self> {
// For zero-length files, we have to read the terminating TOK_PAR
// immediately, since FileReader::read may never be called; we've
// already reached semantic EOF by definition.
if len == 0 {
read::token(reader.inner, &wire::TOK_PAR)?;
#[cfg(debug_assertions)]
{
reader.status.ready_parent();
}
}
Ok(Self {
reader,
@ -175,7 +156,7 @@ impl FileReader<'_, '_> {
let mut buf = try_or_poison!(self.reader, self.reader.inner.fill_buf());
if buf.is_empty() {
poison!(self.reader);
self.reader.status.poison();
return Err(UnexpectedEof.into());
}
@ -237,7 +218,7 @@ impl Read for FileReader<'_, '_> {
self.len -= n as u64;
if n == 0 {
poison!(self.reader);
self.reader.status.poison();
return Err(UnexpectedEof.into());
}
@ -260,18 +241,15 @@ impl FileReader<'_, '_> {
try_or_poison!(self.reader, self.reader.inner.read_exact(&mut buf[pad..]));
if buf != [0; 8] {
poison!(self.reader);
self.reader.status.poison();
return Err(InvalidData.into());
}
}
try_or_poison!(self.reader, read::token(self.reader.inner, &wire::TOK_PAR));
#[cfg(debug_assertions)]
{
// Done with reading this file, allow going back up the chain of readers
self.reader.status.ready_parent();
}
Ok(())
}
@ -321,10 +299,7 @@ impl<'a, 'r> DirReader<'a, 'r> {
// Determine if there are more entries to follow
if let wire::Entry::None = try_or_poison!(self.reader, read::tag(self.reader.inner)) {
// We've reached the end of this directory.
#[cfg(debug_assertions)]
{
self.reader.status.ready_parent();
}
return Ok(None);
}
@ -339,7 +314,7 @@ impl<'a, 'r> DirReader<'a, 'r> {
|| name == b"."
|| name == b".."
{
poison!(self.reader);
self.reader.status.poison();
return Err(InvalidData.into());
}
@ -350,7 +325,7 @@ impl<'a, 'r> DirReader<'a, 'r> {
}
Some(prev_name) => {
if *prev_name >= name {
poison!(self.reader);
self.reader.status.poison();
return Err(InvalidData.into());
}
@ -373,12 +348,12 @@ impl<'a, 'r> DirReader<'a, 'r> {
/// so we can check they are abandoned when an error occurs
/// * Make sure only the most recently created object is read from, and is fully exhausted
/// before anything it was created from is used again.
#[cfg(debug_assertions)]
enum ArchiveReaderStatus<'a> {
StackTop {
poisoned: bool,
ready: bool,
},
#[cfg(not(debug_assertions))]
None(PhantomData<&'a ()>),
#[cfg(debug_assertions)]
StackTop { poisoned: bool, ready: bool },
#[cfg(debug_assertions)]
StackChild {
poisoned: &'a mut bool,
parent_ready: &'a mut bool,
@ -386,12 +361,28 @@ enum ArchiveReaderStatus<'a> {
},
}
#[cfg(debug_assertions)]
impl ArchiveReaderStatus<'_> {
fn top() -> Self {
#[cfg(debug_assertions)]
{
ArchiveReaderStatus::StackTop {
poisoned: false,
ready: true,
}
}
#[cfg(not(debug_assertions))]
ArchiveReaderStatus::None(PhantomData)
}
/// Poison all the objects sharing the same reader, to be used when an error occurs
fn poison(&mut self) {
match self {
#[cfg(not(debug_assertions))]
ArchiveReaderStatus::None(_) => {}
#[cfg(debug_assertions)]
ArchiveReaderStatus::StackTop { poisoned: x, .. } => *x = true,
#[cfg(debug_assertions)]
ArchiveReaderStatus::StackChild { poisoned: x, .. } => **x = true,
}
}
@ -399,10 +390,14 @@ impl ArchiveReaderStatus<'_> {
/// Mark the parent as ready, allowing it to be used again and preventing this reference to the reader being used again.
fn ready_parent(&mut self) {
match self {
Self::StackTop { ready, .. } => {
#[cfg(not(debug_assertions))]
ArchiveReaderStatus::None(_) => {}
#[cfg(debug_assertions)]
ArchiveReaderStatus::StackTop { ready, .. } => {
*ready = false;
}
Self::StackChild {
#[cfg(debug_assertions)]
ArchiveReaderStatus::StackChild {
ready,
parent_ready,
..
@ -415,15 +410,23 @@ impl ArchiveReaderStatus<'_> {
fn poisoned(&self) -> bool {
match self {
Self::StackTop { poisoned, .. } => *poisoned,
Self::StackChild { poisoned, .. } => **poisoned,
#[cfg(not(debug_assertions))]
ArchiveReaderStatus::None(_) => false,
#[cfg(debug_assertions)]
ArchiveReaderStatus::StackTop { poisoned, .. } => *poisoned,
#[cfg(debug_assertions)]
ArchiveReaderStatus::StackChild { poisoned, .. } => **poisoned,
}
}
fn ready(&self) -> bool {
match self {
Self::StackTop { ready, .. } => *ready,
Self::StackChild { ready, .. } => *ready,
#[cfg(not(debug_assertions))]
ArchiveReaderStatus::None(_) => true,
#[cfg(debug_assertions)]
ArchiveReaderStatus::StackTop { ready, .. } => *ready,
#[cfg(debug_assertions)]
ArchiveReaderStatus::StackChild { ready, .. } => *ready,
}
}
}
@ -434,6 +437,8 @@ impl<'a, 'r> ArchiveReader<'a, 'r> {
fn child(&mut self) -> ArchiveReader<'_, 'r> {
ArchiveReader {
inner: self.inner,
#[cfg(not(debug_assertions))]
status: ArchiveReaderStatus::None(PhantomData),
#[cfg(debug_assertions)]
status: match &mut self.status {
ArchiveReaderStatus::StackTop { poisoned, ready } => {
@ -462,16 +467,13 @@ impl<'a, 'r> ArchiveReader<'a, 'r> {
/// Only does anything when debug assertions are on.
#[inline(always)]
fn check_correct(&self) {
#[cfg(debug_assertions)]
{
debug_assert!(
assert!(
!self.status.poisoned(),
"Archive reader used after it was meant to be abandoned!"
);
debug_assert!(
assert!(
self.status.ready(),
"Non-ready archive reader used! (Should've been reading from something else)"
)
}
);
}
}