nix_compat/nix_daemon/
handler.rs

1use std::{future::Future, ops::DerefMut, sync::Arc};
2
3use bytes::Bytes;
4use tokio::{
5    io::{AsyncReadExt, AsyncWriteExt, ReadHalf, WriteHalf, copy_buf, split},
6    sync::Mutex,
7};
8use tracing::{debug, warn};
9
10use super::{
11    NixDaemonIO,
12    framing::{NixFramedReader, StderrReadFramedReader},
13    types::{QueryValidPaths, ValidPathInfo},
14    worker_protocol::{ClientSettings, Operation, STDERR_LAST, Trust, server_handshake_client},
15};
16
17use crate::{
18    nix_daemon::types::BuildPaths,
19    store_path::StorePath,
20    wire::{
21        ProtocolVersion,
22        de::{NixRead, NixReader},
23        ser::{NixSerialize, NixWrite, NixWriter, NixWriterBuilder},
24    },
25};
26
27use crate::{nix_daemon::types::NixError, worker_protocol::STDERR_ERROR};
28
29/// Handles a single connection with a nix client.
30///
31/// As part of its [`initialization`] it performs the handshake with the client
32/// and determines the [ProtocolVersion] and [ClientSettings] to use for the remainder of the session.
33///
34/// Once initialized, [NixDaemon::handle_client] needs to be called to handle
35/// the rest of the session, it delegates all operation handling to an instance
36/// of [NixDaemonIO].
37///
38/// [`initialization`]: NixDaemon::initialize
39#[allow(dead_code)]
40pub struct NixDaemon<IO, R, W> {
41    io: Arc<IO>,
42    protocol_version: ProtocolVersion,
43    client_settings: ClientSettings,
44    reader: NixReader<R>,
45    writer: Arc<Mutex<NixWriter<W>>>,
46}
47
48impl<IO, R, W> NixDaemon<IO, R, W>
49where
50    IO: NixDaemonIO + Sync + Send,
51{
52    pub fn new(
53        io: Arc<IO>,
54        protocol_version: ProtocolVersion,
55        client_settings: ClientSettings,
56        reader: NixReader<R>,
57        writer: NixWriter<W>,
58    ) -> Self {
59        Self {
60            io,
61            protocol_version,
62            client_settings,
63            reader,
64            writer: Arc::new(Mutex::new(writer)),
65        }
66    }
67}
68
69impl<IO, RW> NixDaemon<IO, ReadHalf<RW>, WriteHalf<RW>>
70where
71    RW: AsyncReadExt + AsyncWriteExt + Send + Unpin + 'static,
72    IO: NixDaemonIO + Sync + Send,
73{
74    /// Async constructor for NixDaemon.
75    ///
76    /// Performs the initial handshake with the client and retrieves the client's preferred
77    /// settings.
78    ///
79    /// The resulting daemon can handle the client session by calling [NixDaemon::handle_client].
80    pub async fn initialize(io: Arc<IO>, mut connection: RW) -> Result<Self, std::io::Error>
81    where
82        RW: AsyncReadExt + AsyncWriteExt + Send + Unpin,
83    {
84        let protocol_version =
85            server_handshake_client(&mut connection, "2.18.2", Trust::Trusted).await?;
86
87        connection.write_u64_le(STDERR_LAST).await?;
88        let (reader, writer) = split(connection);
89        let mut reader = NixReader::builder()
90            .set_version(protocol_version)
91            .build(reader);
92        let mut writer = NixWriterBuilder::default()
93            .set_version(protocol_version)
94            .build(writer);
95
96        // The first op is always SetOptions
97        let operation: Operation = reader.read_value().await?;
98        if operation != Operation::SetOptions {
99            return Err(std::io::Error::other(
100                "Expected SetOptions operation, but got {operation}",
101            ));
102        }
103        let client_settings: ClientSettings = reader.read_value().await?;
104        writer.write_number(STDERR_LAST).await?;
105        writer.flush().await?;
106
107        Ok(Self::new(
108            io,
109            protocol_version,
110            client_settings,
111            reader,
112            writer,
113        ))
114    }
115
116    /// Main client connection loop, reads client's requests and responds to them accordingly.
117    pub async fn handle_client(&mut self) -> Result<(), std::io::Error> {
118        let io = self.io.clone();
119        loop {
120            let op_code = self.reader.read_number().await?;
121            let op = TryInto::<Operation>::try_into(op_code);
122            debug!(?op, "Received operation");
123            match op {
124                // Note: please keep operations sorted in ascending order of their numerical op number.
125                Ok(operation) => match operation {
126                    Operation::IsValidPath => {
127                        let path: StorePath = self.reader.read_value().await?;
128                        Self::handle(&self.writer, io.is_valid_path(&path)).await?
129                    }
130                    // Note this operation does not currently delegate to NixDaemonIO,
131                    // The general idea is that we will pass relevant ClientSettings
132                    // into individual NixDaemonIO method calls if the need arises.
133                    // For now we just store the settings in the NixDaemon for future use.
134                    Operation::SetOptions => {
135                        self.client_settings = self.reader.read_value().await?;
136                        Self::handle(&self.writer, async { Ok(()) }).await?
137                    }
138                    Operation::QueryPathInfo => {
139                        let path: StorePath = self.reader.read_value().await?;
140                        Self::handle(&self.writer, io.query_path_info(&path)).await?
141                    }
142                    Operation::QueryPathFromHashPart => {
143                        let hash: Bytes = self.reader.read_value().await?;
144                        Self::handle(&self.writer, io.query_path_from_hash_part(&hash)).await?
145                    }
146                    Operation::QueryValidPaths => {
147                        let query: QueryValidPaths = self.reader.read_value().await?;
148                        Self::handle(&self.writer, io.query_valid_paths(&query)).await?
149                    }
150                    Operation::QueryValidDerivers => {
151                        let path: StorePath = self.reader.read_value().await?;
152                        Self::handle(&self.writer, io.query_valid_derivers(&path)).await?
153                    }
154                    // FUTUREWORK: These are just stubs that return an empty list.
155                    // It's important not to return an error for the local-overlay:// store
156                    // to work properly. While it will not see certain referrers and realizations
157                    // it will not fail on various operations like gc and optimize store. At the
158                    // same time, returning an empty list here shouldn't break any of local-overlay store's
159                    // invariants.
160                    Operation::QueryReferrers | Operation::QueryRealisation => {
161                        let _: String = self.reader.read_value().await?;
162                        Self::handle(&self.writer, async move {
163                            warn!(
164                                ?operation,
165                                "This operation is not implemented. Returning empty result..."
166                            );
167                            Ok(Vec::<StorePath>::new())
168                        })
169                        .await?
170                    }
171                    Operation::AddMultipleToStore => {
172                        let repair = self.reader.read_value::<bool>().await?;
173                        let dont_check_sigs = self.reader.read_value::<bool>().await?;
174
175                        let builder = NixReader::builder().set_version(self.reader.version());
176                        let mut framed = NixFramedReader::new(&mut self.reader);
177                        Self::handle(&self.writer, async {
178                            let mut source = builder.build(&mut framed);
179                            let count = source.read_number().await?;
180                            for _ in 0..count {
181                                let info = source.read_value::<ValidPathInfo>().await?;
182                                self.io
183                                    .add_to_store_nar(info, &mut source, repair, dont_check_sigs)
184                                    .await?;
185                            }
186                            Ok(())
187                        })
188                        .await?;
189
190                        // framing desynchronisation
191                        // this MUST kill the connection
192                        if !framed.is_eof_unpin().await? {
193                            return Err(std::io::Error::new(
194                                std::io::ErrorKind::InvalidData,
195                                "payload was not fully consumed",
196                            ));
197                        }
198                    }
199                    Operation::AddToStoreNar => {
200                        let info = self.reader.read_value::<ValidPathInfo>().await?;
201                        let repair = self.reader.read_value::<bool>().await?;
202                        let dont_check_sigs = self.reader.read_value::<bool>().await?;
203
204                        let minor_version = self.protocol_version.minor();
205                        match minor_version {
206                            ..21 => {
207                                // Before protocol version 1.21, the nar is sent unframed, so we just
208                                // pass the reader directly to the operation.
209                                Self::handle(
210                                    &self.writer,
211                                    self.io.add_to_store_nar(
212                                        info,
213                                        &mut self.reader,
214                                        repair,
215                                        dont_check_sigs,
216                                    ),
217                                )
218                                .await?
219                            }
220                            21..23 => {
221                                // Protocol versions 1.21 .. 1.23 use STDERR_READ protocol, see logging.md#stderr_read.
222                                Self::handle(&self.writer, async {
223                                    let mut writer = self.writer.lock().await;
224                                    let mut reader = StderrReadFramedReader::new(
225                                        &mut self.reader,
226                                        writer.deref_mut(),
227                                    );
228                                    self.io
229                                        .add_to_store_nar(
230                                            info,
231                                            &mut reader,
232                                            repair,
233                                            dont_check_sigs,
234                                        )
235                                        .await
236                                    // TODO(edef): enforce framing synchronisation
237                                })
238                                .await?
239                            }
240                            23.. => {
241                                // Starting at protocol version 1.23, the framed protocol is used, see serialization.md#framed
242                                let mut framed = NixFramedReader::new(&mut self.reader);
243
244                                Self::handle(&self.writer, async {
245                                    self.io
246                                        .add_to_store_nar(
247                                            info,
248                                            &mut framed,
249                                            repair,
250                                            dont_check_sigs,
251                                        )
252                                        .await
253                                })
254                                .await?;
255
256                                // framing desynchronisation
257                                // this MUST kill the connection
258                                if !framed.is_eof_unpin().await? {
259                                    return Err(std::io::Error::new(
260                                        std::io::ErrorKind::InvalidData,
261                                        "payload was not fully consumed",
262                                    ));
263                                }
264                            }
265                        }
266                    }
267                    Operation::BuildPaths => {
268                        let args: BuildPaths = self.reader.read_value().await?;
269                        Self::handle(&self.writer, self.io.build_paths(args.paths, args.mode))
270                            .await?
271                    }
272                    Operation::BuildPathsWithResults => {
273                        let args: BuildPaths = self.reader.read_value().await?;
274                        Self::handle(
275                            &self.writer,
276                            self.io.build_paths_with_results(args.paths, args.mode),
277                        )
278                        .await?
279                    }
280                    Operation::QueryMissing => {
281                        let derived_paths = self.reader.read_value().await?;
282                        Self::handle(&self.writer, self.io.query_missing(derived_paths)).await?
283                    }
284                    Operation::EnsurePath => {
285                        let store_path = self.reader.read_value().await?;
286                        Self::handle(&self.writer, self.io.ensure_path(&store_path)).await?
287                    }
288                    Operation::QueryDerivationOutputMap => {
289                        let store_path = self.reader.read_value().await?;
290                        Self::handle(
291                            &self.writer,
292                            self.io.query_derivation_output_map(&store_path),
293                        )
294                        .await?
295                    }
296                    Operation::NarFromPath => {
297                        let store_path = self.reader.read_value().await?;
298                        let result = self.io.nar_from_path(&store_path).await;
299                        let mut writer = self.writer.lock().await;
300
301                        match result {
302                            Ok(mut reader) => {
303                                // the protocol requires that we first indicate that we are done sending logs
304                                // by sending STDERR_LAST and then the response.
305                                writer.write_number(STDERR_LAST).await?;
306                                copy_buf(&mut reader, &mut *writer).await?;
307                                writer.flush().await?;
308                            }
309                            Err(err) => {
310                                debug!(%err, "IO error");
311                                writer.write_number(STDERR_ERROR).await?;
312                                writer.write_value(&NixError::new(format!("{err}"))).await?;
313                                writer.flush().await?;
314                            }
315                        }
316                    }
317                    _ => {
318                        return Err(std::io::Error::other(format!(
319                            "Operation {operation:?} is not implemented"
320                        )));
321                    }
322                },
323                _ => {
324                    return Err(std::io::Error::other(format!(
325                        "Unknown operation code received: {op_code}"
326                    )));
327                }
328            }
329        }
330    }
331
332    /// Handles the operation and sends the response or error to the client.
333    ///
334    /// As per nix daemon protocol, after sending the request, the client expects zero or more
335    /// log lines/activities followed by either
336    /// * STDERR_LAST and the response bytes
337    /// * STDERR_ERROR and the error
338    ///
339    /// This is a helper method, awaiting on the passed in future and then
340    /// handling log lines/activities as described above.
341    async fn handle<T>(
342        writer: &Arc<Mutex<NixWriter<WriteHalf<RW>>>>,
343        future: impl Future<Output = std::io::Result<T>>,
344    ) -> Result<(), std::io::Error>
345    where
346        T: NixSerialize + Send,
347    {
348        let result = future.await;
349        let mut writer = writer.lock().await;
350
351        match result {
352            Ok(r) => {
353                // the protocol requires that we first indicate that we are done sending logs
354                // by sending STDERR_LAST and then the response.
355                writer.write_number(STDERR_LAST).await?;
356                writer.write_value(&r).await?;
357                writer.flush().await
358            }
359            Err(e) => {
360                debug!(err = ?e, "IO error");
361                writer.write_number(STDERR_ERROR).await?;
362                writer.write_value(&NixError::new(format!("{e:?}"))).await?;
363                writer.flush().await
364            }
365        }
366    }
367}
368
369#[cfg(test)]
370mod tests {
371    use super::*;
372    use std::{io::ErrorKind, sync::Arc};
373
374    use mockall::predicate;
375    use tokio::io::AsyncWriteExt;
376
377    use crate::{
378        nix_daemon::MockNixDaemonIO,
379        wire::ProtocolVersion,
380        worker_protocol::{ClientSettings, WORKER_MAGIC_1, WORKER_MAGIC_2},
381    };
382
383    #[tokio::test]
384    async fn test_daemon_initialization() {
385        let mut builder = tokio_test::io::Builder::new();
386        let test_conn = builder
387            .read(&WORKER_MAGIC_1.to_le_bytes())
388            .write(&WORKER_MAGIC_2.to_le_bytes())
389            // Our version is 1.37
390            .write(&[37, 1, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00])
391            // The client's versin is 1.35
392            .read(&[35, 1, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00])
393            // cpu affinity
394            .read(&[0; 8])
395            // reservespace
396            .read(&[0; 8])
397            // version (size)
398            .write(&[0x06, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00])
399            // version (data == 2.18.2 + padding)
400            .write(&[50, 46, 49, 56, 46, 50, 0, 0])
401            // Trusted (1 == client trusted)
402            .write(&[1, 0, 0, 0, 0, 0, 0, 0])
403            // STDERR_LAST
404            .write(&[115, 116, 108, 97, 0, 0, 0, 0]);
405
406        let mut bytes = Vec::new();
407        let mut writer = NixWriter::new(&mut bytes);
408        writer
409            .write_value(&ClientSettings::default())
410            .await
411            .unwrap();
412        writer.flush().await.unwrap();
413
414        let test_conn = test_conn
415            // SetOptions op
416            .read(&[19, 0, 0, 0, 0, 0, 0, 0])
417            .read(&bytes)
418            // STDERR_LAST
419            .write(&[115, 116, 108, 97, 0, 0, 0, 0])
420            .build();
421
422        let mock = MockNixDaemonIO::new();
423        let daemon = NixDaemon::initialize(Arc::new(mock), test_conn)
424            .await
425            .unwrap();
426        assert_eq!(daemon.client_settings, ClientSettings::default());
427        assert_eq!(daemon.protocol_version, ProtocolVersion::from_parts(1, 35));
428    }
429
430    async fn serialize<T>(req: &T, protocol_version: ProtocolVersion) -> Vec<u8>
431    where
432        T: NixSerialize + Send,
433    {
434        let mut result: Vec<u8> = Vec::new();
435        let mut w = NixWriter::builder()
436            .set_version(protocol_version)
437            .build(&mut result);
438        w.write_value(req).await.unwrap();
439        w.flush().await.unwrap();
440        result
441    }
442
443    async fn respond<T>(
444        resp: &Result<T, std::io::Error>,
445        protocol_version: ProtocolVersion,
446    ) -> Vec<u8>
447    where
448        T: NixSerialize + Send,
449    {
450        let mut result: Vec<u8> = Vec::new();
451        let mut w = NixWriter::builder()
452            .set_version(protocol_version)
453            .build(&mut result);
454        match resp {
455            Ok(value) => {
456                w.write_value(&STDERR_LAST).await.unwrap();
457                w.write_value(value).await.unwrap();
458            }
459            Err(e) => {
460                w.write_value(&STDERR_ERROR).await.unwrap();
461                w.write_value(&NixError::new(format!("{e:?}")))
462                    .await
463                    .unwrap();
464            }
465        }
466        w.flush().await.unwrap();
467        result
468    }
469
470    #[tokio::test]
471    async fn test_handle_is_valid_path_ok() {
472        let version = ProtocolVersion::from_parts(1, 37);
473        let (io, mut handle) = tokio_test::io::Builder::new().build_with_handle();
474        let mut mock = MockNixDaemonIO::new();
475        let (reader, writer) = split(io);
476        let path = StorePath::from_absolute_path(
477            "/nix/store/33l4p0pn0mybmqzaxfkpppyh7vx1c74p-hello-2.12.1".as_bytes(),
478        )
479        .unwrap();
480        mock.expect_is_valid_path()
481            .with(predicate::eq(path.clone()))
482            .times(1)
483            .returning(|_| Box::pin(async { Ok(true) }));
484
485        handle.read(&Into::<u64>::into(Operation::IsValidPath).to_le_bytes());
486        handle.read(&serialize(&path, version).await);
487        handle.write(&respond(&Ok(true), version).await);
488        drop(handle);
489
490        let mut daemon = NixDaemon::new(
491            Arc::new(mock),
492            version,
493            ClientSettings::default(),
494            NixReader::new(reader),
495            NixWriter::new(writer),
496        );
497        assert_eq!(
498            ErrorKind::UnexpectedEof,
499            daemon
500                .handle_client()
501                .await
502                .expect_err("Expecting eof")
503                .kind()
504        );
505    }
506
507    #[tokio::test]
508    async fn test_handle_is_valid_path_err() {
509        let version = ProtocolVersion::from_parts(1, 37);
510        let (io, mut handle) = tokio_test::io::Builder::new().build_with_handle();
511        let mut mock = MockNixDaemonIO::new();
512        let (reader, writer) = split(io);
513        let path: StorePath = StorePath::from_absolute_path(
514            "/nix/store/33l4p0pn0mybmqzaxfkpppyh7vx1c74p-hello-2.12.1".as_bytes(),
515        )
516        .unwrap();
517        mock.expect_is_valid_path()
518            .with(predicate::eq(path.clone()))
519            .times(1)
520            .returning(|_| Box::pin(async { Err(std::io::Error::other("hello")) }));
521
522        handle.read(&Into::<u64>::into(Operation::IsValidPath).to_le_bytes());
523        handle.read(&serialize(&path, version).await);
524        handle.write(&respond::<bool>(&Err(std::io::Error::other("hello")), version).await);
525        drop(handle);
526
527        let mut daemon = NixDaemon::new(
528            Arc::new(mock),
529            version,
530            ClientSettings::default(),
531            NixReader::new(reader),
532            NixWriter::new(writer),
533        );
534        assert_eq!(
535            ErrorKind::UnexpectedEof,
536            daemon
537                .handle_client()
538                .await
539                .expect_err("Expecting eof")
540                .kind()
541        );
542    }
543}
nix_compat/nix_daemon/handler.rs

nix_compat/nix_daemon/
handler.rs
