main.rs

extern crate thrussh;
extern crate thrussh_keys;
extern crate futures;
extern crate tokio;
extern crate env_logger;
use std::sync::Arc;
use thrussh::*;
use thrussh::server::{Auth, Session};
use thrussh_keys::*;

const SERVE_MAGIC_1: u64 = 0x390c9deb;
const SERVE_MAGIC_2: u64 = 0x5452eecb;
const SERVE_PROTOCOL_VERSION: u64 = 0x205;
const OTHER_PROTOCOL_VERSION: u32 = 0x205;

#[derive(Clone)]
struct Server {
    client_pubkey: Arc<thrussh_keys::key::PublicKey>,
    stage: NixStage,
}

#[derive(Clone, Debug)]
enum NixStage {
    Start,
    Ready,
}

// https://github.com/NixOS/nix/blob/013dd28b15c33ff00b4e66d97dd00cb7e2b7d863/src/libstore/serve-protocol.hh#L6
#[derive(Debug)]
enum ServeCommand {
    QueryValidPaths,
    QueryPathInfos,
    DumpStorePath,
    ImportPaths,
    ExportPaths,
    BuildPaths,
    QueryClosure,
    BuildDerivation,
    AddToStoreNar,
}
impl From<ServeCommand> for u64 {
    fn from(command: ServeCommand) -> u64 {
        match command {
            ServeCommand::QueryValidPaths => 1,
            ServeCommand::QueryPathInfos => 2,
            ServeCommand::DumpStorePath => 3,
            ServeCommand::ImportPaths => 4,
            ServeCommand::ExportPaths => 5,
            ServeCommand::BuildPaths => 6,
            ServeCommand::QueryClosure => 7,
            ServeCommand::BuildDerivation => 8,
            ServeCommand::AddToStoreNar => 9,
        }
    }
}
impl From<u64> for ServeCommand {
    fn from(command: u64) -> ServeCommand {
        match command {
            1 => ServeCommand::QueryValidPaths,
            2 => ServeCommand::QueryPathInfos,
            3 => ServeCommand::DumpStorePath,
            4 => ServeCommand::ImportPaths,
            5 => ServeCommand::ExportPaths,
            6 => ServeCommand::BuildPaths,
            7 => ServeCommand::QueryClosure,
            8 => ServeCommand::BuildDerivation,
            9 => ServeCommand::AddToStoreNar,
            _ => panic!("Unknown u64 when decoding ServeCommand: {:?}", command),
        }
    }
}


impl server::Server for Server {
    type Handler = Self;
    fn new(&self) -> Self {
        self.clone()
    }
}

impl Server {
    fn ready(&self) -> Self {
        Server {
            client_pubkey: self.client_pubkey.clone(),
            stage: NixStage::Ready,
        }
    }
}

fn u64_from_u8(bytes: [u8; 8]) -> u64 {
    unsafe { std::mem::transmute(bytes) }
}

fn u32_from_u8(bytes: [u8; 4]) -> u32{
    unsafe { std::mem::transmute(bytes) }
}

fn u16_from_u8(bytes: [u8; 2]) -> u16 {
    unsafe { std::mem::transmute(bytes) }
}

fn u32_to_u8(num: u32) -> [u8; 4] {
    unsafe { std::mem::transmute(num) }
}

fn u64_to_u8(num: u64) -> [u8; 8] {
    unsafe { std::mem::transmute(num) }
}

fn take_u64(bytes: &mut std::slice::Iter<'_, u8>) -> u64 {
    let mut data: [u8; 8] = [
        *bytes.next().unwrap(),
        *bytes.next().unwrap(),
        *bytes.next().unwrap(),
        *bytes.next().unwrap(),
        *bytes.next().unwrap(),
        *bytes.next().unwrap(),
        *bytes.next().unwrap(),
        *bytes.next().unwrap(),
    ];

    u64_from_u8(data)
}

fn take_u32(bytes: &mut std::slice::Iter<'_, u8>) -> u32 {
    let mut data: [u8; 4] = [
        *bytes.next().unwrap(),
        *bytes.next().unwrap(),
        *bytes.next().unwrap(),
        *bytes.next().unwrap(),
    ];

    unsafe { std::mem::transmute(data) }
}

fn take_u16(bytes: &mut std::slice::Iter<'_, u8>) -> u16 {
    let mut data: [u8; 2] = [
        *bytes.next().unwrap(),
        *bytes.next().unwrap(),
    ];

    u16_from_u8(data)
}

impl server::Handler for Server {
   type Error = std::io::Error;
   type FutureAuth = futures::Finished<(Self, server::Auth), Self::Error>;
   type FutureUnit = futures::Finished<(Self, server::Session), Self::Error>;
   type FutureBool = futures::Finished<(Self, server::Session, bool), Self::Error>;

   fn finished_auth(self, auth: Auth) -> Self::FutureAuth {
       futures::finished((self, auth))
   }
   fn finished_bool(self, session: Session, b: bool) -> Self::FutureBool {
       futures::finished((self, session, b))
   }
   fn finished(self, session: Session) -> Self::FutureUnit {
       futures::finished((self, session))
   }

   fn auth_publickey(self, _: &str, _: &key::PublicKey) -> Self::FutureAuth {
       futures::finished((self, server::Auth::Accept))
   }

    fn data(self, channel: ChannelId, data: &[u8], mut session: server::Session) -> Self::FutureUnit {
        // println!("bytes: {:?}", data);
        let mut bytes = data.into_iter();

        match self.stage {
            NixStage::Start => {
                let serve_magic = take_u64(&mut bytes);
                if serve_magic != SERVE_MAGIC_1 {
                    println!("Bad serve magic: {:?}", serve_magic);
                    return futures::finished((self, session));
                } else {
                    println!("OK serve magic");
                }

                let serve_protocol_version = take_u64(&mut bytes);
                if serve_protocol_version != SERVE_PROTOCOL_VERSION {
                    println!("Bad serve protocol version: {:x}", serve_protocol_version);
                    return futures::finished((self, session));
                } else {
                    println!("OK serve protocol version");
                }

                println!("Sending serve magic 2");
                session.data(channel, None, &u64_to_u8(SERVE_MAGIC_2));
                println!("Sending protocol version");
                session.data(channel, None, &u64_to_u8(SERVE_PROTOCOL_VERSION));

                // session.data(channel, None, data);
                futures::finished((self.ready(), session))
            }

            NixStage::Ready => {
                let cmd: ServeCommand = take_u64(&mut bytes).into();
                println!("Serve command: {:?}", cmd);

                match cmd {
                    ServeCommand::QueryPathInfos => {
                        let paths = take_u64(&mut bytes);

                        for path in 1..paths {
                            let len = take_u64(&mut bytes);
                            println!("{:?}", path);
                            let path_len = take_u64(&mut bytes);
                            println!("len: {:?}, str: {:?}",
                                     len,
                                     std::str::from_utf8(&bytes.cloned().collect::<Vec<u8>>()),
                            );
                        }
                    }
                    _ => {
                        panic!("oenth");
                    }
                }

                futures::finished((self, session))
            }
        }
   }
}


fn main() {
    println!("Hello, world!");


   env_logger::init();
   // Starting the server thread.
   let client_key = thrussh_keys::key::KeyPair::generate_ed25519().unwrap();
   let client_pubkey = Arc::new(client_key.clone_public_key());
   let t = std::thread::spawn(|| {
       let mut config = thrussh::server::Config::default();
       config.connection_timeout = Some(std::time::Duration::from_secs(600));
       config.auth_rejection_time = std::time::Duration::from_secs(3);
       config.keys.push(thrussh_keys::key::KeyPair::generate_ed25519().unwrap());
       let config = Arc::new(config);
       let sh = Server{ client_pubkey, stage: NixStage::Start };
       thrussh::server::run(config, "0.0.0.0:2222", sh);
   });

    t.join().unwrap();
}