refactor(tvix/nix-compat): move handshake function to nix-compat r/7876

Moving out the handshake route out of tvix-daemon to nix-compat.

We're bumping the protocol to version 37. It does not change anything
for us. Version 37 introduced a change in the build result
serialization. We do not implement that for now.

Tested tvix-daemon against Nix via the integration test to make sure
we did not break anything in the process.

Taking advantage of this refactor to remove the unused anyhow
dependency in tvix-daemon.

Change-Id: I8a2aad0b6dbb0aaa997446a612a365fd7d70336c
Reviewed-on: https://cl.tvl.fyi/c/depot/+/11379
Tested-by: BuildkiteCI
Reviewed-by: flokli <flokli@flokli.de>

1 files changed, 19 insertions, 119 deletions

diff --git a/users/picnoir/tvix-daemon/src/main.rs b/users/picnoir/tvix-daemon/src/main.rs
index 398c8af01f..82220637c6 100644
--- a/users/picnoir/tvix-daemon/src/main.rs
+++ b/users/picnoir/tvix-daemon/src/main.rs
@@ -1,12 +1,11 @@
-use anyhow::anyhow;
 use clap::Parser;
 use tokio::io::{AsyncReadExt, AsyncWriteExt};
 use tokio_listener::{self, SystemOptions, UserOptions};
 use tracing::{debug, error, info, instrument, Level};
 
 use nix_compat::wire::{
-    bytes, primitive,
-    worker_protocol::{self, ClientSettings},
+    primitive,
+    worker_protocol::{self, server_handshake_client, ClientSettings, Trust},
 };
 
 #[derive(Parser, Debug)]
@@ -19,15 +18,6 @@ struct Cli {
     verbosity: Option<Level>,
 }
 
-/// Structure used to hold the client socket connection and some
-/// metadata about the connection.
-#[derive(Debug)]
-struct ClientConnection<R: AsyncReadExt + AsyncWriteExt + Unpin> {
-    conn: R,
-    version_minor: u64,
-    client_settings: Option<ClientSettings>,
-}
-
 #[tokio::main]
 #[instrument()]
 async fn main() {
@@ -62,6 +52,15 @@ async fn main() {
     }
 }
 
+/// Structure used to hold the client socket connection and some
+/// metadata about the connection.
+#[derive(Debug)]
+struct ClientConnection<R: AsyncReadExt + AsyncWriteExt + Unpin> {
+    pub conn: R,
+    pub version_minor: u64,
+    pub client_settings: Option<ClientSettings>,
+}
+
 /// Worker in charge to respond a Nix client using the Nix wire
 /// protocol.
 #[instrument()]
@@ -69,9 +68,15 @@ async fn worker<R>(mut conn: R)
 where
     R: AsyncReadExt + AsyncWriteExt + Unpin + std::fmt::Debug,
 {
-    match perform_init_handshake(&mut conn).await {
-        Ok(mut client_connection) => {
+    match server_handshake_client(&mut conn, "2.18.2", Trust::Trusted).await {
+        Ok(client_protocol_version) => {
+            let mut client_connection = ClientConnection {
+                conn,
+                version_minor: client_protocol_version,
+                client_settings: None,
+            };
             debug!("Client hanshake succeeded");
+            debug!(client_protocol_version = ?client_protocol_version);
             // TODO: implement logging. For now, we'll just send
             // STDERR_LAST, which is good enough to get Nix respond to
             // us.
@@ -110,108 +115,3 @@ where
     primitive::write_u64(&mut conn.conn, worker_protocol::STDERR_LAST).await?;
     Ok(settings)
 }
-
-/// Performs the initial handshake. During the handshake, the client
-/// will first send a magic u64, to which the daemon needs to respond
-/// with another magic u64.
-///
-/// We then retrieve the client version, and discard a bunch of now
-/// obsolete data.
-#[instrument()]
-async fn perform_init_handshake<'a, R: 'a>(
-    mut conn: &'a mut R,
-) -> anyhow::Result<ClientConnection<&'a mut R>>
-where
-    &'a mut R: AsyncReadExt + AsyncWriteExt + Unpin + std::fmt::Debug,
-{
-    let worker_magic_1 = primitive::read_u64(&mut conn).await?;
-    debug!("Hello read");
-    if worker_magic_1 != worker_protocol::WORKER_MAGIC_1 {
-        Err(anyhow!(
-            "Invalid client hello received: {:?}, expected {:?}",
-            worker_magic_1,
-            worker_protocol::WORKER_MAGIC_1
-        ))
-    } else {
-        primitive::write_u64(&mut conn, worker_protocol::WORKER_MAGIC_2).await?;
-        primitive::write_u64(&mut conn, worker_protocol::PROTOCOL_VERSION).await?;
-        conn.flush().await?;
-        debug!("Hello responded");
-        let client_version = primitive::read_u64(&mut conn).await?;
-        debug!("Version read");
-        if client_version < 0x10a {
-            return Err(anyhow!("The nix client version is too old"));
-        }
-        let protocol_minor = client_version & 0x00ff;
-        let protocol_major = client_version & 0xff00;
-        debug!(client.version = %client_version, client.minor = %protocol_minor, client.major = %protocol_major);
-        if protocol_minor >= 14 {
-            debug!("read cpu affinity");
-            // Obsolete CPU affinity.
-            let read_affinity = primitive::read_u64(&mut conn).await?;
-            if read_affinity != 0 {
-                skip_8_bytes(&mut conn).await?;
-            };
-        }
-        if protocol_minor >= 11 {
-            // Obsolete reserveSpace
-            debug!("read reservespace");
-            skip_8_bytes(&mut conn).await?;
-        }
-        if protocol_minor >= 33 {
-            // Nix version. We're plain lying, we're not Nix, but eh…
-            // Setting it to the 2.3 lineage. Not 100% sure this is a
-            // good idea.
-            debug!("write version");
-            // Plain str padded to 64 bits.
-            bytes::write_bytes(&mut conn, "2.3.17").await?;
-            conn.flush().await?;
-        }
-        if protocol_minor >= 35 {
-            worker_protocol::write_worker_trust_level(&mut conn, worker_protocol::Trust::Trusted)
-                .await?;
-            info!("Trust sent");
-        }
-        Ok(ClientConnection {
-            conn,
-            version_minor: protocol_minor,
-            client_settings: None,
-        })
-    }
-}
-
-async fn skip_8_bytes<R>(conn: &mut R) -> anyhow::Result<()>
-where
-    R: AsyncReadExt + Unpin + std::fmt::Debug,
-{
-    let mut _discard_buffer = [0; 8];
-    conn.read_exact(&mut _discard_buffer).await?;
-    Ok(())
-}
-
-#[cfg(test)]
-mod integration_tests {
-    use nix_compat::wire::worker_protocol;
-    #[tokio::test]
-    async fn test_init_handshake() {
-        let mut test_conn = tokio_test::io::Builder::new()
-            .read(&worker_protocol::WORKER_MAGIC_1.to_le_bytes())
-            .write(&worker_protocol::WORKER_MAGIC_2.to_le_bytes())
-            .write(&worker_protocol::PROTOCOL_VERSION.to_le_bytes())
-            // Let's say the client is in sync with the daemon
-            // protocol-wise
-            .read(&worker_protocol::PROTOCOL_VERSION.to_le_bytes())
-            // cpu affinity
-            .read(&vec![0; 8])
-            // reservespace
-            .read(&vec![0; 8])
-            // version (size)
-            .write(&vec![0x06, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00])
-            // version (data == 2.2.17 + padding)
-            .write(&vec![50, 46, 51, 46, 49, 55, 0, 0])
-            // Trusted (1 == client trusted
-            .write(&vec![1, 0, 0, 0, 0, 0, 0, 0])
-            .build();
-        crate::perform_init_handshake(&mut test_conn).await.unwrap();
-    }
-}