<commit_before>#![deny(warnings)]

extern crate coreutils;
extern crate extra;

use std::env;
use std::fs;
use std::fs::{Metadata, FileType};
use std::path::Path;
use std::io::{stdout, StdoutLock, stderr, Stderr, Write};
use std::os::unix::fs::MetadataExt;

use std::process::exit;

use coreutils::{ArgParser, to_human_readable_string};
use extra::option::OptionalExt;

const MAN_PAGE: &'static str = /* @MANSTART{ls} */ r#"
NAME
    ls - list directory contents

SYNOPSIS
    ls [ -h | --help | -l ][FILE]...

DESCRIPTION
    List information about the FILE(s), or the current directory

OPTIONS
    -a, --all
        do not ignore entries starting with .
    -h, --human-readable
        with -l, print human readable sizes
    --help
        display this help and exit
    -l
        use a long listing format
    -r, --reverse
        reverse order while sorting
    -R, --recursive
        list subdirectories recursively
"#; /* @MANEND */

fn mode_to_human_readable(file_type: &FileType, mode: u32) -> String {

    let mut result = String::from("");
    //TODO is a symlink? => push 'l'
    if file_type.is_dir() {
        result.push('d');
    }else{
        result.push('-');
    }

    let mode_str = format!("{:>6o}", mode);
    let mode_chars = mode_str[3..].chars();
    for i in mode_chars {
        match i {
            '7' => result.push_str("rwx"),
            '6' => result.push_str("rw-"),
            '5' => result.push_str("r-x"),
            '4' => result.push_str("r--"),
            '3' => result.push_str("-wx"),
            '2' => result.push_str("-w-"),
            '1' => result.push_str("--x"),
            _   => result.push_str("---")
        }
    }

    return result;
}

fn print_item(item_path: &str, metadata: &Metadata, parser: &ArgParser, stdout: &mut StdoutLock, stderr: &mut Stderr){
    if parser.found("long-format") {
    stdout.write(&format!("{} {:>5} {:>5} ",
            mode_to_human_readable(&(metadata.file_type()), metadata.mode()),
            metadata.uid(),
            metadata.gid()).as_bytes()).try(stderr);
        if parser.found("human-readable") {
            stdout.write(&format!("{:>6} ",
                    to_human_readable_string(metadata.size())).as_bytes()).try(stderr);
        } else {
            stdout.write(&format!("{:>8} ", metadata.size()).as_bytes()).try(stderr);
        }
    }
    if item_path.starts_with("./") {
        stdout.write(&item_path[2..].as_bytes()).try(stderr);
    }else{
        stdout.write(item_path.as_bytes()).try(stderr);
    }
    stdout.write("\n".as_bytes()).try(stderr);
    stdout.flush().try(stderr);
}

fn list_dir(path: &str, parser: &ArgParser, stdout: &mut StdoutLock, stderr: &mut Stderr) {
    let mut show_hidden = false;
    if parser.found("all") {
        show_hidden = true;
    }

    let metadata = fs::metadata(path).try(stderr);
    if metadata.is_dir() {
        let read_dir = Path::new(path).read_dir().try(stderr);

        let mut entries: Vec<String> = read_dir
                .filter_map(|x| x.ok())
                .map(|x| {
                    let file_name = x.file_name().to_string_lossy().into_owned();
                    file_name
                })
                .filter(|x| {
                    match show_hidden {
                        true => true,
                        false => !x.starts_with(".")
                    }
                })
                .collect();

        if parser.found("reverse") {
            entries.sort_by(|a, b| b.cmp(a));
        } else {
            entries.sort_by(|a, b| a.cmp(b));
        }

        for entry in entries.iter() {
            let mut entry_path = path.to_owned();
            if !entry_path.ends_with('/') {
                entry_path.push('/');
            }
            entry_path.push_str(&entry);
            let metadata = fs::metadata(&entry_path).try(stderr);
            print_item(&entry_path, &metadata, &parser, stdout, stderr);
            if parser.found("recursive") && metadata.is_dir() {
                list_dir(&entry_path, parser, stdout, stderr);
            }
        }
    } else {
        print_item(&path, &metadata, &parser, stdout, stderr);
    }
}

fn main() {
    let stdout = stdout();
    let mut stdout = stdout.lock();
    let mut stderr = stderr();

    let mut parser = ArgParser::new(6)
        .add_flag(&["a", "all"])
        .add_flag(&["l", "long-format"])
        .add_flag(&["h", "human-readable"])
        .add_flag(&["r", "reverse"])
        .add_flag(&["R", "recursive"])
        .add_flag(&["", "help"]);
    parser.parse(env::args());

    if parser.found("help") {
        stdout.write(MAN_PAGE.as_bytes()).try(&mut stderr);
        stdout.flush().try(&mut stderr);
        exit(0);
    }

    if parser.args.is_empty() {
        list_dir(".", &parser, &mut stdout, &mut stderr);
    } else {
        for dir in parser.args.iter() {
            list_dir(&dir, &parser, &mut stdout, &mut stderr);
        }
    }
}

#[test]
fn test_human_readable() {
    assert_eq!(to_human_readable_string(0), "0");
    assert_eq!(to_human_readable_string(1023), "1023");
    assert_eq!(to_human_readable_string(1024), "1.0K");
    assert_eq!(to_human_readable_string(1024 + 100), "1.1K");
    assert_eq!(to_human_readable_string(1024u64.pow(2) * 2), "2.0M");
    assert_eq!(to_human_readable_string(1024u64.pow(3) * 3), "3.0G");
    assert_eq!(to_human_readable_string(1024u64.pow(4) * 4), "4.0T");
    assert_eq!(to_human_readable_string(1024u64.pow(5) * 5), "5.0P");
    assert_eq!(to_human_readable_string(1024u64.pow(6) * 6), "6.0E");
}
<commit_msg>added basic support for symbolic links to ls command ls shows 'l' for link now<commit_after>#![deny(warnings)]

extern crate coreutils;
extern crate extra;

use std::env;
use std::fs;
use std::fs::{Metadata, FileType};
use std::path::Path;
use std::io::{stdout, StdoutLock, stderr, Stderr, Write};
use std::os::unix::fs::MetadataExt;

use std::process::exit;

use coreutils::{ArgParser, to_human_readable_string};
use extra::option::OptionalExt;

const MAN_PAGE: &'static str = /* @MANSTART{ls} */ r#"
NAME
    ls - list directory contents

SYNOPSIS
    ls [ -h | --help | -l ][FILE]...

DESCRIPTION
    List information about the FILE(s), or the current directory

OPTIONS
    -a, --all
        do not ignore entries starting with .
    -h, --human-readable
        with -l, print human readable sizes
    --help
        display this help and exit
    -l
        use a long listing format
    -r, --reverse
        reverse order while sorting
    -R, --recursive
        list subdirectories recursively
"#; /* @MANEND */

fn mode_to_human_readable(file_type: &FileType, symlink_file_type: &FileType, mode: u32) -> String {

    let mut result = String::from("");
    if symlink_file_type.is_symlink(){
            result.push('l')
    }else if file_type.is_dir() {
        result.push('d');
    }else{
        result.push('-');
    }

    let mode_str = format!("{:>6o}", mode);
    let mode_chars = mode_str[3..].chars();
    for i in mode_chars {
        match i {
            '7' => result.push_str("rwx"),
            '6' => result.push_str("rw-"),
            '5' => result.push_str("r-x"),
            '4' => result.push_str("r--"),
            '3' => result.push_str("-wx"),
            '2' => result.push_str("-w-"),
            '1' => result.push_str("--x"),
            _   => result.push_str("---")
        }
    }

    return result;
}

fn print_item(item_path: &str, metadata: &Metadata, symlink_metadata: &Metadata, parser: &ArgParser, stdout: &mut StdoutLock, stderr: &mut Stderr){
    if parser.found("long-format") {
    stdout.write(&format!("{} {:>5} {:>5} ",
            mode_to_human_readable(&(metadata.file_type()), &(symlink_metadata.file_type()), metadata.mode()),
            metadata.uid(),
            metadata.gid()).as_bytes()).try(stderr);
        if parser.found("human-readable") {
            stdout.write(&format!("{:>6} ",
                    to_human_readable_string(metadata.size())).as_bytes()).try(stderr);
        } else {
            stdout.write(&format!("{:>8} ", metadata.size()).as_bytes()).try(stderr);
        }
    }
    if item_path.starts_with("./") {
        stdout.write(&item_path[2..].as_bytes()).try(stderr);
    }else{
        stdout.write(item_path.as_bytes()).try(stderr);
    }
    stdout.write("\n".as_bytes()).try(stderr);
    stdout.flush().try(stderr);
}

fn list_dir(path: &str, parser: &ArgParser, stdout: &mut StdoutLock, stderr: &mut Stderr) {
    let mut show_hidden = false;
    if parser.found("all") {
        show_hidden = true;
    }

    let symlink_metadata = fs::symlink_metadata(path).try(stderr);
    let metadata = fs::metadata(path).try(stderr);
    if metadata.is_dir() {
        let read_dir = Path::new(path).read_dir().try(stderr);

        let mut entries: Vec<String> = read_dir
                .filter_map(|x| x.ok())
                .map(|x| {
                    let file_name = x.file_name().to_string_lossy().into_owned();
                    file_name
                })
                .filter(|x| {
                    match show_hidden {
                        true => true,
                        false => !x.starts_with(".")
                    }
                })
                .collect();

        if parser.found("reverse") {
            entries.sort_by(|a, b| b.cmp(a));
        } else {
            entries.sort_by(|a, b| a.cmp(b));
        }

        for entry in entries.iter() {
            let mut entry_path = path.to_owned();
            if !entry_path.ends_with('/') {
                entry_path.push('/');
            }
            entry_path.push_str(&entry);
            let symlink_metadata = fs::symlink_metadata(&entry_path).try(stderr);
            let metadata = fs::metadata(&entry_path).try(stderr);
            print_item(&entry_path, &metadata, &symlink_metadata, &parser, stdout, stderr);
            if parser.found("recursive") && metadata.is_dir() {
                list_dir(&entry_path, parser, stdout, stderr);
            }
        }
    } else {
        print_item(&path, &metadata, &symlink_metadata, &parser, stdout, stderr);
    }
}

fn main() {
    let stdout = stdout();
    let mut stdout = stdout.lock();
    let mut stderr = stderr();

    let mut parser = ArgParser::new(6)
        .add_flag(&["a", "all"])
        .add_flag(&["l", "long-format"])
        .add_flag(&["h", "human-readable"])
        .add_flag(&["r", "reverse"])
        .add_flag(&["R", "recursive"])
        .add_flag(&["", "help"]);
    parser.parse(env::args());

    if parser.found("help") {
        stdout.write(MAN_PAGE.as_bytes()).try(&mut stderr);
        stdout.flush().try(&mut stderr);
        exit(0);
    }

    if parser.args.is_empty() {
        list_dir(".", &parser, &mut stdout, &mut stderr);
    } else {
        for dir in parser.args.iter() {
            list_dir(&dir, &parser, &mut stdout, &mut stderr);
        }
    }
}

#[test]
fn test_human_readable() {
    assert_eq!(to_human_readable_string(0), "0");
    assert_eq!(to_human_readable_string(1023), "1023");
    assert_eq!(to_human_readable_string(1024), "1.0K");
    assert_eq!(to_human_readable_string(1024 + 100), "1.1K");
    assert_eq!(to_human_readable_string(1024u64.pow(2) * 2), "2.0M");
    assert_eq!(to_human_readable_string(1024u64.pow(3) * 3), "3.0G");
    assert_eq!(to_human_readable_string(1024u64.pow(4) * 4), "4.0T");
    assert_eq!(to_human_readable_string(1024u64.pow(5) * 5), "5.0P");
    assert_eq!(to_human_readable_string(1024u64.pow(6) * 6), "6.0E");
}
<|endoftext|>


<commit_before>//! Buffers to safely work with audio samples.

use std::iter::{Zip, IntoIterator};
use std::vec::IntoIter;
use std::slice;

use num::Float;

/// A buffer containing `ChannelBuffer` buffers for each input/output.
#[allow(dead_code)]
pub struct AudioBuffer<'a, T: 'a + Float> {
    inputs: Vec<ChannelBuffer<'a, T>>,
    outputs: Vec<ChannelBuffer<'a, T>>
}

/// Iterator over channel buffers for either inputs or outputs.
pub type ChannelBufferIter<'a, T> = IntoIter<ChannelBuffer<'a, T>>;

#[allow(dead_code)]
impl<'a, T: 'a + Float> AudioBuffer<'a, T> {
    /// Create an `AudioBuffer` from vectors of slices. Each vector represents either an input or
    /// output, and contains an array of samples.
    /// Eg if inputs was a vector of size 2 containing slices of size 512, it would hold 2 inputs
    /// where each input holds 512 samples.
    pub fn new(inputs: Vec<&'a mut [T]>, outputs: Vec<&'a mut [T]>) -> AudioBuffer<'a, T> {
        // Creats an input / output for each array in respective Vec
        AudioBuffer {
            inputs: inputs.into_iter().map(|input| ChannelBuffer::new(input)).collect(),
            outputs: outputs.into_iter().map(|output| ChannelBuffer::new(output)).collect()
        }
    }

    /// Create an `AudioBuffer` from raw pointers. Only really useful for interacting with the VST
    /// API.
    pub unsafe fn from_raw(inputs_raw: *mut *mut T, outputs_raw: *mut *mut T, num_inputs: usize, num_outputs: usize, samples: usize) -> AudioBuffer<'a, T> {
        // Allocate an array size of vst input count
        let mut inputs: Vec<&mut [T]> = Vec::with_capacity(num_inputs);
        for i in 0 .. inputs.capacity() {
            // Push samples for each input to `inputs` array
            inputs.push(slice::from_raw_parts_mut(*inputs_raw.offset(i as isize), samples));
        }

        // Allocate an array size of vst output count
        let mut outputs: Vec<&mut [T]> = Vec::with_capacity(num_outputs);
        for i in 0 .. outputs.capacity() {
            // Push samples for each output to `outputs` array
            outputs.push(slice::from_raw_parts_mut(*outputs_raw.offset(i as isize), samples));
        }

        // Call constructor with slices
        AudioBuffer::new(inputs, outputs)
    }

    /// Retreieve input at specified index.
    pub fn input(&'a mut self, index: usize) -> Option<&'a mut ChannelBuffer<'a, T>> {
        self.inputs.get_mut(index)
    }

    /// Retreieve output at specified index.
    pub fn output(&'a mut self, index: usize) -> Option<&'a mut ChannelBuffer<'a, T>> {
        self.outputs.get_mut(index)
    }

    /// Create an iterator over all inputs.
    pub fn inputs(self) -> ChannelBufferIter<'a, T> {
        self.inputs.into_iter()
    }

    /// Create an iterator over all outputs.
    pub fn outputs(self) -> ChannelBufferIter<'a, T> {
        self.outputs.into_iter()
    }

    /// Zip together buffers.
    pub fn zip(self) -> Zip<ChannelBufferIter<'a, T>, ChannelBufferIter<'a, T>> {
        self.inputs.into_iter().zip(self.outputs.into_iter())
    }
}

/// Buffer samples for one channel.
pub struct ChannelBuffer<'a, T: 'a + Float> {
    data: &'a mut [T]
}

impl<'a, T: 'a + Float> ChannelBuffer<'a, T> {
    /// Construct a new `ChannelBuffer` from a slice.
    pub fn new(data: &'a mut [T]) -> ChannelBuffer<'a, T> {
        ChannelBuffer {
            data: data
        }
    }
}

impl<'a, T: 'a + Float> IntoIterator for ChannelBuffer<'a, T> {
    type Item = &'a mut T;
    type IntoIter = slice::IterMut<'a, T>;

    fn into_iter(self) -> slice::IterMut<'a, T> {
        self.data.iter_mut()
    }
}
<commit_msg>Added basic buffer tests.<commit_after>//! Buffers to safely work with audio samples.

use std::iter::{Zip, IntoIterator};
use std::vec::IntoIter;
use std::slice;

use num::Float;

/// A buffer containing `ChannelBuffer` buffers for each input/output.
#[allow(dead_code)]
pub struct AudioBuffer<'a, T: 'a + Float> {
    inputs: Vec<ChannelBuffer<'a, T>>,
    outputs: Vec<ChannelBuffer<'a, T>>
}

/// Iterator over channel buffers for either inputs or outputs.
pub type ChannelBufferIter<'a, T> = IntoIter<ChannelBuffer<'a, T>>;

#[allow(dead_code)]
impl<'a, T: 'a + Float> AudioBuffer<'a, T> {
    /// Create an `AudioBuffer` from vectors of slices. Each vector represents either an input or
    /// output, and contains an array of samples.
    /// Eg if inputs was a vector of size 2 containing slices of size 512, it would hold 2 inputs
    /// where each input holds 512 samples.
    pub fn new(inputs: Vec<&'a mut [T]>, outputs: Vec<&'a mut [T]>) -> AudioBuffer<'a, T> {
        // Creats an input / output for each array in respective Vec
        AudioBuffer {
            inputs: inputs.into_iter().map(|input| ChannelBuffer::new(input)).collect(),
            outputs: outputs.into_iter().map(|output| ChannelBuffer::new(output)).collect()
        }
    }

    /// Create an `AudioBuffer` from raw pointers. Only really useful for interacting with the VST
    /// API.
    pub unsafe fn from_raw(inputs_raw: *mut *mut T, outputs_raw: *mut *mut T, num_inputs: usize, num_outputs: usize, samples: usize) -> AudioBuffer<'a, T> {
        // Allocate an array size of vst input count
        let mut inputs: Vec<&mut [T]> = Vec::with_capacity(num_inputs);
        for i in 0 .. inputs.capacity() {
            // Push samples for each input to `inputs` array
            inputs.push(slice::from_raw_parts_mut(*inputs_raw.offset(i as isize), samples));
        }

        // Allocate an array size of vst output count
        let mut outputs: Vec<&mut [T]> = Vec::with_capacity(num_outputs);
        for i in 0 .. outputs.capacity() {
            // Push samples for each output to `outputs` array
            outputs.push(slice::from_raw_parts_mut(*outputs_raw.offset(i as isize), samples));
        }

        // Call constructor with slices
        AudioBuffer::new(inputs, outputs)
    }

    /// Retreieve input at specified index.
    pub fn input(&'a mut self, index: usize) -> Option<&'a mut ChannelBuffer<'a, T>> {
        self.inputs.get_mut(index)
    }

    /// Retreieve output at specified index.
    pub fn output(&'a mut self, index: usize) -> Option<&'a mut ChannelBuffer<'a, T>> {
        self.outputs.get_mut(index)
    }

    /// Create an iterator over all inputs.
    pub fn inputs(self) -> ChannelBufferIter<'a, T> {
        self.inputs.into_iter()
    }

    /// Create an iterator over all outputs.
    pub fn outputs(self) -> ChannelBufferIter<'a, T> {
        self.outputs.into_iter()
    }

    /// Zip together buffers.
    pub fn zip(self) -> Zip<ChannelBufferIter<'a, T>, ChannelBufferIter<'a, T>> {
        self.inputs.into_iter().zip(self.outputs.into_iter())
    }
}

/// Buffer samples for one channel.
pub struct ChannelBuffer<'a, T: 'a + Float> {
    data: &'a mut [T]
}

impl<'a, T: 'a + Float> ChannelBuffer<'a, T> {
    /// Construct a new `ChannelBuffer` from a slice.
    pub fn new(data: &'a mut [T]) -> ChannelBuffer<'a, T> {
        ChannelBuffer {
            data: data
        }
    }
}

impl<'a, T: 'a + Float> IntoIterator for ChannelBuffer<'a, T> {
    type Item = &'a mut T;
    type IntoIter = slice::IterMut<'a, T>;

    fn into_iter(self) -> slice::IterMut<'a, T> {
        self.data.iter_mut()
    }
}

#[cfg(test)]
mod tests {
    use AudioBuffer;

    #[test]
    fn buffer_input() {
        // This test creates a channel for 2 inputs and 2 outputs. The input channels are simply
        // values from 0 to `SIZE` (e.g. [0, 1, 2, 3, 4, .. , SIZE - 1]) and the output channels
        // are just 0. This test assures that when the buffers are zipped together, the input
        // values do not change.
        const SIZE: usize = 1024;
        let vec: Vec<f32> = (0..SIZE).collect::<Vec<usize>>().iter().map(|&x| x as f32).collect();
        let mut in1 = vec.clone();
        let mut in2 = vec.clone();
        let mut out1 = vec![0.0; SIZE];
        let mut out2 = vec![0.0; SIZE];

        let buffer = AudioBuffer::new(vec![&mut in1, &mut in2],
                                      vec![&mut out1, &mut out2]);

        for (input, output) in buffer.zip() {
            input.into_iter().zip(output.into_iter())

            .fold(0.0, |acc, (input, output)| {
                assert!((*input - acc).abs() <= 0.001);
                assert_eq!(*output, 0.0);
                acc + 1.0
            });
        }
    }
}
<|endoftext|>


<commit_before>// Copyright 2020 Google LLC
//
// Use of this source code is governed by an MIT-style license that can be found
// in the LICENSE file or at https://opensource.org/licenses/MIT.

use byteorder::{LittleEndian, ReadBytesExt, WriteBytesExt};
use fleetspeak_proto::common::{Message, Address};
use fleetspeak_proto::channel::{StartupData};
use lazy_static::lazy_static;
use prost;
use prost_types;
use std::io::{Read, Write, Result};
use std::marker::{Send, Sync};
use std::sync::Mutex;

pub struct Connection<R, W> {
    pub input: R,
    pub output: W,
}

impl<R: Read, W: Write> Connection<R, W> {

    pub fn new(input: R, output: W) -> Result<Self> {
        let mut conn = Connection {
            input: input,
            output: output,
        };
        conn.handshake()?;

        Ok(conn)
    }

    pub fn heartbeat(&mut self) -> Result<()> {
        let msg = Message {
            message_type: "Heartbeat".to_string(),
            destination: Some(Address {
                service_name: "system".to_string(),
                ..Default::default()
            }),
            ..Default::default()
        };

        self.emit(msg)
    }

    pub fn startup(&mut self, version: &str) -> Result<()> {
        let data = StartupData {
            pid: std::process::id() as i64,
            version: version.to_string(),
        };

        let mut buf = Vec::new();
        prost::Message::encode(&data, &mut buf).map_err(invalid_data_error)?;

        let msg = Message {
            message_type: "StartupData".to_string(),
            destination: Some(Address {
                service_name: "system".to_string(),
                ..Default::default()
            }),
            data: Some(prost_types::Any {
                value: buf,
                type_url: "type.googleapis.com/fleetspeak.channel.StartupData".to_string(),
            }),
            ..Default::default()
        };

        self.emit(msg)
    }

    pub fn send<M>(&mut self, service: &str, kind: &str, data: M) -> Result<()>
    where
        M: prost::Message,
    {
        let mut buf = Vec::new();
        prost::Message::encode(&data, &mut buf).map_err(invalid_data_error)?;

        let msg = Message {
            message_type: kind.to_string(),
            destination: Some(Address {
                service_name: service.to_string(),
                ..Default::default()
            }),
            data: Some(prost_types::Any {
                value: buf.to_vec(),
                ..Default::default()
            }),
            ..Default::default()
        };

        self.emit(msg)
    }

    pub fn receive<M>(&mut self) -> Result<M>
    where
        M: prost::Message + Default,
    {
        let msg = self.collect()?;

        // It is not clear what is the best approach here. If there is no data,
        // should we error-out or return a default value? For the time being we
        // stick to the default approach, but if this proves to be not working
        // well in practice, it might be reconsidered.
        let data = match msg.data {
            Some(data) => data,
            None => return Ok(Default::default()),
        };

        prost::Message::decode(&data.value[..]).map_err(invalid_data_error)
    }

    fn emit(&mut self, msg: Message) -> Result<()> {
        let mut buf = Vec::new();
        prost::Message::encode(&msg, &mut buf).map_err(invalid_data_error)?;

        self.output.write_u32::<LittleEndian>(buf.len() as u32)?;
        self.output.write(&buf)?;
        self.output.write_u32::<LittleEndian>(MAGIC)?;
        self.output.flush()?;

        Ok(())
    }

    fn collect(&mut self) -> Result<Message> {
        let len = self.input.read_u32::<LittleEndian>()? as usize;
        let mut buf = vec!(0; len);
        self.input.read_exact(&mut buf[..])?;

        let magic = self.input.read_u32::<LittleEndian>()?;
        if magic != MAGIC {
            let err = invalid_data_error(format!("invalid magic: `{}`", magic));
            return Err(err);
        }

        prost::Message::decode(&buf[..]).map_err(invalid_data_error)
    }

    fn handshake(&mut self) -> Result<()> {
        self.output.write_u32::<LittleEndian>(MAGIC)?;
        self.output.flush()?;

        let magic = self.input.read_u32::<LittleEndian>()?;
        if magic != MAGIC {
            let err = invalid_data_error(format!("invalid magic `{}`", magic));
            return Err(err);
        }

        Ok(())
    }

}

fn invalid_data_error<E>(err: E) -> std::io::Error
where
    E: Into<Box<dyn std::error::Error + Send + Sync>>,
{
    return std::io::Error::new(std::io::ErrorKind::InvalidData, err);
}

const MAGIC: u32 = 0xf1ee1001;

fn open(var: &str) -> std::fs::File {
    let fd = match std::env::var(var) {
        Ok(fd) => fd,
        Err(err) => panic!("invalid variable `{}`: {}", var, err),
    };

    let fd = match fd.parse() {
        Ok(fd) => fd,
        Err(err) => panic!("failed to parse a file descriptor: {}", err),
    };

    // TODO: Add support for Windows.
    unsafe {
        std::os::unix::io::FromRawFd::from_raw_fd(fd)
    }
}

lazy_static! {
    static ref CONNECTION: Mutex<Connection<std::fs::File, std::fs::File>> = {
        let input = open("FLEETSPEAK_COMMS_CHANNEL_INFD");
        let output = open("FLEETSPEAK_COMMS_CHANNEL_OUTFD");

        let conn = Connection::new(input, output).expect("handshake failure");
        Mutex::new(conn)
    };
}

pub fn heartbeat() -> Result<()> {
    connected(|conn| conn.heartbeat())
}

pub fn startup(version: &str) -> Result<()> {
    connected(|conn| conn.startup(version))
}

pub fn send<M>(service: &str, kind: &str, data: M) -> Result<()>
where
    M: prost::Message,
{
    connected(|conn| conn.send(service, kind, data))
}

pub fn receive<M>() -> Result<M>
where
    M: prost::Message + Default,
{
    connected(|conn| conn.receive())
}

fn connected<F, T>(f: F) -> Result<T>
where
    F: FnOnce(&mut Connection<std::fs::File, std::fs::File>) -> Result<T>
{
    let mut conn = CONNECTION.lock().expect("poisoned connection mutex");
    f(&mut conn)
}

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

    #[test]
    fn handshake_good_magic() {
        let mut buf_in = [0; 1024];
        let mut buf_out = [0; 1024];

        let mut cur = Cursor::new(&mut buf_in[..]);
        assert!(cur.write_u32::<LittleEndian>(MAGIC).is_ok());

        let cur_in = Cursor::new(&mut buf_in[..]);
        let cur_out = Cursor::new(&mut buf_out[..]);
        assert!(Connection::new(cur_in, cur_out).is_ok());

        let mut cur = Cursor::new(&mut buf_out[..]);
        assert_eq!(cur.read_u32::<LittleEndian>().unwrap(), MAGIC);
    }

    #[test]
    fn handshake_bad_magic() {
        let mut buf_in = [0; 1024];
        let mut buf_out = [0; 1024];

        let mut cur = Cursor::new(&mut buf_in[..]);
        assert!(cur.write_u32::<LittleEndian>(0xf1ee1337).is_ok());

        let cur_in = Cursor::new(&mut buf_in[..]);
        let cur_out = Cursor::new(&mut buf_out[..]);
        assert!(Connection::new(cur_in, cur_out).is_err());
    }
}
<commit_msg>Refactor magic-handling code to separate methods.<commit_after>// Copyright 2020 Google LLC
//
// Use of this source code is governed by an MIT-style license that can be found
// in the LICENSE file or at https://opensource.org/licenses/MIT.

use byteorder::{LittleEndian, ReadBytesExt, WriteBytesExt};
use fleetspeak_proto::common::{Message, Address};
use fleetspeak_proto::channel::{StartupData};
use lazy_static::lazy_static;
use prost;
use prost_types;
use std::io::{Read, Write, Result};
use std::marker::{Send, Sync};
use std::sync::Mutex;

pub struct Connection<R, W> {
    pub input: R,
    pub output: W,
}

impl<R: Read, W: Write> Connection<R, W> {

    pub fn new(input: R, output: W) -> Result<Self> {
        let mut conn = Connection {
            input: input,
            output: output,
        };
        conn.handshake()?;

        Ok(conn)
    }

    pub fn heartbeat(&mut self) -> Result<()> {
        let msg = Message {
            message_type: "Heartbeat".to_string(),
            destination: Some(Address {
                service_name: "system".to_string(),
                ..Default::default()
            }),
            ..Default::default()
        };

        self.emit(msg)
    }

    pub fn startup(&mut self, version: &str) -> Result<()> {
        let data = StartupData {
            pid: std::process::id() as i64,
            version: version.to_string(),
        };

        let mut buf = Vec::new();
        prost::Message::encode(&data, &mut buf).map_err(invalid_data_error)?;

        let msg = Message {
            message_type: "StartupData".to_string(),
            destination: Some(Address {
                service_name: "system".to_string(),
                ..Default::default()
            }),
            data: Some(prost_types::Any {
                value: buf,
                type_url: "type.googleapis.com/fleetspeak.channel.StartupData".to_string(),
            }),
            ..Default::default()
        };

        self.emit(msg)
    }

    pub fn send<M>(&mut self, service: &str, kind: &str, data: M) -> Result<()>
    where
        M: prost::Message,
    {
        let mut buf = Vec::new();
        prost::Message::encode(&data, &mut buf).map_err(invalid_data_error)?;

        let msg = Message {
            message_type: kind.to_string(),
            destination: Some(Address {
                service_name: service.to_string(),
                ..Default::default()
            }),
            data: Some(prost_types::Any {
                value: buf.to_vec(),
                ..Default::default()
            }),
            ..Default::default()
        };

        self.emit(msg)
    }

    pub fn receive<M>(&mut self) -> Result<M>
    where
        M: prost::Message + Default,
    {
        let msg = self.collect()?;

        // It is not clear what is the best approach here. If there is no data,
        // should we error-out or return a default value? For the time being we
        // stick to the default approach, but if this proves to be not working
        // well in practice, it might be reconsidered.
        let data = match msg.data {
            Some(data) => data,
            None => return Ok(Default::default()),
        };

        prost::Message::decode(&data.value[..]).map_err(invalid_data_error)
    }

    fn emit(&mut self, msg: Message) -> Result<()> {
        let mut buf = Vec::new();
        prost::Message::encode(&msg, &mut buf).map_err(invalid_data_error)?;

        self.output.write_u32::<LittleEndian>(buf.len() as u32)?;
        self.output.write(&buf)?;
        self.write_magic()?;
        self.output.flush()?;

        Ok(())
    }

    fn collect(&mut self) -> Result<Message> {
        let len = self.input.read_u32::<LittleEndian>()? as usize;
        let mut buf = vec!(0; len);
        self.input.read_exact(&mut buf[..])?;
        self.read_magic()?;

        prost::Message::decode(&buf[..]).map_err(invalid_data_error)
    }

    fn handshake(&mut self) -> Result<()> {
        self.write_magic()?;
        self.output.flush()?;
        self.read_magic()?;

        Ok(())
    }

    fn write_magic(&mut self) -> Result<()> {
        self.output.write_u32::<LittleEndian>(MAGIC)?;

        Ok(())
    }

    fn read_magic(&mut self) -> Result<()> {
        let magic = self.input.read_u32::<LittleEndian>()?;
        if magic != MAGIC {
            let err = invalid_data_error(format!("invalid magic `{}`", magic));
            return Err(err);
        }

        Ok(())
    }
}

fn invalid_data_error<E>(err: E) -> std::io::Error
where
    E: Into<Box<dyn std::error::Error + Send + Sync>>,
{
    return std::io::Error::new(std::io::ErrorKind::InvalidData, err);
}

const MAGIC: u32 = 0xf1ee1001;

fn open(var: &str) -> std::fs::File {
    let fd = match std::env::var(var) {
        Ok(fd) => fd,
        Err(err) => panic!("invalid variable `{}`: {}", var, err),
    };

    let fd = match fd.parse() {
        Ok(fd) => fd,
        Err(err) => panic!("failed to parse a file descriptor: {}", err),
    };

    // TODO: Add support for Windows.
    unsafe {
        std::os::unix::io::FromRawFd::from_raw_fd(fd)
    }
}

lazy_static! {
    static ref CONNECTION: Mutex<Connection<std::fs::File, std::fs::File>> = {
        let input = open("FLEETSPEAK_COMMS_CHANNEL_INFD");
        let output = open("FLEETSPEAK_COMMS_CHANNEL_OUTFD");

        let conn = Connection::new(input, output).expect("handshake failure");
        Mutex::new(conn)
    };
}

pub fn heartbeat() -> Result<()> {
    connected(|conn| conn.heartbeat())
}

pub fn startup(version: &str) -> Result<()> {
    connected(|conn| conn.startup(version))
}

pub fn send<M>(service: &str, kind: &str, data: M) -> Result<()>
where
    M: prost::Message,
{
    connected(|conn| conn.send(service, kind, data))
}

pub fn receive<M>() -> Result<M>
where
    M: prost::Message + Default,
{
    connected(|conn| conn.receive())
}

fn connected<F, T>(f: F) -> Result<T>
where
    F: FnOnce(&mut Connection<std::fs::File, std::fs::File>) -> Result<T>
{
    let mut conn = CONNECTION.lock().expect("poisoned connection mutex");
    f(&mut conn)
}

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

    #[test]
    fn handshake_good_magic() {
        let mut buf_in = [0; 1024];
        let mut buf_out = [0; 1024];

        let mut cur = Cursor::new(&mut buf_in[..]);
        assert!(cur.write_u32::<LittleEndian>(MAGIC).is_ok());

        let cur_in = Cursor::new(&mut buf_in[..]);
        let cur_out = Cursor::new(&mut buf_out[..]);
        assert!(Connection::new(cur_in, cur_out).is_ok());

        let mut cur = Cursor::new(&mut buf_out[..]);
        assert_eq!(cur.read_u32::<LittleEndian>().unwrap(), MAGIC);
    }

    #[test]
    fn handshake_bad_magic() {
        let mut buf_in = [0; 1024];
        let mut buf_out = [0; 1024];

        let mut cur = Cursor::new(&mut buf_in[..]);
        assert!(cur.write_u32::<LittleEndian>(0xf1ee1337).is_ok());

        let cur_in = Cursor::new(&mut buf_in[..]);
        let cur_out = Cursor::new(&mut buf_out[..]);
        assert!(Connection::new(cur_in, cur_out).is_err());
    }
}
<|endoftext|>


<commit_before>use std::fmt;
use std::str;
use std::sync::Arc;
use std::time::Instant;

use futures::Future;
use futures::future::{
    self,
    Loop,
};
use parking_lot::Mutex;
use tokio_timer::Delay;

use pb::master::{
    DeleteTableRequestPb,
    IsAlterTableDoneRequestPb,
    IsCreateTableDoneRequestPb,
    ListTablesRequestPb,
    MasterService,
    TableIdentifierPb,
};
use pb::ExpectField;

use Error;
use IntoMasterAddrs;
use MasterInfo;
use Options;
use TableId;
use TabletServerInfo;
use backoff::Backoff;
use meta_cache::MetaCache;
use table::AlterTableBuilder;
use table::Table;
use table::TableBuilder;

/// A Kudu database client.
///
/// Encapsulates the connection to a Kudu cluster. Only a single instance should be used per
/// application per cluster.
#[derive(Clone)]
pub struct Client {
    meta_cache: MetaCache,
    latest_observed_timestamp: Arc<Mutex<u64>>, // Replace with AtomicU64 when stable.
}

impl Client {

    /// Creates a new client with the provided configuration.
    pub fn new<Addrs>(master_addresses: Addrs, options: Options) -> impl Future<Item=Client, Error=Error>
    where Addrs: IntoMasterAddrs {
        future::result(master_addresses.into_master_addrs())
               .and_then(|master_addresses| MetaCache::new(master_addresses, options))
               .map(move |meta_cache| Client {
                   meta_cache,
                   latest_observed_timestamp: Arc::new(Mutex::new(0)),
               })
    }

    /// Creates a new Kudu table with the schema and options specified by `builder`. Returns the
    /// new table's ID, or an error on failure.
    pub fn create_table(&mut self, builder: TableBuilder) -> impl Future<Item=TableId, Error=Error> {
        let mut client = self.clone();
        future::result(
            builder.into_pb()
                   .map(|request| {
                       self.meta_cache.master_rpc(
                           MasterService::create_table(Arc::new(request), self.deadline()))
                   }))
            .flatten()
            .and_then(|resp| TableId::parse_bytes(&resp.table_id.expect_field("CreateTableResponsePb", "table_id")?))
            .and_then(move |table_id| {
                client.wait_for_table_creation(table_id)
                      .map(move |_| table_id)
            })
    }

    /// Returns a future which completes when the table is created.
    ///
    /// Not on timeout: this method will not timeout if the master is reachable and responsive.
    fn wait_for_table_creation(&mut self, table: TableId) -> impl Future<Item=(), Error=Error> {
        struct State {
            client: Client,
            table: TableId,
            backoff: Backoff,
        }

        let state = State {
            client: self.clone(),
            table,
            backoff: Backoff::with_duration_range(32, 2048),
        };

        future::loop_fn(state, |mut state: State| {
            Delay::new(Instant::now() + state.backoff.next_backoff())
                 .map_err(|error| -> Error { panic!("timer failed: {}", error); })
                 .and_then(move |_| {
                    let call = MasterService::is_create_table_done(
                        Arc::new(IsCreateTableDoneRequestPb { table: state.table.into() }),
                        state.client.deadline());

                    state.client
                         .meta_cache
                         .master_rpc(call)
                         .map(move |resp| if resp.done() { Loop::Break(()) } else { Loop::Continue(state) })
                 })
        })
    }

    /// Deletes the table.
    pub fn delete_table<S>(&mut self, table: S) -> impl Future<Item=(), Error=Error>
    where S: Into<String> {
        self.do_delete_table(table.into().into())
    }

    /// Deletes the table.
    pub fn delete_table_by_id(&mut self, id: TableId) -> impl Future<Item=(), Error=Error> {
        self.do_delete_table(id.into())
    }

    fn do_delete_table(&mut self, table: TableIdentifierPb) -> impl Future<Item=(), Error=Error>{
        let call = MasterService::delete_table(Arc::new(DeleteTableRequestPb { table }),
                                               self.deadline());
        self.meta_cache.master_rpc(call).map(|_| ())
    }

    pub fn alter_table<S>(&mut self, table: S, alter: AlterTableBuilder) -> impl Future<Item=TableId, Error=Error>
    where S: Into<String> {
        self.do_alter_table(table.into().into(), alter)
    }

    pub fn alter_table_by_id(&mut self, id: TableId, alter: AlterTableBuilder) -> impl Future<Item=(), Error=Error> {
        self.do_alter_table(id.into(), alter).map(|_| ())
    }

    pub fn do_alter_table(&mut self,
                          identifier: TableIdentifierPb,
                          alter: AlterTableBuilder)
                          -> impl Future<Item=TableId, Error=Error> {
        let AlterTableBuilder { result, mut pb, schema } = alter;
        let client = self.clone();
        future::result(
            result.map(move |_| {
                pb.table = identifier;
                let call = MasterService::alter_table(Arc::new(pb), self.deadline());
                self.meta_cache.master_rpc(call)
            }))
            .flatten()
            .and_then(move |resp| {
                let table_id = str::from_utf8(resp.table_id())
                                .map_err(|error| Error::Serialization(format!("{}", error)))
                                .and_then(TableId::parse)?;

                // If the table partitioning was altered and there is an existing meta cache for the
                // table, clear it.
                if schema.is_some() {
                    // TODO
                    // client.meta_cache.clear_table_locations(table_id);
                }

                Ok((table_id, client))

            })
            .and_then(|(table_id, mut client): (TableId, Client)| {
                client.wait_for_table_alteration(table_id)
                      .map(move |_| table_id)
            })
    }

    /// Returns a future which completes when the table is altered.
    ///
    /// Not on timeout: this method will not timeout if the master is reachable and responsive.
    fn wait_for_table_alteration(&mut self, table: TableId) -> impl Future<Item=(), Error=Error> {
        struct State {
            client: Client,
            table: TableId,
            backoff: Backoff,
        }

        let state = State {
            client: self.clone(),
            table,
            backoff: Backoff::with_duration_range(32, 2048),
        };

        future::loop_fn(state, |mut state: State| {
            Delay::new(Instant::now() + state.backoff.next_backoff())
                 .map_err(|error| -> Error { panic!("timer failed: {}", error); })
                 .and_then(move |_| {

                    let call = MasterService::is_alter_table_done(
                        Arc::new(IsAlterTableDoneRequestPb { table: state.table.into() }),
                        state.client.deadline());

                    state.client
                         .meta_cache
                         .master_rpc(call)
                         .map(move |resp| if resp.done() { Loop::Break(()) } else { Loop::Continue(state) })
                    })
        })
    }

    /// Lists all tables and their associated table ID.
    pub fn tables(&mut self) -> impl Future<Item=Vec<(String, TableId)>, Error=Error> {
        self.do_list_tables(Default::default())
    }

    /// Lists all tables with the a name matching the provided prefix, and their associated table ID.
    pub fn tables_with_prefix<S>(&mut self, name_prefix: S) -> impl Future<Item=Vec<(String, TableId)>, Error=Error>
    where S: Into<String> {
        self.do_list_tables(Arc::new(ListTablesRequestPb { name_filter: Some(name_prefix.into()) }))
    }

    fn do_list_tables(&mut self, request: Arc<ListTablesRequestPb>) -> impl Future<Item=Vec<(String, TableId)>, Error=Error> {
        let call = MasterService::list_tables(request, self.deadline());

        self.meta_cache.master_rpc(call).and_then(|resp| {
            let mut tables = Vec::with_capacity(resp.tables.len());
            for table in resp.tables {
                tables.push((table.name, TableId::parse_bytes(&table.id)?));
            }
            Ok(tables)
        })
    }

    pub fn masters(&mut self) -> impl Future<Item=Vec<MasterInfo>, Error=Error> {
        let call = MasterService::list_masters(Default::default(), self.deadline());

        self.meta_cache.master_rpc(call).and_then(|resp| {
            let mut servers = Vec::with_capacity(resp.masters.len());
            for server in resp.masters {
                servers.push(MasterInfo::from_pb(server)?);
            }
            Ok(servers)
        })
    }

    pub fn tablet_servers(&mut self) -> impl Future<Item=Vec<TabletServerInfo>, Error=Error> {
        let call = MasterService::list_tablet_servers(Default::default(), self.deadline());

        self.meta_cache.master_rpc(call).and_then(|resp| {
            let mut servers = Vec::with_capacity(resp.servers.len());
            for server in resp.servers {
                servers.push(TabletServerInfo::from_pb(server)?);
            }
            Ok(servers)
        })
    }

    /// Returns an open table.
    pub fn open_table<S>(&mut self, table: S) -> impl Future<Item=Table, Error=Error> + 'static
    where S: Into<String> {
        self.meta_cache.open_table(TableIdentifierPb::from(table.into()), self.deadline())
    }

    /// Returns an open table.
    pub fn open_table_by_id(&mut self, id: TableId) -> impl Future<Item=Table, Error=Error> + 'static {
        self.meta_cache.open_table(id.into(), self.deadline())
    }

    pub fn latest_observed_timestamp(&self) -> u64 {
        *self.latest_observed_timestamp.lock()
    }

    pub fn observe_timestamp(&self, timestamp: u64) {
        let mut latest = self.latest_observed_timestamp.lock();
        if timestamp > *latest {
            *latest = timestamp;
        }
    }

    fn deadline(&self) -> Instant {
        Instant::now() + self.meta_cache.options().admin_timeout
    }
}

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

#[cfg(test)]
mod tests {

    use AlterTableBuilder;
    use Column;
    use DataType;
    use RangePartitionBound;
    use SchemaBuilder;
    use TableBuilder;
    use mini_cluster::{MiniCluster, MiniClusterConfig};
    use schema::tests::simple_schema;
    use super::*;

    use env_logger;
    use tokio::runtime::current_thread::Runtime;

    #[test]
    fn table_lifecycle() {
        let _ = env_logger::try_init();
        let mut cluster = MiniCluster::default();
        let mut runtime = Runtime::new().unwrap();

        let mut client = runtime.block_on(Client::new(cluster.master_addrs(), Options::default()))
                                .expect("client");

        let schema = SchemaBuilder::new()
            .add_column(Column::builder("key", DataType::Int32).set_not_null())
            .add_column(Column::builder("val", DataType::Int32))
            .set_primary_key(vec!["key"])
            .build()
            .unwrap();

        let mut table_builder = TableBuilder::new("table_lifecycle", schema.clone());
        table_builder.add_hash_partitions(vec!["key"], 4);
        table_builder.set_num_replicas(1);

        let table_id = runtime.block_on(client.create_table(table_builder)).expect("create_table");
        let mut alter_builder = AlterTableBuilder::new();
        alter_builder.rename_table("table_lifecycle_renamed");
        runtime.block_on(client.alter_table_by_id(table_id, alter_builder)).expect("alter_table_by_id");

        let table = runtime.block_on(client.open_table("table_lifecycle_renamed")).expect("open_table");
        assert_eq!(table_id, table.id());

        let tables = runtime.block_on(client.tables()).expect("tables");
        assert_eq!(vec![("table_lifecycle_renamed".to_string(), table_id)], tables);

        runtime.block_on(client.delete_table_by_id(table_id)).expect("delete_table");
        assert!(runtime.block_on(client.tables()).expect("tables").is_empty());
    }

    #[test]
    fn list_servers() {
        let _ = env_logger::try_init();
        let mut cluster = MiniCluster::new(MiniClusterConfig::default()
                                                             .num_masters(3)
                                                             .num_tservers(3));
        let mut runtime = Runtime::new().unwrap();

        let mut client = runtime.block_on(Client::new(cluster.master_addrs(), Options::default()))
                                .expect("client");

        let tablet_servers = runtime.block_on(client.tablet_servers()).expect("list_table_servers");
        assert_eq!(3, tablet_servers.len());

        let masters = runtime.block_on(client.masters()).expect("masters");
        assert_eq!(3, masters.len());
    }

    #[test]
    fn alter_table() {
        let _ = env_logger::try_init();
        let mut cluster = MiniCluster::new(MiniClusterConfig::default()
                                                             .num_masters(1)
                                                             .num_tservers(1));
        let mut runtime = Runtime::new().unwrap();

        let mut client = runtime.block_on(Client::new(cluster.master_addrs(), Options::default()))
                                .expect("client");

        let mut table_builder = TableBuilder::new("t", simple_schema());
        table_builder.set_num_replicas(1);
        table_builder.set_range_partition_columns(vec!["key"]);
        let table_id = runtime.block_on(client.create_table(table_builder)).expect("create_table");

        let mut alter_builder = AlterTableBuilder::new();
        alter_builder.add_column(Column::builder("c0", DataType::Int32));

        let _ = runtime.block_on(client.alter_table("t", alter_builder)).expect("add column");

        let schema = runtime.block_on(client.open_table("t")).expect("open_table").schema().clone();
        assert_eq!(3, schema.columns().len());

        let mut alter_builder = AlterTableBuilder::new();
        alter_builder.drop_range_partition(&RangePartitionBound::Inclusive(schema.new_row()),
                                           &RangePartitionBound::Exclusive(schema.new_row()));

        runtime.block_on(client.alter_table_by_id(table_id, alter_builder)).expect("drop range partition");

        let mut lower_bound = schema.new_row();
        let mut upper_bound = schema.new_row();

        lower_bound.set_by_name("key", "a").unwrap();
        upper_bound.set_by_name("key", "z").unwrap();

        let mut alter_builder = AlterTableBuilder::new();
        alter_builder.add_range_partition(&RangePartitionBound::Inclusive(lower_bound),
                                          &RangePartitionBound::Exclusive(upper_bound))
                     .rename_table("u")
                     .drop_column("c0");
        runtime.block_on(client.alter_table_by_id(table_id, alter_builder)).unwrap();

        let schema = runtime.block_on(client.open_table("u")).unwrap().schema().clone();
        assert_eq!(2, schema.columns().len());
    }
}
<commit_msg>Revert "open_table should be 'static"<commit_after>use std::fmt;
use std::str;
use std::sync::Arc;
use std::time::Instant;

use futures::Future;
use futures::future::{
    self,
    Loop,
};
use parking_lot::Mutex;
use tokio_timer::Delay;

use pb::master::{
    DeleteTableRequestPb,
    IsAlterTableDoneRequestPb,
    IsCreateTableDoneRequestPb,
    ListTablesRequestPb,
    MasterService,
    TableIdentifierPb,
};
use pb::ExpectField;

use Error;
use IntoMasterAddrs;
use MasterInfo;
use Options;
use TableId;
use TabletServerInfo;
use backoff::Backoff;
use meta_cache::MetaCache;
use table::AlterTableBuilder;
use table::Table;
use table::TableBuilder;

/// A Kudu database client.
///
/// Encapsulates the connection to a Kudu cluster. Only a single instance should be used per
/// application per cluster.
#[derive(Clone)]
pub struct Client {
    meta_cache: MetaCache,
    latest_observed_timestamp: Arc<Mutex<u64>>, // Replace with AtomicU64 when stable.
}

impl Client {

    /// Creates a new client with the provided configuration.
    pub fn new<Addrs>(master_addresses: Addrs, options: Options) -> impl Future<Item=Client, Error=Error>
    where Addrs: IntoMasterAddrs {
        future::result(master_addresses.into_master_addrs())
               .and_then(|master_addresses| MetaCache::new(master_addresses, options))
               .map(move |meta_cache| Client {
                   meta_cache,
                   latest_observed_timestamp: Arc::new(Mutex::new(0)),
               })
    }

    /// Creates a new Kudu table with the schema and options specified by `builder`. Returns the
    /// new table's ID, or an error on failure.
    pub fn create_table(&mut self, builder: TableBuilder) -> impl Future<Item=TableId, Error=Error> {
        let mut client = self.clone();
        future::result(
            builder.into_pb()
                   .map(|request| {
                       self.meta_cache.master_rpc(
                           MasterService::create_table(Arc::new(request), self.deadline()))
                   }))
            .flatten()
            .and_then(|resp| TableId::parse_bytes(&resp.table_id.expect_field("CreateTableResponsePb", "table_id")?))
            .and_then(move |table_id| {
                client.wait_for_table_creation(table_id)
                      .map(move |_| table_id)
            })
    }

    /// Returns a future which completes when the table is created.
    ///
    /// Not on timeout: this method will not timeout if the master is reachable and responsive.
    fn wait_for_table_creation(&mut self, table: TableId) -> impl Future<Item=(), Error=Error> {
        struct State {
            client: Client,
            table: TableId,
            backoff: Backoff,
        }

        let state = State {
            client: self.clone(),
            table,
            backoff: Backoff::with_duration_range(32, 2048),
        };

        future::loop_fn(state, |mut state: State| {
            Delay::new(Instant::now() + state.backoff.next_backoff())
                 .map_err(|error| -> Error { panic!("timer failed: {}", error); })
                 .and_then(move |_| {
                    let call = MasterService::is_create_table_done(
                        Arc::new(IsCreateTableDoneRequestPb { table: state.table.into() }),
                        state.client.deadline());

                    state.client
                         .meta_cache
                         .master_rpc(call)
                         .map(move |resp| if resp.done() { Loop::Break(()) } else { Loop::Continue(state) })
                 })
        })
    }

    /// Deletes the table.
    pub fn delete_table<S>(&mut self, table: S) -> impl Future<Item=(), Error=Error>
    where S: Into<String> {
        self.do_delete_table(table.into().into())
    }

    /// Deletes the table.
    pub fn delete_table_by_id(&mut self, id: TableId) -> impl Future<Item=(), Error=Error> {
        self.do_delete_table(id.into())
    }

    fn do_delete_table(&mut self, table: TableIdentifierPb) -> impl Future<Item=(), Error=Error>{
        let call = MasterService::delete_table(Arc::new(DeleteTableRequestPb { table }),
                                               self.deadline());
        self.meta_cache.master_rpc(call).map(|_| ())
    }

    pub fn alter_table<S>(&mut self, table: S, alter: AlterTableBuilder) -> impl Future<Item=TableId, Error=Error>
    where S: Into<String> {
        self.do_alter_table(table.into().into(), alter)
    }

    pub fn alter_table_by_id(&mut self, id: TableId, alter: AlterTableBuilder) -> impl Future<Item=(), Error=Error> {
        self.do_alter_table(id.into(), alter).map(|_| ())
    }

    pub fn do_alter_table(&mut self,
                          identifier: TableIdentifierPb,
                          alter: AlterTableBuilder)
                          -> impl Future<Item=TableId, Error=Error> {
        let AlterTableBuilder { result, mut pb, schema } = alter;
        let client = self.clone();
        future::result(
            result.map(move |_| {
                pb.table = identifier;
                let call = MasterService::alter_table(Arc::new(pb), self.deadline());
                self.meta_cache.master_rpc(call)
            }))
            .flatten()
            .and_then(move |resp| {
                let table_id = str::from_utf8(resp.table_id())
                                .map_err(|error| Error::Serialization(format!("{}", error)))
                                .and_then(TableId::parse)?;

                // If the table partitioning was altered and there is an existing meta cache for the
                // table, clear it.
                if schema.is_some() {
                    // TODO
                    // client.meta_cache.clear_table_locations(table_id);
                }

                Ok((table_id, client))

            })
            .and_then(|(table_id, mut client): (TableId, Client)| {
                client.wait_for_table_alteration(table_id)
                      .map(move |_| table_id)
            })
    }

    /// Returns a future which completes when the table is altered.
    ///
    /// Not on timeout: this method will not timeout if the master is reachable and responsive.
    fn wait_for_table_alteration(&mut self, table: TableId) -> impl Future<Item=(), Error=Error> {
        struct State {
            client: Client,
            table: TableId,
            backoff: Backoff,
        }

        let state = State {
            client: self.clone(),
            table,
            backoff: Backoff::with_duration_range(32, 2048),
        };

        future::loop_fn(state, |mut state: State| {
            Delay::new(Instant::now() + state.backoff.next_backoff())
                 .map_err(|error| -> Error { panic!("timer failed: {}", error); })
                 .and_then(move |_| {

                    let call = MasterService::is_alter_table_done(
                        Arc::new(IsAlterTableDoneRequestPb { table: state.table.into() }),
                        state.client.deadline());

                    state.client
                         .meta_cache
                         .master_rpc(call)
                         .map(move |resp| if resp.done() { Loop::Break(()) } else { Loop::Continue(state) })
                    })
        })
    }

    /// Lists all tables and their associated table ID.
    pub fn tables(&mut self) -> impl Future<Item=Vec<(String, TableId)>, Error=Error> {
        self.do_list_tables(Default::default())
    }

    /// Lists all tables with the a name matching the provided prefix, and their associated table ID.
    pub fn tables_with_prefix<S>(&mut self, name_prefix: S) -> impl Future<Item=Vec<(String, TableId)>, Error=Error>
    where S: Into<String> {
        self.do_list_tables(Arc::new(ListTablesRequestPb { name_filter: Some(name_prefix.into()) }))
    }

    fn do_list_tables(&mut self, request: Arc<ListTablesRequestPb>) -> impl Future<Item=Vec<(String, TableId)>, Error=Error> {
        let call = MasterService::list_tables(request, self.deadline());

        self.meta_cache.master_rpc(call).and_then(|resp| {
            let mut tables = Vec::with_capacity(resp.tables.len());
            for table in resp.tables {
                tables.push((table.name, TableId::parse_bytes(&table.id)?));
            }
            Ok(tables)
        })
    }

    pub fn masters(&mut self) -> impl Future<Item=Vec<MasterInfo>, Error=Error> {
        let call = MasterService::list_masters(Default::default(), self.deadline());

        self.meta_cache.master_rpc(call).and_then(|resp| {
            let mut servers = Vec::with_capacity(resp.masters.len());
            for server in resp.masters {
                servers.push(MasterInfo::from_pb(server)?);
            }
            Ok(servers)
        })
    }

    pub fn tablet_servers(&mut self) -> impl Future<Item=Vec<TabletServerInfo>, Error=Error> {
        let call = MasterService::list_tablet_servers(Default::default(), self.deadline());

        self.meta_cache.master_rpc(call).and_then(|resp| {
            let mut servers = Vec::with_capacity(resp.servers.len());
            for server in resp.servers {
                servers.push(TabletServerInfo::from_pb(server)?);
            }
            Ok(servers)
        })
    }

    /// Returns an open table.
    pub fn open_table<S>(&mut self, table: S) -> impl Future<Item=Table, Error=Error>
    where S: Into<String> {
        self.meta_cache.open_table(TableIdentifierPb::from(table.into()), self.deadline())
    }

    /// Returns an open table.
    pub fn open_table_by_id(&mut self, id: TableId) -> impl Future<Item=Table, Error=Error> {
        self.meta_cache.open_table(id.into(), self.deadline())
    }

    pub fn latest_observed_timestamp(&self) -> u64 {
        *self.latest_observed_timestamp.lock()
    }

    pub fn observe_timestamp(&self, timestamp: u64) {
        let mut latest = self.latest_observed_timestamp.lock();
        if timestamp > *latest {
            *latest = timestamp;
        }
    }

    fn deadline(&self) -> Instant {
        Instant::now() + self.meta_cache.options().admin_timeout
    }
}

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

#[cfg(test)]
mod tests {

    use AlterTableBuilder;
    use Column;
    use DataType;
    use RangePartitionBound;
    use SchemaBuilder;
    use TableBuilder;
    use mini_cluster::{MiniCluster, MiniClusterConfig};
    use schema::tests::simple_schema;
    use super::*;

    use env_logger;
    use tokio::runtime::current_thread::Runtime;

    #[test]
    fn table_lifecycle() {
        let _ = env_logger::try_init();
        let mut cluster = MiniCluster::default();
        let mut runtime = Runtime::new().unwrap();

        let mut client = runtime.block_on(Client::new(cluster.master_addrs(), Options::default()))
                                .expect("client");

        let schema = SchemaBuilder::new()
            .add_column(Column::builder("key", DataType::Int32).set_not_null())
            .add_column(Column::builder("val", DataType::Int32))
            .set_primary_key(vec!["key"])
            .build()
            .unwrap();

        let mut table_builder = TableBuilder::new("table_lifecycle", schema.clone());
        table_builder.add_hash_partitions(vec!["key"], 4);
        table_builder.set_num_replicas(1);

        let table_id = runtime.block_on(client.create_table(table_builder)).expect("create_table");
        let mut alter_builder = AlterTableBuilder::new();
        alter_builder.rename_table("table_lifecycle_renamed");
        runtime.block_on(client.alter_table_by_id(table_id, alter_builder)).expect("alter_table_by_id");

        let table = runtime.block_on(client.open_table("table_lifecycle_renamed")).expect("open_table");
        assert_eq!(table_id, table.id());

        let tables = runtime.block_on(client.tables()).expect("tables");
        assert_eq!(vec![("table_lifecycle_renamed".to_string(), table_id)], tables);

        runtime.block_on(client.delete_table_by_id(table_id)).expect("delete_table");
        assert!(runtime.block_on(client.tables()).expect("tables").is_empty());
    }

    #[test]
    fn list_servers() {
        let _ = env_logger::try_init();
        let mut cluster = MiniCluster::new(MiniClusterConfig::default()
                                                             .num_masters(3)
                                                             .num_tservers(3));
        let mut runtime = Runtime::new().unwrap();

        let mut client = runtime.block_on(Client::new(cluster.master_addrs(), Options::default()))
                                .expect("client");

        let tablet_servers = runtime.block_on(client.tablet_servers()).expect("list_table_servers");
        assert_eq!(3, tablet_servers.len());

        let masters = runtime.block_on(client.masters()).expect("masters");
        assert_eq!(3, masters.len());
    }

    #[test]
    fn alter_table() {
        let _ = env_logger::try_init();
        let mut cluster = MiniCluster::new(MiniClusterConfig::default()
                                                             .num_masters(1)
                                                             .num_tservers(1));
        let mut runtime = Runtime::new().unwrap();

        let mut client = runtime.block_on(Client::new(cluster.master_addrs(), Options::default()))
                                .expect("client");

        let mut table_builder = TableBuilder::new("t", simple_schema());
        table_builder.set_num_replicas(1);
        table_builder.set_range_partition_columns(vec!["key"]);
        let table_id = runtime.block_on(client.create_table(table_builder)).expect("create_table");

        let mut alter_builder = AlterTableBuilder::new();
        alter_builder.add_column(Column::builder("c0", DataType::Int32));

        let _ = runtime.block_on(client.alter_table("t", alter_builder)).expect("add column");

        let schema = runtime.block_on(client.open_table("t")).expect("open_table").schema().clone();
        assert_eq!(3, schema.columns().len());

        let mut alter_builder = AlterTableBuilder::new();
        alter_builder.drop_range_partition(&RangePartitionBound::Inclusive(schema.new_row()),
                                           &RangePartitionBound::Exclusive(schema.new_row()));

        runtime.block_on(client.alter_table_by_id(table_id, alter_builder)).expect("drop range partition");

        let mut lower_bound = schema.new_row();
        let mut upper_bound = schema.new_row();

        lower_bound.set_by_name("key", "a").unwrap();
        upper_bound.set_by_name("key", "z").unwrap();

        let mut alter_builder = AlterTableBuilder::new();
        alter_builder.add_range_partition(&RangePartitionBound::Inclusive(lower_bound),
                                          &RangePartitionBound::Exclusive(upper_bound))
                     .rename_table("u")
                     .drop_column("c0");
        runtime.block_on(client.alter_table_by_id(table_id, alter_builder)).unwrap();

        let schema = runtime.block_on(client.open_table("u")).unwrap().schema().clone();
        assert_eq!(2, schema.columns().len());
    }
}
<|endoftext|>


<commit_before>use std::io::Read;

use hyper::Client as HyperClient;
use hyper::HttpError;
use hyper::client::Response as HyperResponse;
use hyper::header::ContentType;
use hyper::status::StatusCode;
use rustc_serialize::json;
use url::{ParseError, Url};
use url::form_urlencoded::serialize_owned;

#[derive(Debug)]
#[derive(RustcDecodable)]
pub struct Response {
    pub action: String,
    pub node: Node,
    pub prev_node: Option<Node>,
}

#[derive(Debug)]
#[derive(RustcDecodable)]
pub struct Node {
    pub created_index: Option<u64>,
    pub dir: Option<bool>,
    pub expiration: Option<String>,
    pub key: Option<String>,
    pub modified_index: Option<u64>,
    pub nodes: Option<Vec<Node>>,
    pub ttl: Option<i64>,
    pub value: Option<String>,
}

#[derive(Debug)]
pub enum Error {
    EtcdError(EtcdError),
    HttpError(HttpError),
}

#[derive(Debug)]
#[derive(RustcDecodable)]
#[allow(non_snake_case)]
pub struct EtcdError {
    pub cause: Option<String>,
    pub errorCode: u64,
    pub index: u64,
    pub message: String,
}

#[derive(Debug)]
pub struct Client {
    root_url: String,
}

impl Client {
    pub fn new(root_url: &str) -> Result<Client, ParseError> {
        let url = try!(Url::parse(root_url));
        let client = Client {
            root_url: format!("{}", url.serialize()),
        };

        Ok(client)
    }

    pub fn default() -> Client {
        Client {
            root_url: "http://127.0.0.1:2379/".to_string(),
        }
    }

    pub fn get(&self, key: &str) -> Result<Response, EtcdError> {
        Err(EtcdError)
    }

    pub fn mk(&self, key: &str, value: &str, ttl: Option<u64>) -> Result<Response, Error> {
        let url = self.build_url(key);
        let mut options = vec![];

        options.push(("value".to_string(), value.to_string()));
        options.push(("prevExist".to_string(), "false".to_string()));

        if ttl.is_some() {
            options.push(("ttl".to_string(), format!("{}", ttl.unwrap())));
        }

        let body = serialize_owned(&options);

        match self.put(url, body) {
            Ok(mut response) => {
                let mut response_body = String::new();

                response.read_to_string(&mut response_body).unwrap();

                println!("{:?} - {:?} - {:?}", response.status, response.headers, response_body);

                match response.status {
                    StatusCode::Created => Ok(json::decode(&response_body).unwrap()),
                    _ => Err(Error::EtcdError(json::decode(&response_body).unwrap())),
                }
            },
            Err(error) => Err(Error::HttpError(error)),
        }
    }

    pub fn mkdir(&self, key: &str, ttl: Option<u64>) -> Result<Response, EtcdError> {
        Err(EtcdError)
    }

    pub fn set(
        &self,
        key: &str,
        value: &str,
        ttl: Option<u64>,
        prev_value: Option<&str>,
        prev_index: Option<u64>
    ) -> Result<Response, EtcdError> {
        Err(EtcdError)
    }

    // private

    fn build_url(&self, path: &str) -> String {
        format!("{}v2/keys{}", self.root_url, path)
    }

    fn put(&self, url: String, body: String) -> Result<HyperResponse, HttpError> {
        let mut client = HyperClient::new();
        let content_type: ContentType = ContentType(
            "application/x-www-form-urlencoded".parse().unwrap()
        );

        client.put(&url[..]).header(content_type).body(&body[..]).send()
    }
}

#[cfg(test)]
mod mk_tests {
    use super::Client;

    #[test]
    fn mk_key() {
        let client = Client::new("http://etcd:2379").unwrap();

        let response = client.mk("/foo", "bar", Some(100)).ok().unwrap();

        assert_eq!(response.action, "create".to_string());
        assert_eq!(response.node.value.unwrap(), "bar".to_string());
        assert_eq!(response.node.ttl.unwrap(), 100);
    }

    #[test]
    fn mk_key_failure() {
        let client = Client::default();

        assert!(client.mk("/foo", "bar", None).is_ok());
        assert!(client.mk("/foo", "bar", None).is_err());
    }
}

#[cfg(test)]
mod mkdir_tests {
    use super::Client;

    #[test]
    fn mkdir() {
        let client = Client::default();

        assert!(client.mkdir("/foo", None).ok().unwrap().node.dir.unwrap());
    }

    #[test]
    fn mkdir_failure() {
        let client = Client::default();

        assert!(client.mkdir("/foo", None).is_ok());
        assert!(client.mkdir("/foo", None).is_err());
    }
}

#[cfg(test)]
mod get_tests {
    use super::Client;

    #[test]
    fn set_and_get_key() {
        let client = Client::default();

        assert!(client.set("/foo", "bar", None, None, None).is_ok());
        assert_eq!(client.get("/foo").ok().unwrap().node.value.unwrap(), "bar");
    }
}
<commit_msg>Test mk failure.<commit_after>use std::io::Read;

use hyper::Client as HyperClient;
use hyper::HttpError;
use hyper::client::Response as HyperResponse;
use hyper::header::ContentType;
use hyper::status::StatusCode;
use rustc_serialize::json;
use url::{ParseError, Url};
use url::form_urlencoded::serialize_owned;

#[derive(Debug)]
#[derive(RustcDecodable)]
pub struct Response {
    pub action: String,
    pub node: Node,
    pub prev_node: Option<Node>,
}

#[derive(Debug)]
#[derive(RustcDecodable)]
pub struct Node {
    pub created_index: Option<u64>,
    pub dir: Option<bool>,
    pub expiration: Option<String>,
    pub key: Option<String>,
    pub modified_index: Option<u64>,
    pub nodes: Option<Vec<Node>>,
    pub ttl: Option<i64>,
    pub value: Option<String>,
}

#[derive(Debug)]
pub enum Error {
    EtcdError(EtcdError),
    HttpError(HttpError),
}

#[derive(Debug)]
#[derive(RustcDecodable)]
#[allow(non_snake_case)]
pub struct EtcdError {
    pub cause: Option<String>,
    pub errorCode: u64,
    pub index: u64,
    pub message: String,
}

#[derive(Debug)]
pub struct Client {
    root_url: String,
}

impl Client {
    pub fn new(root_url: &str) -> Result<Client, ParseError> {
        let url = try!(Url::parse(root_url));
        let client = Client {
            root_url: format!("{}", url.serialize()),
        };

        Ok(client)
    }

    pub fn default() -> Client {
        Client {
            root_url: "http://127.0.0.1:2379/".to_string(),
        }
    }

    pub fn get(&self, key: &str) -> Result<Response, EtcdError> {
        Err(EtcdError)
    }

    pub fn mk(&self, key: &str, value: &str, ttl: Option<u64>) -> Result<Response, Error> {
        let url = self.build_url(key);
        let mut options = vec![];

        options.push(("value".to_string(), value.to_string()));
        options.push(("prevExist".to_string(), "false".to_string()));

        if ttl.is_some() {
            options.push(("ttl".to_string(), format!("{}", ttl.unwrap())));
        }

        let body = serialize_owned(&options);

        match self.put(url, body) {
            Ok(mut response) => {
                let mut response_body = String::new();

                response.read_to_string(&mut response_body).unwrap();

                println!("{:?} - {:?} - {:?}", response.status, response.headers, response_body);

                match response.status {
                    StatusCode::Created => Ok(json::decode(&response_body).unwrap()),
                    _ => Err(Error::EtcdError(json::decode(&response_body).unwrap())),
                }
            },
            Err(error) => Err(Error::HttpError(error)),
        }
    }

    pub fn mkdir(&self, key: &str, ttl: Option<u64>) -> Result<Response, EtcdError> {
        Err(EtcdError)
    }

    pub fn set(
        &self,
        key: &str,
        value: &str,
        ttl: Option<u64>,
        prev_value: Option<&str>,
        prev_index: Option<u64>
    ) -> Result<Response, EtcdError> {
        Err(EtcdError)
    }

    // private

    fn build_url(&self, path: &str) -> String {
        format!("{}v2/keys{}", self.root_url, path)
    }

    fn put(&self, url: String, body: String) -> Result<HyperResponse, HttpError> {
        let mut client = HyperClient::new();
        let content_type: ContentType = ContentType(
            "application/x-www-form-urlencoded".parse().unwrap()
        );

        client.put(&url[..]).header(content_type).body(&body[..]).send()
    }
}

#[cfg(test)]
mod mk_tests {
    use super::{Client, Error};

    #[test]
    fn mk_key() {
        let client = Client::new("http://etcd:2379").unwrap();

        let response = client.mk("/foo", "bar", Some(100)).ok().unwrap();

        assert_eq!(response.action, "create".to_string());
        assert_eq!(response.node.value.unwrap(), "bar".to_string());
        assert_eq!(response.node.ttl.unwrap(), 100);
    }

    #[test]
    fn mk_key_failure() {
        let client = Client::new("http://etcd:2379").unwrap();

        assert!(client.mk("/foo", "bar", None).is_ok());

        match client.mk("/foo", "bar", None).err().unwrap() {
            Error::EtcdError(error) => assert_eq!(error.message, "Key already exists".to_string()),
            _ => panic!("expected EtcdError due to pre-existing key"),
        };
    }
}

#[cfg(test)]
mod mkdir_tests {
    use super::Client;

    #[test]
    fn mkdir() {
        let client = Client::default();

        assert!(client.mkdir("/foo", None).ok().unwrap().node.dir.unwrap());
    }

    #[test]
    fn mkdir_failure() {
        let client = Client::default();

        assert!(client.mkdir("/foo", None).is_ok());
        assert!(client.mkdir("/foo", None).is_err());
    }
}

#[cfg(test)]
mod get_tests {
    use super::Client;

    #[test]
    fn set_and_get_key() {
        let client = Client::default();

        assert!(client.set("/foo", "bar", None, None, None).is_ok());
        assert_eq!(client.get("/foo").ok().unwrap().node.value.unwrap(), "bar");
    }
}
<|endoftext|>


<commit_before>//! # Client
//!
//! Main module to handle HTTP request signing
//! and Rest calls
//!

extern crate serde;
extern crate serde_json;

use config::Credential;
use std::io::Read;

use hyper::header::{Headers, UserAgent, Accept, qitem, ContentType};
use hyper::mime::{Value, Mime, TopLevel, SubLevel, Attr};
use hyper::client::Response;
use hyper::client::Client;
use hyper;

use chrono::*;

use crypto::digest::Digest;
use crypto::sha1::Sha1;

// Required headers for auth
header! { (XOvhApplication, "X-Ovh-Application") => [String] }
header! { (XOvhTimestamp, "X-Ovh-Timestamp") => [String] }
header! { (XOvhSignature, "X-Ovh-Signature") => [String] }
header! { (XOvhConsumer, "X-Ovh-Consumer") => [String] }

#[derive(Debug,Clone)]
pub struct OVHClient {
    pub credential: Credential,
}

impl OVHClient {
    /// Initialize a new `Credential` from default path a App Key, App secret, Consumer token.
    pub fn new() -> OVHClient {
        OVHClient { credential: Credential::new() }
    }

    /// Compute signature for OVH.
    fn build_sig(method: &str,
                 query: &str,
                 body: &str,
                 timestamp: &str,
                 aas: &str,
                 ck: &str)
                 -> String {
        let sep = "+";
        let prefix = "$1$".to_string();

        let capacity = 1 + &aas.len() + &sep.len() + &ck.len() + &method.len() + &sep.len() +
                       &query.len() + &sep.len() +
                       &body.len() + &sep.len() + &timestamp.len();
        let mut signature = String::with_capacity(capacity);
        signature.push_str(&aas);
        signature.push_str(&sep);
        signature.push_str(&ck);
        signature.push_str(&sep);
        signature.push_str(&method);
        signature.push_str(&sep);
        signature.push_str(&query);
        signature.push_str(&sep);
        signature.push_str(&body);
        signature.push_str(&sep);
        signature.push_str(&timestamp);

        debug!("Signature: {}", &signature);
        let mut hasher = Sha1::new();
        hasher.input_str(&signature);
        let hex = hasher.result_str();
        debug!("hex: {}", &hex);

        let sign = prefix + &hex;
        sign
    }

    /// Ask time to OVH API server to compute delta time
    fn remote_time() -> u64 {
        let query = "https://eu.api.ovh.com/1.0/auth/time".to_string();
        // Create a client.
        let client = Client::new();

        // Creating an outgoing request.
        let mut res = client.get(&query)
            .send()
            .unwrap();
        let mut body = String::new();
        res.read_to_string(&mut body).unwrap();
        body.parse::<u64>().unwrap()
    }

    /// compute delta time
    fn compute_time_delta() -> u64 {
        let localtime = Local::now()
            .format("%s")
            .to_string()
            .parse::<u64>()
            .unwrap();
        let remotetime = OVHClient::remote_time();
        let deltatime = remotetime - localtime;
        info!("Delta time: {:?}", deltatime);
        deltatime
    }

    /// Start a client request with given method
    /// Use Hyper client, maybe curl-rs in later version
    pub fn request(credential: &Credential, method: &str, query: &str, body: &str) -> String {
        let localtime = Local::now().format("%s").to_string().parse::<u64>().unwrap();
        let computed_time = localtime + OVHClient::compute_time_delta();
        let timestamp = computed_time.to_string();

        let protocol = "https://".to_string();
        let base_path = "/1.0";
        let url = protocol + &credential.host + &base_path + &query;
        let sign = OVHClient::build_sig(&method,
                                        &url,
                                        &body,
                                        &timestamp,
                                        credential.application_secret.as_str(),
                                        credential.consumer_key.as_str());

        // build headers
        let mut headers = Headers::new();
        headers.set(XOvhApplication(credential.application_key.to_string()));
        headers.set(XOvhTimestamp(timestamp.to_string()));
        headers.set(XOvhSignature(sign.to_string()));
        headers.set(XOvhConsumer(credential.consumer_key.to_string()));
        headers.set(Accept(vec![
                qitem(Mime(
                        TopLevel::Application,
                        SubLevel::Json,
                        vec![(Attr::Charset,
                              Value::Utf8)])
                ),
            ]));
        headers.set(ContentType(Mime(TopLevel::Application,
                                     SubLevel::Json,
                                     vec![(Attr::Charset, Value::Utf8)])));
        headers.set(UserAgent("hyper/0.9/ovh-rs".to_owned()));

        // Create a client.
        let client = Client::new();

        debug!("Signature: {}", sign.to_string());

        // Creating an outgoing request.
        let res: Result<Response, hyper::Error> = match method {
            "HEAD" => {
                client.head(&url)
                    .headers(headers)
                    .send()
            }
            "GET" => {
                client.get(&url)
                    .headers(headers)
                    .send()
            }
            "POST" => {
                client.post(&url)
                    .headers(headers)
                    .body(body)
                    .send()
            }
            "PUT" => {
                client.put(&url)
                    .headers(headers)
                    .body(body)
                    .send()
            }
            "PATCH" => {
                client.patch(&url)
                    .headers(headers)
                    .body(body)
                    .send()
            }
            "DELETE" => {
                client.delete(&url)
                    .headers(headers)
                    .body(body)
                    .send()
            }
            _ => Err(hyper::Error::Method),
        };
        let mut body = String::new();
        res.unwrap().read_to_string(&mut body).unwrap();
        body
    }
}

#[cfg(test)]
mod tests {
    use super::OVHClient;
    extern crate serde;
    extern crate serde_json;

    #[test]
    fn test_build_sig() {
        let _ovh = OVHClient::new();
        let method = "GET";
        let query = "https://eu.api.ovh.com/1.0/ipLoadbalancing";
        let body = "";
        let timestamp = "1466716163";
        let aas = "somesecret";
        let ck = "fakeconsumerkey";
        let signature = OVHClient::build_sig(&method, &query, &body, &timestamp, &aas, &ck);
        assert_eq!(&signature, "$1$7ff04a6c8610e4f96a1c0a04dff50ed760a6b724");
    }

    #[test]
    fn test_remote_time() {
        let remote_time = OVHClient::remote_time();
        assert_eq!(true, remote_time > 0);
    }

    #[test]
    fn test_get() {
        let ovh = OVHClient::new();
        let cred = ovh.credential;

        let response = OVHClient::request(&cred, "GET", "/ipLoadbalancing", "");
        // should assert after json parse
        let deser_value: self::serde_json::Value = serde_json::from_str(&response).unwrap();
        assert_eq!(true, deser_value.is_array());
    }

    #[test]
    fn test_post() {
        let ovh = OVHClient::new();
        let cred = ovh.credential;

        let mut body = "{\"ovhSubsidiary\": \"FR\"}";
        let mut response = OVHClient::request(&cred, "POST", "/order/cart", &body);
        // should assert after json parse
        let deser_value: self::serde_json::Value = serde_json::from_str(&response).unwrap();
        assert_eq!(true, deser_value.is_object());
        let obj = deser_value.as_object().unwrap();
        assert_eq!(true, obj.get("cartId").unwrap().is_string());
        assert_eq!(true, obj.get("expire").unwrap().is_string());
        assert_eq!(true, obj.get("description").unwrap().is_string());
        assert_eq!("Default cart",
                   obj.get("description").unwrap().as_str().unwrap());
        assert_eq!(true, obj.get("readOnly").unwrap().is_boolean());
        assert_eq!(true, obj.get("items").unwrap().is_array());

        // test_get_with_query
        let cart_id = obj.get("cartId").unwrap().as_str().unwrap();
        let mut url = "/order/cart/".to_string() + cart_id + "/domain?domain=rustyrust.fr";

        response = OVHClient::request(&cred, "GET", &url, "");
        let deser_value: self::serde_json::Value = serde_json::from_str(&response).unwrap();
        // should assert after json parse
        assert_eq!(true, deser_value.is_array());

        // test_put
        url = "/order/cart/".to_string() + cart_id;
        body = "{\"description\": \"a new rust cart description\"}";
        response = OVHClient::request(&cred, "PUT", &url, &body);
        // should assert after json parse
        let deser_value: self::serde_json::Value = serde_json::from_str(&response).unwrap();
        assert_eq!(true, deser_value.is_object());
        let obj = deser_value.as_object().unwrap();
        assert_eq!(true, obj.get("cartId").unwrap().is_string());
        assert_eq!(true, obj.get("expire").unwrap().is_string());
        assert_eq!(true, obj.get("description").unwrap().is_string());
        assert_eq!("a new rust cart description",
                   obj.get("description").unwrap().as_str().unwrap());
        assert_eq!(true, obj.get("readOnly").unwrap().is_boolean());
        assert_eq!(true, obj.get("items").unwrap().is_array());

        // test assign
        url = "/order/cart/".to_string() + cart_id + "/assign";
        body = "";
        response = OVHClient::request(&cred, "POST", &url, &body);
        assert_eq!("null", response);

        // test_delete
        url = "/order/cart/".to_string() + cart_id;
        body = "";
        response = OVHClient::request(&cred, "DELETE", &url, &body);
        assert_eq!("null", response);
    }

}
<commit_msg>fix(client): when remote time parsing fail<commit_after>//! # Client
//!
//! Main module to handle HTTP request signing
//! and Rest calls
//!

extern crate serde;
extern crate serde_json;

use config::Credential;
use std::io::Read;

use hyper::header::{Headers, UserAgent, Accept, qitem, ContentType};
use hyper::mime::{Value, Mime, TopLevel, SubLevel, Attr};
use hyper::client::Response;
use hyper::client::Client;
use hyper;

use chrono::*;

use crypto::digest::Digest;
use crypto::sha1::Sha1;

// Required headers for auth
header! { (XOvhApplication, "X-Ovh-Application") => [String] }
header! { (XOvhTimestamp, "X-Ovh-Timestamp") => [String] }
header! { (XOvhSignature, "X-Ovh-Signature") => [String] }
header! { (XOvhConsumer, "X-Ovh-Consumer") => [String] }

#[derive(Debug,Clone)]
pub struct OVHClient {
    pub credential: Credential,
}

impl OVHClient {
    /// Initialize a new `Credential` from default path a App Key, App secret, Consumer token.
    pub fn new() -> OVHClient {
        OVHClient { credential: Credential::new() }
    }

    /// Compute signature for OVH.
    fn build_sig(method: &str,
                 query: &str,
                 body: &str,
                 timestamp: &str,
                 aas: &str,
                 ck: &str)
                 -> String {
        let sep = "+";
        let prefix = "$1$".to_string();

        let capacity = 1 + &aas.len() + &sep.len() + &ck.len() + &method.len() + &sep.len() +
                       &query.len() + &sep.len() +
                       &body.len() + &sep.len() + &timestamp.len();
        let mut signature = String::with_capacity(capacity);
        signature.push_str(&aas);
        signature.push_str(&sep);
        signature.push_str(&ck);
        signature.push_str(&sep);
        signature.push_str(&method);
        signature.push_str(&sep);
        signature.push_str(&query);
        signature.push_str(&sep);
        signature.push_str(&body);
        signature.push_str(&sep);
        signature.push_str(&timestamp);

        debug!("Signature: {}", &signature);
        let mut hasher = Sha1::new();
        hasher.input_str(&signature);
        let hex = hasher.result_str();
        debug!("hex: {}", &hex);

        let sign = prefix + &hex;
        sign
    }

    /// Ask time to OVH API server to compute delta time
    fn remote_time() -> u64 {
        let query = "https://eu.api.ovh.com/1.0/auth/time".to_string();
        // Create a client.
        let client = Client::new();

        // Creating an outgoing request.
        let mut res = client.get(&query)
            .send()
            .unwrap();
        let mut body = String::new();
        res.read_to_string(&mut body).unwrap();
        match body.parse::<u64>() {
            Ok(time) => time,
            Err(_) => 1,
        }
    }

    /// compute delta time
    fn compute_time_delta() -> u64 {
        let localtime = Local::now()
            .format("%s")
            .to_string()
            .parse::<u64>()
            .unwrap();
        let remotetime = OVHClient::remote_time();
        if remotetime <= localtime {
            info!("fail to fetch remote time");
            0
        } else {
            let deltatime = remotetime - localtime;
            info!("Delta time: {:?}", deltatime);
            deltatime
        }
    }

    /// Start a client request with given method
    /// Use Hyper client, maybe curl-rs in later version
    pub fn request(credential: &Credential, method: &str, query: &str, body: &str) -> String {
        let localtime = Local::now().format("%s").to_string().parse::<u64>().unwrap();
        let computed_time = localtime + OVHClient::compute_time_delta();
        let timestamp = computed_time.to_string();

        let protocol = "https://".to_string();
        let base_path = "/1.0";
        let url = protocol + &credential.host + &base_path + &query;
        let sign = OVHClient::build_sig(&method,
                                        &url,
                                        &body,
                                        &timestamp,
                                        credential.application_secret.as_str(),
                                        credential.consumer_key.as_str());

        // build headers
        let mut headers = Headers::new();
        headers.set(XOvhApplication(credential.application_key.to_string()));
        headers.set(XOvhTimestamp(timestamp.to_string()));
        headers.set(XOvhSignature(sign.to_string()));
        headers.set(XOvhConsumer(credential.consumer_key.to_string()));
        headers.set(Accept(vec![
                qitem(Mime(
                        TopLevel::Application,
                        SubLevel::Json,
                        vec![(Attr::Charset,
                              Value::Utf8)])
                ),
            ]));
        headers.set(ContentType(Mime(TopLevel::Application,
                                     SubLevel::Json,
                                     vec![(Attr::Charset, Value::Utf8)])));
        headers.set(UserAgent("hyper/0.9/ovh-rs".to_owned()));

        // Create a client.
        let client = Client::new();

        debug!("Signature: {}", sign.to_string());

        // Creating an outgoing request.
        let res: Result<Response, hyper::Error> = match method {
            "HEAD" => {
                client.head(&url)
                    .headers(headers)
                    .send()
            }
            "GET" => {
                client.get(&url)
                    .headers(headers)
                    .send()
            }
            "POST" => {
                client.post(&url)
                    .headers(headers)
                    .body(body)
                    .send()
            }
            "PUT" => {
                client.put(&url)
                    .headers(headers)
                    .body(body)
                    .send()
            }
            "PATCH" => {
                client.patch(&url)
                    .headers(headers)
                    .body(body)
                    .send()
            }
            "DELETE" => {
                client.delete(&url)
                    .headers(headers)
                    .body(body)
                    .send()
            }
            _ => Err(hyper::Error::Method),
        };
        let mut body = String::new();
        res.unwrap().read_to_string(&mut body).unwrap();
        body
    }
}

#[cfg(test)]
mod tests {
    use super::OVHClient;
    extern crate serde;
    extern crate serde_json;

    #[test]
    fn test_build_sig() {
        let method = "GET";
        let query = "https://eu.api.ovh.com/1.0/ipLoadbalancing";
        let body = "";
        let timestamp = "1466716163";
        let aas = "somesecret";
        let ck = "fakeconsumerkey";
        let signature = OVHClient::build_sig(&method, &query, &body, &timestamp, &aas, &ck);
        assert_eq!(&signature, "$1$7ff04a6c8610e4f96a1c0a04dff50ed760a6b724");
    }

    #[test]
    fn test_remote_time() {
        let remote_time = OVHClient::remote_time();
        assert_eq!(true, remote_time > 0);
    }

    #[test]
    fn test_get() {
        let ovh = OVHClient::new();
        let cred = ovh.credential;

        let response = OVHClient::request(&cred, "GET", "/ipLoadbalancing", "");
        // should assert after json parse
        let deser_value: self::serde_json::Value = serde_json::from_str(&response).unwrap();
        assert_eq!(true, deser_value.is_array());
    }

    #[test]
    fn test_post() {
        let ovh = OVHClient::new();
        let cred = ovh.credential;

        let mut body = "{\"ovhSubsidiary\": \"FR\"}";
        let mut response = OVHClient::request(&cred, "POST", "/order/cart", &body);
        // should assert after json parse
        let deser_value: self::serde_json::Value = serde_json::from_str(&response).unwrap();
        assert_eq!(true, deser_value.is_object());
        let obj = deser_value.as_object().unwrap();
        assert_eq!(true, obj.get("cartId").unwrap().is_string());
        assert_eq!(true, obj.get("expire").unwrap().is_string());
        assert_eq!(true, obj.get("description").unwrap().is_string());
        assert_eq!("Default cart",
                   obj.get("description").unwrap().as_str().unwrap());
        assert_eq!(true, obj.get("readOnly").unwrap().is_boolean());
        assert_eq!(true, obj.get("items").unwrap().is_array());

        // test_get_with_query
        let cart_id = obj.get("cartId").unwrap().as_str().unwrap();
        let mut url = "/order/cart/".to_string() + cart_id + "/domain?domain=rustyrust.fr";

        response = OVHClient::request(&cred, "GET", &url, "");
        let deser_value: self::serde_json::Value = serde_json::from_str(&response).unwrap();
        // should assert after json parse
        assert_eq!(true, deser_value.is_array());

        // test_put
        url = "/order/cart/".to_string() + cart_id;
        body = "{\"description\": \"a new rust cart description\"}";
        response = OVHClient::request(&cred, "PUT", &url, &body);
        // should assert after json parse
        let deser_value: self::serde_json::Value = serde_json::from_str(&response).unwrap();
        assert_eq!(true, deser_value.is_object());
        let obj = deser_value.as_object().unwrap();
        assert_eq!(true, obj.get("cartId").unwrap().is_string());
        assert_eq!(true, obj.get("expire").unwrap().is_string());
        assert_eq!(true, obj.get("description").unwrap().is_string());
        assert_eq!("a new rust cart description",
                   obj.get("description").unwrap().as_str().unwrap());
        assert_eq!(true, obj.get("readOnly").unwrap().is_boolean());
        assert_eq!(true, obj.get("items").unwrap().is_array());

        // test assign
        url = "/order/cart/".to_string() + cart_id + "/assign";
        body = "";
        response = OVHClient::request(&cred, "POST", &url, &body);
        assert_eq!("null", response);

        // test_delete
        url = "/order/cart/".to_string() + cart_id;
        body = "";
        response = OVHClient::request(&cred, "DELETE", &url, &body);
        assert_eq!("null", response);
    }

}
<|endoftext|>