<commit_before>extern crate futures;
extern crate text_ui;
extern crate tokio;
extern crate tokio_serde_msgpack;
extern crate tokio_uds;
extern crate weaver;

use text_ui::app::App;
use text_ui::backend::{color, Backend, Color};
use text_ui::pane::Pane;
//use text_ui::widget::DbgDump;
use text_ui::widget::Widget;
use text_ui::widget::{shared, Line, Linear, Readline, Shared, Text};
use text_ui::{text_to_lines, Event, Input, Key, Position, Size};

use std::sync::{Arc, RwLock};
use std::thread;

use tokio::prelude::Future;
use weaver::{WeaverClient, WeaverCommand, WeaverNotification, WeaverState};

struct WeaverStateWidget {
    state: Shared<WeaverState>,
}

fn render_command_summary(cmd: &WeaverCommand, width: usize, maxlines: usize) -> Pane {
    let mut pane = Pane::new_width(width);
    let (icon, style) = match cmd.status {
        None => ('…', "command.running"),
        Some(0) => ('✔', "command.success"),
        _ => ('X', "command.failed"),
    };
    let status_pane = Pane::new_styled(
        Position::new(0, 0),
        Size::new(1, 1),
        vec![icon.to_string()],
        style,
    );
    pane.push_child(status_pane);
    let subwidth = width - 1;
    let command_line = text_to_lines(cmd.cmd.clone(), subwidth);
    let pos = Position::new(1, 0);
    let textlen = command_line.len();
    let mut offset = textlen;
    let size = Size::new(subwidth as u16, textlen as u16);
    pane.push_child(Pane::new_styled(pos, size, command_line, "command"));
    if cmd.stdout.len() > 0 {
        let mut stdout = text_to_lines(cmd.stdout.clone(), subwidth);
        let pos = Position::new(1, offset as u16);
        let mut textlen = stdout.len();
        if textlen > maxlines {
            stdout = stdout.split_off(textlen - maxlines);
            textlen = maxlines;
        }
        offset += textlen;
        let size = Size::new(subwidth as u16, textlen as u16);
        pane.push_child(Pane::new_styled(pos, size, stdout, "stdout"));
    }
    if cmd.stderr.len() > 0 {
        let mut stderr = text_to_lines(cmd.stderr.clone(), subwidth);
        let pos = Position::new(1, offset as u16);
        let mut textlen = stderr.len();
        if textlen > maxlines {
            stderr = stderr.split_off(textlen - maxlines);
            textlen = maxlines;
        }
        let size = Size::new(subwidth as u16, textlen as u16);
        pane.push_child(Pane::new_styled(pos, size, stderr, "stderr"));
    }
    pane
}

impl Widget for WeaverStateWidget {
    fn render_children(&self, size: Size) -> Option<Vec<Pane>> {
        let height = size.height as usize;
        let mut ctr: usize = 0;
        let state = self.state.read().unwrap();
        let mut children: Vec<Pane> = vec![];
        for (_i, cmd) in state.command_history.iter().rev() {
            let mut child = render_command_summary(cmd, size.width as usize, 10);
            let offset = child.size.height as usize;

            ctr += offset;
            if ctr >= height {
                let spillover = (ctr - height) as u16;
                let clip_pos = Position::new(0, spillover);
                let clip_size = Size::new(size.width, child.size.height - spillover);
                child = child.clip(clip_pos, clip_size).unwrap();
                child.position = Position::new(0, 0);
                ctr = height;
            }
            let child = child.offset(Position::new(0, (height - ctr) as u16));
            children.push(child);
            if ctr == height {
                break;
            }
        }
        Some(children)
    }
}

struct WeaverTui {
    log: Shared<Text>,
    input: Shared<Readline>,
    vbox: Shared<Linear>,
    content: Shared<Linear>,
    state: Shared<WeaverState>,
    statew: Shared<WeaverStateWidget>,
    show_debug: bool,
}

impl WeaverTui {
    fn new(state: Arc<RwLock<WeaverState>>) -> WeaverTui {
        let log = shared(Text::new(vec![]));
        let input = shared(Readline::new());
        let state: Shared<WeaverState> = state.into();
        let statew = shared(WeaverStateWidget {
            state: state.clone(),
        });
        //let dbgdump = shared(DbgDump::new(&state));
        let show_debug = false;
        let mut contentbox = Linear::hbox();
        //contentbox.push(&dbgdump);
        contentbox.push(&statew);
        let content = shared(contentbox);
        let mut mainbox = Linear::vbox();
        mainbox.push(&content);
        mainbox.push(&input);
        let vbox = shared(mainbox);
        WeaverTui {
            log,
            input,
            vbox,
            content,
            state,
            statew,
            show_debug,
        }
    }

    fn toggle_debug(&mut self) {
        let mut content = self.content.write().unwrap();
        match self.show_debug {
            true => {
                self.show_debug = false;
                content.contents.truncate(0);
                content.push(&self.statew);
            }
            false => {
                self.show_debug = true;
                content.contents.truncate(0);
                content.push(&self.statew);
                content.push(&shared(Line::vertical()));
                content.push(&self.log);
            }
        }
    }

    fn submit_input(&mut self) {
        let text = self.input.write().unwrap().finalize();
        self.state
            .write()
            .unwrap()
            .run_command(text.clone())
            .unwrap();
    }

    fn log_msg(&mut self, msg: &str) {
        let lines: Vec<String> = msg.lines().map(|l| l.to_owned()).collect();
        self.log.write().unwrap().lines.extend(lines);
    }

    fn input(&mut self, key: Key) {
        match key {
            Key::Char('\n') => self.submit_input(),
            Key::Alt('\r') => self.input.write().unwrap().process_key(Key::Char('\n')),
            k => self.input.write().unwrap().process_key(k),
        }
    }
}

impl App for WeaverTui {
    type UI = Shared<Linear>;
    type MyEvent = WeaverNotification;
    fn widget(&self) -> Self::UI {
        self.vbox.clone()
    }
    fn handle_event(&mut self, event: Event<Self::MyEvent>) -> Result<(), Option<String>> {
        match event {
            Event::InputEvent(i) => match i {
                Input::Key(Key::Esc) => Err(None),
                Input::Key(Key::Alt('d')) => {
                    self.toggle_debug();
                    Ok(())
                }
                Input::Key(k) => {
                    self.input(k);
                    Ok(())
                }
                _ => Ok(()),
            },
            Event::AppEvent(_) => {
                self.log_msg(&format!("{:?}", event));
                Ok(())
            }
        }
    }
    fn style(&self, name: &str) -> (Option<Box<Color>>, Option<Box<Color>>) {
        match name {
            "command" => (None, Some(Box::new(color::Rgb(16, 16, 32)))),
            "stderr" => (None, Some(Box::new(color::Rgb(32, 16, 16)))),
            "input" => (None, Some(Box::new(color::Rgb(32, 32, 32)))),
            "command.running" => (
                Some(Box::new(color::LightYellow)),
                Some(Box::new(color::Rgb(16, 16, 32))),
            ),
            "command.success" => (
                Some(Box::new(color::LightGreen)),
                Some(Box::new(color::Rgb(16, 16, 32))),
            ),
            "command.failed" => (Some(Box::new(color::LightRed)), None),
            _ => (None, None),
        }
    }
}

fn main() {
    let be = Backend::new();
    let sender = be.sender.clone();
    let weaver = WeaverClient::new(sender);
    let mut app = WeaverTui::new(weaver.state.clone());
    thread::spawn(move || {
        tokio::run(weaver.map_err(|e| panic!("Client Error: {:#?}", e)));
    });
    app.log_msg("Esc to exit");
    be.run_app(&mut app);
}
<commit_msg>Start interacting with history<commit_after>extern crate futures;
extern crate text_ui;
extern crate tokio;
extern crate tokio_serde_msgpack;
extern crate tokio_uds;
extern crate weaver;

use text_ui::app::App;
use text_ui::backend::{color, Backend, Color};
use text_ui::pane::Pane;
//use text_ui::widget::DbgDump;
use text_ui::widget::Widget;
use text_ui::widget::{shared, Line, Linear, Readline, Shared, Text};
use text_ui::{text_to_lines, Event, Input, Key, Position, Size};

use std::sync::{Arc, RwLock};
use std::thread;

use tokio::prelude::Future;
use weaver::{WeaverClient, WeaverCommand, WeaverNotification, WeaverState};

struct WeaverStateWidget {
    state: Shared<WeaverState>,
    selected: Option<u32>,
}

impl WeaverStateWidget {
    pub fn new(state: Shared<WeaverState>) -> Self {
        let selected = None;
        WeaverStateWidget { state, selected }
    }
}

fn render_command_summary(
    cmd: &WeaverCommand,
    width: usize,
    maxlines: usize,
    selected: bool,
) -> Pane {
    let mut pane = Pane::new_width(width);
    let (icon, style) = match cmd.status {
        None => ('…', "command.running"),
        Some(0) => ('✔', "command.success"),
        _ => ('X', "command.failed"),
    };
    let status_pane = Pane::new_styled(
        Position::new(0, 0),
        Size::new(1, 1),
        vec![icon.to_string()],
        style,
    );
    pane.push_child(status_pane);
    let subwidth = width - 1;
    let command_line = text_to_lines(cmd.cmd.clone(), subwidth);
    let pos = Position::new(1, 0);
    let textlen = command_line.len();
    let mut offset = textlen;
    let size = Size::new(subwidth as u16, textlen as u16);
    let prefix = match selected {
        true => "selected.",
        false => "",
    };
    pane.push_child(Pane::new_styled(
        pos,
        size,
        command_line,
        &format!("{}command", prefix),
    ));
    if cmd.stdout.len() > 0 {
        let mut stdout = text_to_lines(cmd.stdout.clone(), subwidth);
        let pos = Position::new(1, offset as u16);
        let mut textlen = stdout.len();
        if textlen > maxlines {
            stdout = stdout.split_off(textlen - maxlines);
            textlen = maxlines;
        }
        offset += textlen;
        let size = Size::new(subwidth as u16, textlen as u16);
        pane.push_child(Pane::new_styled(
            pos,
            size,
            stdout,
            &format!("{}stdout", prefix),
        ));
    }
    if cmd.stderr.len() > 0 {
        let mut stderr = text_to_lines(cmd.stderr.clone(), subwidth);
        let pos = Position::new(1, offset as u16);
        let mut textlen = stderr.len();
        if textlen > maxlines {
            stderr = stderr.split_off(textlen - maxlines);
            textlen = maxlines;
        }
        let size = Size::new(subwidth as u16, textlen as u16);
        pane.push_child(Pane::new_styled(
            pos,
            size,
            stderr,
            &format!("{}stderr", prefix),
        ));
    }
    pane
}

fn render_command_detail(cmd: &WeaverCommand, size: Size) -> Pane {
    let mut pane = Pane::new_width(size.width as usize);
    let (icon, style) = match cmd.status {
        None => ('…', "command.running"),
        Some(0) => ('✔', "command.success"),
        _ => ('X', "command.failed"),
    };
    let status_pane = Pane::new_styled(
        Position::new(0, 0),
        Size::new(1, 1),
        vec![icon.to_string()],
        style,
    );
    pane.push_child(status_pane);
    let subwidth: usize = size.width as usize - 1;
    let command_line = text_to_lines(cmd.cmd.clone(), subwidth);
    let pos = Position::new(1, 0);
    let textlen = command_line.len();
    let mut offset = textlen;
    let pane_size = Size::new(subwidth as u16, textlen as u16);
    let prefix = "selected.";
    let mut maxlines = size.height as usize - textlen;
    pane.push_child(Pane::new_styled(
        pos,
        pane_size,
        command_line,
        &format!("{}command", prefix),
    ));
    if cmd.stdout.len() > 0 {
        let mut stdout = text_to_lines(cmd.stdout.clone(), subwidth);
        let pos = Position::new(1, offset as u16);
        let mut textlen = stdout.len();
        if textlen > maxlines {
            stdout = stdout.split_off(textlen - maxlines);
            textlen = maxlines;
        }
        maxlines -= textlen;
        offset += textlen;
        let pane_size = Size::new(subwidth as u16, textlen as u16);
        pane.push_child(Pane::new_styled(
            pos,
            pane_size,
            stdout,
            &format!("{}stdout", prefix),
        ));
    }
    if cmd.stderr.len() > 0 {
        let mut stderr = text_to_lines(cmd.stderr.clone(), subwidth);
        let pos = Position::new(1, offset as u16);
        let mut textlen = stderr.len();
        if textlen > maxlines {
            stderr = stderr.split_off(textlen - maxlines);
            textlen = maxlines;
        }
        let pane_size = Size::new(subwidth as u16, textlen as u16);
        pane.push_child(Pane::new_styled(
            pos,
            pane_size,
            stderr,
            &format!("{}stderr", prefix),
        ));
    }
    pane
}

impl Widget for WeaverStateWidget {
    fn render_children(&self, size: Size) -> Option<Vec<Pane>> {
        let height = size.height as usize;
        let mut ctr: usize = 0;
        let state = self.state.read().unwrap();
        let mut children: Vec<Pane> = vec![];
        let mut i = 0;
        let child_width: usize = match self.selected {
            None => size.width as usize,
            Some(_) => size.width as usize / 2,
        };
        for (_i, cmd) in state.command_history.iter().rev() {
            let selected: bool = match self.selected {
                None => false,
                Some(idx) => idx == i,
            };
            let mut child = render_command_summary(cmd, child_width, 10, selected);
            let offset = child.size.height as usize;

            ctr += offset;
            if ctr >= height {
                let spillover = (ctr - height) as u16;
                let clip_pos = Position::new(0, spillover);
                let clip_size = Size::new(child_width as u16, child.size.height - spillover);
                child = child.clip(clip_pos, clip_size).unwrap();
                child.position = Position::new(0, 0);
                ctr = height;
            }
            let child = child.offset(Position::new(0, (height - ctr) as u16));
            children.push(child);
            if selected {
                let child_pos = Position::new(child_width as u16, 0);
                let child_size = Size::new(size.width - (child_width as u16), size.height);
                let child = render_command_detail(cmd, child_size).offset(child_pos);
                children.push(child);
            }
            if ctr == height {
                break;
            }
            i += 1;
        }
        Some(children)
    }
}

struct WeaverTui {
    log: Shared<Text>,
    input: Shared<Readline>,
    vbox: Shared<Linear>,
    content: Shared<Linear>,
    state: Shared<WeaverState>,
    statew: Shared<WeaverStateWidget>,
    show_debug: bool,
}

impl WeaverTui {
    fn new(state: Arc<RwLock<WeaverState>>) -> WeaverTui {
        let log = shared(Text::new(vec![]));
        let input = shared(Readline::new());
        let state: Shared<WeaverState> = state.into();
        let statew = shared(WeaverStateWidget::new(state.clone()));
        //let dbgdump = shared(DbgDump::new(&state));
        let show_debug = false;
        let mut contentbox = Linear::hbox();
        //contentbox.push(&dbgdump);
        contentbox.push(&statew);
        let content = shared(contentbox);
        let mut mainbox = Linear::vbox();
        mainbox.push(&content);
        mainbox.push(&input);
        let vbox = shared(mainbox);
        WeaverTui {
            log,
            input,
            vbox,
            content,
            state,
            statew,
            show_debug,
        }
    }

    fn toggle_debug(&mut self) {
        let mut content = self.content.write().unwrap();
        match self.show_debug {
            true => {
                self.show_debug = false;
                content.contents.truncate(0);
                content.push(&self.statew);
            }
            false => {
                self.show_debug = true;
                content.contents.truncate(0);
                content.push(&self.statew);
                content.push(&shared(Line::vertical()));
                content.push(&self.log);
            }
        }
    }

    fn submit_input(&mut self) {
        let text = self.input.write().unwrap().finalize();
        self.state
            .write()
            .unwrap()
            .run_command(text.clone())
            .unwrap();
        self.statew.write().unwrap().selected = Some(0);
    }

    fn log_msg(&mut self, msg: &str) {
        let lines: Vec<String> = msg.lines().map(|l| l.to_owned()).collect();
        self.log.write().unwrap().lines.extend(lines);
    }

    fn input(&mut self, key: Key) {
        match key {
            Key::Char('\n') => self.submit_input(),
            Key::Alt('\r') => self.input.write().unwrap().process_key(Key::Char('\n')),
            Key::Up => {
                let mut statew = self.statew.write().unwrap();
                statew.selected = match statew.selected.take() {
                    None => Some(0),
                    Some(i) => Some(i + 1),
                }
            }
            Key::Down => {
                let mut statew = self.statew.write().unwrap();
                statew.selected = match statew.selected.take() {
                    None => None,
                    Some(0) => None,
                    Some(i) => Some(i - 1),
                }
            }
            k => self.input.write().unwrap().process_key(k),
        }
    }
}

impl App for WeaverTui {
    type UI = Shared<Linear>;
    type MyEvent = WeaverNotification;
    fn widget(&self) -> Self::UI {
        self.vbox.clone()
    }
    fn handle_event(&mut self, event: Event<Self::MyEvent>) -> Result<(), Option<String>> {
        match event {
            Event::InputEvent(i) => match i {
                Input::Key(Key::Esc) => Err(None),
                Input::Key(Key::Alt('d')) => {
                    self.toggle_debug();
                    Ok(())
                }
                Input::Key(k) => {
                    self.input(k);
                    Ok(())
                }
                _ => Ok(()),
            },
            Event::AppEvent(_) => {
                self.log_msg(&format!("{:?}", event));
                Ok(())
            }
        }
    }
    fn style(&self, name: &str) -> (Option<Box<Color>>, Option<Box<Color>>) {
        match name {
            "command" => (None, Some(Box::new(color::Rgb(16, 16, 32)))),
            "stderr" => (None, Some(Box::new(color::Rgb(32, 16, 16)))),
            "stdout" => (None, Some(Box::new(color::Rgb(16, 32, 16)))),
            "input" => (None, Some(Box::new(color::Rgb(32, 32, 32)))),
            "command.running" => (
                Some(Box::new(color::LightYellow)),
                Some(Box::new(color::Rgb(16, 16, 32))),
            ),
            "command.success" => (
                Some(Box::new(color::LightGreen)),
                Some(Box::new(color::Rgb(16, 16, 32))),
            ),
            "command.failed" => (Some(Box::new(color::LightRed)), None),
            "selected.command" => (
                Some(Box::new(color::LightWhite)),
                Some(Box::new(color::Rgb(32, 32, 128))),
            ),
            "selected.stderr" => (
                Some(Box::new(color::LightWhite)),
                Some(Box::new(color::Rgb(64, 16, 16))),
            ),
            "selected.stdout" => (
                Some(Box::new(color::LightWhite)),
                Some(Box::new(color::Rgb(16, 64, 16))),
            ),
            "selected.command.running" => (
                Some(Box::new(color::LightYellow)),
                Some(Box::new(color::Rgb(16, 16, 32))),
            ),
            "selected.command.success" => (
                Some(Box::new(color::LightGreen)),
                Some(Box::new(color::Rgb(16, 16, 32))),
            ),
            "selected.command.failed" => (Some(Box::new(color::LightRed)), None),
            _ => (None, None),
        }
    }
}

fn main() {
    let be = Backend::new();
    let sender = be.sender.clone();
    let weaver = WeaverClient::new(sender);
    let mut app = WeaverTui::new(weaver.state.clone());
    thread::spawn(move || {
        tokio::run(weaver.map_err(|e| panic!("Client Error: {:#?}", e)));
    });
    app.log_msg("Esc to exit");
    be.run_app(&mut app);
}
<|endoftext|>


<commit_before>use super::*;

#[test]
fn no_proxies() {
    let client = new_client();
    let proxies = client.proxies().unwrap();
    assert_eq!(proxies.len(), 0);
}

#[test]
fn create_proxy() {
    let client = new_client();
    let proxy = test_proxy();
    client.create_proxy(&proxy).unwrap();

    let proxies = client.proxies().unwrap();
    assert_eq!(1, proxies.len());
    assert_eq!(Some(&proxy), proxies.get(&proxy.name));
}

#[test]
fn update_proxy() {
    let client = new_client();
    let mut proxy = test_proxy();
    client.create_proxy(&proxy).unwrap();

    proxy.enabled = false;
    proxy.upstream = "example.org:80".to_string();
    client.update_proxy(&proxy).unwrap();

    let proxies = client.proxies().unwrap();
    assert_eq!(1, proxies.len());
    assert_eq!(Some(&proxy), proxies.get(&proxy.name));
}

#[test]
fn delete_proxy() {
    let client = new_client();
    let proxy = test_proxy();
    client.create_proxy(&proxy).unwrap();
    client.delete_proxy(&proxy.name).unwrap();

    let proxies = client.proxies().unwrap();
    assert_eq!(0, proxies.len());
}

#[test]
fn unknown_proxy() {
    let client = new_client();
    let proxy = test_proxy();

    assert!(client.update_proxy(&proxy).is_err());
    assert!(client.delete_proxy(&proxy.name).is_err());
}

#[test]
fn duplicate_proxy() {
    let client = new_client();
    let proxy = test_proxy();
    client.create_proxy(&proxy).unwrap();
    assert!(client.create_proxy(&proxy).is_err());
}

#[test]
fn no_toxics() {
    let client = new_client();
    let proxy = test_proxy();

    client.create_proxy(&proxy).unwrap();
    let toxics = client.toxics(&proxy.name).unwrap();
    assert_eq!(0, toxics.len());
}

#[test]
fn create_toxic() {
    let client = new_client();
    let proxy = test_proxy();
    client.create_proxy(&proxy).unwrap();

    let toxic = test_toxic();
    client.create_toxic(&proxy.name, &toxic).unwrap();

    let toxics = client.toxics(&proxy.name).unwrap();
    assert_eq!(1, toxics.len());
    assert_eq!(toxic, toxics[0]);
}

#[test]
fn update_toxic() {
    let client = new_client();
    let proxy = test_proxy();
    client.create_proxy(&proxy).unwrap();

    let mut toxic = test_toxic();
    client.create_toxic(&proxy.name, &toxic).unwrap();

    toxic.toxicity = 0.5;
    client.update_toxic(&proxy.name, &toxic).unwrap();

    let toxics = client.toxics(&proxy.name).unwrap();
    assert_eq!(1, toxics.len());
    assert_eq!(toxic, toxics[0]);
}

#[test]
fn delete_toxic() {
    let client = new_client();
    let proxy = test_proxy();
    client.create_proxy(&proxy).unwrap();

    let toxic = test_toxic();
    client.create_toxic(&proxy.name, &toxic).unwrap();
    client.delete_toxic(&proxy.name, &toxic.name).unwrap();

    let toxics = client.toxics(&proxy.name).unwrap();
    assert_eq!(0, toxics.len());
}

#[test]
fn unknown_toxic() {
    let client = new_client();
    let proxy = test_proxy();
    client.create_proxy(&proxy).unwrap();

    let toxic = test_toxic();
    assert!(client.update_toxic(&proxy.name, &toxic).is_err());
    assert!(client.delete_toxic(&proxy.name, &toxic.name).is_err());
}

#[test]
fn duplicate_toxic() {
    let client = new_client();
    let proxy = test_proxy();
    client.create_proxy(&proxy).unwrap();

    let toxic = test_toxic();
    client.create_toxic(&proxy.name, &toxic).unwrap();
    assert!(client.create_toxic(&proxy.name, &toxic).is_err());
}

fn new_client() -> Client {
    let client = Client::new("localhost:8474");
    let proxies = client.proxies().unwrap();

    for name in proxies.keys() {
        client.delete_proxy(name).unwrap();
    }

    client
}

fn test_proxy() -> Proxy {
    Proxy {
        name: "testproxy".to_string(),
        listen: "127.0.0.1:13306".to_string(),
        upstream: "127.0.0.1:3306".to_string(),
        enabled: true,
    }
}

fn test_toxic() -> Toxic {
    Toxic {
        name: "testtoxic".to_string(),
        stream: Stream::Downstream,
        toxicity: 1.0,
        typ: ToxicType::Latency {
            latency: 100,
            jitter: 0,
        },
    }
}
<commit_msg>Rename new_client to test_client<commit_after>use super::*;

#[test]
fn no_proxies() {
    let client = test_client();
    let proxies = client.proxies().unwrap();
    assert_eq!(proxies.len(), 0);
}

#[test]
fn create_proxy() {
    let client = test_client();
    let proxy = test_proxy();
    client.create_proxy(&proxy).unwrap();

    let proxies = client.proxies().unwrap();
    assert_eq!(1, proxies.len());
    assert_eq!(Some(&proxy), proxies.get(&proxy.name));
}

#[test]
fn update_proxy() {
    let client = test_client();
    let mut proxy = test_proxy();
    client.create_proxy(&proxy).unwrap();

    proxy.enabled = false;
    proxy.upstream = "example.org:80".to_string();
    client.update_proxy(&proxy).unwrap();

    let proxies = client.proxies().unwrap();
    assert_eq!(1, proxies.len());
    assert_eq!(Some(&proxy), proxies.get(&proxy.name));
}

#[test]
fn delete_proxy() {
    let client = test_client();
    let proxy = test_proxy();
    client.create_proxy(&proxy).unwrap();
    client.delete_proxy(&proxy.name).unwrap();

    let proxies = client.proxies().unwrap();
    assert_eq!(0, proxies.len());
}

#[test]
fn unknown_proxy() {
    let client = test_client();
    let proxy = test_proxy();

    assert!(client.update_proxy(&proxy).is_err());
    assert!(client.delete_proxy(&proxy.name).is_err());
}

#[test]
fn duplicate_proxy() {
    let client = test_client();
    let proxy = test_proxy();
    client.create_proxy(&proxy).unwrap();
    assert!(client.create_proxy(&proxy).is_err());
}

#[test]
fn no_toxics() {
    let client = test_client();
    let proxy = test_proxy();

    client.create_proxy(&proxy).unwrap();
    let toxics = client.toxics(&proxy.name).unwrap();
    assert_eq!(0, toxics.len());
}

#[test]
fn create_toxic() {
    let client = test_client();
    let proxy = test_proxy();
    client.create_proxy(&proxy).unwrap();

    let toxic = test_toxic();
    client.create_toxic(&proxy.name, &toxic).unwrap();

    let toxics = client.toxics(&proxy.name).unwrap();
    assert_eq!(1, toxics.len());
    assert_eq!(toxic, toxics[0]);
}

#[test]
fn update_toxic() {
    let client = test_client();
    let proxy = test_proxy();
    client.create_proxy(&proxy).unwrap();

    let mut toxic = test_toxic();
    client.create_toxic(&proxy.name, &toxic).unwrap();

    toxic.toxicity = 0.5;
    client.update_toxic(&proxy.name, &toxic).unwrap();

    let toxics = client.toxics(&proxy.name).unwrap();
    assert_eq!(1, toxics.len());
    assert_eq!(toxic, toxics[0]);
}

#[test]
fn delete_toxic() {
    let client = test_client();
    let proxy = test_proxy();
    client.create_proxy(&proxy).unwrap();

    let toxic = test_toxic();
    client.create_toxic(&proxy.name, &toxic).unwrap();
    client.delete_toxic(&proxy.name, &toxic.name).unwrap();

    let toxics = client.toxics(&proxy.name).unwrap();
    assert_eq!(0, toxics.len());
}

#[test]
fn unknown_toxic() {
    let client = test_client();
    let proxy = test_proxy();
    client.create_proxy(&proxy).unwrap();

    let toxic = test_toxic();
    assert!(client.update_toxic(&proxy.name, &toxic).is_err());
    assert!(client.delete_toxic(&proxy.name, &toxic.name).is_err());
}

#[test]
fn duplicate_toxic() {
    let client = test_client();
    let proxy = test_proxy();
    client.create_proxy(&proxy).unwrap();

    let toxic = test_toxic();
    client.create_toxic(&proxy.name, &toxic).unwrap();
    assert!(client.create_toxic(&proxy.name, &toxic).is_err());
}

fn test_client() -> Client {
    let client = Client::new("localhost:8474");
    let proxies = client.proxies().unwrap();

    for name in proxies.keys() {
        client.delete_proxy(name).unwrap();
    }

    client
}

fn test_proxy() -> Proxy {
    Proxy {
        name: "testproxy".to_string(),
        listen: "127.0.0.1:13306".to_string(),
        upstream: "127.0.0.1:3306".to_string(),
        enabled: true,
    }
}

fn test_toxic() -> Toxic {
    Toxic {
        name: "testtoxic".to_string(),
        stream: Stream::Downstream,
        toxicity: 1.0,
        typ: ToxicType::Latency {
            latency: 100,
            jitter: 0,
        },
    }
}
<|endoftext|>


<commit_before>use num::{Float, Num, FromPrimitive};
use std::ops::{Neg, AddAssign, MulAssign, SubAssign, DivAssign};
use cmp::ApproxEq;
use structure::FieldApprox;

#[allow(missing_docs)]

/// Trait shared by all approximate reals.
///
/// Approximate reals are approximate fields aquipped with funtinos that are commonly used on
/// reals. The results of those functions only have to be approximately equal to the actual
/// theoretical values.
pub trait RealApprox: Copy + Num + FieldApprox + FromPrimitive +
                      Neg<Output = Self> + AddAssign + MulAssign + SubAssign + DivAssign +
                      ApproxEq<Eps = Self> + Ord {
    fn floor(self) -> Self;
    fn ceil(self) -> Self;
    fn round(self) -> Self;
    fn trunc(self) -> Self;
    fn fract(self) -> Self;
    fn abs(self) -> Self;
    fn signum(self) -> Self;
    fn is_sign_positive(self) -> bool;
    fn is_sign_negative(self) -> bool;
    fn mul_add(self, a: Self, b: Self) -> Self;
    fn recip(self) -> Self;
    fn powi(self, n: i32) -> Self;
    fn powf(self, n: Self) -> Self;
    fn sqrt(self) -> Self;
    fn exp(self) -> Self;
    fn exp2(self) -> Self;
    fn ln(self) -> Self;
    fn log(self, base: Self) -> Self;
    fn log2(self) -> Self;
    fn log10(self) -> Self;
    fn max(self, other: Self) -> Self;
    fn min(self, other: Self) -> Self;
    fn abs_sub(self, other: Self) -> Self;
    fn cbrt(self) -> Self;
    fn hypot(self, other: Self) -> Self;
    fn sin(self) -> Self;
    fn cos(self) -> Self;
    fn tan(self) -> Self;
    fn asin(self) -> Self;
    fn acos(self) -> Self;
    fn atan(self) -> Self;
    fn atan2(self, other: Self) -> Self;
    fn sin_cos(self) -> (Self, Self);
    fn exp_m1(self) -> Self;
    fn ln_1p(self) -> Self;
    fn sinh(self) -> Self;
    fn cosh(self) -> Self;
    fn tanh(self) -> Self;
    fn asinh(self) -> Self;
    fn acosh(self) -> Self;
    fn atanh(self) -> Self;
}

impl<S> RealApprox for S
where S: Float + FieldApprox + FromPrimitive + Neg +
         AddAssign + MulAssign + SubAssign + DivAssign +
         ApproxEq<Eps = S> + Ord {
    fn floor(self) -> Self {
        self.floor()
    }

    fn ceil(self) -> Self {
        self.ceil()
    }

    fn round(self) -> Self {
        self.round()
    }

    fn trunc(self) -> Self {
        self.trunc()
    }

    fn fract(self) -> Self {
        self.fract()
    }

    fn abs(self) -> Self {
        self.abs()
    }

    fn signum(self) -> Self {
        self.signum()
    }

    fn is_sign_positive(self) -> bool {
        self.is_sign_positive()
    }

    fn is_sign_negative(self) -> bool {
        self.is_sign_negative()
    }

    fn mul_add(self, a: Self, b: Self) -> Self {
        self.mul_add(a, b)
    }

    fn recip(self) -> Self {
        self.recip()
    }

    fn powi(self, n: i32) -> Self {
        self.powi(n)
    }

    fn powf(self, n: Self) -> Self {
        self.powf(n)
    }

    fn sqrt(self) -> Self {
        self.sqrt()
    }

    fn exp(self) -> Self {
        self.exp()
    }

    fn exp2(self) -> Self {
        self.exp2()
    }

    fn ln(self) -> Self {
        self.ln()
    }

    fn log(self, base: Self) -> Self {
        self.log(base)
    }

    fn log2(self) -> Self {
        self.log2()
    }

    fn log10(self) -> Self {
        self.log10()
    }

    fn max(self, other: Self) -> Self {
        self.max(other)
    }

    fn min(self, other: Self) -> Self {
        self.min(other)
    }

    fn abs_sub(self, other: Self) -> Self {
        self.abs_sub(other)
    }

    fn cbrt(self) -> Self {
        self.cbrt()
    }

    fn hypot(self, other: Self) -> Self {
        self.hypot(other)
    }

    fn sin(self) -> Self {
        self.sin()
    }

    fn cos(self) -> Self {
        self.cos()
    }

    fn tan(self) -> Self {
        self.tan()
    }

    fn asin(self) -> Self {
        self.asin()
    }

    fn acos(self) -> Self {
        self.acos()
    }

    fn atan(self) -> Self {
        self.atan()
    }

    fn atan2(self, other: Self) -> Self {
        self.atan2(other)
    }

    fn sin_cos(self) -> (Self, Self) {
        self.sin_cos()
    }

    fn exp_m1(self) -> Self {
        self.exp_m1()
    }

    fn ln_1p(self) -> Self {
        self.ln_1p()
    }

    fn sinh(self) -> Self {
        self.sinh()
    }

    fn cosh(self) -> Self {
        self.cosh()
    }

    fn tanh(self) -> Self {
        self.tanh()
    }

    fn asinh(self) -> Self {
        self.asinh()
    }

    fn acosh(self) -> Self {
        self.acosh()
    }

    fn atanh(self) -> Self {
        self.atanh()
    }
}
<commit_msg>Inline the implementation of `RealApprox` for floats.<commit_after>use num::{Float, Num, FromPrimitive};
use std::ops::{Neg, AddAssign, MulAssign, SubAssign, DivAssign};
use cmp::ApproxEq;
use structure::FieldApprox;

#[allow(missing_docs)]

/// Trait shared by all approximate reals.
///
/// Approximate reals are approximate fields aquipped with funtinos that are commonly used on
/// reals. The results of those functions only have to be approximately equal to the actual
/// theoretical values.
pub trait RealApprox: Copy + Num + FieldApprox + FromPrimitive +
                      Neg<Output = Self> + AddAssign + MulAssign + SubAssign + DivAssign +
                      ApproxEq<Eps = Self> + Ord {
    fn floor(self) -> Self;
    fn ceil(self) -> Self;
    fn round(self) -> Self;
    fn trunc(self) -> Self;
    fn fract(self) -> Self;
    fn abs(self) -> Self;
    fn signum(self) -> Self;
    fn is_sign_positive(self) -> bool;
    fn is_sign_negative(self) -> bool;
    fn mul_add(self, a: Self, b: Self) -> Self;
    fn recip(self) -> Self;
    fn powi(self, n: i32) -> Self;
    fn powf(self, n: Self) -> Self;
    fn sqrt(self) -> Self;
    fn exp(self) -> Self;
    fn exp2(self) -> Self;
    fn ln(self) -> Self;
    fn log(self, base: Self) -> Self;
    fn log2(self) -> Self;
    fn log10(self) -> Self;
    fn max(self, other: Self) -> Self;
    fn min(self, other: Self) -> Self;
    fn abs_sub(self, other: Self) -> Self;
    fn cbrt(self) -> Self;
    fn hypot(self, other: Self) -> Self;
    fn sin(self) -> Self;
    fn cos(self) -> Self;
    fn tan(self) -> Self;
    fn asin(self) -> Self;
    fn acos(self) -> Self;
    fn atan(self) -> Self;
    fn atan2(self, other: Self) -> Self;
    fn sin_cos(self) -> (Self, Self);
    fn exp_m1(self) -> Self;
    fn ln_1p(self) -> Self;
    fn sinh(self) -> Self;
    fn cosh(self) -> Self;
    fn tanh(self) -> Self;
    fn asinh(self) -> Self;
    fn acosh(self) -> Self;
    fn atanh(self) -> Self;
}

impl<S> RealApprox for S
where S: Float + FieldApprox + FromPrimitive + Neg +
         AddAssign + MulAssign + SubAssign + DivAssign +
         ApproxEq<Eps = S> + Ord {
    #[inline]
    fn floor(self) -> Self {
        self.floor()
    }

    #[inline]
    fn ceil(self) -> Self {
        self.ceil()
    }

    #[inline]
    fn round(self) -> Self {
        self.round()
    }

    #[inline]
    fn trunc(self) -> Self {
        self.trunc()
    }

    #[inline]
    fn fract(self) -> Self {
        self.fract()
    }

    #[inline]
    fn abs(self) -> Self {
        self.abs()
    }

    #[inline]
    fn signum(self) -> Self {
        self.signum()
    }

    #[inline]
    fn is_sign_positive(self) -> bool {
        self.is_sign_positive()
    }

    #[inline]
    fn is_sign_negative(self) -> bool {
        self.is_sign_negative()
    }

    #[inline]
    fn mul_add(self, a: Self, b: Self) -> Self {
        self.mul_add(a, b)
    }

    #[inline]
    fn recip(self) -> Self {
        self.recip()
    }

    #[inline]
    fn powi(self, n: i32) -> Self {
        self.powi(n)
    }

    #[inline]
    fn powf(self, n: Self) -> Self {
        self.powf(n)
    }

    #[inline]
    fn sqrt(self) -> Self {
        self.sqrt()
    }

    #[inline]
    fn exp(self) -> Self {
        self.exp()
    }

    #[inline]
    fn exp2(self) -> Self {
        self.exp2()
    }

    #[inline]
    fn ln(self) -> Self {
        self.ln()
    }

    #[inline]
    fn log(self, base: Self) -> Self {
        self.log(base)
    }

    #[inline]
    fn log2(self) -> Self {
        self.log2()
    }

    #[inline]
    fn log10(self) -> Self {
        self.log10()
    }

    #[inline]
    fn max(self, other: Self) -> Self {
        self.max(other)
    }

    #[inline]
    fn min(self, other: Self) -> Self {
        self.min(other)
    }

    #[inline]
    fn abs_sub(self, other: Self) -> Self {
        self.abs_sub(other)
    }

    #[inline]
    fn cbrt(self) -> Self {
        self.cbrt()
    }

    #[inline]
    fn hypot(self, other: Self) -> Self {
        self.hypot(other)
    }

    #[inline]
    fn sin(self) -> Self {
        self.sin()
    }

    #[inline]
    fn cos(self) -> Self {
        self.cos()
    }

    #[inline]
    fn tan(self) -> Self {
        self.tan()
    }

    #[inline]
    fn asin(self) -> Self {
        self.asin()
    }

    #[inline]
    fn acos(self) -> Self {
        self.acos()
    }

    #[inline]
    fn atan(self) -> Self {
        self.atan()
    }

    #[inline]
    fn atan2(self, other: Self) -> Self {
        self.atan2(other)
    }

    #[inline]
    fn sin_cos(self) -> (Self, Self) {
        self.sin_cos()
    }

    #[inline]
    fn exp_m1(self) -> Self {
        self.exp_m1()
    }

    #[inline]
    fn ln_1p(self) -> Self {
        self.ln_1p()
    }

    #[inline]
    fn sinh(self) -> Self {
        self.sinh()
    }

    #[inline]
    fn cosh(self) -> Self {
        self.cosh()
    }

    #[inline]
    fn tanh(self) -> Self {
        self.tanh()
    }

    #[inline]
    fn asinh(self) -> Self {
        self.asinh()
    }

    #[inline]
    fn acosh(self) -> Self {
        self.acosh()
    }

    #[inline]
    fn atanh(self) -> Self {
        self.atanh()
    }
}
<|endoftext|>


<commit_before>// check-pass

// This test exhibits undefined behavior, but it is impossible to prevent generally during
// const eval, even if possible to prevent in the cases here. The reason it's impossible in general
// is that we run into query cycles even *without* UB, just because we're checking for UB.
// We do not detect it if you create references to statics
// in ways that are UB.

enum Foo {
    A = 5,
    B = 42,
}
enum Bar {
    C = 42,
    D = 99,
}
#[repr(C)]
union Union {
    foo: &'static Foo,
    bar: &'static Bar,
    u8: &'static u8,
}
static BAR: u8 = 5;
static FOO: (&Foo, &Bar) = unsafe {
    (
        // undefined behavior
        Union { u8: &BAR }.foo,
        Union { u8: &BAR }.bar,
    )
};
static FOO2: (&Foo, &Bar) = unsafe { (std::mem::transmute(&BAR), std::mem::transmute(&BAR)) };
//^ undefined behavior

fn main() {}
<commit_msg>Clarify a statement in UB test<commit_after>// check-pass

// This test exhibits undefined behavior, but it is impossible to prevent generally during
// const eval, even if it could be prevented in the cases here if we added expensive and
// complex checks in rustc.
// The reason it's impossible in general
// is that we run into query cycles even *without* UB, just because we're checking for UB.
// We do not detect it if you create references to statics
// in ways that are UB.

enum Foo {
    A = 5,
    B = 42,
}
enum Bar {
    C = 42,
    D = 99,
}
#[repr(C)]
union Union {
    foo: &'static Foo,
    bar: &'static Bar,
    u8: &'static u8,
}
static BAR: u8 = 5;
static FOO: (&Foo, &Bar) = unsafe {
    (
        // undefined behavior
        Union { u8: &BAR }.foo,
        Union { u8: &BAR }.bar,
    )
};
static FOO2: (&Foo, &Bar) = unsafe { (std::mem::transmute(&BAR), std::mem::transmute(&BAR)) };
//^ undefined behavior

fn main() {}
<|endoftext|>


<commit_before>use std::{
    ops::Deref,
    sync::{
        atomic::{
            AtomicUsize,
            Ordering::{Acquire, Relaxed, Release},
        },
        Arc, Mutex, RwLock,
    },
};

use pagecache::{FastMap8, PagePtr};

use super::*;

/// The `sled` embedded database!
#[derive(Clone)]
pub struct Db {
    config: Config,
    pages: Arc<PageCache<BLinkMaterializer, Frag, Recovery>>,
    idgen: Arc<AtomicUsize>,
    idgen_persists: Arc<AtomicUsize>,
    idgen_persist_mu: Arc<Mutex<()>>,
    tenants: Arc<RwLock<FastMap8<Vec<u8>, Arc<Tree>>>>,
    default: Arc<Tree>,
    transactions: Arc<Tree>,
    was_recovered: bool,
}

unsafe impl Send for Db {}

unsafe impl Sync for Db {}

#[cfg(feature = "event_log")]
impl Drop for Db {
    fn drop(&mut self) {
        let guard = pin();

        self.config.event_log.meta_before_restart(
            meta::meta(&*self.pages, &guard)
                .expect("should get meta under test")
                .clone(),
        );
    }
}

impl Deref for Db {
    type Target = Tree;

    fn deref(&self) -> &Tree {
        &self.default
    }
}

impl Db {
    /// Load existing or create a new `Db` with a default configuration.
    pub fn start_default<P: AsRef<std::path::Path>>(
        path: P,
    ) -> Result<Db, ()> {
        let config = ConfigBuilder::new().path(path).build();
        Self::start(config)
    }

    /// Load existing or create a new `Db`.
    pub fn start(config: Config) -> Result<Db, ()> {
        let _measure = Measure::new(&M.tree_start);

        #[cfg(any(test, feature = "check_snapshot_integrity"))]
        match config
            .verify_snapshot::<BLinkMaterializer, Frag, Recovery>()
        {
            Ok(_) => {}
            #[cfg(feature = "failpoints")]
            Err(Error::FailPoint) => {}
            other => panic!("failed to verify snapshot: {:?}", other),
        }

        let pages = Arc::new(PageCache::start(config.clone())?);

        let recovery: Recovery =
            pages.recovered_state().unwrap_or_default();
        let idgen_recovery =
            recovery.counter + (2 * config.idgen_persist_interval);
        let idgen_persists = recovery.counter
            / config.idgen_persist_interval
            * config.idgen_persist_interval;

        let guard = pin();

        let was_recovered: bool;

        if pages
            .get(META_PID, &guard)
            .map_err(|e| e.danger_cast())?
            .is_unallocated()
        {
            was_recovered = false;

            // set up meta
            let meta_id = pages.allocate(&guard)?;

            assert_eq!(
                meta_id,
                META_PID,
                "we expect the meta page to have pid {}, but it had pid {} instead",
                META_PID,
                meta_id,
            );

            let meta = Frag::Meta(Meta::default());
            pages
                .replace(meta_id, PagePtr::allocated(), meta, &guard)
                .map_err(|e| e.danger_cast())?;

            // set up idgen
            let counter_id = pages.allocate(&guard)?;

            assert_eq!(
                counter_id,
                COUNTER_PID,
                "we expect the counter to have pid {}, but it had pid {} instead",
                COUNTER_PID,
                counter_id,
            );

            let counter = Frag::Counter(0);
            pages
                .replace(
                    counter_id,
                    PagePtr::allocated(),
                    counter,
                    &guard,
                )
                .map_err(|e| e.danger_cast())?;
        } else {
            was_recovered = true;
        };

        let default_tree = meta::open_tree(
            pages.clone(),
            config.clone(),
            DEFAULT_TREE_ID.to_vec(),
            &guard,
        )?;
        let tx_tree = meta::open_tree(
            pages.clone(),
            config.clone(),
            TX_TREE_ID.to_vec(),
            &guard,
        )?;

        let ret = Db {
            pages: pages,
            config,
            idgen: Arc::new(AtomicUsize::new(idgen_recovery)),
            idgen_persists: Arc::new(AtomicUsize::new(
                idgen_persists,
            )),
            idgen_persist_mu: Arc::new(Mutex::new(())),
            tenants: Arc::new(RwLock::new(FastMap8::default())),
            default: Arc::new(default_tree),
            transactions: Arc::new(tx_tree),
            was_recovered,
        };

        #[cfg(feature = "event_log")]
        ret.config.event_log.meta_after_restart(
            meta::meta(&*ret.pages, &guard)
                .expect("should be able to get meta under test")
                .clone(),
        );

        let mut tenants = ret.tenants.write().unwrap();

        for (id, _root) in
            meta::meta(&*ret.pages, &guard)?.tenants().into_iter()
        {
            let tree = Tree {
                tree_id: id.clone(),
                subscriptions: Arc::new(Subscriptions::default()),
                config: ret.config.clone(),
                pages: ret.pages.clone(),
            };
            tenants.insert(id, Arc::new(tree));
        }

        drop(tenants);

        let mut ret = ret;

        ret.default = ret.open_tree(DEFAULT_TREE_ID.to_vec())?;
        ret.transactions = ret.open_tree(TX_TREE_ID.to_vec())?;

        Ok(ret)
    }

    /// Generate a monotonic ID. Not guaranteed to be
    /// contiguous. Written to disk every `idgen_persist_interval`
    /// operations, followed by a blocking flush. During recovery, we
    /// take the last recovered generated ID and add 2x
    /// the `idgen_persist_interval` to it. While persisting, if the
    /// previous persisted counter wasn't synced to disk yet, we will do
    /// a blocking flush to fsync the latest counter, ensuring
    /// that we will never give out the same counter twice.
    pub fn generate_id(&self) -> Result<usize, ()> {
        let ret = self.idgen.fetch_add(1, Relaxed);

        let interval = self.config.idgen_persist_interval;
        let necessary_persists = ret / interval * interval;
        let mut persisted = self.idgen_persists.load(Acquire);

        while persisted < necessary_persists {
            let _mu = self.idgen_persist_mu.lock().unwrap();
            persisted = self.idgen_persists.load(Acquire);
            if persisted < necessary_persists {
                // it's our responsibility to persist up to our ID
                let guard = pin();
                let (current, key) = self
                    .pages
                    .get(COUNTER_PID, &guard)
                    .map_err(|e| e.danger_cast())?
                    .unwrap();

                if let Frag::Counter(current) = current {
                    assert_eq!(*current, persisted);
                } else {
                    panic!(
                        "counter pid contained non-Counter: {:?}",
                        current
                    );
                }

                let counter_frag = Frag::Counter(necessary_persists);

                let old = self
                    .idgen_persists
                    .swap(necessary_persists, Release);
                assert_eq!(old, persisted);

                self.pages
                    .replace(
                        COUNTER_PID,
                        key.clone(),
                        counter_frag,
                        &guard,
                    )
                    .map_err(|e| e.danger_cast())?;

                // during recovery we add 2x the interval. we only
                // need to block if the last one wasn't stable yet.
                if key.last_lsn() > self.pages.stable_lsn() {
                    self.pages.make_stable(key.last_lsn())?;
                }

                guard.flush();
            }
        }

        Ok(ret)
    }

    /// Open or create a new disk-backed Tree with its own keyspace,
    /// accessible from the `Db` via the provided identifier.
    pub fn open_tree(&self, name: Vec<u8>) -> Result<Arc<Tree>, ()> {
        let guard = pin();

        let tenants = self.tenants.read().unwrap();
        if let Some(tree) = tenants.get(&name) {
            return Ok(tree.clone());
        }
        // drop reader lock
        drop(tenants);

        // set up empty leaf
        let leaf_id = self.pages.allocate(&guard)?;
        trace!(
            "allocated pid {} for leaf in open_tree for namespace {:?}",
            leaf_id,
            name,
        );

        let leaf = Frag::Base(Node {
            id: leaf_id,
            data: Data::Leaf(vec![]),
            next: None,
            lo: vec![],
            hi: vec![],
        });

        self.pages
            .replace(leaf_id, PagePtr::allocated(), leaf, &guard)
            .map_err(|e| e.danger_cast())?;

        // set up root index
        let root_id = self.pages.allocate(&guard)?;

        debug!(
            "allocated pid {} for root of new_tree {:?}",
            root_id, name
        );

        // vec![0] represents a prefix-encoded empty prefix
        let root_index_vec = vec![(vec![0], leaf_id)];

        let root = Frag::Base(Node {
            id: root_id,
            data: Data::Index(root_index_vec),
            next: None,
            lo: vec![],
            hi: vec![],
        });

        self.pages
            .replace(root_id, PagePtr::allocated(), root, &guard)
            .map_err(|e| e.danger_cast())?;

        meta::cas_root(
            &*self.pages,
            name.clone(),
            None,
            Some(root_id),
            &guard,
        )
        .map_err(|e| e.danger_cast())?;

        let mut tenants = self.tenants.write().unwrap();
        let tree = Arc::new(Tree {
            tree_id: name.clone(),
            subscriptions: Arc::new(Subscriptions::default()),
            config: self.config.clone(),
            pages: self.pages.clone(),
        });
        tenants.insert(name, tree.clone());

        Ok(tree)
    }

    /// Remove a disk-backed collection.
    pub fn drop_tree(&self, name: &[u8]) -> Result<bool, ()> {
        if name == DEFAULT_TREE_ID || name == TX_TREE_ID {
            return Err(Error::Unsupported(
                "cannot remove the core structures".into(),
            ));
        }
        trace!("dropping tree {:?}", name,);

        let mut tenants = self.tenants.write().unwrap();

        let tree = if let Some(tree) = tenants.remove(&*name) {
            tree
        } else {
            return Ok(false);
        };

        let guard = pin();

        let mut root_id =
            meta::pid_for_name(&*self.pages, &name, &guard)?;

        let leftmost_chain: Vec<PageId> = tree
            .path_for_key(b"", &guard)?
            .into_iter()
            .map(|(frag, _tp)| frag.unwrap_base().id)
            .collect();

        loop {
            let res = meta::cas_root(
                &*self.pages,
                name.to_vec(),
                Some(root_id),
                None,
                &guard,
            )
            .map_err(|e| e.danger_cast());

            match res {
                Ok(_) => break,
                Err(Error::CasFailed(actual)) => root_id = actual,
                Err(other) => return Err(other.danger_cast()),
            }
        }

        // drop writer lock
        drop(tenants);

        guard.defer(move || tree.gc_pages(leftmost_chain));

        guard.flush();

        Ok(true)
    }

    /// Returns `true` if the database was
    /// recovered from a previous process.
    /// Note that database state is only
    /// guaranteed to be present up to the
    /// last call to `flush`! Otherwise state
    /// is synced to disk periodically if the
    /// `sync_every_ms` configuration option
    /// is set to `Some(number_of_ms_between_syncs)`
    /// or if the IO buffer gets filled to
    /// capacity before being rotated.
    pub fn was_recovered(&self) -> bool {
        self.was_recovered
    }

    /// Record a set of pages as being involved in
    /// a transaction that is about to be written to
    /// one of the structures in the Db. Do this
    /// before performing transactions in case
    /// the system crashes before all changes
    /// can be written. During recovery,
    /// you can call `possibly_dirty_pages`
    /// to get the set of pages that may have
    /// transactional state written to them.
    pub fn mark_pending_tx(
        &self,
        _pages: Vec<PageId>,
        _txid: u64,
    ) -> Result<(), ()> {
        unimplemented!()
    }

    /// Mark a transaction as being aborted. This
    /// ensures that during recovery we know that
    /// this transaction is not the last one that
    /// should be recovered, because the process
    /// finished the transaction, but we may need
    /// to clean up its pending state still.
    pub fn abort_tx(&self, _txid: u64) -> Result<(), ()> {
        unimplemented!()
    }

    /// Mark a transaction as being completed. This
    /// ensures that during recovery we know that
    /// this transaction completed, and that we
    /// can safely remove its entry in the
    /// transaction table.
    pub fn commit_tx(&self, _txid: u64) -> Result<(), ()> {
        unimplemented!()
    }

    /// Remove entries associated with this transaction
    /// in the transaction table, because we have
    /// completed the transaction and finished
    /// applying or removing pending state, and do
    /// not need to know about it during crash recovery.
    pub fn cleanup_tx(&self, _txid: u64) -> Result<(), ()> {
        unimplemented!()
    }

    /// Call this during recovery to retrieve the set of
    /// pages that were involved in a transaction
    /// that was in a pending or aborted (but not
    /// yet fully cleaned up) state during process
    /// termination.
    pub fn possibly_dirty_pages(
        &self,
    ) -> Result<Vec<(u64, Vec<PageId>)>, ()> {
        unimplemented!()
    }
}
<commit_msg>Always call PageCache::flush on Db Drop<commit_after>use std::{
    ops::Deref,
    sync::{
        atomic::{
            AtomicUsize,
            Ordering::{Acquire, Relaxed, Release},
        },
        Arc, Mutex, RwLock,
    },
};

use pagecache::{FastMap8, PagePtr};

use super::*;

/// The `sled` embedded database!
#[derive(Clone)]
pub struct Db {
    config: Config,
    pages: Arc<PageCache<BLinkMaterializer, Frag, Recovery>>,
    idgen: Arc<AtomicUsize>,
    idgen_persists: Arc<AtomicUsize>,
    idgen_persist_mu: Arc<Mutex<()>>,
    tenants: Arc<RwLock<FastMap8<Vec<u8>, Arc<Tree>>>>,
    default: Arc<Tree>,
    transactions: Arc<Tree>,
    was_recovered: bool,
}

unsafe impl Send for Db {}

unsafe impl Sync for Db {}

impl Drop for Db {
    fn drop(&mut self) {
        if let Err(e) = self.pages.flush() {
            error!(
                "failed to flush underlying \
                 pagecache during drop: {:?}",
                e
            );
        }

        #[cfg(feature = "event_log")]
        {
            let guard = pin();

            self.config.event_log.meta_before_restart(
                meta::meta(&*self.pages, &guard)
                    .expect("should get meta under test")
                    .clone(),
            );
        }
    }
}

impl Deref for Db {
    type Target = Tree;

    fn deref(&self) -> &Tree {
        &self.default
    }
}

impl Db {
    /// Load existing or create a new `Db` with a default configuration.
    pub fn start_default<P: AsRef<std::path::Path>>(
        path: P,
    ) -> Result<Db, ()> {
        let config = ConfigBuilder::new().path(path).build();
        Self::start(config)
    }

    /// Load existing or create a new `Db`.
    pub fn start(config: Config) -> Result<Db, ()> {
        let _measure = Measure::new(&M.tree_start);

        #[cfg(any(test, feature = "check_snapshot_integrity"))]
        match config
            .verify_snapshot::<BLinkMaterializer, Frag, Recovery>()
        {
            Ok(_) => {}
            #[cfg(feature = "failpoints")]
            Err(Error::FailPoint) => {}
            other => panic!("failed to verify snapshot: {:?}", other),
        }

        let pages = Arc::new(PageCache::start(config.clone())?);

        let recovery: Recovery =
            pages.recovered_state().unwrap_or_default();
        let idgen_recovery =
            recovery.counter + (2 * config.idgen_persist_interval);
        let idgen_persists = recovery.counter
            / config.idgen_persist_interval
            * config.idgen_persist_interval;

        let guard = pin();

        let was_recovered: bool;

        if pages
            .get(META_PID, &guard)
            .map_err(|e| e.danger_cast())?
            .is_unallocated()
        {
            was_recovered = false;

            // set up meta
            let meta_id = pages.allocate(&guard)?;

            assert_eq!(
                meta_id,
                META_PID,
                "we expect the meta page to have pid {}, but it had pid {} instead",
                META_PID,
                meta_id,
            );

            let meta = Frag::Meta(Meta::default());
            pages
                .replace(meta_id, PagePtr::allocated(), meta, &guard)
                .map_err(|e| e.danger_cast())?;

            // set up idgen
            let counter_id = pages.allocate(&guard)?;

            assert_eq!(
                counter_id,
                COUNTER_PID,
                "we expect the counter to have pid {}, but it had pid {} instead",
                COUNTER_PID,
                counter_id,
            );

            let counter = Frag::Counter(0);
            pages
                .replace(
                    counter_id,
                    PagePtr::allocated(),
                    counter,
                    &guard,
                )
                .map_err(|e| e.danger_cast())?;
        } else {
            was_recovered = true;
        };

        let default_tree = meta::open_tree(
            pages.clone(),
            config.clone(),
            DEFAULT_TREE_ID.to_vec(),
            &guard,
        )?;
        let tx_tree = meta::open_tree(
            pages.clone(),
            config.clone(),
            TX_TREE_ID.to_vec(),
            &guard,
        )?;

        let ret = Db {
            pages: pages,
            config,
            idgen: Arc::new(AtomicUsize::new(idgen_recovery)),
            idgen_persists: Arc::new(AtomicUsize::new(
                idgen_persists,
            )),
            idgen_persist_mu: Arc::new(Mutex::new(())),
            tenants: Arc::new(RwLock::new(FastMap8::default())),
            default: Arc::new(default_tree),
            transactions: Arc::new(tx_tree),
            was_recovered,
        };

        #[cfg(feature = "event_log")]
        ret.config.event_log.meta_after_restart(
            meta::meta(&*ret.pages, &guard)
                .expect("should be able to get meta under test")
                .clone(),
        );

        let mut tenants = ret.tenants.write().unwrap();

        for (id, _root) in
            meta::meta(&*ret.pages, &guard)?.tenants().into_iter()
        {
            let tree = Tree {
                tree_id: id.clone(),
                subscriptions: Arc::new(Subscriptions::default()),
                config: ret.config.clone(),
                pages: ret.pages.clone(),
            };
            tenants.insert(id, Arc::new(tree));
        }

        drop(tenants);

        let mut ret = ret;

        ret.default = ret.open_tree(DEFAULT_TREE_ID.to_vec())?;
        ret.transactions = ret.open_tree(TX_TREE_ID.to_vec())?;

        Ok(ret)
    }

    /// Generate a monotonic ID. Not guaranteed to be
    /// contiguous. Written to disk every `idgen_persist_interval`
    /// operations, followed by a blocking flush. During recovery, we
    /// take the last recovered generated ID and add 2x
    /// the `idgen_persist_interval` to it. While persisting, if the
    /// previous persisted counter wasn't synced to disk yet, we will do
    /// a blocking flush to fsync the latest counter, ensuring
    /// that we will never give out the same counter twice.
    pub fn generate_id(&self) -> Result<usize, ()> {
        let ret = self.idgen.fetch_add(1, Relaxed);

        let interval = self.config.idgen_persist_interval;
        let necessary_persists = ret / interval * interval;
        let mut persisted = self.idgen_persists.load(Acquire);

        while persisted < necessary_persists {
            let _mu = self.idgen_persist_mu.lock().unwrap();
            persisted = self.idgen_persists.load(Acquire);
            if persisted < necessary_persists {
                // it's our responsibility to persist up to our ID
                let guard = pin();
                let (current, key) = self
                    .pages
                    .get(COUNTER_PID, &guard)
                    .map_err(|e| e.danger_cast())?
                    .unwrap();

                if let Frag::Counter(current) = current {
                    assert_eq!(*current, persisted);
                } else {
                    panic!(
                        "counter pid contained non-Counter: {:?}",
                        current
                    );
                }

                let counter_frag = Frag::Counter(necessary_persists);

                let old = self
                    .idgen_persists
                    .swap(necessary_persists, Release);
                assert_eq!(old, persisted);

                self.pages
                    .replace(
                        COUNTER_PID,
                        key.clone(),
                        counter_frag,
                        &guard,
                    )
                    .map_err(|e| e.danger_cast())?;

                // during recovery we add 2x the interval. we only
                // need to block if the last one wasn't stable yet.
                if key.last_lsn() > self.pages.stable_lsn() {
                    self.pages.make_stable(key.last_lsn())?;
                }

                guard.flush();
            }
        }

        Ok(ret)
    }

    /// Open or create a new disk-backed Tree with its own keyspace,
    /// accessible from the `Db` via the provided identifier.
    pub fn open_tree(&self, name: Vec<u8>) -> Result<Arc<Tree>, ()> {
        let guard = pin();

        let tenants = self.tenants.read().unwrap();
        if let Some(tree) = tenants.get(&name) {
            return Ok(tree.clone());
        }
        // drop reader lock
        drop(tenants);

        // set up empty leaf
        let leaf_id = self.pages.allocate(&guard)?;
        trace!(
            "allocated pid {} for leaf in open_tree for namespace {:?}",
            leaf_id,
            name,
        );

        let leaf = Frag::Base(Node {
            id: leaf_id,
            data: Data::Leaf(vec![]),
            next: None,
            lo: vec![],
            hi: vec![],
        });

        self.pages
            .replace(leaf_id, PagePtr::allocated(), leaf, &guard)
            .map_err(|e| e.danger_cast())?;

        // set up root index
        let root_id = self.pages.allocate(&guard)?;

        debug!(
            "allocated pid {} for root of new_tree {:?}",
            root_id, name
        );

        // vec![0] represents a prefix-encoded empty prefix
        let root_index_vec = vec![(vec![0], leaf_id)];

        let root = Frag::Base(Node {
            id: root_id,
            data: Data::Index(root_index_vec),
            next: None,
            lo: vec![],
            hi: vec![],
        });

        self.pages
            .replace(root_id, PagePtr::allocated(), root, &guard)
            .map_err(|e| e.danger_cast())?;

        meta::cas_root(
            &*self.pages,
            name.clone(),
            None,
            Some(root_id),
            &guard,
        )
        .map_err(|e| e.danger_cast())?;

        let mut tenants = self.tenants.write().unwrap();
        let tree = Arc::new(Tree {
            tree_id: name.clone(),
            subscriptions: Arc::new(Subscriptions::default()),
            config: self.config.clone(),
            pages: self.pages.clone(),
        });
        tenants.insert(name, tree.clone());

        Ok(tree)
    }

    /// Remove a disk-backed collection.
    pub fn drop_tree(&self, name: &[u8]) -> Result<bool, ()> {
        if name == DEFAULT_TREE_ID || name == TX_TREE_ID {
            return Err(Error::Unsupported(
                "cannot remove the core structures".into(),
            ));
        }
        trace!("dropping tree {:?}", name,);

        let mut tenants = self.tenants.write().unwrap();

        let tree = if let Some(tree) = tenants.remove(&*name) {
            tree
        } else {
            return Ok(false);
        };

        let guard = pin();

        let mut root_id =
            meta::pid_for_name(&*self.pages, &name, &guard)?;

        let leftmost_chain: Vec<PageId> = tree
            .path_for_key(b"", &guard)?
            .into_iter()
            .map(|(frag, _tp)| frag.unwrap_base().id)
            .collect();

        loop {
            let res = meta::cas_root(
                &*self.pages,
                name.to_vec(),
                Some(root_id),
                None,
                &guard,
            )
            .map_err(|e| e.danger_cast());

            match res {
                Ok(_) => break,
                Err(Error::CasFailed(actual)) => root_id = actual,
                Err(other) => return Err(other.danger_cast()),
            }
        }

        // drop writer lock
        drop(tenants);

        guard.defer(move || tree.gc_pages(leftmost_chain));

        guard.flush();

        Ok(true)
    }

    /// Returns `true` if the database was
    /// recovered from a previous process.
    /// Note that database state is only
    /// guaranteed to be present up to the
    /// last call to `flush`! Otherwise state
    /// is synced to disk periodically if the
    /// `sync_every_ms` configuration option
    /// is set to `Some(number_of_ms_between_syncs)`
    /// or if the IO buffer gets filled to
    /// capacity before being rotated.
    pub fn was_recovered(&self) -> bool {
        self.was_recovered
    }

    /// Record a set of pages as being involved in
    /// a transaction that is about to be written to
    /// one of the structures in the Db. Do this
    /// before performing transactions in case
    /// the system crashes before all changes
    /// can be written. During recovery,
    /// you can call `possibly_dirty_pages`
    /// to get the set of pages that may have
    /// transactional state written to them.
    pub fn mark_pending_tx(
        &self,
        _pages: Vec<PageId>,
        _txid: u64,
    ) -> Result<(), ()> {
        unimplemented!()
    }

    /// Mark a transaction as being aborted. This
    /// ensures that during recovery we know that
    /// this transaction is not the last one that
    /// should be recovered, because the process
    /// finished the transaction, but we may need
    /// to clean up its pending state still.
    pub fn abort_tx(&self, _txid: u64) -> Result<(), ()> {
        unimplemented!()
    }

    /// Mark a transaction as being completed. This
    /// ensures that during recovery we know that
    /// this transaction completed, and that we
    /// can safely remove its entry in the
    /// transaction table.
    pub fn commit_tx(&self, _txid: u64) -> Result<(), ()> {
        unimplemented!()
    }

    /// Remove entries associated with this transaction
    /// in the transaction table, because we have
    /// completed the transaction and finished
    /// applying or removing pending state, and do
    /// not need to know about it during crash recovery.
    pub fn cleanup_tx(&self, _txid: u64) -> Result<(), ()> {
        unimplemented!()
    }

    /// Call this during recovery to retrieve the set of
    /// pages that were involved in a transaction
    /// that was in a pending or aborted (but not
    /// yet fully cleaned up) state during process
    /// termination.
    pub fn possibly_dirty_pages(
        &self,
    ) -> Result<Vec<(u64, Vec<PageId>)>, ()> {
        unimplemented!()
    }
}
<|endoftext|>


<commit_before>use data::*;
use util::*;
use secstr::SecStr;
use rusterpassword::*;
use rustc_serialize::base64::{ToBase64, STANDARD};
use sodiumoxide::crypto::sign::ed25519;
use byteorder::{BigEndian, WriteBytesExt};
#[cfg(all(unix, not(target_os = "android"), not(target_os = "ios")))] use unix_socket::UnixStream;
#[cfg(all(unix, not(target_os = "android"), not(target_os = "ios")))] use std::env;
#[cfg(all(unix, not(target_os = "android"), not(target_os = "ios")))] use std::io::{Read, Write};
#[cfg(all(unix, not(target_os = "android"), not(target_os = "ios")))] use std::net::Shutdown;

pub enum Output {
    PrivateText(SecStr),
    OpenText(String),
    PrivateBinary(SecStr),
    Ed25519Keypair(Ed25519Usage, ed25519::PublicKey, ed25519::SecretKey),
}

fn pick_tpl(tpl: &PasswordTemplate) -> &'static [&'static str] {
    match *tpl {
       PasswordTemplate::Maximum  => TEMPLATES_MAXIMUM,
       PasswordTemplate::Long     => TEMPLATES_LONG,
       PasswordTemplate::Medium   => TEMPLATES_MEDIUM,
       PasswordTemplate::Short    => TEMPLATES_SHORT,
       PasswordTemplate::Basic    => TEMPLATES_BASIC,
       PasswordTemplate::Pin      => TEMPLATES_PIN,
    }
}

pub fn process_output(entry_name: &str, master_key: &SecStr, field: &Field) -> Result<Output> {
    match *field {
        Field::Derived { counter, ref site_name, ref usage } => {
            let site_seed = try!(gen_site_seed(master_key, &site_name.clone().unwrap_or(entry_name.to_string()), counter)
                                 .map_err(|_| Error::SeedGenerationError));
            match *usage {
                DerivedUsage::Password(ref tpl) =>
                    Ok(Output::PrivateText(gen_site_password(&site_seed, pick_tpl(tpl)))),
                DerivedUsage::Ed25519Key(ref keyusage) => {
                    let edseed = try!(ed25519::Seed::from_slice(site_seed.unsecure()).ok_or(Error::WrongDerivedKeyLength));
                    let (pubkey, seckey) = ed25519::keypair_from_seed(&edseed);
                    Ok(Output::Ed25519Keypair(*keyusage, pubkey, seckey))
                },
                DerivedUsage::RawKey =>
                    Ok(Output::PrivateBinary(site_seed)),
            }
        },
        Field::Stored { ref data, ref usage } =>
            match *usage {
                StoredUsage::Password =>
                    Ok(Output::PrivateText(SecStr::new(data.unsecure().to_vec()))),
                StoredUsage::Text =>
                    Ok(Output::OpenText(try!(String::from_utf8(data.unsecure().to_vec()))))
            }
    }
}

pub fn ssh_public_key_output(keypair: &Output, comment: &str) -> Result<String> {
    if let &Output::Ed25519Keypair(Ed25519Usage::SSH, ed25519::PublicKey(pubkey_bytes), _) = keypair {
        let mut raw = vec![];
        try!(raw.write_u32::<BigEndian>(11));
        raw.extend(b"ssh-ed25519");
        try!(raw.write_u32::<BigEndian>(ed25519::PUBLICKEYBYTES as u32));
        raw.extend(&pubkey_bytes);
        Ok("ssh-ed25519 ".to_string() + &raw.to_base64(STANDARD) + " " + comment)
    } else { Err(Error::InappropriateFormat) }
}

pub fn ssh_private_key_agent_message(keypair: &Output, comment: &str) -> Result<SecStr> {
    if let &Output::Ed25519Keypair(Ed25519Usage::SSH, ed25519::PublicKey(pubkey_bytes), ed25519::SecretKey(seckey_bytes)) = keypair {
        let mut msg = vec![17u8];
        try!(msg.write_u32::<BigEndian>(11));
        msg.extend(b"ssh-ed25519");
        try!(msg.write_u32::<BigEndian>(ed25519::PUBLICKEYBYTES as u32));
        msg.extend(&pubkey_bytes);
        try!(msg.write_u32::<BigEndian>(ed25519::SECRETKEYBYTES as u32));
        msg.extend(seckey_bytes.iter()); // LOL, there's no iterator for &[u8, 64] because 64 is a lot
        try!(msg.write_u32::<BigEndian>(comment.as_bytes().len() as u32));
        msg.extend(comment.as_bytes());
        Ok(SecStr::new(msg))
    } else { Err(Error::InappropriateFormat) }
}

#[cfg(all(unix, not(target_os = "android"), not(target_os = "ios")))]
pub fn ssh_agent_send_message(msg: SecStr) -> Result<()> {
    if let Some(sock_path) = env::var_os("SSH_AUTH_SOCK") {
        let mut stream = try!(UnixStream::connect(sock_path));
        try!(stream.write_u32::<BigEndian>(msg.unsecure().len() as u32));
        try!(stream.write(msg.unsecure()));
        let mut response = vec![0; 5];
        try!(stream.read(&mut response));
        try!(stream.shutdown(Shutdown::Both));
        Ok(())
    } else { Err(Error::SSHAgentSocketNotFound) }
}

#[cfg(any(not(unix), target_os = "android", target_os = "ios"))]
pub fn ssh_agent_send_message(_: SecStr) -> Result<()> {
    Err(Error::NotAvailableOnPlatform)
}

pub fn signify_keynum(pubkey_bytes: &[u8]) -> Vec<u8> {
    blake2b(pubkey_bytes, b"freepass.signify", 8)
}

pub fn signify_public_key_output(keypair: &Output, comment: &str) -> Result<String> {
    if let &Output::Ed25519Keypair(Ed25519Usage::Signify, ed25519::PublicKey(pubkey_bytes), _) = keypair {
        let mut raw = vec![];
        raw.extend(b"Ed");
        raw.extend(&signify_keynum(&pubkey_bytes));
        raw.extend(&pubkey_bytes);
        Ok("untrusted comment: ".to_string() + comment + "\n" + &raw.to_base64(STANDARD))
    } else { Err(Error::InappropriateFormat) }
}

pub fn signify_sign(keypair: &Output, comment: &str, data: &[u8]) -> Result<String> {
    if let &Output::Ed25519Keypair(Ed25519Usage::Signify, ed25519::PublicKey(pubkey_bytes), ref seckey) = keypair {
        let ed25519::Signature(sig_bytes) = ed25519::sign_detached(data, &seckey);
        let mut raw = vec![];
        raw.extend(b"Ed");
        raw.extend(&signify_keynum(&pubkey_bytes));
        raw.extend(sig_bytes.iter());
        Ok("untrusted comment: ".to_string() + comment + "\n" + &raw.to_base64(STANDARD))
    } else { Err(Error::InappropriateFormat) }
}

#[cfg(test)]
mod tests {
    use super::*;
    use data::*;
    use sodiumoxide::crypto::sign::ed25519;

    fn keypair() -> Output {
        let (pubkey, seckey) = ed25519::keypair_from_seed(&ed25519::Seed::from_slice(b"aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa").unwrap());
        Output::Ed25519Keypair(Ed25519Usage::SSH, pubkey, seckey)
    }

    #[test]
    fn test_ssh_public_key_output() {
        assert_eq!(ssh_public_key_output(&keypair(), "myComment").unwrap(), "ssh-ed25519 AAAAC3NzaC1lZDI1NTE5AAAAIK8Go+MpFxTk81bBnJsVzRlR7G5mYqp3vgdUfyiTgzQd myComment");
    }

    #[test]
    fn test_ssh_private_key_agent_message() {
        let msg = ssh_private_key_agent_message(&keypair(), "myComment").unwrap();
        assert_eq!(Vec::from(msg.unsecure()), vec![17, 0, 0, 0, 11, 115, 115, 104, 45, 101, 100, 50, 53, 53, 49, 57, 0, 0, 0, 32, 175, 6, 163, 227, 41, 23, 20, 228, 243, 86, 193, 156, 155, 21, 205, 25, 81, 236, 110, 102, 98, 170, 119, 190, 7, 84, 127, 40, 147, 131, 52, 29, 0, 0, 0, 64, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 175, 6, 163, 227, 41, 23, 20, 228, 243, 86, 193, 156, 155, 21, 205, 25, 81, 236, 110, 102, 98, 170, 119, 190, 7, 84, 127, 40, 147, 131, 52, 29, 0, 0, 0, 9, 109, 121, 67, 111, 109, 109, 101, 110, 116]);
    }
}
<commit_msg>signify tests<commit_after>use data::*;
use util::*;
use secstr::SecStr;
use rusterpassword::*;
use rustc_serialize::base64::{ToBase64, STANDARD};
use sodiumoxide::crypto::sign::ed25519;
use byteorder::{BigEndian, WriteBytesExt};
#[cfg(all(unix, not(target_os = "android"), not(target_os = "ios")))] use unix_socket::UnixStream;
#[cfg(all(unix, not(target_os = "android"), not(target_os = "ios")))] use std::env;
#[cfg(all(unix, not(target_os = "android"), not(target_os = "ios")))] use std::io::{Read, Write};
#[cfg(all(unix, not(target_os = "android"), not(target_os = "ios")))] use std::net::Shutdown;

pub enum Output {
    PrivateText(SecStr),
    OpenText(String),
    PrivateBinary(SecStr),
    Ed25519Keypair(Ed25519Usage, ed25519::PublicKey, ed25519::SecretKey),
}

fn pick_tpl(tpl: &PasswordTemplate) -> &'static [&'static str] {
    match *tpl {
       PasswordTemplate::Maximum  => TEMPLATES_MAXIMUM,
       PasswordTemplate::Long     => TEMPLATES_LONG,
       PasswordTemplate::Medium   => TEMPLATES_MEDIUM,
       PasswordTemplate::Short    => TEMPLATES_SHORT,
       PasswordTemplate::Basic    => TEMPLATES_BASIC,
       PasswordTemplate::Pin      => TEMPLATES_PIN,
    }
}

pub fn process_output(entry_name: &str, master_key: &SecStr, field: &Field) -> Result<Output> {
    match *field {
        Field::Derived { counter, ref site_name, ref usage } => {
            let site_seed = try!(gen_site_seed(master_key, &site_name.clone().unwrap_or(entry_name.to_string()), counter)
                                 .map_err(|_| Error::SeedGenerationError));
            match *usage {
                DerivedUsage::Password(ref tpl) =>
                    Ok(Output::PrivateText(gen_site_password(&site_seed, pick_tpl(tpl)))),
                DerivedUsage::Ed25519Key(ref keyusage) => {
                    let edseed = try!(ed25519::Seed::from_slice(site_seed.unsecure()).ok_or(Error::WrongDerivedKeyLength));
                    let (pubkey, seckey) = ed25519::keypair_from_seed(&edseed);
                    Ok(Output::Ed25519Keypair(*keyusage, pubkey, seckey))
                },
                DerivedUsage::RawKey =>
                    Ok(Output::PrivateBinary(site_seed)),
            }
        },
        Field::Stored { ref data, ref usage } =>
            match *usage {
                StoredUsage::Password =>
                    Ok(Output::PrivateText(SecStr::new(data.unsecure().to_vec()))),
                StoredUsage::Text =>
                    Ok(Output::OpenText(try!(String::from_utf8(data.unsecure().to_vec()))))
            }
    }
}

pub fn ssh_public_key_output(keypair: &Output, comment: &str) -> Result<String> {
    if let &Output::Ed25519Keypair(Ed25519Usage::SSH, ed25519::PublicKey(pubkey_bytes), _) = keypair {
        let mut raw = vec![];
        try!(raw.write_u32::<BigEndian>(11));
        raw.extend(b"ssh-ed25519");
        try!(raw.write_u32::<BigEndian>(ed25519::PUBLICKEYBYTES as u32));
        raw.extend(&pubkey_bytes);
        Ok("ssh-ed25519 ".to_string() + &raw.to_base64(STANDARD) + " " + comment)
    } else { Err(Error::InappropriateFormat) }
}

pub fn ssh_private_key_agent_message(keypair: &Output, comment: &str) -> Result<SecStr> {
    if let &Output::Ed25519Keypair(Ed25519Usage::SSH, ed25519::PublicKey(pubkey_bytes), ed25519::SecretKey(seckey_bytes)) = keypair {
        let mut msg = vec![17u8];
        try!(msg.write_u32::<BigEndian>(11));
        msg.extend(b"ssh-ed25519");
        try!(msg.write_u32::<BigEndian>(ed25519::PUBLICKEYBYTES as u32));
        msg.extend(&pubkey_bytes);
        try!(msg.write_u32::<BigEndian>(ed25519::SECRETKEYBYTES as u32));
        msg.extend(seckey_bytes.iter()); // LOL, there's no iterator for &[u8, 64] because 64 is a lot
        try!(msg.write_u32::<BigEndian>(comment.as_bytes().len() as u32));
        msg.extend(comment.as_bytes());
        Ok(SecStr::new(msg))
    } else { Err(Error::InappropriateFormat) }
}

#[cfg(all(unix, not(target_os = "android"), not(target_os = "ios")))]
pub fn ssh_agent_send_message(msg: SecStr) -> Result<()> {
    if let Some(sock_path) = env::var_os("SSH_AUTH_SOCK") {
        let mut stream = try!(UnixStream::connect(sock_path));
        try!(stream.write_u32::<BigEndian>(msg.unsecure().len() as u32));
        try!(stream.write(msg.unsecure()));
        let mut response = vec![0; 5];
        try!(stream.read(&mut response));
        try!(stream.shutdown(Shutdown::Both));
        Ok(())
    } else { Err(Error::SSHAgentSocketNotFound) }
}

#[cfg(any(not(unix), target_os = "android", target_os = "ios"))]
pub fn ssh_agent_send_message(_: SecStr) -> Result<()> {
    Err(Error::NotAvailableOnPlatform)
}

pub fn signify_keynum(pubkey_bytes: &[u8]) -> Vec<u8> {
    blake2b(pubkey_bytes, b"freepass.signify", 8)
}

pub fn signify_public_key_output(keypair: &Output, comment: &str) -> Result<String> {
    if let &Output::Ed25519Keypair(Ed25519Usage::Signify, ed25519::PublicKey(pubkey_bytes), _) = keypair {
        let mut raw = vec![];
        raw.extend(b"Ed");
        raw.extend(&signify_keynum(&pubkey_bytes));
        raw.extend(&pubkey_bytes);
        Ok("untrusted comment: ".to_string() + comment + "\n" + &raw.to_base64(STANDARD))
    } else { Err(Error::InappropriateFormat) }
}

pub fn signify_sign(keypair: &Output, comment: &str, data: &[u8]) -> Result<String> {
    if let &Output::Ed25519Keypair(Ed25519Usage::Signify, ed25519::PublicKey(pubkey_bytes), ref seckey) = keypair {
        let ed25519::Signature(sig_bytes) = ed25519::sign_detached(data, &seckey);
        let mut raw = vec![];
        raw.extend(b"Ed");
        raw.extend(&signify_keynum(&pubkey_bytes));
        raw.extend(sig_bytes.iter());
        Ok("untrusted comment: ".to_string() + comment + "\n" + &raw.to_base64(STANDARD))
    } else { Err(Error::InappropriateFormat) }
}

#[cfg(test)]
mod tests {
    use super::*;
    use data::*;
    use sodiumoxide::crypto::sign::ed25519;

    fn keypair(usage: Ed25519Usage) -> Output {
        let (pubkey, seckey) = ed25519::keypair_from_seed(&ed25519::Seed::from_slice(b"aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa").unwrap());
        Output::Ed25519Keypair(usage, pubkey, seckey)
    }

    #[test]
    fn test_ssh_public_key_output() {
        assert_eq!(ssh_public_key_output(&keypair(Ed25519Usage::SSH), "myComment").unwrap(), "ssh-ed25519 AAAAC3NzaC1lZDI1NTE5AAAAIK8Go+MpFxTk81bBnJsVzRlR7G5mYqp3vgdUfyiTgzQd myComment");
    }

    #[test]
    fn test_ssh_private_key_agent_message() {
        let msg = ssh_private_key_agent_message(&keypair(Ed25519Usage::SSH), "myComment").unwrap();
        assert_eq!(Vec::from(msg.unsecure()), vec![17, 0, 0, 0, 11, 115, 115, 104, 45, 101, 100, 50, 53, 53, 49, 57, 0, 0, 0, 32, 175, 6, 163, 227, 41, 23, 20, 228, 243, 86, 193, 156, 155, 21, 205, 25, 81, 236, 110, 102, 98, 170, 119, 190, 7, 84, 127, 40, 147, 131, 52, 29, 0, 0, 0, 64, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 175, 6, 163, 227, 41, 23, 20, 228, 243, 86, 193, 156, 155, 21, 205, 25, 81, 236, 110, 102, 98, 170, 119, 190, 7, 84, 127, 40, 147, 131, 52, 29, 0, 0, 0, 9, 109, 121, 67, 111, 109, 109, 101, 110, 116]);
    }

    #[test]
    fn test_signify_public_key_output() {
        assert_eq!(signify_public_key_output(&keypair(Ed25519Usage::Signify), "myComment").unwrap(), "untrusted comment: myComment\nRWTvU+MamKpB6a8Go+MpFxTk81bBnJsVzRlR7G5mYqp3vgdUfyiTgzQd");
    }

    #[test]
    fn test_signify_sign() {
        assert_eq!(signify_sign(&keypair(Ed25519Usage::Signify), "myComment", b"hello world\n").unwrap(), "untrusted comment: myComment\nRWTvU+MamKpB6auI/gtU4NWNehjKQjhXsPp15mrjGTC1TfZ4SHESosQOOFkOpd+UPSUsoMac2pA5NEi+u5oHwLkKV9UeLm69JAI=");
    }
}
<|endoftext|>