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
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
use std::error::Error as StdError;
use std::future::Future;
use std::ops::RangeInclusive;
use std::{fmt, io};

use ::bytes::Bytes;

use super::ProtocolVersion;

mod bytes;
mod collections;
mod int;
#[cfg(any(test, feature = "test"))]
pub mod mock;
mod reader;

pub use reader::{NixReader, NixReaderBuilder};

/// Like serde the `Error` trait allows `NixRead` implementations to add
/// custom error handling for `NixDeserialize`.
pub trait Error: Sized + StdError {
    /// A totally custom non-specific error.
    fn custom<T: fmt::Display>(msg: T) -> Self;

    /// Some kind of std::io::Error occured.
    fn io_error(err: std::io::Error) -> Self {
        Self::custom(format_args!("There was an I/O error {}", err))
    }

    /// The data read from `NixRead` is invalid.
    /// This could be that some bytes were supposed to be valid UFT-8 but weren't.
    fn invalid_data<T: fmt::Display>(msg: T) -> Self {
        Self::custom(msg)
    }

    /// Required data is missing. This is mostly like an EOF
    fn missing_data<T: fmt::Display>(msg: T) -> Self {
        Self::custom(msg)
    }
}

impl Error for io::Error {
    fn custom<T: fmt::Display>(msg: T) -> Self {
        io::Error::new(io::ErrorKind::Other, msg.to_string())
    }

    fn io_error(err: std::io::Error) -> Self {
        err
    }

    fn invalid_data<T: fmt::Display>(msg: T) -> Self {
        io::Error::new(io::ErrorKind::InvalidData, msg.to_string())
    }

    fn missing_data<T: fmt::Display>(msg: T) -> Self {
        io::Error::new(io::ErrorKind::UnexpectedEof, msg.to_string())
    }
}

/// A reader of data from the Nix daemon protocol.
/// Basically there are two basic types in the Nix daemon protocol
/// u64 and a bytes buffer. Everything else is more or less built on
/// top of these two types.
pub trait NixRead: Send {
    type Error: Error + Send;

    /// Some types are serialized differently depending on the version
    /// of the protocol and so this can be used for implementing that.
    fn version(&self) -> ProtocolVersion;

    /// Read a single u64 from the protocol.
    /// This returns an Option to support graceful shutdown.
    fn try_read_number(
        &mut self,
    ) -> impl Future<Output = Result<Option<u64>, Self::Error>> + Send + '_;

    /// Read bytes from the protocol.
    /// A size limit on the returned bytes has to be specified.
    /// This returns an Option to support graceful shutdown.
    fn try_read_bytes_limited(
        &mut self,
        limit: RangeInclusive<usize>,
    ) -> impl Future<Output = Result<Option<Bytes>, Self::Error>> + Send + '_;

    /// Read bytes from the protocol without a limit.
    /// The default implementation just calls `try_read_bytes_limited` with a
    /// limit of `0..=usize::MAX` but other implementations are free to have a
    /// reader wide limit.
    /// This returns an Option to support graceful shutdown.
    fn try_read_bytes(
        &mut self,
    ) -> impl Future<Output = Result<Option<Bytes>, Self::Error>> + Send + '_ {
        self.try_read_bytes_limited(0..=usize::MAX)
    }

    /// Read a single u64 from the protocol.
    /// This will return an error if the number could not be read.
    fn read_number(&mut self) -> impl Future<Output = Result<u64, Self::Error>> + Send + '_ {
        async move {
            match self.try_read_number().await? {
                Some(v) => Ok(v),
                None => Err(Self::Error::missing_data("unexpected end-of-file")),
            }
        }
    }

    /// Read bytes from the protocol.
    /// A size limit on the returned bytes has to be specified.
    /// This will return an error if the number could not be read.
    fn read_bytes_limited(
        &mut self,
        limit: RangeInclusive<usize>,
    ) -> impl Future<Output = Result<Bytes, Self::Error>> + Send + '_ {
        async move {
            match self.try_read_bytes_limited(limit).await? {
                Some(v) => Ok(v),
                None => Err(Self::Error::missing_data("unexpected end-of-file")),
            }
        }
    }

    /// Read bytes from the protocol.
    /// The default implementation just calls `read_bytes_limited` with a
    /// limit of `0..=usize::MAX` but other implementations are free to have a
    /// reader wide limit.
    /// This will return an error if the bytes could not be read.
    fn read_bytes(&mut self) -> impl Future<Output = Result<Bytes, Self::Error>> + Send + '_ {
        self.read_bytes_limited(0..=usize::MAX)
    }

    /// Read a value from the protocol.
    /// Uses `NixDeserialize::deserialize` to read a value.
    fn read_value<V: NixDeserialize>(
        &mut self,
    ) -> impl Future<Output = Result<V, Self::Error>> + Send + '_ {
        V::deserialize(self)
    }

    /// Read a value from the protocol.
    /// Uses `NixDeserialize::try_deserialize` to read a value.
    /// This returns an Option to support graceful shutdown.
    fn try_read_value<V: NixDeserialize>(
        &mut self,
    ) -> impl Future<Output = Result<Option<V>, Self::Error>> + Send + '_ {
        V::try_deserialize(self)
    }
}

impl<T: ?Sized + NixRead> NixRead for &mut T {
    type Error = T::Error;

    fn version(&self) -> ProtocolVersion {
        (**self).version()
    }

    fn try_read_number(
        &mut self,
    ) -> impl Future<Output = Result<Option<u64>, Self::Error>> + Send + '_ {
        (**self).try_read_number()
    }

    fn try_read_bytes_limited(
        &mut self,
        limit: RangeInclusive<usize>,
    ) -> impl Future<Output = Result<Option<Bytes>, Self::Error>> + Send + '_ {
        (**self).try_read_bytes_limited(limit)
    }

    fn try_read_bytes(
        &mut self,
    ) -> impl Future<Output = Result<Option<Bytes>, Self::Error>> + Send + '_ {
        (**self).try_read_bytes()
    }

    fn read_number(&mut self) -> impl Future<Output = Result<u64, Self::Error>> + Send + '_ {
        (**self).read_number()
    }

    fn read_bytes_limited(
        &mut self,
        limit: RangeInclusive<usize>,
    ) -> impl Future<Output = Result<Bytes, Self::Error>> + Send + '_ {
        (**self).read_bytes_limited(limit)
    }

    fn read_bytes(&mut self) -> impl Future<Output = Result<Bytes, Self::Error>> + Send + '_ {
        (**self).read_bytes()
    }

    fn try_read_value<V: NixDeserialize>(
        &mut self,
    ) -> impl Future<Output = Result<Option<V>, Self::Error>> + Send + '_ {
        (**self).try_read_value()
    }

    fn read_value<V: NixDeserialize>(
        &mut self,
    ) -> impl Future<Output = Result<V, Self::Error>> + Send + '_ {
        (**self).read_value()
    }
}

/// A data structure that can be deserialized from the Nix daemon
/// worker protocol.
pub trait NixDeserialize: Sized {
    /// Read a value from the reader.
    /// This returns an Option to support gracefull shutdown.
    fn try_deserialize<R>(
        reader: &mut R,
    ) -> impl Future<Output = Result<Option<Self>, R::Error>> + Send + '_
    where
        R: ?Sized + NixRead + Send;

    fn deserialize<R>(reader: &mut R) -> impl Future<Output = Result<Self, R::Error>> + Send + '_
    where
        R: ?Sized + NixRead + Send,
    {
        async move {
            match Self::try_deserialize(reader).await? {
                Some(v) => Ok(v),
                None => Err(R::Error::missing_data("unexpected end-of-file")),
            }
        }
    }
}