about summary refs log tree commit diff
path: root/tvix/nix-compat/src/nar/reader/read.rs
blob: 9938581f2a2e41a9dd4baca7e21ee30c151390ae (plain) (blame)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
//! Helpers for reading [crate::nar::wire] format.

use std::io::{
    self,
    ErrorKind::{Interrupted, InvalidData, UnexpectedEof},
};

use super::Reader;
use crate::nar::wire::Tag;

/// Consume a little-endian [prim@u64] from the reader.
pub fn u64(reader: &mut Reader) -> io::Result<u64> {
    let mut buf = [0; 8];
    reader.read_exact(&mut buf)?;
    Ok(u64::from_le_bytes(buf))
}

/// Consume a byte string from the reader into a provided buffer,
/// returning the data bytes.
pub fn bytes_buf<'a, const N: usize>(
    reader: &mut Reader,
    buf: &'a mut [u8; N],
    max_len: usize,
) -> io::Result<&'a [u8]> {
    assert_eq!(N % 8, 0);
    assert!(max_len <= N);

    // read the length, and reject excessively large values
    let len = self::u64(reader)?;
    if len > max_len as u64 {
        return Err(InvalidData.into());
    }
    // we know the length fits in a usize now
    let len = len as usize;

    // read the data and padding into a buffer
    let buf_len = (len + 7) & !7;
    reader.read_exact(&mut buf[..buf_len])?;

    // verify that the padding is all zeroes
    for &b in &buf[len..buf_len] {
        if b != 0 {
            return Err(InvalidData.into());
        }
    }

    Ok(&buf[..len])
}

/// Consume a byte string of up to `max_len` bytes from the reader.
pub fn bytes(reader: &mut Reader, max_len: usize) -> io::Result<Vec<u8>> {
    assert!(max_len <= isize::MAX as usize);

    // read the length, and reject excessively large values
    let len = self::u64(reader)?;
    if len > max_len as u64 {
        return Err(InvalidData.into());
    }
    // we know the length fits in a usize now
    let len = len as usize;

    // read the data and padding into a buffer
    let buf_len = (len + 7) & !7;
    let mut buf = vec![0; buf_len];
    reader.read_exact(&mut buf)?;

    // verify that the padding is all zeroes
    for b in buf.drain(len..) {
        if b != 0 {
            return Err(InvalidData.into());
        }
    }

    Ok(buf)
}

/// Consume a known token from the reader.
pub fn token<const N: usize>(reader: &mut Reader, token: &[u8; N]) -> io::Result<()> {
    let mut buf = [0u8; N];

    // This implements something similar to [Read::read_exact], but verifies that
    // the input data matches the token while we read it. These two slices respectively
    // represent the remaining token to be verified, and the remaining input buffer.
    let mut token = &token[..];
    let mut buf = &mut buf[..];

    while !token.is_empty() {
        match reader.read(buf) {
            Ok(0) => {
                return Err(UnexpectedEof.into());
            }
            Ok(n) => {
                let (t, b);
                (t, token) = token.split_at(n);
                (b, buf) = buf.split_at_mut(n);

                if t != b {
                    return Err(InvalidData.into());
                }
            }
            Err(e) => {
                if e.kind() != Interrupted {
                    return Err(e);
                }
            }
        }
    }

    Ok(())
}

/// Consume a [Tag] from the reader.
pub fn tag<T: Tag>(reader: &mut Reader) -> io::Result<T> {
    let mut buf = T::make_buf();
    let buf = buf.as_mut();

    // first read the known minimum length…
    reader.read_exact(&mut buf[..T::MIN])?;

    // then decide which tag we're expecting
    let tag = T::from_u8(buf[T::OFF]).ok_or(InvalidData)?;
    let (head, tail) = tag.as_bytes().split_at(T::MIN);

    // make sure what we've read so far is valid
    if buf[..T::MIN] != *head {
        return Err(InvalidData.into());
    }

    // …then read the rest, if any
    if !tail.is_empty() {
        let rest = tail.len();
        reader.read_exact(&mut buf[..rest])?;

        // and make sure it's what we expect
        if buf[..rest] != *tail {
            return Err(InvalidData.into());
        }
    }

    Ok(tag)
}