<commit_before>extern crate libc;
extern crate contain;

use libc::{c_void, c_int, c_ulong, c_char};
use std::ptr;
use libc::{MS_REC, MS_PRIVATE, MS_NOSUID, MS_NODEV, MS_NOEXEC, MNT_DETACH};
use libc::{CLONE_NEWUSER, CLONE_NEWNS, CLONE_NEWPID, SIGCHLD};
use std::fs;
use std::fs::{File, OpenOptions};
use std::path::Path;
use std::io::{Read, Write};
use std::thread;
use std::ffi::{CString, CStr};
use std::os::unix::ffi::OsStrExt;
use std::process::Command;
use contain::Stack;

fn has_busybox() -> Option<String> {
    let o = Command::new("which")
            .arg("busybox")
            .output().expect("failed to start `which' program");
    if o.stderr.len() != 0 {
        println!("{}", String::from_utf8_lossy(&o.stderr));
    }
    if !o.status.success() {
        return None
    }
    //strip newline
    let path = {
        let p = String::from_utf8(o.stdout).unwrap();
        let mut l = p.lines();
        let path = String::from(l.next().unwrap());
        assert_eq!(l.next(), None);
        path
    };

    //check that is statically linked
    let dynamic_cmd = format!("test -z \"$(objdump -p {} | grep NEEDED)\"", path);
    println!("{}", dynamic_cmd);
    let o = Command::new("sh").arg("-c").arg(dynamic_cmd).status().expect("failed to check whether busybox is static linked");
    if !o.success() {
        println!("busybox found at {} but it does not seem to be a static executable", path);
        None
    } else {
        Some(path)
    }

}

//path to cstr
fn p2cstr<P: AsRef<Path>>(p: P) -> CString {
    CString::new(p.as_ref().as_os_str().as_bytes()).expect("string contains NULLs")
}

// CString Option to pointer
fn cstr2p(o: Option<CString>) -> *const c_char {
    o.map_or(ptr::null(), |x| x.as_ptr())
}

mod ffi{
    use libc::{c_int, c_char};
    extern {
        pub fn pivot_root(new_root: *const c_char, put_old: *const c_char) -> c_int;
    }
}
fn pivot_root<P: AsRef<Path>>(new_root: P, put_old: P) {
    let new_root = p2cstr(new_root).as_ptr();
    let put_old = p2cstr(put_old).as_ptr();
    let r = unsafe { ffi::pivot_root(new_root, put_old) };
    assert_eq!(r, 0);
}

fn domount<P: AsRef<Path>>(source: Option<P>, target: Option<P>, filesystemtype: Option<&str>, mountflags: c_ulong) -> () {
    let source = source.map(p2cstr);
    let target = target.map(p2cstr);
    let fstype = match filesystemtype {
            None => ptr::null(),
            Some(fst) => CString::new(fst).expect("fstypes has NULLs").as_ptr(),
        };
    let r = unsafe { libc::mount(cstr2p(source), cstr2p(target), fstype, mountflags, ptr::null()) };
    assert_eq!(r, 0);
}

// a lot copied from man USER_NAMESPACES(7)
extern "C" fn child_func(args: *mut c_void) -> c_int {
    println!("I'm called from child_func");

    let args = args as *mut ChildArgs;
    let r_pipe_fd = unsafe { (*args).r_pipe_fd };
    let w_pipe_fd = unsafe { (*args).w_pipe_fd };
    println!("r_pipe_fd: {}", r_pipe_fd);

    // run everything in a new thread so exceptions bubble up to rust
    let h = thread::spawn(move || {
        // wait for parent
        println!("waiting for parent to set up mapping, ...");
        assert_eq!(unsafe { libc::close(w_pipe_fd) } , 0);
        unsafe {
            let mut buf = 0u8;
            assert_eq!(libc::read(r_pipe_fd, (&mut buf) as *mut u8 as *mut c_void, 1), 0);
        }

        {
            println!("Child status");
            let mut fstatus = File::open("/proc/self/status").expect("open status");
            let mut contents = String::new();
            fstatus.read_to_string(&mut contents).expect("read status");
            for l in contents.lines(){
                if l.starts_with("Seccomp") || l.starts_with("Cap"){
                    println!("{}", l);
                }
            }
        }

        let container_root = Path::new("./mntcont");
        if container_root.exists() {
            assert!(container_root.is_dir());
        } else {
            println!("creating {:?}", container_root.as_os_str());
            fs::create_dir(container_root).expect("mkdir mntcont");
        }

        println!("mounting container root");
        domount(None, Some(container_root), Some("tmpfs"), 0);
        domount(Some("none"), Some("/"), None, MS_REC|MS_PRIVATE);
        {
            let proc_dir = container_root.join("proc");
            println!("setting up {:?}", proc_dir.as_os_str());
            fs::create_dir(proc_dir.as_path()).expect("mkdir proc");
            domount(None, Some(proc_dir.as_path()), Some("proc"), MS_NOSUID|MS_NODEV|MS_NOEXEC);
        }

        // TODO populate root. copy busybox to it
        match has_busybox() {
            None => println!("not populating root (no busybox found)"),
            Some(busyboxpath) => {
                println!("populating root with a busybox image.");
                assert!(fs::copy(busyboxpath, container_root.join("busybox")).expect("copy busybox") > 0);
            }
        }

        let old_root_name = "oldroot";
        {
            println!("Entering new root");
            let old_root = container_root.join(old_root_name);
            fs::create_dir(old_root.as_path()).expect("mkdir oldroot");
            pivot_root(container_root, old_root.as_path());
        }
        // we moved to the new root
        let container_root = Path::new("/");
        let old_root = container_root.join(old_root_name);

        assert_eq!( unsafe { libc::umount2(p2cstr(old_root).as_ptr(), MNT_DETACH) }, 0);

        {
            println!("mountinfo");
            let mut fstatus = File::open("/proc/self/mountinfo").expect("open mountinfo");
            let mut contents = String::new();
            fstatus.read_to_string(&mut contents).expect("read status");
            for l in contents.lines(){
                println!("{}", l);
            }
        }
        0
    });

    match h.join() {
        Ok(_) => {
            let prog = CString::new("/busybox").unwrap();
            let arg0 = CString::new("busybox").unwrap();
            let arg1 = CString::new("sh").unwrap();
            let argv = vec![arg0.as_ptr(), arg1.as_ptr(), ptr::null()];
            unsafe { libc::execv(prog.as_ptr(), argv.as_ptr()) };
            panic!("exec failed!"); //stack unwinding will also panic :D
        }
        Err(_) => 1,
    }
}

struct ChildArgs {
    r_pipe_fd: c_int,
    w_pipe_fd: c_int,
}

fn write(path: &String, content: &[u8]) {
    println!("Writing to {}", path);
    let mut file = OpenOptions::new().read(false).write(true).create(false).open(path).expect("open");
    // needs to happen in one write call
    assert_eq!(file.write(content).expect("write"), content.len());
    //file gets dropped
}
fn main() {
    println!("Hello, world!");

    let mut pipe_fd = vec![0;2];  /* Pipe used to synchronize parent and child */
    assert_eq!(unsafe { libc::pipe(pipe_fd.as_mut_ptr()) }, 0);
    println!("r pipe: {} w pipe: {}", pipe_fd[0], pipe_fd[1]);
    let ptr_child_args = {
        let child_args = Box::new(ChildArgs { r_pipe_fd: pipe_fd[0], w_pipe_fd: pipe_fd[1]});
        Box::into_raw(child_args) as *mut c_void
    };

    let child_stack = Stack::new(512);

    println!("cloning");
    let child_pid = unsafe {
        libc::clone(child_func, child_stack.top, CLONE_NEWUSER|CLONE_NEWNS|CLONE_NEWPID|SIGCHLD, ptr_child_args)
    };
    assert!(child_pid != -1);

    let uid = unsafe { libc::getuid() };
    write(&format!("/proc/{}/uid_map", child_pid), &format!("0 {} 1", uid).into_bytes());

    //proc_setgroups_write man user_namespaces
    write(&format!("/proc/{}/setgroups", child_pid), &format!("deny").into_bytes());

    let gid = unsafe { libc::getgid() };
    write(&format!("/proc/{}/gid_map", child_pid), &format!("0 {} 1", gid).into_bytes());

    //unleash child
    assert_eq!(unsafe { libc::close(pipe_fd[1]) }, 0);

    let mut child_status: c_int = 0;
    assert!(unsafe { libc::waitpid(child_pid, &mut child_status, 0) } != -1);
    println!("Child terminated: {}", child_status);

    unsafe {
        drop(Box::from_raw(ptr_child_args));
    }

    //child_stack gets unmapped
}
<commit_msg>better directory setup<commit_after>extern crate libc;
extern crate contain;

use libc::{c_void, c_int, c_ulong, c_char};
use std::ptr;
use libc::{MS_REC, MS_PRIVATE, MS_NOSUID, MS_NODEV, MS_NOEXEC, MNT_DETACH};
use libc::{CLONE_NEWUSER, CLONE_NEWNS, CLONE_NEWPID, SIGCHLD};
use std::fs;
use std::fs::{File, OpenOptions};
use std::path::Path;
use std::io::{Read, Write};
use std::thread;
use std::ffi::{CString, CStr};
use std::os::unix::ffi::OsStrExt;
use std::process::Command;
use contain::Stack;

fn has_busybox() -> Option<String> {
    let o = Command::new("which")
            .arg("busybox")
            .output().expect("failed to start `which' program");
    if o.stderr.len() != 0 {
        println!("{}", String::from_utf8_lossy(&o.stderr));
    }
    if !o.status.success() {
        return None
    }
    //strip newline
    let path = {
        let p = String::from_utf8(o.stdout).unwrap();
        let mut l = p.lines();
        let path = String::from(l.next().unwrap());
        assert_eq!(l.next(), None);
        path
    };

    //check that is statically linked
    let dynamic_cmd = format!("test -z \"$(objdump -p {} | grep NEEDED)\"", path);
    println!("{}", dynamic_cmd);
    let o = Command::new("sh").arg("-c").arg(dynamic_cmd).status().expect("failed to check whether busybox is static linked");
    if !o.success() {
        println!("busybox found at {} but it does not seem to be a static executable", path);
        None
    } else {
        Some(path)
    }

}

//path to cstr
fn p2cstr<P: AsRef<Path>>(p: P) -> CString {
    CString::new(p.as_ref().as_os_str().as_bytes()).expect("string contains NULLs")
}

// CString Option to pointer
fn cstr2p(o: Option<CString>) -> *const c_char {
    o.map_or(ptr::null(), |x| x.as_ptr())
}

mod ffi{
    use libc::{c_int, c_char};
    extern {
        pub fn pivot_root(new_root: *const c_char, put_old: *const c_char) -> c_int;
    }
}
fn pivot_root<P: AsRef<Path>>(new_root: P, put_old: P) {
    let new_root = p2cstr(new_root).as_ptr();
    let put_old = p2cstr(put_old).as_ptr();
    let r = unsafe { ffi::pivot_root(new_root, put_old) };
    assert_eq!(r, 0);
}

fn domount<P: AsRef<Path>>(source: Option<P>, target: Option<P>, filesystemtype: Option<&str>, mountflags: c_ulong) -> () {
    let source = source.map(p2cstr);
    let target = target.map(p2cstr);
    let fstype = match filesystemtype {
            None => ptr::null(),
            Some(fst) => CString::new(fst).expect("fstypes has NULLs").as_ptr(),
        };
    let r = unsafe { libc::mount(cstr2p(source), cstr2p(target), fstype, mountflags, ptr::null()) };
    assert_eq!(r, 0);
}

// a lot copied from man USER_NAMESPACES(7)
extern "C" fn child_func(args: *mut c_void) -> c_int {
    println!("I'm called from child_func");

    let args = args as *mut ChildArgs;
    let r_pipe_fd = unsafe { (*args).r_pipe_fd };
    let w_pipe_fd = unsafe { (*args).w_pipe_fd };
    println!("r_pipe_fd: {}", r_pipe_fd);

    // run everything in a new thread so exceptions bubble up to rust
    let h = thread::spawn(move || {
        // wait for parent
        println!("waiting for parent to set up mapping, ...");
        assert_eq!(unsafe { libc::close(w_pipe_fd) } , 0);
        unsafe {
            let mut buf = 0u8;
            assert_eq!(libc::read(r_pipe_fd, (&mut buf) as *mut u8 as *mut c_void, 1), 0);
        }

        {
            println!("Child status");
            let mut fstatus = File::open("/proc/self/status").expect("open status");
            let mut contents = String::new();
            fstatus.read_to_string(&mut contents).expect("read status");
            for l in contents.lines(){
                if l.starts_with("Seccomp") || l.starts_with("Cap"){
                    println!("{}", l);
                }
            }
        }

        let container_root = Path::new("./mntcont");
        if container_root.exists() {
            assert!(container_root.is_dir());
        } else {
            println!("creating {:?}", container_root.as_os_str());
            fs::create_dir(container_root).expect("mkdir mntcont");
        }

        println!("mounting container root");
        domount(None, Some(container_root), Some("tmpfs"), 0);
        domount(Some("none"), Some("/"), None, MS_REC|MS_PRIVATE);
        {
            let proc_dir = container_root.join("proc");
            println!("setting up {:?}", proc_dir.as_os_str());
            fs::create_dir(proc_dir.as_path()).expect("mkdir proc");
            domount(None, Some(proc_dir.as_path()), Some("proc"), MS_NOSUID|MS_NODEV|MS_NOEXEC);
        }

        // TODO populate root. copy busybox to it
        match has_busybox() {
            None => println!("not populating root (no busybox found)"),
            Some(busyboxpath) => {
                println!("populating root with a busybox image.");
                assert!(fs::copy(busyboxpath, container_root.join("busybox")).expect("copy busybox") > 0);
            }
        }

        let old_root_name = "oldroot";
        {
            println!("Entering new root");
            let old_root = container_root.join(old_root_name);
            fs::create_dir(old_root.as_path()).expect("mkdir oldroot");
            pivot_root(container_root, old_root.as_path());
        }

        unsafe { let root = CString::new("/").unwrap();
            assert_eq!(libc::chdir(root.as_ptr()), 0);
        };

        // we moved to the new root
        let container_root = Path::new("/");
        let old_root = container_root.join(old_root_name);

        {
            println!("getting rid of old root");
            assert_eq!( unsafe { libc::umount2(p2cstr(old_root.as_path()).as_ptr(), MNT_DETACH) }, 0);
            assert!(fs::read_dir(old_root.as_path()).expect("oldroot dir").next().is_none());
            fs::remove_dir(old_root.as_path()).expect("rmdir oldroot");
        }

        {
            println!("mountinfo");
            let mut fstatus = File::open("/proc/self/mountinfo").expect("open mountinfo");
            let mut contents = String::new();
            fstatus.read_to_string(&mut contents).expect("read status");
            for l in contents.lines(){
                println!("{}", l);
            }
        }
        0
    });

    match h.join() {
        Ok(_) => {
            let prog = CString::new("/busybox").unwrap();
            let arg0 = CString::new("busybox").unwrap();
            let arg1 = CString::new("sh").unwrap();
            let argv = vec![arg0.as_ptr(), arg1.as_ptr(), ptr::null()];
            unsafe { libc::execv(prog.as_ptr(), argv.as_ptr()) };
            panic!("exec failed!"); //stack unwinding will also panic :D
        }
        Err(_) => 1,
    }
}

struct ChildArgs {
    r_pipe_fd: c_int,
    w_pipe_fd: c_int,
}

fn write(path: &String, content: &[u8]) {
    println!("Writing to {}", path);
    let mut file = OpenOptions::new().read(false).write(true).create(false).open(path).expect("open");
    // needs to happen in one write call
    assert_eq!(file.write(content).expect("write"), content.len());
    //file gets dropped
}
fn main() {
    println!("Hello, world!");

    let mut pipe_fd = vec![0;2];  /* Pipe used to synchronize parent and child */
    assert_eq!(unsafe { libc::pipe(pipe_fd.as_mut_ptr()) }, 0);
    println!("r pipe: {} w pipe: {}", pipe_fd[0], pipe_fd[1]);
    let ptr_child_args = {
        let child_args = Box::new(ChildArgs { r_pipe_fd: pipe_fd[0], w_pipe_fd: pipe_fd[1]});
        Box::into_raw(child_args) as *mut c_void
    };

    let child_stack = Stack::new(512);

    println!("cloning");
    let child_pid = unsafe {
        libc::clone(child_func, child_stack.top, CLONE_NEWUSER|CLONE_NEWNS|CLONE_NEWPID|SIGCHLD, ptr_child_args)
    };
    assert!(child_pid != -1);

    let uid = unsafe { libc::getuid() };
    write(&format!("/proc/{}/uid_map", child_pid), &format!("0 {} 1", uid).into_bytes());

    //proc_setgroups_write man user_namespaces
    write(&format!("/proc/{}/setgroups", child_pid), &format!("deny").into_bytes());

    let gid = unsafe { libc::getgid() };
    write(&format!("/proc/{}/gid_map", child_pid), &format!("0 {} 1", gid).into_bytes());

    //unleash child
    assert_eq!(unsafe { libc::close(pipe_fd[1]) }, 0);

    let mut child_status: c_int = 0;
    assert!(unsafe { libc::waitpid(child_pid, &mut child_status, 0) } != -1);
    println!("Child terminated: {}", child_status);

    unsafe {
        drop(Box::from_raw(ptr_child_args));
    }

    //child_stack gets unmapped
}
<|endoftext|>


<commit_before>use rustc_serialize::hex::{ToHex, FromHex};
use rustc_serialize::base64::{ToBase64, FromBase64, STANDARD};
use std::str;
use std::io::BufWriter;
use std::fs::File;
use std::path::Path;
use std::io::prelude::*;
use std::cmp;
use std::fs::OpenOptions;
use std::fs;

pub struct CipherText(Vec<u8>);

impl CipherText {
    pub fn new(bytes: &[u8]) -> CipherText {
        CipherText(bytes.to_vec())
    }

    pub fn from_hex(hex: &str) -> Result<CipherText, Error> {
        let cipher_bytes: Vec<u8> = try!(hex.from_hex()
            .map_err(|e| Error::Hex(e.to_string())));
        Ok(CipherText(cipher_bytes))
    }

    pub fn to_hex(&self) -> String {
        let &CipherText(ref vec_bytes) = self;
        vec_bytes.to_hex()
    }

    pub fn from_b64<S: AsRef<str>>(b64: S) -> Result<CipherText, Error> {
        let b64 = b64.as_ref();
        let cipher_bytes: Vec<u8> = try!(b64.from_base64()
            .map_err(|e| Error::Base64(e.to_string())));
        Ok(CipherText(cipher_bytes))
    }

    pub fn to_b64(&self) -> String {
        let &CipherText(ref vec_bytes) = self;
        vec_bytes.to_base64(STANDARD)
    }

    fn from_file<P: AsRef<Path>>(path: P) -> Result<CipherText, Error> {
        buffer_file(path.as_ref()).and_then(CipherText::from_b64)
    }

    pub fn to_file<P: AsRef<Path>>(&self, path: P) -> Result<(), Error> {
        let path = path.as_ref();
        let text_b64 = self.to_b64();
        if path.exists() {
            try!(fs::remove_file(path).map_err(|e| Error::File(e.to_string())));
        }

        let mut file = try!(OpenOptions::new()
            .write(true)
            .append(true)
            .create(true)
            .open(path)
            .map_err(|e| Error::File(e.to_string())));

        let mut i = 0;
        let mut j = cmp::min(text_b64.len(), 80);
        while i != j {
            try!(file.write_all(&text_b64[i..j].as_bytes())
                .map_err(|e| Error::File(e.to_string())));

            try!(file.write_all("\n".as_bytes()).map_err(|e| Error::File(e.to_string())));

            i = j;
            j = cmp::min(text_b64.len(), j + 80);
        }
        Ok(())
    }
}
pub struct PlainText(Vec<u8>);

impl PlainText {
    pub fn from_string<S: AsRef<str>>(string: S) -> PlainText {
        PlainText(string.as_ref().as_bytes().to_vec())
    }

    pub fn from_bytes(bytes: &[u8]) -> PlainText {
        PlainText(bytes.to_vec())
    }

    fn from_file<P: AsRef<Path>>(path: P) -> Result<PlainText, Error> {
        buffer_file(path.as_ref()).map(PlainText::from_string)
    }

    pub fn to_file<P: AsRef<Path>>(&self, path: P) -> Result<(), Error> {
        let path = path.as_ref();
        let &PlainText(ref vec_bytes) = self;
        let file = try!(OpenOptions::new()
            .write(true)
            .truncate(true)
            .create(true)
            .open(path)
            .map_err(|e| Error::File(e.to_string())));

        let mut out_file = BufWriter::new(file);
        try!(out_file.write_all(&vec_bytes).map_err(|e| Error::File(e.to_string())));

        Ok(())
    }

    pub fn to_utf8(&self) -> Result<String, Error> {
        let &PlainText(ref vec_bytes) = self;
        let plain = try!(str::from_utf8(&vec_bytes).map_err(|e| Error::UTF8(e.to_string())));
        Ok(plain.to_string())
    }
}

fn buffer_file(path: &Path) -> Result<String, Error> {
    let mut file = try!(File::open(Path::new(path)).map_err(|e| Error::File(e.to_string())));
    let capacity = file.metadata().ok().map_or(0, |x| x.len());
    let mut buffer = String::with_capacity(capacity as usize);
    try!(file.read_to_string(&mut buffer).map_err(|e| Error::File(e.to_string())));
    Ok(buffer)
}

pub fn encrypt(text: &PlainText, key: &str) -> CipherText {
    let &PlainText(ref bytes) = text;
    let xored: Vec<u8> = repeating_xor(bytes, key.as_bytes());
    CipherText::new(&xored)
}

pub fn decrypt(cipher_text: &CipherText, key: &str) -> Result<PlainText, Error> {
    let &CipherText(ref bytes) = cipher_text;
    let xored: Vec<u8> = repeating_xor(bytes, key.as_bytes());
    Ok(PlainText::from_bytes(&xored))
}

pub fn decrypt_file(input_path: &str, output_path: &str, key: &str) -> Result<(), Error> {
    let cipher = try!(CipherText::from_file(input_path));
    let plain = try!(decrypt(&cipher, key));
    try!(plain.to_file(output_path));
    Ok(())
}

pub fn encrypt_file(input_path: &str, output_path: &str, key: &str) -> Result<(), Error> {
    let plain = try!(PlainText::from_file(input_path));
    let cipher = encrypt(&plain, key);
    try!(cipher.to_file(output_path));
    Ok(())
}

#[derive(Debug)]
pub enum Error {
    Hex(String),
    Base64(String),
    UTF8(String),
    File(String),
}

fn repeating_xor(input: &[u8], key: &[u8]) -> Vec<u8> {
    let mut out: Vec<u8> = Vec::with_capacity(input.len());
    unsafe {
        out.set_len(input.len());
    }
    for (i, in_val) in input.iter().enumerate() {
        out[i] = in_val ^ key[i % key.len()];
    }
    out
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_encrypt_and_decrypt() {
        check_text("This is the plain text\n");
        check_text("123, 456");
        check_text("!£$%^&*@~:");
        check_text("日本語");
    }

    #[test]
    fn test_encrypt_and_decrypt_files() {
        check_text_file("This is the plain text\n");
        check_text_file("123, 456");
        check_text_file("!£$%^&*@~:");
        check_text_file("日本語");
        check_text_file("This is the plain text\nThis is the plain text\nThis is the plain text\n");
    }

    #[test]
    fn test_wrong_key() {
        let text = "This is the plain text";
        let plain_text = PlainText::from_string(text);
        let key: &str = "toy";
        let decoded_text = decrypt(&encrypt(&plain_text, key), "wrong_key").unwrap();
        assert!(text != decoded_text.to_utf8().unwrap());
    }

    fn check_text(text: &str) {
        let plain_text = PlainText::from_string(text);
        let key: &str = "toy";
        let decoded_text = decrypt(&encrypt(&plain_text, key), key).unwrap();
        assert_eq!(text, decoded_text.to_utf8().unwrap());
    }

    fn check_text_file(text: &str) {
        let plain_text = PlainText::from_string(text);
        let key: &str = "toy";
        let plain_path = "./target/test-plain.txt";
        let cipher_path = "./target/test-cipher.txt";
        let decoded_path = "./target/test-decoded.txt";

        plain_text.to_file(plain_path).unwrap();
        encrypt_file(plain_path, cipher_path, key).unwrap();
        decrypt_file(cipher_path, decoded_path, key).unwrap();
        let decoded_text = PlainText::from_file(decoded_path).unwrap();

        assert_eq!(text, decoded_text.to_utf8().unwrap());
    }
}
<commit_msg>Simply to_file methods<commit_after>use rustc_serialize::hex::{ToHex, FromHex};
use rustc_serialize::base64::{ToBase64, FromBase64, STANDARD};
use std::str;
use std::io::BufWriter;
use std::fs::File;
use std::path::Path;
use std::io::prelude::*;
use std::cmp;
use std::fs::OpenOptions;
use std::fs;

enum WriteMethod { Truncate, Append }

pub struct CipherText(Vec<u8>);

impl CipherText {
    pub fn new(bytes: &[u8]) -> CipherText {
        CipherText(bytes.to_vec())
    }

    pub fn from_hex(hex: &str) -> Result<CipherText, Error> {
        let cipher_bytes: Vec<u8> = try!(hex.from_hex()
            .map_err(|e| Error::Hex(e.to_string())));
        Ok(CipherText(cipher_bytes))
    }

    pub fn to_hex(&self) -> String {
        let &CipherText(ref vec_bytes) = self;
        vec_bytes.to_hex()
    }

    pub fn from_b64<S: AsRef<str>>(b64: S) -> Result<CipherText, Error> {
        let cipher_bytes: Vec<u8> = try!(b64.as_ref().from_base64()
            .map_err(|e| Error::Base64(e.to_string())));
        Ok(CipherText(cipher_bytes))
    }

    pub fn to_b64(&self) -> String {
        let &CipherText(ref vec_bytes) = self;
        vec_bytes.to_base64(STANDARD)
    }

    fn from_file<P: AsRef<Path>>(path: P) -> Result<CipherText, Error> {
        buffer_file(path.as_ref()).and_then(CipherText::from_b64)
    }

    pub fn to_file<P: AsRef<Path>>(&self, path: P) -> Result<(), Error> {
        let text_b64 = self.to_b64();
        let mut file = try!(create_file(path.as_ref(), WriteMethod::Append));

        let mut i = 0;
        let mut j = cmp::min(text_b64.len(), 80);
        while i != j {
            try!(file.write_all(&text_b64[i..j].as_bytes())
                .map_err(|e| Error::File(e.to_string())));

            try!(file.write_all("\n".as_bytes()).map_err(|e| Error::File(e.to_string())));

            i = j;
            j = cmp::min(text_b64.len(), j + 80);
        }
        Ok(())
    }
}
pub struct PlainText(Vec<u8>);

impl PlainText {
    pub fn from_string<S: AsRef<str>>(string: S) -> PlainText {
        PlainText(string.as_ref().as_bytes().to_vec())
    }

    pub fn from_bytes(bytes: &[u8]) -> PlainText {
        PlainText(bytes.to_vec())
    }

    fn from_file<P: AsRef<Path>>(path: P) -> Result<PlainText, Error> {
        buffer_file(path.as_ref()).map(PlainText::from_string)
    }

    pub fn to_file<P: AsRef<Path>>(&self, path: P) -> Result<(), Error> {
        let &PlainText(ref vec_bytes) = self;
        let file = try!(create_file(path.as_ref(), WriteMethod::Truncate));
        let mut out_file = BufWriter::new(file);
        try!(out_file.write_all(&vec_bytes).map_err(|e| Error::File(e.to_string())));
        Ok(())
    }

    pub fn to_utf8(&self) -> Result<String, Error> {
        let &PlainText(ref vec_bytes) = self;
        let plain = try!(str::from_utf8(&vec_bytes).map_err(|e| Error::UTF8(e.to_string())));
        Ok(plain.to_string())
    }
}

fn buffer_file(path: &Path) -> Result<String, Error> {
    let mut file = try!(File::open(Path::new(path)).map_err(|e| Error::File(e.to_string())));
    let capacity = file.metadata().ok().map_or(0, |x| x.len());
    let mut buffer = String::with_capacity(capacity as usize);
    try!(file.read_to_string(&mut buffer).map_err(|e| Error::File(e.to_string())));
    Ok(buffer)
}

fn create_file(path: &Path, method: WriteMethod) -> Result<File, Error> {
    if path.exists() {
        try!(fs::remove_file(path).map_err(|e| Error::File(e.to_string())));
    }

    match method {
        WriteMethod::Append => {
            OpenOptions::new()
                .write(true)
                .append(true)
                .create(true)
                .open(path)
                .map_err(|e| Error::File(e.to_string()))
        },
        WriteMethod::Truncate => {
            OpenOptions::new()
                .write(true)
                .truncate(true)
                .create(true)
                .open(path)
                .map_err(|e| Error::File(e.to_string()))
        }
    }

}

pub fn encrypt(text: &PlainText, key: &str) -> CipherText {
    let &PlainText(ref bytes) = text;
    let xored: Vec<u8> = repeating_xor(bytes, key.as_bytes());
    CipherText::new(&xored)
}

pub fn decrypt(cipher_text: &CipherText, key: &str) -> Result<PlainText, Error> {
    let &CipherText(ref bytes) = cipher_text;
    let xored: Vec<u8> = repeating_xor(bytes, key.as_bytes());
    Ok(PlainText::from_bytes(&xored))
}

pub fn decrypt_file(input_path: &str, output_path: &str, key: &str) -> Result<(), Error> {
    let cipher = try!(CipherText::from_file(input_path));
    let plain = try!(decrypt(&cipher, key));
    try!(plain.to_file(output_path));
    Ok(())
}

pub fn encrypt_file(input_path: &str, output_path: &str, key: &str) -> Result<(), Error> {
    let plain = try!(PlainText::from_file(input_path));
    let cipher = encrypt(&plain, key);
    try!(cipher.to_file(output_path));
    Ok(())
}

#[derive(Debug)]
pub enum Error {
    Hex(String),
    Base64(String),
    UTF8(String),
    File(String),
}

fn repeating_xor(input: &[u8], key: &[u8]) -> Vec<u8> {
    let mut out: Vec<u8> = Vec::with_capacity(input.len());
    unsafe {
        out.set_len(input.len());
    }
    for (i, in_val) in input.iter().enumerate() {
        out[i] = in_val ^ key[i % key.len()];
    }
    out
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_encrypt_and_decrypt() {
        check_text("This is the plain text\n");
        check_text("123, 456");
        check_text("!£$%^&*@~:");
        check_text("日本語");
    }

    #[test]
    fn test_encrypt_and_decrypt_files() {
        check_text_file("This is the plain text\n");
        check_text_file("123, 456");
        check_text_file("!£$%^&*@~:");
        check_text_file("日本語");
        check_text_file("This is the plain text\nThis is the plain text\nThis is the plain text\n");
    }

    #[test]
    fn test_wrong_key() {
        let text = "This is the plain text";
        let plain_text = PlainText::from_string(text);
        let key: &str = "toy";
        let decoded_text = decrypt(&encrypt(&plain_text, key), "wrong_key").unwrap();
        assert!(text != decoded_text.to_utf8().unwrap());
    }

    fn check_text(text: &str) {
        let plain_text = PlainText::from_string(text);
        let key: &str = "toy";
        let decoded_text = decrypt(&encrypt(&plain_text, key), key).unwrap();
        assert_eq!(text, decoded_text.to_utf8().unwrap());
    }

    fn check_text_file(text: &str) {
        let plain_text = PlainText::from_string(text);
        let key: &str = "toy";
        let plain_path = "./target/test-plain.txt";
        let cipher_path = "./target/test-cipher.txt";
        let decoded_path = "./target/test-decoded.txt";

        plain_text.to_file(plain_path).unwrap();
        encrypt_file(plain_path, cipher_path, key).unwrap();
        decrypt_file(cipher_path, decoded_path, key).unwrap();
        let decoded_text = PlainText::from_file(decoded_path).unwrap();

        assert_eq!(text, decoded_text.to_utf8().unwrap());
    }
}
<|endoftext|>


<commit_before>//! The `Client` allows users of the `raft` library to connect to remote `Server` instances and
//! issue commands to be applied to the `StateMachine`.

use std::collections::HashSet;
use std::io::Write;
use std::net::SocketAddr;
use std::net::TcpStream;
use std::fmt;

use bufstream::BufStream;
use capnp::{serialize, MessageReader, ReaderOptions};

use messages_capnp::{client_response, proposal_response};
use messages;
use ClientId;
use Result;
use RaftError;

/// The representation of a Client connection to the cluster.
pub struct Client {
    /// The `Uuid` of the client, should be unique in the cluster.
    pub id: ClientId,
    /// The current connection to the current leader.
    /// If it is none it may mean that there is no estabished leader or that there has been
    /// a disconnection.
    leader_connection: Option<BufStream<TcpStream>>,
    /// A lookup for the cluster's nodes.
    cluster: HashSet<SocketAddr>,
}

impl Client {

    /// Creates a new client.
    pub fn new(cluster: HashSet<SocketAddr>) -> Client {
        Client {
            id: ClientId::new(),
            leader_connection: None,
            cluster: cluster,
        }
    }

    /// Proposes an entry to be appended to the replicated log. This will only
    /// return once the entry has been durably committed.
    /// Returns `Error` when the entire cluster has an unknown leader. Try proposing again later.
    pub fn propose(&mut self, entry: &[u8]) -> Result<()> {
        scoped_trace!("{:?}: propose", self);
        let mut message = messages::proposal_request(entry);

        let mut members = self.cluster.iter().cloned();

        loop {
            let mut connection = match self.leader_connection.take() {
                Some(cxn) => {
                    scoped_debug!("had existing connection {:?}", cxn.get_ref().peer_addr());
                    cxn
                },
                None => {
                    let leader = try!(members.next().ok_or(RaftError::LeaderSearchExhausted));
                    scoped_debug!("connecting to potential leader {}", leader);
                    // Send the preamble.
                    let preamble = messages::client_connection_preamble(self.id);
                    let mut stream = BufStream::new(try!(TcpStream::connect(leader)));
                    try!(serialize::write_message(&mut stream, &*preamble));
                    stream
                }
            };
            try!(serialize::write_message(&mut connection, &mut message));
            try!(connection.flush());
            let response = try!(serialize::read_message(&mut connection, ReaderOptions::new()));
            match try!(response.get_root::<client_response::Reader>()).which().unwrap() {
                client_response::Which::Proposal(Ok(response)) => {
                    match response.which().unwrap() {
                        proposal_response::Which::Success(()) => {
                            scoped_debug!("recieved response Success");
                            self.leader_connection = Some(connection);
                            return Ok(())
                        },
                        proposal_response::Which::UnknownLeader(()) => {
                            scoped_debug!("recieved response UnknownLeader");
                            ()
                        },
                        proposal_response::Which::NotLeader(leader) => {
                            scoped_debug!("recieved response NotLeader");
                            let mut connection: TcpStream = try!(TcpStream::connect(try!(leader)));
                            let preamble = messages::client_connection_preamble(self.id);
                            try!(serialize::write_message(&mut connection, &*preamble));
                            self.leader_connection = Some(BufStream::new(connection));
                        }
                    }
                },
                _ => panic!("Unexpected message type"), // TODO: return a proper error
            }
        }
    }
}

impl fmt::Debug for Client {
    fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
        write!(fmt, "Client")
    }
}

#[cfg(test)]
mod test {
    extern crate env_logger;
    use Client;
    use uuid::Uuid;
    use std::net::{SocketAddr, TcpListener};
    use std::collections::HashSet;
    use std::str::FromStr;
    use std::thread;
    use std::io::{Read, Write};
    use capnp::{serialize, ReaderOptions};
    use capnp::message::MessageReader;
    use messages;
    use messages_capnp::{connection_preamble, client_request};

    #[test]
    fn test_proposal_standalone() {
        setup_test!("test_proposal_standalone");
        let mut cluster = HashSet::new();
        let test_server = TcpListener::bind(SocketAddr::from_str("127.0.0.1:0").unwrap()).unwrap();
        let test_addr = test_server.local_addr().unwrap();
        cluster.insert(test_addr);

        // TODO: Test if the second server is not in the set.
        let second_server = TcpListener::bind(SocketAddr::from_str("127.0.0.1:0").unwrap()).unwrap();
        let second_addr = second_server.local_addr().unwrap();
        // cluster.insert(second_addr);

        let mut client = Client::new(cluster);
        let client_id = client.id.0.clone();
        let to_propose = b"Bears";

        // The client connects on the first proposal.
        // Wait for it.
        let child = thread::spawn(move || {
            let (mut connection, _)  = test_server.accept().unwrap();

            // First proposal should be fine, no errors.
            scoped_debug!("Should get preamble and proposal. Responds Success");

            // Expect Preamble.
            let message = serialize::read_message(&mut connection, ReaderOptions::new()).unwrap();
            let preamble = message.get_root::<connection_preamble::Reader>().unwrap();
            // Test to make sure preamble has the right id.
            if let connection_preamble::id::Which::Client(Ok(id)) = preamble.get_id().which().unwrap() {
                scoped_debug!("got preamble");
                assert_eq!(Uuid::from_bytes(id).unwrap(), client_id);
            } else { panic!("Invalid preamble."); }

            // Expect first proposal! (success!)
            let message = serialize::read_message(&mut connection, ReaderOptions::new()).unwrap();
            let request = message.get_root::<client_request::Reader>().unwrap();
            // Test to make sure request has the right value.
            if let client_request::Which::Proposal(Ok(proposal)) = request.which().unwrap() {
                scoped_debug!("got proposal");
                assert_eq!(proposal.get_entry().unwrap(), to_propose);
            } else { panic!("Invalid request."); }

            // Send first response! (success!)
            let response = messages::proposal_response_success();
            serialize::write_message(&mut connection, &*response).unwrap();
            connection.flush();

            // Second proposal should report unknown leader, and have the client return error.
            scoped_debug!("Should get proposal. Responds UnknownLeader");

            // Expect proposal! (unknown leader!)
            let message = serialize::read_message(&mut connection, ReaderOptions::new()).unwrap();
            let request = message.get_root::<client_request::Reader>().unwrap();
            // Test to make sure request has the right value.
            if let client_request::Which::Proposal(Ok(proposal)) = request.which().unwrap() {
                scoped_debug!("got proposal");
                assert_eq!(proposal.get_entry().unwrap(), to_propose);
            } else { panic!("Invalid request."); }

            // Send response! (unknown leader!) Client should drop connection.
            let response = messages::proposal_response_unknown_leader();
            serialize::write_message(&mut connection, &*response).unwrap();
            connection.flush();

            let (mut connection, _)  = test_server.accept().unwrap();
            serialize::write_message(&mut connection, &*response).unwrap();
            connection.flush();

            // Third Proposal should report NotLeader. Client should choose the server we direct it to.
            scoped_debug!("Should get preamble and proposal. Responds NotLeader.");
            let (mut connection, _)  = test_server.accept().unwrap();

            // Expect Preamble.
            let message = serialize::read_message(&mut connection, ReaderOptions::new()).unwrap();
            let preamble = message.get_root::<connection_preamble::Reader>().unwrap();
            // Test to make sure preamble has the right id.
            if let connection_preamble::id::Which::Client(Ok(id)) = preamble.get_id().which().unwrap() {
                scoped_debug!("got third preamble");
                assert_eq!(Uuid::from_bytes(id).unwrap(), client_id);
            } else { panic!("Invalid preamble."); }

            // Expect proposal! (not leader!)
            let message = serialize::read_message(&mut connection, ReaderOptions::new()).unwrap();
            let request = message.get_root::<client_request::Reader>().unwrap();
            // Test to make sure request has the right value.
            if let client_request::Which::Proposal(Ok(proposal)) = request.which().unwrap() {
                scoped_debug!("got second proposal");
                assert_eq!(proposal.get_entry().unwrap(), to_propose);
            } else { panic!("Invalid request."); }

            // Send response! (not leader!)
            let response = messages::proposal_response_not_leader(&format!("{}", second_addr));
            serialize::write_message(&mut connection, &*response).unwrap();
            connection.flush();

            // Test that it seeks out other server and proposes.
            scoped_debug!("Second server should get preamble and proposal. Responds Success.");

            // Accept on the second server.
            let (mut connection, _)  = second_server.accept().unwrap();
            // Expect Preamble.
            let message = serialize::read_message(&mut connection, ReaderOptions::new()).unwrap();
            let preamble = message.get_root::<connection_preamble::Reader>().unwrap();
            // Test to make sure preamble has the right id.
            if let connection_preamble::id::Which::Client(Ok(id)) = preamble.get_id().which().unwrap() {
                scoped_debug!("got fourth preamble");
                assert_eq!(Uuid::from_bytes(id).unwrap(), client_id);
            } else { panic!("Invalid preamble."); }

            // Expect proposal! (again!)
            let message = serialize::read_message(&mut connection, ReaderOptions::new()).unwrap();
            let request = message.get_root::<client_request::Reader>().unwrap();
            // Test to make sure request has the right value.
            if let client_request::Which::Proposal(Ok(proposal)) = request.which().unwrap() {
                scoped_debug!("got third proposal");
                assert_eq!(proposal.get_entry().unwrap(), to_propose);
            } else { panic!("Invalid request."); }

            // Send final response! (Success!)
            let response = messages::proposal_response_success();
            serialize::write_message(&mut connection, &*response).unwrap();

        });
        // Propose. It's a marriage made in heaven! :)
        // Should be ok
        scoped_debug!("first starting");
        client.propose(to_propose).unwrap();
        assert!(client.leader_connection.is_some());
        scoped_debug!("first done");
        // Should be err
        scoped_debug!("second starting");
        assert!(client.propose(to_propose).is_err());
        scoped_debug!("second done");
        // Should be ok, change leader connection.
        scoped_debug!("third starting");
        client.propose(to_propose).unwrap();
        assert!(client.leader_connection.is_some());
        scoped_debug!("third done");

        child.join().unwrap();
    }
}
<commit_msg>Update regarding @danburkert's nits<commit_after>//! The `Client` allows users of the `raft` library to connect to remote `Server` instances and
//! issue commands to be applied to the `StateMachine`.

use std::collections::HashSet;
use std::fmt;
use std::io::Write;
use std::net::SocketAddr;
use std::net::TcpStream;

use bufstream::BufStream;
use capnp::{serialize, MessageReader, ReaderOptions};

use messages_capnp::{client_response, proposal_response};
use messages;
use ClientId;
use Result;
use RaftError;

/// The representation of a Client connection to the cluster.
pub struct Client {
    /// The `Uuid` of the client, should be unique in the cluster.
    pub id: ClientId,
    /// The current connection to the current leader.
    /// If it is none it may mean that there is no estabished leader or that there has been
    /// a disconnection.
    leader_connection: Option<BufStream<TcpStream>>,
    /// A lookup for the cluster's nodes.
    cluster: HashSet<SocketAddr>,
}

impl Client {

    /// Creates a new client.
    pub fn new(cluster: HashSet<SocketAddr>) -> Client {
        Client {
            id: ClientId::new(),
            leader_connection: None,
            cluster: cluster,
        }
    }

    /// Proposes an entry to be appended to the replicated log. This will only
    /// return once the entry has been durably committed.
    /// Returns `Error` when the entire cluster has an unknown leader. Try proposing again later.
    pub fn propose(&mut self, entry: &[u8]) -> Result<()> {
        scoped_trace!("{:?}: propose", self);
        let mut message = messages::proposal_request(entry);

        let mut members = self.cluster.iter().cloned();

        loop {
            let mut connection = match self.leader_connection.take() {
                Some(cxn) => {
                    scoped_debug!("had existing connection {:?}", cxn.get_ref().peer_addr());
                    cxn
                },
                None => {
                    let leader = try!(members.next().ok_or(RaftError::LeaderSearchExhausted));
                    scoped_debug!("connecting to potential leader {}", leader);
                    // Send the preamble.
                    let preamble = messages::client_connection_preamble(self.id);
                    let mut stream = BufStream::new(try!(TcpStream::connect(leader)));
                    try!(serialize::write_message(&mut stream, &*preamble));
                    stream
                }
            };
            try!(serialize::write_message(&mut connection, &mut message));
            try!(connection.flush());
            let response = try!(serialize::read_message(&mut connection, ReaderOptions::new()));
            match try!(response.get_root::<client_response::Reader>()).which().unwrap() {
                client_response::Which::Proposal(Ok(response)) => {
                    match response.which().unwrap() {
                        proposal_response::Which::Success(()) => {
                            scoped_debug!("received response Success");
                            self.leader_connection = Some(connection);
                            return Ok(())
                        },
                        proposal_response::Which::UnknownLeader(()) => {
                            scoped_debug!("received response UnknownLeader");
                            ()
                        },
                        proposal_response::Which::NotLeader(leader) => {
                            scoped_debug!("received response NotLeader");
                            let mut connection: TcpStream = try!(TcpStream::connect(try!(leader)));
                            let preamble = messages::client_connection_preamble(self.id);
                            try!(serialize::write_message(&mut connection, &*preamble));
                            self.leader_connection = Some(BufStream::new(connection));
                        }
                    }
                },
                _ => panic!("Unexpected message type"), // TODO: return a proper error
            }
        }
    }
}

impl fmt::Debug for Client {
    fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
        write!(fmt, "Client({})", self.id)
    }
}

#[cfg(test)]
mod test {
    extern crate env_logger;

    use std::collections::HashSet;
    use std::io::{Read, Write};
    use std::net::{TcpStream, SocketAddr, TcpListener};
    use std::str::FromStr;
    use std::thread;

    use uuid::Uuid;
    use capnp::{serialize, ReaderOptions};
    use capnp::message::MessageReader;

    use {Client, ClientId, messages, Result};
    use messages_capnp::{connection_preamble, client_request};

    fn expect_preamble(connection: &mut TcpStream, client_id: Uuid) -> Result<bool> {
        let message = try!(serialize::read_message(connection, ReaderOptions::new()));
        let preamble = try!(message.get_root::<connection_preamble::Reader>());
        // Test to make sure preamble has the right id.
        if let connection_preamble::id::Which::Client(Ok(id)) = try!(preamble.get_id().which()) {
            Ok(Uuid::from_bytes(id).unwrap() == client_id)
        } else {
            Ok(false)
        }
    }

    fn expect_proposal(connection: &mut TcpStream, value: &[u8]) -> Result<bool> {
        let message = try!(serialize::read_message(connection, ReaderOptions::new()));
        let request = try!(message.get_root::<client_request::Reader>());
        // Test to make sure request has the right value.
        if let client_request::Which::Proposal(Ok(proposal)) = try!(request.which()) {
            Ok(proposal.get_entry().unwrap() == value)
        } else {
            Ok(false)
        }
    }

    #[test]
    fn test_proposal_standalone() {
        setup_test!("test_proposal_standalone");
        let mut cluster = HashSet::new();
        let test_server = TcpListener::bind("127.0.0.1:0").unwrap();
        let test_addr = test_server.local_addr().unwrap();
        cluster.insert(test_addr);

        // TODO: Test if the second server is not in the set.
        let second_server = TcpListener::bind("127.0.0.1:0").unwrap();
        let second_addr = second_server.local_addr().unwrap();
        // cluster.insert(second_addr);

        let mut client = Client::new(cluster);
        let client_id = client.id.0.clone();
        let to_propose = b"Bears";

        // The client connects on the first proposal.
        // Wait for it.
        let child = thread::spawn(move || {
            let (mut connection, _)  = test_server.accept().unwrap();

            // First proposal should be fine, no errors.
            scoped_debug!("Should get preamble and proposal. Responds Success");
            expect_preamble(&mut connection, client_id).unwrap();
            expect_proposal(&mut connection, to_propose).unwrap();
            // Send first response! (success!)
            let response = messages::proposal_response_success();
            serialize::write_message(&mut connection, &*response).unwrap();
            connection.flush();

            // Second proposal should report unknown leader, and have the client return error.
            scoped_debug!("Should get proposal. Responds UnknownLeader");
            expect_proposal(&mut connection, to_propose).unwrap();
            // Send response! (unknown leader!) Client should drop connection.
            let response = messages::proposal_response_unknown_leader();
            serialize::write_message(&mut connection, &*response).unwrap();
            connection.flush();
            // Since the client will iter over the members, we need to accept and drop again.
            let (mut connection, _)  = test_server.accept().unwrap();
            serialize::write_message(&mut connection, &*response).unwrap();
            connection.flush();

            // Third Proposal should report NotLeader. Client should choose the server we direct it to.
            scoped_debug!("Should get preamble and proposal. Responds NotLeader.");
            let (mut connection, _)  = test_server.accept().unwrap();
            expect_preamble(&mut connection, client_id).unwrap();
            expect_proposal(&mut connection, to_propose).unwrap();

            // Send response! (not leader!)
            let response = messages::proposal_response_not_leader(&format!("{}", second_addr));
            serialize::write_message(&mut connection, &*response).unwrap();
            connection.flush();

            // Test that it seeks out other server and proposes.
            scoped_debug!("Second server should get preamble and proposal. Responds Success.");
            let (mut connection, _)  = second_server.accept().unwrap();
            expect_preamble(&mut connection, client_id).unwrap();
            expect_proposal(&mut connection, to_propose).unwrap();

            // Send final response! (Success!)
            let response = messages::proposal_response_success();
            serialize::write_message(&mut connection, &*response).unwrap();

        });
        // Propose. It's a marriage made in heaven! :)
        // Should be ok
        client.propose(to_propose).unwrap();
        assert!(client.leader_connection.is_some());
        // Should be err
        assert!(client.propose(to_propose).is_err());
        // Should be ok, change leader connection.
        client.propose(to_propose).unwrap();
        assert!(client.leader_connection.is_some());

        child.join().unwrap();
    }
}
<|endoftext|>


<commit_before>//! The HTTP Client that perform query
use std::collections::HashMap;
use std::mem;
use std::net::SocketAddr;
use std::pin::Pin;

use futures::future::{BoxFuture, Future};

use async_std::io::{self, Write};
use async_std::task::Context;
use async_std::task::Poll;

use super::constants;
use super::http;
use super::request::Request;
use super::response::{Response, ResponseBuilder};
use super::results::CabotResult;

/// Perform the http query
#[derive(Default)]
pub struct Client {
    verbose: bool,
    ipv4: bool,
    ipv6: bool,
    authorities: HashMap<String, SocketAddr>,
    read_timeout: u64,
    connect_timeout: u64,
    dns_timeout: u64,
    request_timeout: u64,
    max_redir: u8,
}

impl<'a> Client {
    /// Construct a new `Client`
    pub fn new() -> Self {
        Client {
            verbose: false,
            ipv4: true,
            ipv6: true,
            authorities: HashMap::new(),
            dns_timeout: constants::DNS_LOOKUP_TIMEOUT * 1000,
            connect_timeout: constants::CONNECT_TIMEOUT * 1000,
            read_timeout: constants::READ_TIMEOUT * 1000,
            request_timeout: constants::REQUEST_TIMEOUT * 1000,
            max_redir: constants::NUMBER_OF_REDIRECT,
        }
    }

    /// Set Address Type authorized for DNS resolution.
    /// If every version is set to false, the resolution will failed.
    pub fn set_ip_version(&mut self, ipv4: bool, ipv6: bool) {
        self.ipv4 = ipv4;
        self.ipv6 = ipv6;
    }

    /// Avoid DNS resolution, force an address to be resolve to a given endpoint.
    /// authority has the format "host:port".
    pub fn add_authority(&mut self, authority: &str, sock_addr: &SocketAddr) {
        self.authorities
            .insert(authority.to_owned(), sock_addr.clone());
    }

    /// Set the timeout for DNS resolution in seconds.
    pub fn set_dns_timeout(&mut self, timeout: u64) {
        self.dns_timeout = timeout * 1000;
    }
    /// Set the connect socket timeout in seconds.
    pub fn set_connect_timeout(&mut self, timeout: u64) {
        self.connect_timeout = timeout * 1000;
    }
    /// Set the read socket timeout in seconds.
    pub fn set_read_timeout(&mut self, timeout: u64) {
        self.read_timeout = timeout * 1000;
    }
    /// Set the request timeout in milliseconds.
    pub fn set_request_timeout(&mut self, timeout: u64) {
        self.request_timeout = timeout * 1000;
    }

    /// Set the timeout for DNS resolution in milliseconds.
    pub fn set_dns_timeout_ms(&mut self, timeout: u64) {
        self.dns_timeout = timeout;
    }
    /// Set the connect socket timeout in milliseconds.
    pub fn set_connect_timeout_ms(&mut self, timeout: u64) {
        self.connect_timeout = timeout;
    }
    /// Set the read socket timeout in milliseconds.
    pub fn set_read_timeout_ms(&mut self, timeout: u64) {
        self.read_timeout = timeout;
    }

    /// Set the request timeout in milliseconds.
    pub fn set_request_timeout_ms(&mut self, timeout: u64) {
        self.request_timeout = timeout;
    }

    /// Set the number of redirection to follow before giving up and return it.
    pub fn set_max_redir(&mut self, max_redir: u8) {
        self.max_redir = max_redir;
    }

    /// Execute the [Request](../request/struct.Request.html) and
    /// return the associate [Response](../response/struct.Response.html).
    pub async fn execute(
        &self,
        request: &Request,
    ) -> CabotResult<Response> {
        self.execute_fut(request).await
    }

    /// Execute the [Request](../request/struct.Request.html) and
    /// return the associate [Response](../response/struct.Response.html) in a box
    /// in order to user it in a async-std task.
    pub fn execute_box(
        &'a self,
        request: &'a Request,
    ) -> BoxFuture<'a, CabotResult<Response>> {
        let fut = Box::pin(self.execute_fut(request));
        Box::pin(ResponseFuture { fut })
    }

    /// Execute the [Request](../request/struct.Request.html) and
    /// return a Future instance in order to use it Client public api.
    fn execute_fut(
        &'a self,
        request: &'a Request,
    ) -> impl Future<Output = CabotResult<Response>> + 'a {
        async move {
            let mut out = CabotLibWrite::new();
            http::http_query(
                request,
                &mut out,
                &self.authorities,
                self.verbose,
                self.ipv4,
                self.ipv6,
                self.dns_timeout,
                self.connect_timeout,
                self.read_timeout,
                self.request_timeout,
                self.max_redir,
            )
            .await?;
            out.response()
        }
    }
}

/// A Future that implement Send to use Client inside task::spwan
pub struct ResponseFuture<'a> {
    fut: Pin<Box<dyn Future<Output = CabotResult<Response>> + 'a>>,
}

unsafe impl<'a> Send for ResponseFuture<'a> {}

impl<'a> Future for ResponseFuture<'a> {
    type Output = CabotResult<Response>;

    fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
        Pin::new(&mut self.fut).poll(cx)
    }
}

struct CabotLibWrite {
    header_read: bool,
    body_buffer: Vec<u8>,
    response_builder: ResponseBuilder,
}

impl CabotLibWrite {
    pub fn new() -> Self {
        CabotLibWrite {
            response_builder: ResponseBuilder::new(),
            body_buffer: Vec::new(),
            header_read: false,
        }
    }

    fn split_headers(&mut self, buf: &[u8]) {
        let mut builder = ResponseBuilder::new();
        if let Some(pos) = buf.iter().position(|&x| x == b'\n') {
            let (status_line, hdrs) = buf.split_at(pos);
            let status_line = String::from_utf8_lossy(status_line);
            builder = builder.set_status_line(status_line.trim_end());
            let mut header = "".to_owned();
            for hdr in hdrs.split(|&x| x == b'\n') {
                let hdr = String::from_utf8_lossy(hdr);
                if hdr.starts_with(' ') || hdr.starts_with('\t') {
                    debug!("Obsolete line folded header reveived in {}", header);
                    header.push_str(" ");
                    header.push_str(hdr.trim());
                } else {
                    let clean_hdr = header.trim();
                    if clean_hdr.len() > 0 {
                        builder = builder.add_header(clean_hdr.trim());
                        header.clear();
                    }
                    header.push_str(hdr.trim());
                }
            }
            let clean_hdr = header.trim();
            if clean_hdr.len() > 0 {
                builder = builder.add_header(clean_hdr.trim());
            }
            self.response_builder = builder;
        }
    }

    pub fn response(&self) -> CabotResult<Response> {
        self.response_builder.build()
    }
}

impl Write for CabotLibWrite {
    fn poll_write(self: Pin<&mut Self>, _cx: &mut Context, buf: &[u8]) -> Poll<io::Result<usize>> {
        let self_ = Pin::get_mut(self);
        if !self_.header_read {
            self_.split_headers(&buf);
            self_.header_read = true;
            Poll::Ready(Ok(0))
        } else {
            self_.body_buffer.extend_from_slice(&buf);
            Poll::Ready(Ok(buf.len()))
        }
    }

    fn poll_flush(self: Pin<&mut Self>, _cx: &mut Context) -> Poll<io::Result<()>> {
        let self_ = Pin::get_mut(self);
        info!("Adding body {:?}", self_.body_buffer);
        let builder = mem::replace(&mut self_.response_builder, ResponseBuilder::new());
        self_.response_builder = builder.set_body(self_.body_buffer.as_slice());
        Poll::Ready(Ok(()))
    }

    // Don't implemented unused method

    fn poll_close(self: Pin<&mut Self>, _cx: &mut Context) -> Poll<io::Result<()>> {
        Poll::Ready(Err(io::Error::new(io::ErrorKind::Other, "Not Implemented")))
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use async_std;
    use async_std::prelude::*;

    #[async_std::test]
    async fn test_build_http_response_from_string() -> std::io::Result<()> {
        let response = [
            vec![
                "HTTP/1.1 200 Ok",
                "Content-Type: text/plain",
                "Content-Length: 12",
            ]
            .join("\r\n"),
            vec!["Hello World!"].join("\r\n"),
        ];

        let mut out = CabotLibWrite::new();
        out.write(response[0].as_bytes()).await.unwrap();
        out.write(response[1].as_bytes()).await.unwrap();
        out.flush().await.unwrap();
        let response = out.response().unwrap();
        assert_eq!(response.http_version(), "HTTP/1.1");
        assert_eq!(response.status_code(), 200);
        assert_eq!(response.status_line(), "200 Ok");
        let headers: &[&str] = &["Content-Type: text/plain", "Content-Length: 12"];
        assert_eq!(response.headers(), headers);
        assert_eq!(
            response.body_as_string().unwrap(),
            "Hello World!".to_owned()
        );
        Ok(())
    }

    #[async_std::test]
    async fn test_build_http_header_obsolete_line_folding() -> std::io::Result<()> {
        let response = [
            vec![
                "HTTP/1.1 200 Ok",
                "ows: https://tools.ietf.org/html/rfc7230",
                "  #section-3.2.4",
                "Content-Length: 12",
            ]
            .join("\r\n"),
            vec!["Hello World!"].join("\r\n"),
        ];

        let mut out = CabotLibWrite::new();
        out.write(response[0].as_bytes()).await.unwrap();
        out.write(response[1].as_bytes()).await.unwrap();
        out.flush().await.unwrap();
        let response = out.response().unwrap();
        assert_eq!(response.http_version(), "HTTP/1.1");
        assert_eq!(response.status_code(), 200);
        assert_eq!(response.status_line(), "200 Ok");
        let headers: &[&str] = &[
            "ows: https://tools.ietf.org/html/rfc7230 #section-3.2.4",
            "Content-Length: 12",
        ];
        assert_eq!(response.headers(), headers);
        assert_eq!(
            response.body_as_string().unwrap(),
            "Hello World!".to_owned()
        );
        Ok(())
    }

    #[async_std::test]
    async fn test_build_http_header_obsolete_line_folding_tab() -> std::io::Result<()> {
        let response = [
            vec![
                "HTTP/1.1 200 Ok",
                "ows: https://tools.ietf.org/html/rfc7230",
                "\t#section-3.2.4",
                "Content-Length: 12",
            ]
            .join("\r\n"),
            vec!["Hello World!"].join("\r\n"),
        ];

        let mut out = CabotLibWrite::new();
        out.write(response[0].as_bytes()).await.unwrap();
        out.write(response[1].as_bytes()).await.unwrap();
        out.flush().await.unwrap();
        let response = out.response().unwrap();
        assert_eq!(response.http_version(), "HTTP/1.1");
        assert_eq!(response.status_code(), 200);
        assert_eq!(response.status_line(), "200 Ok");
        let headers: &[&str] = &[
            "ows: https://tools.ietf.org/html/rfc7230 #section-3.2.4",
            "Content-Length: 12",
        ];
        assert_eq!(response.headers(), headers);
        assert_eq!(
            response.body_as_string().unwrap(),
            "Hello World!".to_owned()
        );
        Ok(())
    }

    #[async_std::test]
    async fn test_build_http_no_response_body() -> std::io::Result<()> {
        let response = vec![
            "HTTP/1.1 302 Moved",
            "Location: https://tools.ietf.org/html/rfc7230#section-3.3",
        ]
        .join("\r\n");

        let mut out = CabotLibWrite::new();
        out.write(response.as_bytes()).await.unwrap();
        out.flush().await.unwrap();
        let response = out.response().unwrap();
        assert_eq!(response.http_version(), "HTTP/1.1");
        assert_eq!(response.status_code(), 302);
        assert_eq!(response.status_line(), "302 Moved");
        let headers: &[&str] = &["Location: https://tools.ietf.org/html/rfc7230#section-3.3"];
        assert_eq!(response.headers(), headers);
        assert_eq!(response.body_as_string().unwrap(), "");
        Ok(())
    }
}
<commit_msg>Fix set_request_timeout docs<commit_after>//! The HTTP Client that perform query
use std::collections::HashMap;
use std::mem;
use std::net::SocketAddr;
use std::pin::Pin;

use futures::future::{BoxFuture, Future};

use async_std::io::{self, Write};
use async_std::task::Context;
use async_std::task::Poll;

use super::constants;
use super::http;
use super::request::Request;
use super::response::{Response, ResponseBuilder};
use super::results::CabotResult;

/// Perform the http query
#[derive(Default)]
pub struct Client {
    verbose: bool,
    ipv4: bool,
    ipv6: bool,
    authorities: HashMap<String, SocketAddr>,
    read_timeout: u64,
    connect_timeout: u64,
    dns_timeout: u64,
    request_timeout: u64,
    max_redir: u8,
}

impl<'a> Client {
    /// Construct a new `Client`
    pub fn new() -> Self {
        Client {
            verbose: false,
            ipv4: true,
            ipv6: true,
            authorities: HashMap::new(),
            dns_timeout: constants::DNS_LOOKUP_TIMEOUT * 1000,
            connect_timeout: constants::CONNECT_TIMEOUT * 1000,
            read_timeout: constants::READ_TIMEOUT * 1000,
            request_timeout: constants::REQUEST_TIMEOUT * 1000,
            max_redir: constants::NUMBER_OF_REDIRECT,
        }
    }

    /// Set Address Type authorized for DNS resolution.
    /// If every version is set to false, the resolution will failed.
    pub fn set_ip_version(&mut self, ipv4: bool, ipv6: bool) {
        self.ipv4 = ipv4;
        self.ipv6 = ipv6;
    }

    /// Avoid DNS resolution, force an address to be resolve to a given endpoint.
    /// authority has the format "host:port".
    pub fn add_authority(&mut self, authority: &str, sock_addr: &SocketAddr) {
        self.authorities
            .insert(authority.to_owned(), sock_addr.clone());
    }

    /// Set the timeout for DNS resolution in seconds.
    pub fn set_dns_timeout(&mut self, timeout: u64) {
        self.dns_timeout = timeout * 1000;
    }
    /// Set the connect socket timeout in seconds.
    pub fn set_connect_timeout(&mut self, timeout: u64) {
        self.connect_timeout = timeout * 1000;
    }
    /// Set the read socket timeout in seconds.
    pub fn set_read_timeout(&mut self, timeout: u64) {
        self.read_timeout = timeout * 1000;
    }
    /// Set the request timeout in seconds.
    pub fn set_request_timeout(&mut self, timeout: u64) {
        self.request_timeout = timeout * 1000;
    }

    /// Set the timeout for DNS resolution in milliseconds.
    pub fn set_dns_timeout_ms(&mut self, timeout: u64) {
        self.dns_timeout = timeout;
    }
    /// Set the connect socket timeout in milliseconds.
    pub fn set_connect_timeout_ms(&mut self, timeout: u64) {
        self.connect_timeout = timeout;
    }
    /// Set the read socket timeout in milliseconds.
    pub fn set_read_timeout_ms(&mut self, timeout: u64) {
        self.read_timeout = timeout;
    }

    /// Set the request timeout in milliseconds.
    pub fn set_request_timeout_ms(&mut self, timeout: u64) {
        self.request_timeout = timeout;
    }

    /// Set the number of redirection to follow before giving up and return it.
    pub fn set_max_redir(&mut self, max_redir: u8) {
        self.max_redir = max_redir;
    }

    /// Execute the [Request](../request/struct.Request.html) and
    /// return the associate [Response](../response/struct.Response.html).
    pub async fn execute(
        &self,
        request: &Request,
    ) -> CabotResult<Response> {
        self.execute_fut(request).await
    }

    /// Execute the [Request](../request/struct.Request.html) and
    /// return the associate [Response](../response/struct.Response.html) in a box
    /// in order to user it in a async-std task.
    pub fn execute_box(
        &'a self,
        request: &'a Request,
    ) -> BoxFuture<'a, CabotResult<Response>> {
        let fut = Box::pin(self.execute_fut(request));
        Box::pin(ResponseFuture { fut })
    }

    /// Execute the [Request](../request/struct.Request.html) and
    /// return a Future instance in order to use it Client public api.
    fn execute_fut(
        &'a self,
        request: &'a Request,
    ) -> impl Future<Output = CabotResult<Response>> + 'a {
        async move {
            let mut out = CabotLibWrite::new();
            http::http_query(
                request,
                &mut out,
                &self.authorities,
                self.verbose,
                self.ipv4,
                self.ipv6,
                self.dns_timeout,
                self.connect_timeout,
                self.read_timeout,
                self.request_timeout,
                self.max_redir,
            )
            .await?;
            out.response()
        }
    }
}

/// A Future that implement Send to use Client inside task::spwan
pub struct ResponseFuture<'a> {
    fut: Pin<Box<dyn Future<Output = CabotResult<Response>> + 'a>>,
}

unsafe impl<'a> Send for ResponseFuture<'a> {}

impl<'a> Future for ResponseFuture<'a> {
    type Output = CabotResult<Response>;

    fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
        Pin::new(&mut self.fut).poll(cx)
    }
}

struct CabotLibWrite {
    header_read: bool,
    body_buffer: Vec<u8>,
    response_builder: ResponseBuilder,
}

impl CabotLibWrite {
    pub fn new() -> Self {
        CabotLibWrite {
            response_builder: ResponseBuilder::new(),
            body_buffer: Vec::new(),
            header_read: false,
        }
    }

    fn split_headers(&mut self, buf: &[u8]) {
        let mut builder = ResponseBuilder::new();
        if let Some(pos) = buf.iter().position(|&x| x == b'\n') {
            let (status_line, hdrs) = buf.split_at(pos);
            let status_line = String::from_utf8_lossy(status_line);
            builder = builder.set_status_line(status_line.trim_end());
            let mut header = "".to_owned();
            for hdr in hdrs.split(|&x| x == b'\n') {
                let hdr = String::from_utf8_lossy(hdr);
                if hdr.starts_with(' ') || hdr.starts_with('\t') {
                    debug!("Obsolete line folded header reveived in {}", header);
                    header.push_str(" ");
                    header.push_str(hdr.trim());
                } else {
                    let clean_hdr = header.trim();
                    if clean_hdr.len() > 0 {
                        builder = builder.add_header(clean_hdr.trim());
                        header.clear();
                    }
                    header.push_str(hdr.trim());
                }
            }
            let clean_hdr = header.trim();
            if clean_hdr.len() > 0 {
                builder = builder.add_header(clean_hdr.trim());
            }
            self.response_builder = builder;
        }
    }

    pub fn response(&self) -> CabotResult<Response> {
        self.response_builder.build()
    }
}

impl Write for CabotLibWrite {
    fn poll_write(self: Pin<&mut Self>, _cx: &mut Context, buf: &[u8]) -> Poll<io::Result<usize>> {
        let self_ = Pin::get_mut(self);
        if !self_.header_read {
            self_.split_headers(&buf);
            self_.header_read = true;
            Poll::Ready(Ok(0))
        } else {
            self_.body_buffer.extend_from_slice(&buf);
            Poll::Ready(Ok(buf.len()))
        }
    }

    fn poll_flush(self: Pin<&mut Self>, _cx: &mut Context) -> Poll<io::Result<()>> {
        let self_ = Pin::get_mut(self);
        info!("Adding body {:?}", self_.body_buffer);
        let builder = mem::replace(&mut self_.response_builder, ResponseBuilder::new());
        self_.response_builder = builder.set_body(self_.body_buffer.as_slice());
        Poll::Ready(Ok(()))
    }

    // Don't implemented unused method

    fn poll_close(self: Pin<&mut Self>, _cx: &mut Context) -> Poll<io::Result<()>> {
        Poll::Ready(Err(io::Error::new(io::ErrorKind::Other, "Not Implemented")))
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use async_std;
    use async_std::prelude::*;

    #[async_std::test]
    async fn test_build_http_response_from_string() -> std::io::Result<()> {
        let response = [
            vec![
                "HTTP/1.1 200 Ok",
                "Content-Type: text/plain",
                "Content-Length: 12",
            ]
            .join("\r\n"),
            vec!["Hello World!"].join("\r\n"),
        ];

        let mut out = CabotLibWrite::new();
        out.write(response[0].as_bytes()).await.unwrap();
        out.write(response[1].as_bytes()).await.unwrap();
        out.flush().await.unwrap();
        let response = out.response().unwrap();
        assert_eq!(response.http_version(), "HTTP/1.1");
        assert_eq!(response.status_code(), 200);
        assert_eq!(response.status_line(), "200 Ok");
        let headers: &[&str] = &["Content-Type: text/plain", "Content-Length: 12"];
        assert_eq!(response.headers(), headers);
        assert_eq!(
            response.body_as_string().unwrap(),
            "Hello World!".to_owned()
        );
        Ok(())
    }

    #[async_std::test]
    async fn test_build_http_header_obsolete_line_folding() -> std::io::Result<()> {
        let response = [
            vec![
                "HTTP/1.1 200 Ok",
                "ows: https://tools.ietf.org/html/rfc7230",
                "  #section-3.2.4",
                "Content-Length: 12",
            ]
            .join("\r\n"),
            vec!["Hello World!"].join("\r\n"),
        ];

        let mut out = CabotLibWrite::new();
        out.write(response[0].as_bytes()).await.unwrap();
        out.write(response[1].as_bytes()).await.unwrap();
        out.flush().await.unwrap();
        let response = out.response().unwrap();
        assert_eq!(response.http_version(), "HTTP/1.1");
        assert_eq!(response.status_code(), 200);
        assert_eq!(response.status_line(), "200 Ok");
        let headers: &[&str] = &[
            "ows: https://tools.ietf.org/html/rfc7230 #section-3.2.4",
            "Content-Length: 12",
        ];
        assert_eq!(response.headers(), headers);
        assert_eq!(
            response.body_as_string().unwrap(),
            "Hello World!".to_owned()
        );
        Ok(())
    }

    #[async_std::test]
    async fn test_build_http_header_obsolete_line_folding_tab() -> std::io::Result<()> {
        let response = [
            vec![
                "HTTP/1.1 200 Ok",
                "ows: https://tools.ietf.org/html/rfc7230",
                "\t#section-3.2.4",
                "Content-Length: 12",
            ]
            .join("\r\n"),
            vec!["Hello World!"].join("\r\n"),
        ];

        let mut out = CabotLibWrite::new();
        out.write(response[0].as_bytes()).await.unwrap();
        out.write(response[1].as_bytes()).await.unwrap();
        out.flush().await.unwrap();
        let response = out.response().unwrap();
        assert_eq!(response.http_version(), "HTTP/1.1");
        assert_eq!(response.status_code(), 200);
        assert_eq!(response.status_line(), "200 Ok");
        let headers: &[&str] = &[
            "ows: https://tools.ietf.org/html/rfc7230 #section-3.2.4",
            "Content-Length: 12",
        ];
        assert_eq!(response.headers(), headers);
        assert_eq!(
            response.body_as_string().unwrap(),
            "Hello World!".to_owned()
        );
        Ok(())
    }

    #[async_std::test]
    async fn test_build_http_no_response_body() -> std::io::Result<()> {
        let response = vec![
            "HTTP/1.1 302 Moved",
            "Location: https://tools.ietf.org/html/rfc7230#section-3.3",
        ]
        .join("\r\n");

        let mut out = CabotLibWrite::new();
        out.write(response.as_bytes()).await.unwrap();
        out.flush().await.unwrap();
        let response = out.response().unwrap();
        assert_eq!(response.http_version(), "HTTP/1.1");
        assert_eq!(response.status_code(), 302);
        assert_eq!(response.status_line(), "302 Moved");
        let headers: &[&str] = &["Location: https://tools.ietf.org/html/rfc7230#section-3.3"];
        assert_eq!(response.headers(), headers);
        assert_eq!(response.body_as_string().unwrap(), "");
        Ok(())
    }
}
<|endoftext|>


<commit_before>extern crate crypto;
extern crate rand;

use std::iter::repeat;
use std::fs;
use std::io;
use std::fs::File;
use std::io::{Read, Write};
use std::os::unix::fs::PermissionsExt;
use std::process::Command;
use std::env::set_current_dir;
use std::path::Path;

use self::crypto::salsa20::Salsa20;
use self::crypto::symmetriccipher::SynchronousStreamCipher;
use self::rand::{OsRng, Rng};

use base64;

const SALT_LENGTH: usize = 24;
const KEY_LENGTH: usize = 32;

pub fn generate_password(key: &[u8], meat: Vec<u8>, salt: Vec<u8>) -> Vec<u8> {
    let mut cipher = Salsa20::new_xsalsa20(&key, &salt);
    let mut buf: Vec<u8> = repeat(0).take(meat.len()).collect();
    cipher.process(&meat, &mut buf);
    buf
}

pub fn generate_salt() -> Vec<u8> {
    let mut rng = OsRng::new().expect("OsRng init failed");
    rng.gen_iter::<u8>().take(SALT_LENGTH).collect()
}

pub fn generate_meat(i: usize) -> Vec<u8> {
    let mut rng = OsRng::new().expect("OsRng init failed");
    rng.gen_iter::<u8>().take(i).collect()
}

pub fn load_or_create_key(filename: &str) -> Vec<u8> {
    match read_data(filename) {
        Ok(s) => {
            let s = s.lines().next().expect("Reading base64 key failed");
            base64::decode(s).expect("Key base64 decoding failed")
        }
        Err(_) => {
            println!("Creating a new key in {}", filename);
            let mut rng = OsRng::new().expect("OsRng init failed");
            let new_key: Vec<u8> = rng.gen_iter::<u8>().take(KEY_LENGTH).collect();
            let key_base64 = base64::encode(&new_key); // .to_base64(base64::STANDARD);
            write_data(&key_base64, filename);
            set_file_perms(filename, 0o400);
            new_key
        }
    }
}

pub fn ensure_data_dir(data_dir: &str) {
    fs::create_dir_all(data_dir.to_string())
        .expect(&format!("Creating data directory {} failed", data_dir));
    set_file_perms(&data_dir, 0o700);

    if !Path::new(data_dir).join(".git").exists() {
        set_current_dir(data_dir).expect("Failed to set dir to data dir");

        Command::new("git").arg("init").status().expect("Failed to init data git repo");

        let git_args = [
            "config",
            "user.email",
            "chaos"
        ];
        Command::new("git").args(&git_args).status().expect("Failed to set user.email");

        let git_args = [
            "config",
            "user.name",
            "chaos"
        ];
        Command::new("git").args(&git_args).status().expect("Failed to set user.name");
    }
}

pub fn read_data(path: &str) -> Result<String, io::Error> {
    let mut f = try!(File::open(path));
    let mut s = String::new();
    try!(f.read_to_string(&mut s));
    Ok(s)
}

pub fn commit_data(data_dir: &str, data_file: &str, commit_text: &str) {
    set_current_dir(data_dir).expect("Failed to set dir to data dir");

    let git_args = [
        "add",
        data_file,
    ].to_vec();

    Command::new("git").args(&git_args).status().unwrap_or_else(|e| {
        panic!("Failed to add data file: {}", e);
    });

    let git_args = [
        "commit",
        "-m",
        commit_text
    ].to_vec();
    Command::new("git").args(&git_args).status().unwrap_or_else(|e| {
        panic!("Failed to git commit: {}", e);
    });

}

pub fn write_data(data: &str, filename: &str) {
    let mut f = File::create(filename).unwrap();
    f.write_all(data.as_bytes()).expect("Data file write failed");
    f.write_all(b"\n").expect("Newline write failed!?");
    f.sync_all().expect("Sync failed");
}

pub fn set_file_perms(filename: &str, mode: u32) {
    let mut perms = fs::metadata(filename).expect("Gettings perms failed").permissions();
    perms.set_mode(mode);
    fs::set_permissions(filename, perms).expect("Setting permission failed");
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]

    #[test]
    fn test_generate_salt() {
        let salt1 = generate_salt();
        let salt2 = generate_salt();
        assert!(salt1 != salt2);
    }

    #[test]
    fn test_generate_meat() {
        let meat1 = generate_meat(8);
        let meat2 = generate_meat(8);
        assert!(meat1 != meat2);
        assert_eq!(8, meat1.len());
        assert_eq!(8, meat2.len());
    }

    #[test]
    fn test_generate_password() {
        let pass1 = generate_password(&[1,2,3,4,5,6,7,8,9,10,11,12,13,14,15,16,17,18,19,20,21,22,23,24,25,26,27,28,29,30,31,32], vec!(1, 2), vec!(1,2,3,4,5,6,7,8,9,10,11,12,13,14,15,16,17,18,19,20,21,22,23,24));
        assert_eq!(vec!(230, 6), pass1);

    }
}
<commit_msg>Nothing here<commit_after>extern crate crypto;
extern crate rand;

use std::iter::repeat;
use std::fs;
use std::io;
use std::fs::File;
use std::io::{Read, Write};
use std::os::unix::fs::PermissionsExt;
use std::process::Command;
use std::env::set_current_dir;
use std::path::Path;

use self::crypto::salsa20::Salsa20;
use self::crypto::symmetriccipher::SynchronousStreamCipher;
use self::rand::{OsRng, Rng};

use base64;

const SALT_LENGTH: usize = 24;
const KEY_LENGTH: usize = 32;

pub fn generate_password(key: &[u8], meat: Vec<u8>, salt: Vec<u8>) -> Vec<u8> {
    let mut cipher = Salsa20::new_xsalsa20(&key, &salt);
    let mut buf: Vec<u8> = repeat(0).take(meat.len()).collect();
    cipher.process(&meat, &mut buf);
    buf
}

pub fn generate_salt() -> Vec<u8> {
    let mut rng = OsRng::new().expect("OsRng init failed");
    rng.gen_iter::<u8>().take(SALT_LENGTH).collect()
}

pub fn generate_meat(i: usize) -> Vec<u8> {
    let mut rng = OsRng::new().expect("OsRng init failed");
    rng.gen_iter::<u8>().take(i).collect()
}

pub fn load_or_create_key(filename: &str) -> Vec<u8> {
    match read_data(filename) {
        Ok(s) => {
            let s = s.lines().next().expect("Reading base64 key failed");
            base64::decode(s).expect("Key base64 decoding failed")
        }
        Err(_) => {
            println!("Creating a new key in {}", filename);
            let mut rng = OsRng::new().expect("OsRng init failed");
            let new_key: Vec<u8> = rng.gen_iter::<u8>().take(KEY_LENGTH).collect();
            let key_base64 = base64::encode(&new_key); // .to_base64(base64::STANDARD);
            write_data(&key_base64, filename);
            set_file_perms(filename, 0o400);
            new_key
        }
    }
}

pub fn ensure_data_dir(data_dir: &str) {
    fs::create_dir_all(data_dir.to_string())
        .expect(&format!("Creating data directory {} failed", data_dir));
    set_file_perms(&data_dir, 0o700);

    if !Path::new(data_dir).join(".git").exists() {
        set_current_dir(data_dir).expect("Failed to set dir to data dir");

        Command::new("git").arg("init").status().expect("Failed to init data git repo");

        let git_args = [
            "config",
            "user.email",
            "chaos"
        ];
        Command::new("git").args(&git_args).status().expect("Failed to set user.email");

        let git_args = [
            "config",
            "user.name",
            "chaos"
        ];
        Command::new("git").args(&git_args).status().expect("Failed to set user.name");
    }
}

pub fn read_data(path: &str) -> Result<String, io::Error> {
    let mut f = try!(File::open(path));
    let mut s = String::new();
    try!(f.read_to_string(&mut s));
    Ok(s)
}

pub fn commit_data(data_dir: &str, data_file: &str, commit_text: &str) {
    set_current_dir(data_dir).expect("Failed to set dir to data dir");

    let git_args = [
        "add",
        data_file,
    ].to_vec();

    Command::new("git").args(&git_args).status().unwrap_or_else(|e| {
        panic!("Failed to add data file: {}", e);
    });

    let git_args = [
        "commit",
        "-m",
        commit_text
    ].to_vec();
    Command::new("git").args(&git_args).status().unwrap_or_else(|e| {
        panic!("Failed to git commit: {}", e);
    });

}

pub fn write_data(data: &str, filename: &str) {
    let mut f = File::create(filename).unwrap();
    f.write_all(data.as_bytes()).expect("Data file write failed");
    f.write_all(b"\n").expect("Newline write failed!?");
    f.sync_all().expect("Sync failed");
}

pub fn set_file_perms(filename: &str, mode: u32) {
    let mut perms = fs::metadata(filename).expect("Gettings perms failed").permissions();
    perms.set_mode(mode);
    fs::set_permissions(filename, perms).expect("Setting permission failed");
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_generate_salt() {
        let salt1 = generate_salt();
        let salt2 = generate_salt();
        assert!(salt1 != salt2);
    }

    #[test]
    fn test_generate_meat() {
        let meat1 = generate_meat(8);
        let meat2 = generate_meat(8);
        assert!(meat1 != meat2);
        assert_eq!(8, meat1.len());
        assert_eq!(8, meat2.len());
    }

    #[test]
    fn test_generate_password() {
        let pass1 = generate_password(&[1,2,3,4,5,6,7,8,9,10,11,12,13,14,15,16,17,18,19,20,21,22,23,24,25,26,27,28,29,30,31,32], vec!(1, 2), vec!(1,2,3,4,5,6,7,8,9,10,11,12,13,14,15,16,17,18,19,20,21,22,23,24));
        assert_eq!(vec!(230, 6), pass1);

    }
}
<|endoftext|>