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::{BasicDerivation, 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::BuildDerivation => {
281                        let drv_path = self.reader.read_value().await?;
282                        let basic_derivation: BasicDerivation = self.reader.read_value().await?;
283                        let derivation =
284                            basic_derivation.try_into().map_err(std::io::Error::other)?;
285                        let mode = self.reader.read_value().await?;
286                        Self::handle(
287                            &self.writer,
288                            self.io.build_derivation(drv_path, derivation, mode),
289                        )
290                        .await?
291                    }
292                    Operation::QueryMissing => {
293                        let derived_paths = self.reader.read_value().await?;
294                        Self::handle(&self.writer, self.io.query_missing(derived_paths)).await?
295                    }
296                    Operation::EnsurePath => {
297                        let store_path = self.reader.read_value().await?;
298                        Self::handle(&self.writer, self.io.ensure_path(&store_path)).await?
299                    }
300                    Operation::QueryDerivationOutputMap => {
301                        let store_path = self.reader.read_value().await?;
302                        Self::handle(
303                            &self.writer,
304                            self.io.query_derivation_output_map(&store_path),
305                        )
306                        .await?
307                    }
308                    Operation::NarFromPath => {
309                        let store_path = self.reader.read_value().await?;
310                        let result = self.io.nar_from_path(&store_path).await;
311                        let mut writer = self.writer.lock().await;
312
313                        match result {
314                            Ok(mut reader) => {
315                                // the protocol requires that we first indicate that we are done sending logs
316                                // by sending STDERR_LAST and then the response.
317                                writer.write_number(STDERR_LAST).await?;
318                                copy_buf(&mut reader, &mut *writer).await?;
319                                writer.flush().await?;
320                            }
321                            Err(err) => {
322                                debug!(%err, "IO error");
323                                writer.write_number(STDERR_ERROR).await?;
324                                writer.write_value(&NixError::new(format!("{err}"))).await?;
325                                writer.flush().await?;
326                            }
327                        }
328                    }
329
330                    Operation::AddTempRoot => {
331                        let store_path = self.reader.read_value().await?;
332                        Self::handle(&self.writer, self.io.add_temp_root(&store_path)).await?;
333                        let mut writer = self.writer.lock().await;
334                        writer.write_number(1).await?;
335                        writer.flush().await?;
336                    }
337
338                    Operation::AddIndirectRoot => {
339                        let store_path: String = self.reader.read_value().await?;
340                        Self::handle(&self.writer, self.io.add_indirect_root(&store_path)).await?;
341                        let mut writer = self.writer.lock().await;
342                        writer.write_number(1).await?;
343                        writer.flush().await?;
344                    }
345
346                    _ => {
347                        return Err(std::io::Error::other(format!(
348                            "Operation {operation:?} is not implemented"
349                        )));
350                    }
351                },
352                _ => {
353                    return Err(std::io::Error::other(format!(
354                        "Unknown operation code received: {op_code}"
355                    )));
356                }
357            }
358        }
359    }
360
361    /// Handles the operation and sends the response or error to the client.
362    ///
363    /// As per nix daemon protocol, after sending the request, the client expects zero or more
364    /// log lines/activities followed by either
365    /// * STDERR_LAST and the response bytes
366    /// * STDERR_ERROR and the error
367    ///
368    /// This is a helper method, awaiting on the passed in future and then
369    /// handling log lines/activities as described above.
370    async fn handle<T>(
371        writer: &Arc<Mutex<NixWriter<WriteHalf<RW>>>>,
372        future: impl Future<Output = std::io::Result<T>>,
373    ) -> Result<(), std::io::Error>
374    where
375        T: NixSerialize + Send,
376    {
377        let result = future.await;
378        let mut writer = writer.lock().await;
379
380        match result {
381            Ok(r) => {
382                // the protocol requires that we first indicate that we are done sending logs
383                // by sending STDERR_LAST and then the response.
384                writer.write_number(STDERR_LAST).await?;
385                writer.write_value(&r).await?;
386                writer.flush().await
387            }
388            Err(e) => {
389                debug!(err = ?e, "IO error");
390                writer.write_number(STDERR_ERROR).await?;
391                writer.write_value(&NixError::new(format!("{e:?}"))).await?;
392                writer.flush().await
393            }
394        }
395    }
396}
397
398#[cfg(test)]
399mod tests {
400    use super::*;
401    use std::{io::ErrorKind, sync::Arc};
402
403    use mockall::predicate;
404    use tokio::io::AsyncWriteExt;
405
406    use crate::{
407        nix_daemon::MockNixDaemonIO,
408        wire::ProtocolVersion,
409        worker_protocol::{ClientSettings, WORKER_MAGIC_1, WORKER_MAGIC_2},
410    };
411
412    #[tokio::test]
413    async fn test_daemon_initialization() {
414        let mut builder = tokio_test::io::Builder::new();
415        let test_conn = builder
416            .read(&WORKER_MAGIC_1.to_le_bytes())
417            .write(&WORKER_MAGIC_2.to_le_bytes())
418            // Our version is 1.37
419            .write(&[37, 1, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00])
420            // The client's versin is 1.35
421            .read(&[35, 1, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00])
422            // cpu affinity
423            .read(&[0; 8])
424            // reservespace
425            .read(&[0; 8])
426            // version (size)
427            .write(&[0x06, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00])
428            // version (data == 2.18.2 + padding)
429            .write(&[50, 46, 49, 56, 46, 50, 0, 0])
430            // Trusted (1 == client trusted)
431            .write(&[1, 0, 0, 0, 0, 0, 0, 0])
432            // STDERR_LAST
433            .write(&[115, 116, 108, 97, 0, 0, 0, 0]);
434
435        let mut bytes = Vec::new();
436        let mut writer = NixWriter::new(&mut bytes);
437        writer
438            .write_value(&ClientSettings::default())
439            .await
440            .unwrap();
441        writer.flush().await.unwrap();
442
443        let test_conn = test_conn
444            // SetOptions op
445            .read(&[19, 0, 0, 0, 0, 0, 0, 0])
446            .read(&bytes)
447            // STDERR_LAST
448            .write(&[115, 116, 108, 97, 0, 0, 0, 0])
449            .build();
450
451        let mock = MockNixDaemonIO::new();
452        let daemon = NixDaemon::initialize(Arc::new(mock), test_conn)
453            .await
454            .unwrap();
455        assert_eq!(daemon.client_settings, ClientSettings::default());
456        assert_eq!(daemon.protocol_version, ProtocolVersion::from_parts(1, 35));
457    }
458
459    async fn serialize<T>(req: &T, protocol_version: ProtocolVersion) -> Vec<u8>
460    where
461        T: NixSerialize + Send,
462    {
463        let mut result: Vec<u8> = Vec::new();
464        let mut w = NixWriter::builder()
465            .set_version(protocol_version)
466            .build(&mut result);
467        w.write_value(req).await.unwrap();
468        w.flush().await.unwrap();
469        result
470    }
471
472    async fn respond<T>(
473        resp: &Result<T, std::io::Error>,
474        protocol_version: ProtocolVersion,
475    ) -> Vec<u8>
476    where
477        T: NixSerialize + Send,
478    {
479        let mut result: Vec<u8> = Vec::new();
480        let mut w = NixWriter::builder()
481            .set_version(protocol_version)
482            .build(&mut result);
483        match resp {
484            Ok(value) => {
485                w.write_value(&STDERR_LAST).await.unwrap();
486                w.write_value(value).await.unwrap();
487            }
488            Err(e) => {
489                w.write_value(&STDERR_ERROR).await.unwrap();
490                w.write_value(&NixError::new(format!("{e:?}")))
491                    .await
492                    .unwrap();
493            }
494        }
495        w.flush().await.unwrap();
496        result
497    }
498
499    #[tokio::test]
500    async fn test_handle_is_valid_path_ok() {
501        let version = ProtocolVersion::from_parts(1, 37);
502        let (io, mut handle) = tokio_test::io::Builder::new().build_with_handle();
503        let mut mock = MockNixDaemonIO::new();
504        let (reader, writer) = split(io);
505        let path = StorePath::from_absolute_path(
506            "/nix/store/33l4p0pn0mybmqzaxfkpppyh7vx1c74p-hello-2.12.1".as_bytes(),
507        )
508        .unwrap();
509        mock.expect_is_valid_path()
510            .with(predicate::eq(path.clone()))
511            .times(1)
512            .returning(|_| Box::pin(async { Ok(true) }));
513
514        handle.read(&Into::<u64>::into(Operation::IsValidPath).to_le_bytes());
515        handle.read(&serialize(&path, version).await);
516        handle.write(&respond(&Ok(true), version).await);
517        drop(handle);
518
519        let mut daemon = NixDaemon::new(
520            Arc::new(mock),
521            version,
522            ClientSettings::default(),
523            NixReader::new(reader),
524            NixWriter::new(writer),
525        );
526        assert_eq!(
527            ErrorKind::UnexpectedEof,
528            daemon
529                .handle_client()
530                .await
531                .expect_err("Expecting eof")
532                .kind()
533        );
534    }
535
536    #[tokio::test]
537    async fn test_handle_is_valid_path_err() {
538        let version = ProtocolVersion::from_parts(1, 37);
539        let (io, mut handle) = tokio_test::io::Builder::new().build_with_handle();
540        let mut mock = MockNixDaemonIO::new();
541        let (reader, writer) = split(io);
542        let path: StorePath = StorePath::from_absolute_path(
543            "/nix/store/33l4p0pn0mybmqzaxfkpppyh7vx1c74p-hello-2.12.1".as_bytes(),
544        )
545        .unwrap();
546        mock.expect_is_valid_path()
547            .with(predicate::eq(path.clone()))
548            .times(1)
549            .returning(|_| Box::pin(async { Err(std::io::Error::other("hello")) }));
550
551        handle.read(&Into::<u64>::into(Operation::IsValidPath).to_le_bytes());
552        handle.read(&serialize(&path, version).await);
553        handle.write(&respond::<bool>(&Err(std::io::Error::other("hello")), version).await);
554        drop(handle);
555
556        let mut daemon = NixDaemon::new(
557            Arc::new(mock),
558            version,
559            ClientSettings::default(),
560            NixReader::new(reader),
561            NixWriter::new(writer),
562        );
563        assert_eq!(
564            ErrorKind::UnexpectedEof,
565            daemon
566                .handle_client()
567                .await
568                .expect_err("Expecting eof")
569                .kind()
570        );
571    }
572}
nix_compat/nix_daemon/handler.rs

nix_compat/nix_daemon/
handler.rs
