feat(nix-daemon): Implement client handler. r/8907

This change includes only the basic nix handshake protocol handling and
sets up a client session. The only supported operation at this point is
SetOptions.

Additional operations will be implemented in subsequent cls.

Change-Id: I3eccd9e0ceb270c3865929543c702f1491768852
Reviewed-on: https://cl.tvl.fyi/c/depot/+/12743
Autosubmit: Vladimir Kryachko <v.kryachko@gmail.com>
Tested-by: BuildkiteCI
Reviewed-by: flokli <flokli@flokli.de>
Reviewed-by: edef <edef@edef.eu>
Reviewed-by: Brian Olsen <me@griff.name>

1 files changed, 229 insertions, 0 deletions

diff --git a/tvix/nix-compat/src/nix_daemon/handler.rs b/tvix/nix-compat/src/nix_daemon/handler.rs
new file mode 100644
index 000000000000..c0257d979acc
--- /dev/null
+++ b/tvix/nix-compat/src/nix_daemon/handler.rs
@@ -0,0 +1,229 @@
+use std::{future::Future, sync::Arc};
+
+use tokio::{
+    io::{split, AsyncReadExt, AsyncWriteExt, ReadHalf, WriteHalf},
+    sync::Mutex,
+};
+use tracing::debug;
+
+use super::{
+    worker_protocol::{server_handshake_client, ClientSettings, Operation, Trust, STDERR_LAST},
+    NixDaemonIO,
+};
+use crate::wire::{
+    de::{NixRead, NixReader},
+    ser::{NixSerialize, NixWrite, NixWriter, NixWriterBuilder},
+    ProtocolVersion,
+};
+use crate::{nix_daemon::types::NixError, worker_protocol::STDERR_ERROR};
+
+/// Handles a single connection with a nix client.
+///
+/// As part of its [`initialization`] it performs the handshake with the client
+/// and determines the [ProtocolVersion] and [ClientSettings] to use for the remainder of the session.
+///
+/// Once initialized, [`handle_client`] needs to be called to handle the rest of the session,
+/// it delegates all operation handling to an instance of [NixDaemonIO].
+///
+/// [`initialization`]: NixDaemon::initialize
+#[allow(dead_code)]
+pub struct NixDaemon<IO, R, W> {
+    io: Arc<IO>,
+    protocol_version: ProtocolVersion,
+    client_settings: ClientSettings,
+    reader: NixReader<R>,
+    writer: Arc<Mutex<NixWriter<W>>>,
+}
+
+impl<IO, R, W> NixDaemon<IO, R, W>
+where
+    IO: NixDaemonIO + Sync + Send,
+{
+    pub fn new(
+        io: Arc<IO>,
+        protocol_version: ProtocolVersion,
+        client_settings: ClientSettings,
+        reader: NixReader<R>,
+        writer: NixWriter<W>,
+    ) -> Self {
+        Self {
+            io,
+            protocol_version,
+            client_settings,
+            reader,
+            writer: Arc::new(Mutex::new(writer)),
+        }
+    }
+}
+
+impl<IO, RW> NixDaemon<IO, ReadHalf<RW>, WriteHalf<RW>>
+where
+    RW: AsyncReadExt + AsyncWriteExt + Send + Unpin + 'static,
+    IO: NixDaemonIO + Sync + Send,
+{
+    /// Async constructor for NixDaemon.
+    ///
+    /// Performs the initial handshake with the client and retrieves the client's preferred
+    /// settings.
+    ///
+    /// The resulting daemon can handle the client session by calling [NixDaemon::handle_client].
+    pub async fn initialize(io: Arc<IO>, mut connection: RW) -> Result<Self, std::io::Error>
+    where
+        RW: AsyncReadExt + AsyncWriteExt + Send + Unpin,
+    {
+        let protocol_version =
+            server_handshake_client(&mut connection, "2.18.2", Trust::Trusted).await?;
+
+        connection.write_u64_le(STDERR_LAST).await?;
+        let (reader, writer) = split(connection);
+        let mut reader = NixReader::builder()
+            .set_version(protocol_version)
+            .build(reader);
+        let mut writer = NixWriterBuilder::default()
+            .set_version(protocol_version)
+            .build(writer);
+
+        // The first op is always SetOptions
+        let operation: Operation = reader.read_value().await?;
+        if operation != Operation::SetOptions {
+            return Err(std::io::Error::other(
+                "Expected SetOptions operation, but got {operation}",
+            ));
+        }
+        let client_settings: ClientSettings = reader.read_value().await?;
+        writer.write_number(STDERR_LAST).await?;
+        writer.flush().await?;
+
+        Ok(Self::new(
+            io,
+            protocol_version,
+            client_settings,
+            reader,
+            writer,
+        ))
+    }
+
+    /// Main client connection loop, reads client's requests and responds to them accordingly.
+    pub async fn handle_client(&mut self) -> Result<(), std::io::Error> {
+        loop {
+            let op_code = self.reader.read_number().await?;
+            match TryInto::<Operation>::try_into(op_code) {
+                Ok(operation) => match operation {
+                    Operation::SetOptions => {
+                        self.client_settings = self.reader.read_value().await?;
+                        self.handle(async { Ok(()) }).await?
+                    }
+                    _ => {
+                        return Err(std::io::Error::other(format!(
+                            "Operation {operation:?} is not implemented"
+                        )));
+                    }
+                },
+                _ => {
+                    return Err(std::io::Error::other(format!(
+                        "Unknown operation code received: {op_code}"
+                    )));
+                }
+            }
+        }
+    }
+
+    /// Handles the operation and sends the response or error to the client.
+    ///
+    /// As per nix daemon protocol, after sending the request, the client expects zero or more
+    /// log lines/activities followed by either
+    /// * STDERR_LAST and the response bytes
+    /// * STDERR_ERROR and the error
+    ///
+    /// This is a helper method, awaiting on the passed in future and then
+    /// handling log lines/activities as described above.
+    async fn handle<T>(
+        &mut self,
+        future: impl Future<Output = std::io::Result<T>>,
+    ) -> Result<(), std::io::Error>
+    where
+        T: NixSerialize + Send,
+    {
+        let result = future.await;
+        let mut writer = self.writer.lock().await;
+
+        match result {
+            Ok(r) => {
+                // the protocol requires that we first indicate that we are done sending logs
+                // by sending STDERR_LAST and then the response.
+                writer.write_number(STDERR_LAST).await?;
+                writer.write_value(&r).await?;
+                writer.flush().await
+            }
+            Err(e) => {
+                debug!(err = ?e, "IO error");
+                writer.write_number(STDERR_ERROR).await?;
+                writer.write_value(&NixError::new(format!("{e:?}"))).await?;
+                writer.flush().await
+            }
+        }
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use std::sync::Arc;
+
+    use tokio::io::AsyncWriteExt;
+
+    use crate::{
+        wire::ProtocolVersion,
+        worker_protocol::{ClientSettings, WORKER_MAGIC_1, WORKER_MAGIC_2},
+    };
+
+    struct MockDaemonIO {}
+
+    impl NixDaemonIO for MockDaemonIO {}
+
+    #[tokio::test]
+    async fn test_daemon_initialization() {
+        let mut builder = tokio_test::io::Builder::new();
+        let test_conn = builder
+            .read(&WORKER_MAGIC_1.to_le_bytes())
+            .write(&WORKER_MAGIC_2.to_le_bytes())
+            // Our version is 1.37
+            .write(&[37, 1, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00])
+            // The client's versin is 1.35
+            .read(&[35, 1, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00])
+            // cpu affinity
+            .read(&[0; 8])
+            // reservespace
+            .read(&[0; 8])
+            // version (size)
+            .write(&[0x06, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00])
+            // version (data == 2.18.2 + padding)
+            .write(&[50, 46, 49, 56, 46, 50, 0, 0])
+            // Trusted (1 == client trusted)
+            .write(&[1, 0, 0, 0, 0, 0, 0, 0])
+            // STDERR_LAST
+            .write(&[115, 116, 108, 97, 0, 0, 0, 0]);
+
+        let mut bytes = Vec::new();
+        let mut writer = NixWriter::new(&mut bytes);
+        writer
+            .write_value(&ClientSettings::default())
+            .await
+            .unwrap();
+        writer.flush().await.unwrap();
+
+        let test_conn = test_conn
+            // SetOptions op
+            .read(&[19, 0, 0, 0, 0, 0, 0, 0])
+            .read(&bytes)
+            // STDERR_LAST
+            .write(&[115, 116, 108, 97, 0, 0, 0, 0])
+            .build();
+
+        let daemon = NixDaemon::initialize(Arc::new(MockDaemonIO {}), test_conn)
+            .await
+            .unwrap();
+        assert_eq!(daemon.client_settings, ClientSettings::default());
+        assert_eq!(daemon.protocol_version, ProtocolVersion::from_parts(1, 35));
+    }
+}