<commit_before>extern crate actix;
extern crate actix_http;
extern crate actix_net;
extern crate actix_web;
extern crate futures;

use std::{io::Read, io::Write, net, thread, time};

use actix::System;
use actix_net::server::Server;
use actix_web::{client, test};
use futures::future;

use actix_http::{h1, Error, KeepAlive, Response, ServiceConfig};

#[test]
fn test_h1_v2() {
    let addr = test::TestServer::unused_addr();
    thread::spawn(move || {
        Server::new()
            .bind("test", addr, move || {
                let settings = ServiceConfig::build()
                    .keep_alive(KeepAlive::Disabled)
                    .client_timeout(1000)
                    .client_disconnect(1000)
                    .server_hostname("localhost")
                    .server_address(addr)
                    .finish();

                h1::H1Service::new(settings, |_| {
                    future::ok::<_, Error>(Response::Ok().finish())
                })
            }).unwrap()
            .run();
    });

    let mut sys = System::new("test");
    {
        let req = client::ClientRequest::get(format!("http://{}/", addr).as_str())
            .finish()
            .unwrap();
        let response = sys.block_on(req.send()).unwrap();
        assert!(response.status().is_success());
    }
}

#[test]
fn test_slow_request() {
    let addr = test::TestServer::unused_addr();
    thread::spawn(move || {
        Server::new()
            .bind("test", addr, move || {
                let settings = ServiceConfig::build().client_timeout(100).finish();

                h1::H1Service::new(settings, |_| {
                    future::ok::<_, Error>(Response::Ok().finish())
                })
            }).unwrap()
            .run();
    });
    thread::sleep(time::Duration::from_millis(100));

    let mut stream = net::TcpStream::connect(addr).unwrap();
    let _ = stream.write_all(b"GET /test/tests/test HTTP/1.1\r\n");
    let mut data = String::new();
    let _ = stream.read_to_string(&mut data);
    assert!(data.starts_with("HTTP/1.1 408 Request Timeout"));
}

#[test]
fn test_malformed_request() {
    let addr = test::TestServer::unused_addr();
    thread::spawn(move || {
        Server::new()
            .bind("test", addr, move || {
                let settings = ServiceConfig::build().client_timeout(100).finish();
                h1::H1Service::new(settings, |_| {
                    future::ok::<_, Error>(Response::Ok().finish())
                })
            }).unwrap()
            .run();
    });
    thread::sleep(time::Duration::from_millis(100));

    let mut stream = net::TcpStream::connect(addr).unwrap();
    let _ = stream.write_all(b"GET /test/tests/test HTTP1.1\r\n");
    let mut data = String::new();
    let _ = stream.read_to_string(&mut data);
    assert!(data.starts_with("HTTP/1.1 400 Bad Request"));
}
<commit_msg>add content-length test<commit_after>extern crate actix;
extern crate actix_http;
extern crate actix_net;
extern crate actix_web;
extern crate futures;

use std::{io::Read, io::Write, net, thread, time};

use actix::System;
use actix_net::server::Server;
use actix_web::{client, test, HttpMessage};
use futures::future;

use actix_http::{h1, Error, KeepAlive, Request, Response, ServiceConfig};

#[test]
fn test_h1_v2() {
    let addr = test::TestServer::unused_addr();
    thread::spawn(move || {
        Server::new()
            .bind("test", addr, move || {
                let settings = ServiceConfig::build()
                    .keep_alive(KeepAlive::Disabled)
                    .client_timeout(1000)
                    .client_disconnect(1000)
                    .server_hostname("localhost")
                    .server_address(addr)
                    .finish();

                h1::H1Service::new(settings, |_| {
                    future::ok::<_, Error>(Response::Ok().finish())
                })
            }).unwrap()
            .run();
    });

    let mut sys = System::new("test");
    {
        let req = client::ClientRequest::get(format!("http://{}/", addr).as_str())
            .finish()
            .unwrap();
        let response = sys.block_on(req.send()).unwrap();
        assert!(response.status().is_success());
    }
}

#[test]
fn test_slow_request() {
    let addr = test::TestServer::unused_addr();
    thread::spawn(move || {
        Server::new()
            .bind("test", addr, move || {
                let settings = ServiceConfig::build().client_timeout(100).finish();

                h1::H1Service::new(settings, |_| {
                    future::ok::<_, Error>(Response::Ok().finish())
                })
            }).unwrap()
            .run();
    });
    thread::sleep(time::Duration::from_millis(100));

    let mut stream = net::TcpStream::connect(addr).unwrap();
    let _ = stream.write_all(b"GET /test/tests/test HTTP/1.1\r\n");
    let mut data = String::new();
    let _ = stream.read_to_string(&mut data);
    assert!(data.starts_with("HTTP/1.1 408 Request Timeout"));
}

#[test]
fn test_malformed_request() {
    let addr = test::TestServer::unused_addr();
    thread::spawn(move || {
        Server::new()
            .bind("test", addr, move || {
                let settings = ServiceConfig::build().client_timeout(100).finish();
                h1::H1Service::new(settings, |_| {
                    future::ok::<_, Error>(Response::Ok().finish())
                })
            }).unwrap()
            .run();
    });
    thread::sleep(time::Duration::from_millis(100));

    let mut stream = net::TcpStream::connect(addr).unwrap();
    let _ = stream.write_all(b"GET /test/tests/test HTTP1.1\r\n");
    let mut data = String::new();
    let _ = stream.read_to_string(&mut data);
    assert!(data.starts_with("HTTP/1.1 400 Bad Request"));
}

#[test]
fn test_content_length() {
    use actix_http::http::{
        header::{HeaderName, HeaderValue},
        StatusCode,
    };

    let addr = test::TestServer::unused_addr();
    thread::spawn(move || {
        Server::new()
            .bind("test", addr, move || {
                let settings = ServiceConfig::build().client_timeout(100).finish();
                h1::H1Service::new(settings, |req: Request| {
                    let indx: usize = req.uri().path()[1..].parse().unwrap();
                    let statuses = [
                        StatusCode::NO_CONTENT,
                        StatusCode::CONTINUE,
                        StatusCode::SWITCHING_PROTOCOLS,
                        StatusCode::PROCESSING,
                        StatusCode::OK,
                        StatusCode::NOT_FOUND,
                    ];
                    future::ok::<_, Error>(Response::new(statuses[indx]))
                })
            }).unwrap()
            .run();
    });
    thread::sleep(time::Duration::from_millis(100));

    let header = HeaderName::from_static("content-length");
    let value = HeaderValue::from_static("0");

    let mut sys = System::new("test");
    {
        for i in 0..4 {
            let req =
                client::ClientRequest::get(format!("http://{}/{}", addr, i).as_str())
                    .finish()
                    .unwrap();
            let response = sys.block_on(req.send()).unwrap();
            assert_eq!(response.headers().get(&header), None);
        }

        for i in 4..6 {
            let req =
                client::ClientRequest::get(format!("http://{}/{}", addr, i).as_str())
                    .finish()
                    .unwrap();
            let response = sys.block_on(req.send()).unwrap();
            assert_eq!(response.headers().get(&header), Some(&value));
        }
    }
}
<|endoftext|>


<commit_before>#![feature(core)]

// This code is editable and runnable!
fn main() {
    // A simple integer calculator:
    // `+` or `-` means add or subtract by 1
    // `*` or `/` means multiply or divide by 2

    let program = "+ + * - /";
    let mut accumulator = 0;

    for token in program.chars() {
        match token {
            '+' => accumulator += 1,
            '-' => accumulator -= 1,
            '*' => accumulator *= 2,
            '/' => accumulator /= 2,
            _ => { /* ignore everything else */ }
        }
    }

    println!("The program \"{}\" calculates the value {}",
              program, accumulator);
}
<commit_msg>Fix warning about unknown or unused feature<commit_after>// This code is editable and runnable!
fn main() {
    // A simple integer calculator:
    // `+` or `-` means add or subtract by 1
    // `*` or `/` means multiply or divide by 2

    let program = "+ + * - /";
    let mut accumulator = 0;

    for token in program.chars() {
        match token {
            '+' => accumulator += 1,
            '-' => accumulator -= 1,
            '*' => accumulator *= 2,
            '/' => accumulator /= 2,
            _ => { /* ignore everything else */ }
        }
    }

    println!("The program \"{}\" calculates the value {}",
              program, accumulator);
}
<|endoftext|>


<commit_before>// Copyright 2015 The Servo Project Developers. See the COPYRIGHT
// file at the top-level directory of this distribution.
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

use std::sync::mpsc;
use std::sync::{Arc, Mutex, Condvar};
use std::collections::hash_map::HashMap;
use std::cell::{RefCell};
use std::slice;
use std::fmt::{self, Debug, Formatter};
use std::cmp::{PartialEq};
use std::ops::Deref;
use std::mem;

use uuid::Uuid;

struct SenderRecord {
    sender: MpscSender,
    conn_sender: mpsc::Sender<bool>,
    conn_receiver: Mutex<mpsc::Receiver<bool>>,
}

impl SenderRecord {
    fn new(sender: MpscSender) -> SenderRecord {
        let (tx, rx) = mpsc::channel::<bool>();
        SenderRecord {
            sender: sender,
            conn_sender: tx,
            conn_receiver: Mutex::new(rx),
        }
    }

    fn accept(&self) {
        self.conn_receiver.lock().unwrap().recv().unwrap();
    }

    fn connect(&self) {
        self.conn_sender.send(true).unwrap();
    }
}

lazy_static! {
    static ref ONE_SHOT_SENDERS: Mutex<HashMap<String,SenderRecord>> = Mutex::new(HashMap::new());
}

struct MpscChannelMessage(Vec<u8>, Vec<MpscChannel>, Vec<MpscSharedMemory>);

pub fn channel() -> Result<(MpscSender, MpscReceiver),MpscError> {
    let (base_sender, base_receiver) = mpsc::channel::<MpscChannelMessage>();
    Ok((MpscSender::new(base_sender), MpscReceiver::new(base_receiver)))
}

pub struct MpscReceiver {
    receiver: RefCell<Option<mpsc::Receiver<MpscChannelMessage>>>,
}

impl MpscReceiver {
    fn new(receiver: mpsc::Receiver<MpscChannelMessage>) -> MpscReceiver {
        MpscReceiver {
            receiver: RefCell::new(Some(receiver)),
        }
    }

    pub fn consume(&self) -> MpscReceiver {
        let receiver = self.receiver.borrow_mut().take();
        MpscReceiver::new(receiver.unwrap())
    }

    pub fn recv(&self) -> Result<(Vec<u8>, Vec<OpaqueMpscChannel>, Vec<MpscSharedMemory>),MpscError> {
        let r = self.receiver.borrow();
        match r.as_ref().unwrap().recv() {
            Ok(MpscChannelMessage(d,c,s)) => Ok((d,
                                                 c.into_iter().map(OpaqueMpscChannel::new).collect(),
                                                 s)),
            Err(_) => Err(MpscError::ChannelClosedError),
        }
    }

    pub fn try_recv(&self) -> Result<(Vec<u8>, Vec<OpaqueMpscChannel>, Vec<MpscSharedMemory>),MpscError> {
        let r = self.receiver.borrow();
        match r.as_ref().unwrap().try_recv() {
            Ok(MpscChannelMessage(d,c,s)) => Ok((d,
                                                 c.into_iter().map(OpaqueMpscChannel::new).collect(),
                                                 s)),
            Err(_) => Err(MpscError::ChannelClosedError),
        }
    }
}

unsafe impl Send for MpscReceiver { }
unsafe impl Sync for MpscReceiver { }

#[derive(Clone)]
pub struct MpscSender {
    sender: RefCell<mpsc::Sender<MpscChannelMessage>>,
}

unsafe impl Send for MpscSender { }
unsafe impl Sync for MpscSender { }

impl MpscSender {
    fn new(sender: mpsc::Sender<MpscChannelMessage>) -> MpscSender {
        MpscSender {
            sender: RefCell::new(sender),
        }
    }

    pub fn connect(name: String) -> Result<MpscSender,MpscError> {
        let record = ONE_SHOT_SENDERS.lock().unwrap().remove(&name).unwrap();
        record.connect();
        Ok(record.sender)
    }

    pub fn send(&self,
                data: &[u8],
                ports: Vec<MpscChannel>,
                shared_memory_regions: Vec<MpscSharedMemory>)
                -> Result<(),MpscError>
    {
        match self.sender.borrow().send(MpscChannelMessage(data.to_vec(), ports, shared_memory_regions)) {
            Err(_) => Err(MpscError::ChannelClosedError),
            Ok(_) => Ok(()),
        }
    }
}

pub struct MpscReceiverSet {
    last_index: usize,
    receiver_ids: Vec<usize>,
    receivers: Vec<MpscReceiver>,
}

impl MpscReceiverSet {
    pub fn new() -> Result<MpscReceiverSet,MpscError> {
        Ok(MpscReceiverSet {
            last_index: 0,
            receiver_ids: vec![],
            receivers: vec![],
        })
    }

    pub fn add(&mut self, receiver: MpscReceiver) -> Result<i64,MpscError> {
        self.last_index += 1;
        self.receiver_ids.push(self.last_index);
        self.receivers.push(receiver.consume());
        Ok(self.last_index as i64)
    }

    pub fn select(&mut self) -> Result<Vec<MpscSelectionResult>,MpscError> {
        let mut receivers: Vec<Option<mpsc::Receiver<MpscChannelMessage>>> = Vec::with_capacity(self.receivers.len());
        let mut r_id: i64 = -1;
        let mut r_index: usize = 0;

        {
            let select = mpsc::Select::new();
            // we *must* allocate exact capacity for this, because the Handles *can't move*
            let mut handles: Vec<mpsc::Handle<MpscChannelMessage>> = Vec::with_capacity(self.receivers.len());

            for r in &self.receivers {
                let inner_r = mem::replace(&mut *r.receiver.borrow_mut(), None);
                receivers.push(inner_r);
            }

            for r in &receivers {
                unsafe {
                    handles.push(select.handle(r.as_ref().unwrap()));
                    handles.last_mut().unwrap().add();
                }
            }

            let id = select.wait();

            for (index,h) in handles.iter().enumerate() {
                if h.id() == id {
                    r_index = index;
                    r_id = self.receiver_ids[index] as i64;
                    break;
                }
            }
        }

        // put the receivers back
        for (index,r) in self.receivers.iter().enumerate() {
            mem::replace(&mut *r.receiver.borrow_mut(), mem::replace(&mut receivers[index], None));
        }

        if r_id == -1 {
            return Err(MpscError::UnknownError);
        }

        let receivers = &mut self.receivers;
        match receivers[r_index].recv() {
            Ok((data, channels, shmems)) =>
                Ok(vec![MpscSelectionResult::DataReceived(r_id, data, channels, shmems)]),
            Err(MpscError::ChannelClosedError) => {
                receivers.remove(r_index);
                self.receiver_ids.remove(r_index);
                Ok(vec![MpscSelectionResult::ChannelClosed(r_id)])
            },
            Err(err) => Err(err),
        }
    }
}

pub enum MpscSelectionResult {
    DataReceived(i64, Vec<u8>, Vec<OpaqueMpscChannel>, Vec<MpscSharedMemory>),
    ChannelClosed(i64),
}

pub struct MpscOneShotServer {
    receiver: RefCell<Option<MpscReceiver>>,
    name: String,
}

impl MpscOneShotServer {
    pub fn new() -> Result<(MpscOneShotServer, String),MpscError> {
        let (sender, receiver) = match channel() {
            Ok((s,r)) => (s,r),
            Err(err) => return Err(err),
        };

        let name = Uuid::new_v4().to_string();
        let record = SenderRecord::new(sender);
        ONE_SHOT_SENDERS.lock().unwrap().insert(name.clone(), record);
        Ok((MpscOneShotServer {
            receiver: RefCell::new(Some(receiver)),
            name: name.clone(),
        },name.clone()))
    }

    pub fn accept(&self) -> Result<(MpscReceiver,
                                    Vec<u8>,
                                    Vec<OpaqueMpscChannel>,
                                    Vec<MpscSharedMemory>),MpscError>
    {
        ONE_SHOT_SENDERS.lock().unwrap().get(&self.name).unwrap().accept();
        let receiver = self.receiver.borrow_mut().take().unwrap();
        let (data, channels, shmems) = receiver.recv().unwrap();
        Ok((receiver, data, channels, shmems))
    }
}

pub enum MpscChannel {
    Sender(MpscSender),
    Receiver(MpscReceiver),
}

pub struct OpaqueMpscChannel {
    channel: RefCell<Option<MpscChannel>>,
}

impl OpaqueMpscChannel {
    fn new(channel: MpscChannel) -> OpaqueMpscChannel {
        OpaqueMpscChannel {
            channel: RefCell::new(Some(channel))
        }
    }

    pub fn to_receiver(&self) -> MpscReceiver {
        match self.channel.borrow_mut().take().unwrap() {
            MpscChannel::Sender(_) => panic!("Opaque channel is not a receiver!"),
            MpscChannel::Receiver(r) => r
        }
    }

    pub fn to_sender(&self) -> MpscSender {
        match self.channel.borrow_mut().take().unwrap() {
            MpscChannel::Sender(s) => s,
            MpscChannel::Receiver(_) => panic!("Opaque channel is not a sender!"),
        }
    }
}

pub struct MpscSharedMemory {
    ptr: *mut u8,
    length: usize,
    data: Arc<Vec<u8>>,
}

unsafe impl Send for MpscSharedMemory {}
unsafe impl Sync for MpscSharedMemory {}

impl Clone for MpscSharedMemory {
    fn clone(&self) -> MpscSharedMemory {
        MpscSharedMemory {
            ptr: self.ptr,
            length: self.length,
            data: self.data.clone(),
        }
    }
}

impl PartialEq for MpscSharedMemory {
    fn eq(&self, other: &MpscSharedMemory) -> bool {
        **self == **other
    }
}

impl Debug for MpscSharedMemory {
    fn fmt(&self, formatter: &mut Formatter) -> Result<(), fmt::Error> {
        (**self).fmt(formatter)
    }
}

impl Deref for MpscSharedMemory {
    type Target = [u8];

    #[inline]
    fn deref(&self) -> &[u8] {
        if self.ptr.is_null() {
            panic!("attempted to access a consumed `MpscSharedMemory`")
        }
        unsafe {
            slice::from_raw_parts(self.ptr, self.length)
        }
    }
}

impl MpscSharedMemory {
    pub fn from_byte(byte: u8, length: usize) -> MpscSharedMemory {
        let mut v = Arc::new(vec![byte; length]);
        MpscSharedMemory {
            ptr: Arc::get_mut(&mut v).unwrap().as_mut_ptr(),
            length: length,
            data: v
        }
    }

    pub fn from_bytes(bytes: &[u8]) -> MpscSharedMemory {
        let mut v = Arc::new(bytes.to_vec());
        MpscSharedMemory {
            ptr: Arc::get_mut(&mut v).unwrap().as_mut_ptr(),
            length: v.len(),
            data: v
        }
    }
}

#[derive(Debug)]
pub enum MpscError {
    ChannelClosedError,
    UnknownError,
}

<commit_msg>rename ONE_SHOT_SENDERS to ONE_SHOT_SERVERS<commit_after>// Copyright 2015 The Servo Project Developers. See the COPYRIGHT
// file at the top-level directory of this distribution.
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

use std::sync::mpsc;
use std::sync::{Arc, Mutex, Condvar};
use std::collections::hash_map::HashMap;
use std::cell::{RefCell};
use std::slice;
use std::fmt::{self, Debug, Formatter};
use std::cmp::{PartialEq};
use std::ops::Deref;
use std::mem;

use uuid::Uuid;

struct ServerRecord {
    sender: MpscSender,
    conn_sender: mpsc::Sender<bool>,
    conn_receiver: Mutex<mpsc::Receiver<bool>>,
}

impl ServerRecord {
    fn new(sender: MpscSender) -> ServerRecord {
        let (tx, rx) = mpsc::channel::<bool>();
        ServerRecord {
            sender: sender,
            conn_sender: tx,
            conn_receiver: Mutex::new(rx),
        }
    }

    fn accept(&self) {
        self.conn_receiver.lock().unwrap().recv().unwrap();
    }

    fn connect(&self) {
        self.conn_sender.send(true).unwrap();
    }
}

lazy_static! {
    static ref ONE_SHOT_SERVERS: Mutex<HashMap<String,ServerRecord>> = Mutex::new(HashMap::new());
}

struct MpscChannelMessage(Vec<u8>, Vec<MpscChannel>, Vec<MpscSharedMemory>);

pub fn channel() -> Result<(MpscSender, MpscReceiver),MpscError> {
    let (base_sender, base_receiver) = mpsc::channel::<MpscChannelMessage>();
    Ok((MpscSender::new(base_sender), MpscReceiver::new(base_receiver)))
}

pub struct MpscReceiver {
    receiver: RefCell<Option<mpsc::Receiver<MpscChannelMessage>>>,
}

impl MpscReceiver {
    fn new(receiver: mpsc::Receiver<MpscChannelMessage>) -> MpscReceiver {
        MpscReceiver {
            receiver: RefCell::new(Some(receiver)),
        }
    }

    pub fn consume(&self) -> MpscReceiver {
        let receiver = self.receiver.borrow_mut().take();
        MpscReceiver::new(receiver.unwrap())
    }

    pub fn recv(&self) -> Result<(Vec<u8>, Vec<OpaqueMpscChannel>, Vec<MpscSharedMemory>),MpscError> {
        let r = self.receiver.borrow();
        match r.as_ref().unwrap().recv() {
            Ok(MpscChannelMessage(d,c,s)) => Ok((d,
                                                 c.into_iter().map(OpaqueMpscChannel::new).collect(),
                                                 s)),
            Err(_) => Err(MpscError::ChannelClosedError),
        }
    }

    pub fn try_recv(&self) -> Result<(Vec<u8>, Vec<OpaqueMpscChannel>, Vec<MpscSharedMemory>),MpscError> {
        let r = self.receiver.borrow();
        match r.as_ref().unwrap().try_recv() {
            Ok(MpscChannelMessage(d,c,s)) => Ok((d,
                                                 c.into_iter().map(OpaqueMpscChannel::new).collect(),
                                                 s)),
            Err(_) => Err(MpscError::ChannelClosedError),
        }
    }
}

unsafe impl Send for MpscReceiver { }
unsafe impl Sync for MpscReceiver { }

#[derive(Clone)]
pub struct MpscSender {
    sender: RefCell<mpsc::Sender<MpscChannelMessage>>,
}

unsafe impl Send for MpscSender { }
unsafe impl Sync for MpscSender { }

impl MpscSender {
    fn new(sender: mpsc::Sender<MpscChannelMessage>) -> MpscSender {
        MpscSender {
            sender: RefCell::new(sender),
        }
    }

    pub fn connect(name: String) -> Result<MpscSender,MpscError> {
        let record = ONE_SHOT_SERVERS.lock().unwrap().remove(&name).unwrap();
        record.connect();
        Ok(record.sender)
    }

    pub fn send(&self,
                data: &[u8],
                ports: Vec<MpscChannel>,
                shared_memory_regions: Vec<MpscSharedMemory>)
                -> Result<(),MpscError>
    {
        match self.sender.borrow().send(MpscChannelMessage(data.to_vec(), ports, shared_memory_regions)) {
            Err(_) => Err(MpscError::ChannelClosedError),
            Ok(_) => Ok(()),
        }
    }
}

pub struct MpscReceiverSet {
    last_index: usize,
    receiver_ids: Vec<usize>,
    receivers: Vec<MpscReceiver>,
}

impl MpscReceiverSet {
    pub fn new() -> Result<MpscReceiverSet,MpscError> {
        Ok(MpscReceiverSet {
            last_index: 0,
            receiver_ids: vec![],
            receivers: vec![],
        })
    }

    pub fn add(&mut self, receiver: MpscReceiver) -> Result<i64,MpscError> {
        self.last_index += 1;
        self.receiver_ids.push(self.last_index);
        self.receivers.push(receiver.consume());
        Ok(self.last_index as i64)
    }

    pub fn select(&mut self) -> Result<Vec<MpscSelectionResult>,MpscError> {
        let mut receivers: Vec<Option<mpsc::Receiver<MpscChannelMessage>>> = Vec::with_capacity(self.receivers.len());
        let mut r_id: i64 = -1;
        let mut r_index: usize = 0;

        {
            let select = mpsc::Select::new();
            // we *must* allocate exact capacity for this, because the Handles *can't move*
            let mut handles: Vec<mpsc::Handle<MpscChannelMessage>> = Vec::with_capacity(self.receivers.len());

            for r in &self.receivers {
                let inner_r = mem::replace(&mut *r.receiver.borrow_mut(), None);
                receivers.push(inner_r);
            }

            for r in &receivers {
                unsafe {
                    handles.push(select.handle(r.as_ref().unwrap()));
                    handles.last_mut().unwrap().add();
                }
            }

            let id = select.wait();

            for (index,h) in handles.iter().enumerate() {
                if h.id() == id {
                    r_index = index;
                    r_id = self.receiver_ids[index] as i64;
                    break;
                }
            }
        }

        // put the receivers back
        for (index,r) in self.receivers.iter().enumerate() {
            mem::replace(&mut *r.receiver.borrow_mut(), mem::replace(&mut receivers[index], None));
        }

        if r_id == -1 {
            return Err(MpscError::UnknownError);
        }

        let receivers = &mut self.receivers;
        match receivers[r_index].recv() {
            Ok((data, channels, shmems)) =>
                Ok(vec![MpscSelectionResult::DataReceived(r_id, data, channels, shmems)]),
            Err(MpscError::ChannelClosedError) => {
                receivers.remove(r_index);
                self.receiver_ids.remove(r_index);
                Ok(vec![MpscSelectionResult::ChannelClosed(r_id)])
            },
            Err(err) => Err(err),
        }
    }
}

pub enum MpscSelectionResult {
    DataReceived(i64, Vec<u8>, Vec<OpaqueMpscChannel>, Vec<MpscSharedMemory>),
    ChannelClosed(i64),
}

pub struct MpscOneShotServer {
    receiver: RefCell<Option<MpscReceiver>>,
    name: String,
}

impl MpscOneShotServer {
    pub fn new() -> Result<(MpscOneShotServer, String),MpscError> {
        let (sender, receiver) = match channel() {
            Ok((s,r)) => (s,r),
            Err(err) => return Err(err),
        };

        let name = Uuid::new_v4().to_string();
        let record = ServerRecord::new(sender);
        ONE_SHOT_SERVERS.lock().unwrap().insert(name.clone(), record);
        Ok((MpscOneShotServer {
            receiver: RefCell::new(Some(receiver)),
            name: name.clone(),
        },name.clone()))
    }

    pub fn accept(&self) -> Result<(MpscReceiver,
                                    Vec<u8>,
                                    Vec<OpaqueMpscChannel>,
                                    Vec<MpscSharedMemory>),MpscError>
    {
        ONE_SHOT_SERVERS.lock().unwrap().get(&self.name).unwrap().accept();
        let receiver = self.receiver.borrow_mut().take().unwrap();
        let (data, channels, shmems) = receiver.recv().unwrap();
        Ok((receiver, data, channels, shmems))
    }
}

pub enum MpscChannel {
    Sender(MpscSender),
    Receiver(MpscReceiver),
}

pub struct OpaqueMpscChannel {
    channel: RefCell<Option<MpscChannel>>,
}

impl OpaqueMpscChannel {
    fn new(channel: MpscChannel) -> OpaqueMpscChannel {
        OpaqueMpscChannel {
            channel: RefCell::new(Some(channel))
        }
    }

    pub fn to_receiver(&self) -> MpscReceiver {
        match self.channel.borrow_mut().take().unwrap() {
            MpscChannel::Sender(_) => panic!("Opaque channel is not a receiver!"),
            MpscChannel::Receiver(r) => r
        }
    }

    pub fn to_sender(&self) -> MpscSender {
        match self.channel.borrow_mut().take().unwrap() {
            MpscChannel::Sender(s) => s,
            MpscChannel::Receiver(_) => panic!("Opaque channel is not a sender!"),
        }
    }
}

pub struct MpscSharedMemory {
    ptr: *mut u8,
    length: usize,
    data: Arc<Vec<u8>>,
}

unsafe impl Send for MpscSharedMemory {}
unsafe impl Sync for MpscSharedMemory {}

impl Clone for MpscSharedMemory {
    fn clone(&self) -> MpscSharedMemory {
        MpscSharedMemory {
            ptr: self.ptr,
            length: self.length,
            data: self.data.clone(),
        }
    }
}

impl PartialEq for MpscSharedMemory {
    fn eq(&self, other: &MpscSharedMemory) -> bool {
        **self == **other
    }
}

impl Debug for MpscSharedMemory {
    fn fmt(&self, formatter: &mut Formatter) -> Result<(), fmt::Error> {
        (**self).fmt(formatter)
    }
}

impl Deref for MpscSharedMemory {
    type Target = [u8];

    #[inline]
    fn deref(&self) -> &[u8] {
        if self.ptr.is_null() {
            panic!("attempted to access a consumed `MpscSharedMemory`")
        }
        unsafe {
            slice::from_raw_parts(self.ptr, self.length)
        }
    }
}

impl MpscSharedMemory {
    pub fn from_byte(byte: u8, length: usize) -> MpscSharedMemory {
        let mut v = Arc::new(vec![byte; length]);
        MpscSharedMemory {
            ptr: Arc::get_mut(&mut v).unwrap().as_mut_ptr(),
            length: length,
            data: v
        }
    }

    pub fn from_bytes(bytes: &[u8]) -> MpscSharedMemory {
        let mut v = Arc::new(bytes.to_vec());
        MpscSharedMemory {
            ptr: Arc::get_mut(&mut v).unwrap().as_mut_ptr(),
            length: v.len(),
            data: v
        }
    }
}

#[derive(Debug)]
pub enum MpscError {
    ChannelClosedError,
    UnknownError,
}

<|endoftext|>


<commit_before>extern mod std;

extern mod crypto {
  fn SHA1(src: *u8, sz: libc::c_uint, out: *u8) -> *u8;
  fn MD5(src: *u8, sz: libc::c_uint, out: *u8) -> *u8;
  fn SHA224(src: *u8, sz: libc::c_uint, out: *u8) -> *u8;
  fn SHA256(src: *u8, sz: libc::c_uint, out: *u8) -> *u8;
  fn SHA384(src: *u8, sz: libc::c_uint, out: *u8) -> *u8;
  fn SHA512(src: *u8, sz: libc::c_uint, out: *u8) -> *u8;
}

enum HashMethod { MD5, SHA1, SHA224, SHA256, SHA384, SHA512 }

struct HashEngine {
  engine:      (extern unsafe fn(*u8, libc::c_uint, *u8) -> *u8),
  digest_size: uint,
  block_size:  uint
}

impl HashEngine {
  fn new(engine: HashMethod) -> ~HashEngine {
    unsafe {
      match engine {
        MD5    => ~HashEngine{ engine: crypto::MD5,    digest_size: 16, block_size: 64  },
        SHA1   => ~HashEngine{ engine: crypto::SHA1,   digest_size: 20, block_size: 64  },
        SHA224 => ~HashEngine{ engine: crypto::SHA224, digest_size: 28, block_size: 64  },
        SHA256 => ~HashEngine{ engine: crypto::SHA256, digest_size: 32, block_size: 64  },
        SHA384 => ~HashEngine{ engine: crypto::SHA384, digest_size: 48, block_size: 64  },
        SHA512 => ~HashEngine{ engine: crypto::SHA512, digest_size: 64, block_size: 128 }
      }
    }
  }

  fn hash(&self, data: ~[u8]) -> ~Digest {
    let hash_func = self.engine;
    Digest::new(
      unsafe {
        vec::from_buf(
          hash_func(
            vec::raw::to_ptr(data),
            data.len() as libc::c_uint, ptr::null()
          ),
          self.digest_size
        )
      }
    )
  }

  fn hmac(&self, key: ~[u8], message: ~[u8]) -> ~Digest {
    let computed_key = match key.len() {
      _ if key.len() > self.block_size => self.hash(key).hexdigest().to_bytes(),
      _ if key.len() < self.block_size => key + vec::from_elem(self.block_size - key.len(), 0),
      _ => key
    };

    let o_key_pad = HashEngine::xor_with(computed_key, 0x5c);
    let i_key_pad = HashEngine::xor_with(computed_key, 0x36);

    self.hash(o_key_pad + self.hash(i_key_pad + message).digest)
  }

  priv fn xor_with(subject: &[u8], val: u8) -> ~[u8] {
    do subject.map |&c| { val ^ c }
  }
}

struct Digest { digest: ~[u8] }

impl Digest {
  fn new(digest: ~[u8]) -> ~Digest { return ~Digest{ digest: digest } }

  fn hexdigest(&self) -> ~str {
    let mut acc = ~"";
    for self.digest.each |&byte| { acc += fmt!("%02x", byte as uint); }
    acc
  }
}

fn sha1(data: ~[u8]) -> ~Digest {
  HashEngine::new(SHA1).hash(data)
}

fn sha224(data: ~[u8]) -> ~Digest {
  HashEngine::new(SHA224).hash(data)
}

fn sha256(data: ~[u8]) -> ~Digest {
  HashEngine::new(SHA256).hash(data)
}

fn sha384(data: ~[u8]) -> ~Digest {
  HashEngine::new(SHA384).hash(data)
}

fn sha512(data: ~[u8]) -> ~Digest {
  HashEngine::new(SHA512).hash(data)
}

fn md5(data: ~[u8]) -> ~Digest {
  HashEngine::new(MD5).hash(data)
}

fn hmac_md5(key: ~[u8], message: ~[u8]) -> ~Digest {
  HashEngine::new(MD5).hmac(key, message)
}

fn hmac_sha1(key: ~[u8], message: ~[u8]) -> ~Digest {
  HashEngine::new(SHA1).hmac(key, message)
}

fn hmac_sha256(key: ~[u8], message: ~[u8]) -> ~Digest {
  HashEngine::new(SHA256).hmac(key, message)
}

fn hmac_sha512(key: ~[u8], message: ~[u8]) -> ~Digest {
  HashEngine::new(SHA512).hmac(key, message)
}

fn main() {
  if core::os::args().len() < 2 { fail!(~"Nothing to hash: supply an argument") }
  io::println(sha1(core::os::args()[1].clone().to_bytes()).hexdigest());
}

#[test]
fn test_sha1() {
  assert_eq!(sha1("testing SHA1".to_bytes()).hexdigest(), ~"a07bb2f7f56b7b47e54fd9e29b4629c8967dde4c");
}

#[test]
fn test_md5() {
  assert_eq!(md5("testing MD5".to_bytes()).hexdigest(), ~"da35df95d34f166890fa759f5e00e94c");
}

#[test]
fn test_hmac_sha1_with_a_40_byte_key() {
  assert_eq!(
    hmac_sha1("xxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxx".to_bytes(), "some data to hmac".to_bytes()).hexdigest(),
    ~"18f652f1bced11930d55a72ad60bb5e671573dec"
  )
}

#[test]
fn test_hmac_sha1_with_a_too_long_key() {
  assert_eq!(
    hmac_sha1("keyxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxx".to_bytes(), "some data to hmac".to_bytes()).hexdigest(),
    ~"64109bfc49893757afcc3c0c80b8eb61ce0f60a5"
  )
}

#[test]
fn test_hmac_sha1_with_a_too_short_key() {
  assert_eq!(
    hmac_sha1("my key".to_bytes(), "some data to hmac".to_bytes()).hexdigest(),
    ~"331c9bf4d7dc8ad7e9bab2f566c8612042a9f4e2"
  )
}

#[test]
fn test_hmac_sha256() {
  assert_eq!(
    hmac_sha256("my key".to_bytes(), "some data to hmac".to_bytes()).hexdigest(),
    ~"17764a8a2e8c23ae3b9e1781c5ebeb5754c7920806021032e9133051e10118be"
  );
}

#[test]
fn test_hmac_sha512() {
  assert_eq!(
    hmac_sha512("my key".to_bytes(), "some data to hmac".to_bytes()).hexdigest(),
    ~"b54973feba2e436c4f0911855c80e7320d72edd37b359d6474d714718c52775f163832b79214e9fa9d400208d21372a465b2cda2a1d5ab488bbc02b0daea9626"
  );
}

#[test]
fn test_hmac_md5() {
  assert_eq!(hmac_md5("my key".to_bytes(), "some data to hmac".to_bytes()).hexdigest(), ~"befd13e00ee07549520ade402c664a3a");
}
<commit_msg>Convert to Rust 0.7.<commit_after>extern mod std;

use std::{io, vec, ptr, os};
use std::libc::{c_uint};

mod crypto {
  use std::libc::{c_uint};
#[link_args = "-lcrypto"]
  extern {
    fn SHA1(src: *u8, sz: c_uint, out: *u8) -> *u8;
    fn MD5(src: *u8, sz: c_uint, out: *u8) -> *u8;
    fn SHA224(src: *u8, sz: c_uint, out: *u8) -> *u8;
    fn SHA256(src: *u8, sz: c_uint, out: *u8) -> *u8;
    fn SHA384(src: *u8, sz: c_uint, out: *u8) -> *u8;
    fn SHA512(src: *u8, sz: c_uint, out: *u8) -> *u8;
  }
}

enum HashMethod { MD5, SHA1, SHA224, SHA256, SHA384, SHA512 }

struct HashEngine {
  engine:      (extern "C" unsafe fn(*u8, c_uint, *u8) -> *u8),
  digest_size: uint,
  block_size:  uint
}

impl HashEngine {
  fn new(engine: HashMethod) -> ~HashEngine {
    match engine {
      MD5    => ~HashEngine{ engine: crypto::MD5,    digest_size: 16, block_size: 64  },
      SHA1   => ~HashEngine{ engine: crypto::SHA1,   digest_size: 20, block_size: 64  },
      SHA224 => ~HashEngine{ engine: crypto::SHA224, digest_size: 28, block_size: 64  },
      SHA256 => ~HashEngine{ engine: crypto::SHA256, digest_size: 32, block_size: 64  },
      SHA384 => ~HashEngine{ engine: crypto::SHA384, digest_size: 48, block_size: 64  },
      SHA512 => ~HashEngine{ engine: crypto::SHA512, digest_size: 64, block_size: 128 }
    }
  }

  fn hash(&self, data: ~[u8]) -> ~Digest {
    let hash_func = self.engine;
    Digest::new(
      unsafe {
        vec::from_buf(
          hash_func(
            vec::raw::to_ptr(data),
            data.len() as c_uint, ptr::null()
          ),
          self.digest_size
        )
      }
    )
  }

  fn hmac(&self, key: ~[u8], message: ~[u8]) -> ~Digest {
    let computed_key = match key.len() {
      _ if key.len() > self.block_size => self.hash(key).digest,
      _ if key.len() < self.block_size => key + vec::from_elem(self.block_size - key.len(), 0),
      _ => key
    };

    let o_key_pad = HashEngine::xor_with(computed_key, 0x5c);
    let i_key_pad = HashEngine::xor_with(computed_key, 0x36);

    self.hash(o_key_pad + self.hash(i_key_pad + message).digest)
  }

  priv fn xor_with(subject: &[u8], val: u8) -> ~[u8] {
    do subject.map |&c| { val ^ c }
  }
}

struct Digest { digest: ~[u8] }

impl Digest {
  fn new(digest: ~[u8]) -> ~Digest { return ~Digest{ digest: digest } }

  fn hexdigest(&self) -> ~str {
    let mut acc = ~"";
    for self.digest.iter().advance |&byte| { acc = acc.append(fmt!("%02x", byte as uint)); }
    acc
  }
}

fn sha1(data: ~[u8]) -> ~Digest {
  HashEngine::new(SHA1).hash(data)
}

fn sha224(data: ~[u8]) -> ~Digest {
  HashEngine::new(SHA224).hash(data)
}

fn sha256(data: ~[u8]) -> ~Digest {
  HashEngine::new(SHA256).hash(data)
}

fn sha384(data: ~[u8]) -> ~Digest {
  HashEngine::new(SHA384).hash(data)
}

fn sha512(data: ~[u8]) -> ~Digest {
  HashEngine::new(SHA512).hash(data)
}

fn md5(data: ~[u8]) -> ~Digest {
  HashEngine::new(MD5).hash(data)
}

fn hmac_md5(key: ~[u8], message: ~[u8]) -> ~Digest {
  HashEngine::new(MD5).hmac(key, message)
}

fn hmac_sha1(key: ~[u8], message: ~[u8]) -> ~Digest {
  HashEngine::new(SHA1).hmac(key, message)
}

fn hmac_sha256(key: ~[u8], message: ~[u8]) -> ~Digest {
  HashEngine::new(SHA256).hmac(key, message)
}

fn hmac_sha512(key: ~[u8], message: ~[u8]) -> ~Digest {
  HashEngine::new(SHA512).hmac(key, message)
}

fn main() {
  if os::args().len() < 2 { fail!(~"Nothing to hash: supply an argument") }
  let args = ~os::args();
  io::println(sha1(args[1].as_bytes().to_owned()).hexdigest());
}

#[test]
fn test_sha1() {
  assert_eq!(sha1("testing SHA1".as_bytes().to_owned()).hexdigest(), ~"a07bb2f7f56b7b47e54fd9e29b4629c8967dde4c");
}

#[test]
fn test_md5() {
  assert_eq!(md5("testing MD5".as_bytes().to_owned()).hexdigest(), ~"da35df95d34f166890fa759f5e00e94c");
}

#[test]
fn test_hmac_sha1_with_a_64_byte_key() {
  assert_eq!(
    hmac_sha1("xxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxx".as_bytes().to_owned(), "some data to hmac".as_bytes().to_owned()).hexdigest(),
    ~"d6264ce636b593e6c5b43c90ec4ca3550032c99f"
  )
}

#[test]
fn test_hmac_sha1_with_a_too_long_key() {
  assert_eq!(
    hmac_sha1("keyxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxx".as_bytes().to_owned(), "some data to hmac".as_bytes().to_owned()).hexdigest(),
    ~"4361f0a10ae7fc1ce33d2795e998d0e526a38815"
  )
}

#[test]
fn test_hmac_sha1_with_a_too_short_key() {
  assert_eq!(
    hmac_sha1("my key".as_bytes().to_owned(), "some data to hmac".as_bytes().to_owned()).hexdigest(),
    ~"331c9bf4d7dc8ad7e9bab2f566c8612042a9f4e2"
  )
}

#[test]
fn test_hmac_sha256() {
  assert_eq!(
    hmac_sha256("my key".as_bytes().to_owned(), "some data to hmac".as_bytes().to_owned()).hexdigest(),
    ~"17764a8a2e8c23ae3b9e1781c5ebeb5754c7920806021032e9133051e10118be"
  );
}

#[test]
fn test_hmac_sha512() {
  assert_eq!(
    hmac_sha512("my key".as_bytes().to_owned(), "some data to hmac".as_bytes().to_owned()).hexdigest(),
    ~"b54973feba2e436c4f0911855c80e7320d72edd37b359d6474d714718c52775f163832b79214e9fa9d400208d21372a465b2cda2a1d5ab488bbc02b0daea9626"
  );
}

#[test]
fn test_hmac_md5() {
  assert_eq!(hmac_md5("my key".as_bytes().to_owned(), "some data to hmac".as_bytes().to_owned()).hexdigest(), ~"befd13e00ee07549520ade402c664a3a");
}
<|endoftext|>


<commit_before>//! Online backup API
use std::path::{Path};

use super::ffi;
use {SqliteError, SqliteResult, SqliteConnection};

impl SqliteConnection {
    /// Open or create a database file named `dest_path`. Transfer the
    /// content of local database `name` into the `dest_path` database.
    pub fn backup<P: AsRef<Path>>(&self, name: &str, dest_path: &P) -> SqliteResult<()> {
        let dest = try!(SqliteConnection::open(dest_path)); // "cannot open target database: "

        let mut dest_db = dest.db.borrow_mut();
        let dest_name = try!(super::str_to_cstring("main"));
        let src_db = self.db.borrow_mut();
        let src_name = try!(super::str_to_cstring(name));
        let backup = unsafe { ffi::sqlite3_backup_init(dest_db.db(), dest_name.as_ptr(), src_db.db(), src_name.as_ptr()) };
        if backup.is_null() {
            let msg = unsafe { super::errmsg_to_string(ffi::sqlite3_errmsg(dest_db.db())) };
            let err = Err(SqliteError{code: unsafe { ffi::sqlite3_errcode(dest_db.db()) }, message: "backup failed: ".to_string() + &msg });
            let _ = dest_db.close();
            return err;
        }

        let mut rc;
        loop {
            rc = unsafe { ffi::sqlite3_backup_step(backup, 100) };
            if rc != ffi::SQLITE_OK {
                break;
            }
        }
        unsafe { ffi::sqlite3_backup_finish(backup) };
        let result = if rc == ffi::SQLITE_DONE {
            Ok(())
        } else {
            let msg = unsafe { super::errmsg_to_string(ffi::sqlite3_errmsg(dest_db.db())) };
            Err(SqliteError{code: rc, message: "backup failed: ".to_string() + &msg })
        };
        let _ = dest_db.close();
        result
    }

    /// Open a database file named `src_path`. Transfer the content
    /// of `src_path` into the local database `name`.
    pub fn restore<P: AsRef<Path>>(&self, name: &str, src_path: &P) -> SqliteResult<()> {
        let src = try!(SqliteConnection::open_with_flags(src_path, super::SQLITE_OPEN_READ_ONLY | super::SQLITE_OPEN_URI)); // "cannot open source database: "

        let dest_db = self.db.borrow_mut();
        let dest_name = try!(super::str_to_cstring(name));
        let mut src_db = src.db.borrow_mut();
        let src_name = try!(super::str_to_cstring("main"));
        let backup = unsafe { ffi::sqlite3_backup_init(dest_db.db(), dest_name.as_ptr(), src_db.db(), src_name.as_ptr()) };
        if backup.is_null() {
            let msg = unsafe { super::errmsg_to_string(ffi::sqlite3_errmsg(dest_db.db())) };
            let err = Err(SqliteError{code: unsafe { ffi::sqlite3_errcode(dest_db.db()) }, message: "restore failed: ".to_string() + &msg });
            let _ = src_db.close();
            return err;
        }
        let mut rc;
        let mut n_timeout = 0;
        loop {
            rc = unsafe { ffi::sqlite3_backup_step(backup, 100) };
            if rc == ffi::SQLITE_OK {
            } else if rc == ffi::SQLITE_BUSY {
                n_timeout += 1;
                if n_timeout >= 3 {
                    break;
                }
                unsafe { ffi::sqlite3_sleep(100) };
            } else {
                break;
            }
        }
        unsafe { ffi::sqlite3_backup_finish(backup) };
        let result = if rc == ffi::SQLITE_DONE {
            Ok(())
        } else if rc == ffi::SQLITE_BUSY || rc == ffi::SQLITE_LOCKED {
            Err(SqliteError{code: rc, message: "restore failed: source database busy".to_string() })
        } else {
            let msg = unsafe { super::errmsg_to_string(ffi::sqlite3_errmsg(dest_db.db())) };
            Err(SqliteError{code: rc, message: "restore failed: ".to_string() + &msg })
        };
        let _ = src_db.close();
        result
    }
}

#[cfg(test)]
mod test {
    use SqliteConnection;
    use {SQLITE_OPEN_URI, SQLITE_OPEN_CREATE, SQLITE_OPEN_READ_WRITE};

    #[test]
    fn test_backup() {
        let db = SqliteConnection::open_in_memory().unwrap();
        let sql = "BEGIN;
                CREATE TABLE foo(x INTEGER);
                INSERT INTO foo VALUES(42);
                END;";
        db.execute_batch(sql).unwrap();
        db.backup("main", &":memory:").unwrap();
        assert!(db.close().is_ok());
    }

    #[test]
    fn test_restore() {
        let src_path = "file:dummy.db?mode=memory&cache=shared";
        let src = SqliteConnection::open_with_flags(&src_path, SQLITE_OPEN_URI | SQLITE_OPEN_CREATE | SQLITE_OPEN_READ_WRITE).unwrap();
        let sql = "BEGIN;
               CREATE TABLE foo(x INTEGER);
               INSERT INTO foo VALUES(42);
               END;";
        src.execute_batch(sql).unwrap();

        let db = SqliteConnection::open_in_memory().unwrap();
        db.restore("main", &src_path).unwrap();

        let the_answer = db.query_row("SELECT x FROM foo",
                                           &[],
                                           |r| r.get::<i64>(0));
        assert_eq!(42i64, the_answer.unwrap());

        assert!(db.close().is_ok());
        assert!(src.close().is_ok());
    }
}
<commit_msg>Cleanup use statement.<commit_after>//! Online backup API
use std::path::Path;

use super::ffi;
use {SqliteError, SqliteResult, SqliteConnection};

impl SqliteConnection {
    /// Open or create a database file named `dest_path`. Transfer the
    /// content of local database `name` into the `dest_path` database.
    pub fn backup<P: AsRef<Path>>(&self, name: &str, dest_path: &P) -> SqliteResult<()> {
        let dest = try!(SqliteConnection::open(dest_path)); // "cannot open target database: "

        let mut dest_db = dest.db.borrow_mut();
        let dest_name = try!(super::str_to_cstring("main"));
        let src_db = self.db.borrow_mut();
        let src_name = try!(super::str_to_cstring(name));
        let backup = unsafe { ffi::sqlite3_backup_init(dest_db.db(), dest_name.as_ptr(), src_db.db(), src_name.as_ptr()) };
        if backup.is_null() {
            let msg = unsafe { super::errmsg_to_string(ffi::sqlite3_errmsg(dest_db.db())) };
            let err = Err(SqliteError{code: unsafe { ffi::sqlite3_errcode(dest_db.db()) }, message: "backup failed: ".to_string() + &msg });
            let _ = dest_db.close();
            return err;
        }

        let mut rc;
        loop {
            rc = unsafe { ffi::sqlite3_backup_step(backup, 100) };
            if rc != ffi::SQLITE_OK {
                break;
            }
        }
        unsafe { ffi::sqlite3_backup_finish(backup) };
        let result = if rc == ffi::SQLITE_DONE {
            Ok(())
        } else {
            let msg = unsafe { super::errmsg_to_string(ffi::sqlite3_errmsg(dest_db.db())) };
            Err(SqliteError{code: rc, message: "backup failed: ".to_string() + &msg })
        };
        let _ = dest_db.close();
        result
    }

    /// Open a database file named `src_path`. Transfer the content
    /// of `src_path` into the local database `name`.
    pub fn restore<P: AsRef<Path>>(&self, name: &str, src_path: &P) -> SqliteResult<()> {
        let src = try!(SqliteConnection::open_with_flags(src_path, super::SQLITE_OPEN_READ_ONLY | super::SQLITE_OPEN_URI)); // "cannot open source database: "

        let dest_db = self.db.borrow_mut();
        let dest_name = try!(super::str_to_cstring(name));
        let mut src_db = src.db.borrow_mut();
        let src_name = try!(super::str_to_cstring("main"));
        let backup = unsafe { ffi::sqlite3_backup_init(dest_db.db(), dest_name.as_ptr(), src_db.db(), src_name.as_ptr()) };
        if backup.is_null() {
            let msg = unsafe { super::errmsg_to_string(ffi::sqlite3_errmsg(dest_db.db())) };
            let err = Err(SqliteError{code: unsafe { ffi::sqlite3_errcode(dest_db.db()) }, message: "restore failed: ".to_string() + &msg });
            let _ = src_db.close();
            return err;
        }
        let mut rc;
        let mut n_timeout = 0;
        loop {
            rc = unsafe { ffi::sqlite3_backup_step(backup, 100) };
            if rc == ffi::SQLITE_OK {
            } else if rc == ffi::SQLITE_BUSY {
                n_timeout += 1;
                if n_timeout >= 3 {
                    break;
                }
                unsafe { ffi::sqlite3_sleep(100) };
            } else {
                break;
            }
        }
        unsafe { ffi::sqlite3_backup_finish(backup) };
        let result = if rc == ffi::SQLITE_DONE {
            Ok(())
        } else if rc == ffi::SQLITE_BUSY || rc == ffi::SQLITE_LOCKED {
            Err(SqliteError{code: rc, message: "restore failed: source database busy".to_string() })
        } else {
            let msg = unsafe { super::errmsg_to_string(ffi::sqlite3_errmsg(dest_db.db())) };
            Err(SqliteError{code: rc, message: "restore failed: ".to_string() + &msg })
        };
        let _ = src_db.close();
        result
    }
}

#[cfg(test)]
mod test {
    use SqliteConnection;
    use {SQLITE_OPEN_URI, SQLITE_OPEN_CREATE, SQLITE_OPEN_READ_WRITE};

    #[test]
    fn test_backup() {
        let db = SqliteConnection::open_in_memory().unwrap();
        let sql = "BEGIN;
                CREATE TABLE foo(x INTEGER);
                INSERT INTO foo VALUES(42);
                END;";
        db.execute_batch(sql).unwrap();
        db.backup("main", &":memory:").unwrap();
        assert!(db.close().is_ok());
    }

    #[test]
    fn test_restore() {
        let src_path = "file:dummy.db?mode=memory&cache=shared";
        let src = SqliteConnection::open_with_flags(&src_path, SQLITE_OPEN_URI | SQLITE_OPEN_CREATE | SQLITE_OPEN_READ_WRITE).unwrap();
        let sql = "BEGIN;
               CREATE TABLE foo(x INTEGER);
               INSERT INTO foo VALUES(42);
               END;";
        src.execute_batch(sql).unwrap();

        let db = SqliteConnection::open_in_memory().unwrap();
        db.restore("main", &src_path).unwrap();

        let the_answer = db.query_row("SELECT x FROM foo",
                                           &[],
                                           |r| r.get::<i64>(0));
        assert_eq!(42i64, the_answer.unwrap());

        assert!(db.close().is_ok());
        assert!(src.close().is_ok());
    }
}
<|endoftext|>


<commit_before>use sodiumoxide::crypto::{
    hash::sha256::DIGESTBYTES,
    sign::{PUBLICKEYBYTES, SIGNATUREBYTES},
};
use std::io::{Cursor, Error, ErrorKind, Read};

use crate::asset::Asset;
use crate::crypto::{PublicKey, ScriptHash, SigPair, Signature};

macro_rules! read_exact_bytes {
    ($self:expr, $len:expr) => {{
        let mut buf = Vec::with_capacity($len);
        unsafe {
            buf.set_len($len);
        }
        $self.read_exact(&mut buf)?;
        buf
    }};
}

#[inline]
fn zigzag_encode(from: i64) -> u64 {
    ((from << 1) ^ (from >> 63)) as u64
}

#[inline]
fn zigzag_decode(from: u64) -> i64 {
    ((from >> 1) ^ (-((from & 1) as i64)) as u64) as i64
}

pub trait BufWrite {
    fn push_u16(&mut self, num: u16);
    fn push_u32(&mut self, num: u32);
    fn push_i64(&mut self, num: i64);
    fn push_var_i64(&mut self, num: i64);
    fn push_u64(&mut self, num: u64);
    fn push_bytes(&mut self, slice: &[u8]);
    fn push_pub_key(&mut self, key: &PublicKey);
    fn push_script_hash(&mut self, hash: &ScriptHash);
    fn push_sig_pair(&mut self, pair: &SigPair);
    fn push_asset(&mut self, asset: Asset);
}

impl BufWrite for Vec<u8> {
    fn push_u16(&mut self, num: u16) {
        self.push((num >> 8) as u8);
        self.push(num as u8);
    }

    fn push_u32(&mut self, num: u32) {
        self.push((num >> 24) as u8);
        self.push((num >> 16) as u8);
        self.push((num >> 8) as u8);
        self.push(num as u8);
    }

    fn push_i64(&mut self, num: i64) {
        self.push((num >> 56) as u8);
        self.push((num >> 48) as u8);
        self.push((num >> 40) as u8);
        self.push((num >> 32) as u8);
        self.push((num >> 24) as u8);
        self.push((num >> 16) as u8);
        self.push((num >> 8) as u8);
        self.push(num as u8);
    }

    fn push_var_i64(&mut self, num: i64) {
        let mut more = true;
        let mut num = zigzag_encode(num);

        while more {
            let mut byte: u8 = (num & 0x7F) as u8;
            num >>= 7;

            if num == 0 {
                more = false;
            } else {
                byte |= 0x80;
            }

            self.push(byte);
        }
    }

    fn push_u64(&mut self, num: u64) {
        self.push((num >> 56) as u8);
        self.push((num >> 48) as u8);
        self.push((num >> 40) as u8);
        self.push((num >> 32) as u8);
        self.push((num >> 24) as u8);
        self.push((num >> 16) as u8);
        self.push((num >> 8) as u8);
        self.push(num as u8);
    }

    fn push_bytes(&mut self, other: &[u8]) {
        if other.is_empty() {
            self.push_u32(0);
            return;
        }
        self.push_u32(other.len() as u32);
        self.extend_from_slice(other);
    }

    fn push_pub_key(&mut self, key: &PublicKey) {
        self.extend_from_slice(key.as_ref());
    }

    fn push_script_hash(&mut self, hash: &ScriptHash) {
        self.extend_from_slice(hash.as_ref());
    }

    fn push_sig_pair(&mut self, pair: &SigPair) {
        self.push_pub_key(&pair.pub_key);
        self.extend_from_slice(pair.signature.as_ref());
    }

    fn push_asset(&mut self, asset: Asset) {
        self.push_i64(asset.amount);
    }
}

pub trait BufRead {
    fn take_u8(&mut self) -> Result<u8, Error>;
    fn take_u16(&mut self) -> Result<u16, Error>;
    fn take_u32(&mut self) -> Result<u32, Error>;
    fn take_i64(&mut self) -> Result<i64, Error>;
    fn take_var_i64(&mut self) -> Result<i64, Error>;
    fn take_u64(&mut self) -> Result<u64, Error>;
    fn take_bytes(&mut self) -> Result<Vec<u8>, Error>;
    fn take_pub_key(&mut self) -> Result<PublicKey, Error>;
    fn take_script_hash(&mut self) -> Result<ScriptHash, Error>;
    fn take_sig_pair(&mut self) -> Result<SigPair, Error>;
    fn take_asset(&mut self) -> Result<Asset, Error>;
}

impl<T: AsRef<[u8]> + Read> BufRead for Cursor<T> {
    fn take_u8(&mut self) -> Result<u8, Error> {
        let mut buf = [0u8; 1];
        self.read_exact(&mut buf)?;
        Ok(buf[0])
    }

    fn take_u16(&mut self) -> Result<u16, Error> {
        let mut buf = [0u8; 2];
        self.read_exact(&mut buf)?;
        Ok((u16::from(buf[0]) << 8) | u16::from(buf[1]))
    }

    fn take_u32(&mut self) -> Result<u32, Error> {
        let mut buf = [0u8; 4];
        self.read_exact(&mut buf)?;
        Ok(u32_from_buf!(buf))
    }

    fn take_i64(&mut self) -> Result<i64, Error> {
        let mut buf = [0u8; 8];
        self.read_exact(&mut buf)?;
        Ok((i64::from(buf[0]) << 56)
            | (i64::from(buf[1]) << 48)
            | (i64::from(buf[2]) << 40)
            | (i64::from(buf[3]) << 32)
            | (i64::from(buf[4]) << 24)
            | (i64::from(buf[5]) << 16)
            | (i64::from(buf[6]) << 8)
            | i64::from(buf[7]))
    }

    fn take_var_i64(&mut self) -> Result<i64, Error> {
        let mut result: u64 = 0;
        let mut shift = 0;
        let mut buf = [0u8; 1];
        loop {
            if shift > 63 {
                return Err(Error::new(ErrorKind::Other, "overflow taking varint"));
            }

            self.read_exact(&mut buf)?;
            let byte = buf[0];

            result |= ((byte & 0x7F) as u64) << shift;
            if byte & 0x80 == 0 {
                break;
            }

            shift += 7;
        }

        Ok(zigzag_decode(result))
    }

    fn take_u64(&mut self) -> Result<u64, Error> {
        let mut buf = [0u8; 8];
        self.read_exact(&mut buf)?;
        Ok(u64_from_buf!(buf))
    }

    fn take_bytes(&mut self) -> Result<Vec<u8>, Error> {
        let len = self.take_u32()? as usize;
        let buf = read_exact_bytes!(self, len);
        Ok(buf)
    }

    fn take_pub_key(&mut self) -> Result<PublicKey, Error> {
        let buf = read_exact_bytes!(self, PUBLICKEYBYTES);
        PublicKey::from_slice(&buf)
            .ok_or_else(|| Error::new(ErrorKind::Other, "incorrect public key length"))
    }

    fn take_script_hash(&mut self) -> Result<ScriptHash, Error> {
        let buf = read_exact_bytes!(self, DIGESTBYTES);
        ScriptHash::from_slice(&buf)
            .ok_or_else(|| Error::new(ErrorKind::Other, "incorrect script hash length"))
    }

    fn take_sig_pair(&mut self) -> Result<SigPair, Error> {
        let pub_key = self.take_pub_key()?;
        let signature = {
            let buf = read_exact_bytes!(self, SIGNATUREBYTES);
            Signature::from_slice(&buf)
                .ok_or_else(|| Error::new(ErrorKind::Other, "incorrect signature length"))?
        };
        Ok(SigPair { pub_key, signature })
    }

    fn take_asset(&mut self) -> Result<Asset, Error> {
        let amount = self.take_i64()?;
        Ok(Asset::new(amount))
    }
}

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

    #[test]
    fn u32_serialization() {
        let num: u32 = 0x0A0B0C0D;
        let mut v = Vec::with_capacity(4);
        v.push_u32(num);
        assert_eq!(v, [0x0A, 0x0B, 0x0C, 0x0D]);
    }

    #[test]
    fn u64_serialization() {
        let num: u64 = 0x0A0B0C0D_0A0B0C0D;
        let mut v = Vec::with_capacity(8);
        v.push_u64(num);

        assert_eq!(v, [0x0A, 0x0B, 0x0C, 0x0D, 0x0A, 0x0B, 0x0C, 0x0D]);
        let dec = Cursor::<&[u8]>::new(&v).take_u64().unwrap();
        assert_eq!(num, dec);
    }

    #[test]
    fn asset_serialization() {
        let a = "12.34567 GRAEL".parse().unwrap();
        let mut v = vec![];
        v.push_asset(a);

        let mut c = Cursor::<&[u8]>::new(&v);
        let b = c.take_asset().unwrap();
        assert_eq!(a.to_string(), b.to_string());
    }

    #[test]
    fn zigzag() {
        fn cmp(decoded: i64, encoded: u64) {
            assert_eq!(decoded, zigzag_decode(encoded));
            assert_eq!(encoded, zigzag_encode(decoded));
        }

        cmp(0, 0);
        cmp(-1, 1);
        cmp(1, 2);
        cmp(-2, 3);
        cmp(2147483647, 4294967294);
        cmp(-2147483648, 4294967295);
        cmp(9223372036854775807, 18446744073709551614);
        cmp(-9223372036854775808, 18446744073709551615);
    }

    #[test]
    fn var_i64_serialization() {
        let mut buf = vec![];
        buf.push_var_i64(0);
        buf.push_var_i64(300);
        buf.push_var_i64(-300);
        buf.push_var_i64(i64::max_value());
        buf.push_var_i64(i64::min_value());
        // Outputs 0 as only the first bit is checked on the final byte with a shift of 63
        buf.extend(vec![
            0x80, 0x80, 0x80, 0x80, 0x80, 0x80, 0x80, 0x80, 0x80, 0x7E,
        ]);
        // Outputs 1 << 62 as bit 62 is set
        buf.extend(vec![
            0x80, 0x80, 0x80, 0x80, 0x80, 0x80, 0x80, 0x80, 0x80, 0x01,
        ]);

        let mut c = Cursor::<&[u8]>::new(&buf);
        assert_eq!(c.take_var_i64().unwrap(), 0);
        assert_eq!(c.take_var_i64().unwrap(), 300);
        assert_eq!(c.take_var_i64().unwrap(), -300);
        assert_eq!(c.take_var_i64().unwrap(), i64::max_value());
        assert_eq!(c.take_var_i64().unwrap(), i64::min_value());
        assert_eq!(c.take_var_i64().unwrap(), 0);
        assert_eq!(c.take_var_i64().unwrap(), 1 << 62);
    }

    #[test]
    fn var_i64_serialization_overflow() {
        use std::error;

        let buf = vec![
            0x80, 0x80, 0x80, 0x80, 0x80, 0x80, 0x80, 0x80, 0x80, 0x80, 0,
        ];
        let mut c = Cursor::<&[u8]>::new(&buf);
        assert_eq!(
            error::Error::description(&c.take_var_i64().unwrap_err()),
            "overflow taking varint"
        );
    }

    #[test]
    fn var_i64_serialization_eof() {
        let buf = vec![0x80, 0x80, 0x80, 0x80, 0x80, 0x80, 0x80, 0x80, 0x80];
        let mut c = Cursor::<&[u8]>::new(&buf);
        assert_eq!(
            c.take_var_i64().unwrap_err().kind(),
            ErrorKind::UnexpectedEof
        );
    }
}
<commit_msg>Serializing assets now uses variable length i64<commit_after>use sodiumoxide::crypto::{
    hash::sha256::DIGESTBYTES,
    sign::{PUBLICKEYBYTES, SIGNATUREBYTES},
};
use std::io::{Cursor, Error, ErrorKind, Read};

use crate::asset::Asset;
use crate::crypto::{PublicKey, ScriptHash, SigPair, Signature};

macro_rules! read_exact_bytes {
    ($self:expr, $len:expr) => {{
        let mut buf = Vec::with_capacity($len);
        unsafe {
            buf.set_len($len);
        }
        $self.read_exact(&mut buf)?;
        buf
    }};
}

#[inline]
fn zigzag_encode(from: i64) -> u64 {
    ((from << 1) ^ (from >> 63)) as u64
}

#[inline]
fn zigzag_decode(from: u64) -> i64 {
    ((from >> 1) ^ (-((from & 1) as i64)) as u64) as i64
}

pub trait BufWrite {
    fn push_u16(&mut self, num: u16);
    fn push_u32(&mut self, num: u32);
    fn push_i64(&mut self, num: i64);
    fn push_var_i64(&mut self, num: i64);
    fn push_u64(&mut self, num: u64);
    fn push_bytes(&mut self, slice: &[u8]);
    fn push_pub_key(&mut self, key: &PublicKey);
    fn push_script_hash(&mut self, hash: &ScriptHash);
    fn push_sig_pair(&mut self, pair: &SigPair);
    fn push_asset(&mut self, asset: Asset);
}

impl BufWrite for Vec<u8> {
    fn push_u16(&mut self, num: u16) {
        self.push((num >> 8) as u8);
        self.push(num as u8);
    }

    fn push_u32(&mut self, num: u32) {
        self.push((num >> 24) as u8);
        self.push((num >> 16) as u8);
        self.push((num >> 8) as u8);
        self.push(num as u8);
    }

    fn push_i64(&mut self, num: i64) {
        self.push((num >> 56) as u8);
        self.push((num >> 48) as u8);
        self.push((num >> 40) as u8);
        self.push((num >> 32) as u8);
        self.push((num >> 24) as u8);
        self.push((num >> 16) as u8);
        self.push((num >> 8) as u8);
        self.push(num as u8);
    }

    fn push_var_i64(&mut self, num: i64) {
        let mut more = true;
        let mut num = zigzag_encode(num);

        while more {
            let mut byte: u8 = (num & 0x7F) as u8;
            num >>= 7;

            if num == 0 {
                more = false;
            } else {
                byte |= 0x80;
            }

            self.push(byte);
        }
    }

    fn push_u64(&mut self, num: u64) {
        self.push((num >> 56) as u8);
        self.push((num >> 48) as u8);
        self.push((num >> 40) as u8);
        self.push((num >> 32) as u8);
        self.push((num >> 24) as u8);
        self.push((num >> 16) as u8);
        self.push((num >> 8) as u8);
        self.push(num as u8);
    }

    fn push_bytes(&mut self, other: &[u8]) {
        if other.is_empty() {
            self.push_u32(0);
            return;
        }
        self.push_u32(other.len() as u32);
        self.extend_from_slice(other);
    }

    fn push_pub_key(&mut self, key: &PublicKey) {
        self.extend_from_slice(key.as_ref());
    }

    fn push_script_hash(&mut self, hash: &ScriptHash) {
        self.extend_from_slice(hash.as_ref());
    }

    fn push_sig_pair(&mut self, pair: &SigPair) {
        self.push_pub_key(&pair.pub_key);
        self.extend_from_slice(pair.signature.as_ref());
    }

    fn push_asset(&mut self, asset: Asset) {
        self.push_var_i64(asset.amount);
    }
}

pub trait BufRead {
    fn take_u8(&mut self) -> Result<u8, Error>;
    fn take_u16(&mut self) -> Result<u16, Error>;
    fn take_u32(&mut self) -> Result<u32, Error>;
    fn take_i64(&mut self) -> Result<i64, Error>;
    fn take_var_i64(&mut self) -> Result<i64, Error>;
    fn take_u64(&mut self) -> Result<u64, Error>;
    fn take_bytes(&mut self) -> Result<Vec<u8>, Error>;
    fn take_pub_key(&mut self) -> Result<PublicKey, Error>;
    fn take_script_hash(&mut self) -> Result<ScriptHash, Error>;
    fn take_sig_pair(&mut self) -> Result<SigPair, Error>;
    fn take_asset(&mut self) -> Result<Asset, Error>;
}

impl<T: AsRef<[u8]> + Read> BufRead for Cursor<T> {
    fn take_u8(&mut self) -> Result<u8, Error> {
        let mut buf = [0u8; 1];
        self.read_exact(&mut buf)?;
        Ok(buf[0])
    }

    fn take_u16(&mut self) -> Result<u16, Error> {
        let mut buf = [0u8; 2];
        self.read_exact(&mut buf)?;
        Ok((u16::from(buf[0]) << 8) | u16::from(buf[1]))
    }

    fn take_u32(&mut self) -> Result<u32, Error> {
        let mut buf = [0u8; 4];
        self.read_exact(&mut buf)?;
        Ok(u32_from_buf!(buf))
    }

    fn take_i64(&mut self) -> Result<i64, Error> {
        let mut buf = [0u8; 8];
        self.read_exact(&mut buf)?;
        Ok((i64::from(buf[0]) << 56)
            | (i64::from(buf[1]) << 48)
            | (i64::from(buf[2]) << 40)
            | (i64::from(buf[3]) << 32)
            | (i64::from(buf[4]) << 24)
            | (i64::from(buf[5]) << 16)
            | (i64::from(buf[6]) << 8)
            | i64::from(buf[7]))
    }

    fn take_var_i64(&mut self) -> Result<i64, Error> {
        let mut result: u64 = 0;
        let mut shift = 0;
        let mut buf = [0u8; 1];
        loop {
            if shift > 63 {
                return Err(Error::new(ErrorKind::Other, "overflow taking varint"));
            }

            self.read_exact(&mut buf)?;
            let byte = buf[0];

            result |= ((byte & 0x7F) as u64) << shift;
            if byte & 0x80 == 0 {
                break;
            }

            shift += 7;
        }

        Ok(zigzag_decode(result))
    }

    fn take_u64(&mut self) -> Result<u64, Error> {
        let mut buf = [0u8; 8];
        self.read_exact(&mut buf)?;
        Ok(u64_from_buf!(buf))
    }

    fn take_bytes(&mut self) -> Result<Vec<u8>, Error> {
        let len = self.take_u32()? as usize;
        let buf = read_exact_bytes!(self, len);
        Ok(buf)
    }

    fn take_pub_key(&mut self) -> Result<PublicKey, Error> {
        let buf = read_exact_bytes!(self, PUBLICKEYBYTES);
        PublicKey::from_slice(&buf)
            .ok_or_else(|| Error::new(ErrorKind::Other, "incorrect public key length"))
    }

    fn take_script_hash(&mut self) -> Result<ScriptHash, Error> {
        let buf = read_exact_bytes!(self, DIGESTBYTES);
        ScriptHash::from_slice(&buf)
            .ok_or_else(|| Error::new(ErrorKind::Other, "incorrect script hash length"))
    }

    fn take_sig_pair(&mut self) -> Result<SigPair, Error> {
        let pub_key = self.take_pub_key()?;
        let signature = {
            let buf = read_exact_bytes!(self, SIGNATUREBYTES);
            Signature::from_slice(&buf)
                .ok_or_else(|| Error::new(ErrorKind::Other, "incorrect signature length"))?
        };
        Ok(SigPair { pub_key, signature })
    }

    fn take_asset(&mut self) -> Result<Asset, Error> {
        let amount = self.take_var_i64()?;
        Ok(Asset::new(amount))
    }
}

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

    #[test]
    fn u32_serialization() {
        let num: u32 = 0x0A0B0C0D;
        let mut v = Vec::with_capacity(4);
        v.push_u32(num);
        assert_eq!(v, [0x0A, 0x0B, 0x0C, 0x0D]);
    }

    #[test]
    fn u64_serialization() {
        let num: u64 = 0x0A0B0C0D_0A0B0C0D;
        let mut v = Vec::with_capacity(8);
        v.push_u64(num);

        assert_eq!(v, [0x0A, 0x0B, 0x0C, 0x0D, 0x0A, 0x0B, 0x0C, 0x0D]);
        let dec = Cursor::<&[u8]>::new(&v).take_u64().unwrap();
        assert_eq!(num, dec);
    }

    #[test]
    fn asset_serialization() {
        let a = "12.34567 GRAEL".parse().unwrap();
        let mut v = vec![];
        v.push_asset(a);

        let mut c = Cursor::<&[u8]>::new(&v);
        let b = c.take_asset().unwrap();
        assert_eq!(a.to_string(), b.to_string());
    }

    #[test]
    fn zigzag() {
        fn cmp(decoded: i64, encoded: u64) {
            assert_eq!(decoded, zigzag_decode(encoded));
            assert_eq!(encoded, zigzag_encode(decoded));
        }

        cmp(0, 0);
        cmp(-1, 1);
        cmp(1, 2);
        cmp(-2, 3);
        cmp(2147483647, 4294967294);
        cmp(-2147483648, 4294967295);
        cmp(9223372036854775807, 18446744073709551614);
        cmp(-9223372036854775808, 18446744073709551615);
    }

    #[test]
    fn var_i64_serialization() {
        let mut buf = vec![];
        buf.push_var_i64(0);
        buf.push_var_i64(300);
        buf.push_var_i64(-300);
        buf.push_var_i64(i64::max_value());
        buf.push_var_i64(i64::min_value());
        // Outputs 0 as only the first bit is checked on the final byte with a shift of 63
        buf.extend(vec![
            0x80, 0x80, 0x80, 0x80, 0x80, 0x80, 0x80, 0x80, 0x80, 0x7E,
        ]);
        // Outputs 1 << 62 as bit 62 is set
        buf.extend(vec![
            0x80, 0x80, 0x80, 0x80, 0x80, 0x80, 0x80, 0x80, 0x80, 0x01,
        ]);

        let mut c = Cursor::<&[u8]>::new(&buf);
        assert_eq!(c.take_var_i64().unwrap(), 0);
        assert_eq!(c.take_var_i64().unwrap(), 300);
        assert_eq!(c.take_var_i64().unwrap(), -300);
        assert_eq!(c.take_var_i64().unwrap(), i64::max_value());
        assert_eq!(c.take_var_i64().unwrap(), i64::min_value());
        assert_eq!(c.take_var_i64().unwrap(), 0);
        assert_eq!(c.take_var_i64().unwrap(), 1 << 62);
    }

    #[test]
    fn var_i64_serialization_overflow() {
        use std::error;

        let buf = vec![
            0x80, 0x80, 0x80, 0x80, 0x80, 0x80, 0x80, 0x80, 0x80, 0x80, 0,
        ];
        let mut c = Cursor::<&[u8]>::new(&buf);
        assert_eq!(
            error::Error::description(&c.take_var_i64().unwrap_err()),
            "overflow taking varint"
        );
    }

    #[test]
    fn var_i64_serialization_eof() {
        let buf = vec![0x80, 0x80, 0x80, 0x80, 0x80, 0x80, 0x80, 0x80, 0x80];
        let mut c = Cursor::<&[u8]>::new(&buf);
        assert_eq!(
            c.take_var_i64().unwrap_err().kind(),
            ErrorKind::UnexpectedEof
        );
    }
}
<|endoftext|>


<commit_before>use std::rc::{Rc};
use std::cell::{RefCell};

pub struct Broker<T: Subscriber> {
    subscriptions: Subscriptions<T>,
}

impl<T: Subscriber> Broker<T> {
    pub fn new() -> Self {
        Broker{ subscriptions: Subscriptions::<T>::new() }
    }

    pub fn publish(&mut self, topic: &str, payload: &[u8]) {
        self.subscriptions.publish(topic, payload);
    }

    pub fn subscribe(&mut self, subscriber: Rc<RefCell<T>>, topics: &[&str]) {
        let subscriber = self.add_or_find(subscriber);
        subscriber.borrow_mut().append_topics(topics);
    }

    fn add_or_find(&mut self, subscriber: Rc<RefCell<T>>) -> Rc<RefCell<T>> {
        for sub in &self.subscriptions.subscribers {
            let sub1 = sub.clone();
            let sub2 = subscriber.clone();
            //this is horrible
            //borrow doesn't return T, it returns Ref<T>
            //since Ref<T> itself doesn't satisfy the template contraint, we need to deref
            //it to get the to the T inside. And since is_same takes borrows, we reapply
            //the ampersand
            if is_same(&*sub1.borrow(), &*sub2.borrow()) {
                return sub2;
            }
        }

        self.subscriptions.subscribers.push(subscriber);
        self.subscriptions.subscribers[self.subscriptions.subscribers.len() - 1].clone()
    }
}

//is same identity
fn is_same<T>(lhs: &T, rhs: &T) -> bool {
    return lhs as *const T as i64 == rhs as *const T as i64
}

pub struct Subscriptions<T: Subscriber> {
    subscribers: Vec<Rc<RefCell<T>>>,
}

impl<T: Subscriber> Subscriptions<T> {
    fn new() -> Self {
        return Subscriptions { subscribers: vec![], }
    }

    fn publish(&mut self, topic: &str, payload: &[u8]) {
        for s in &mut self.subscribers {
            let sub_clone = s.clone();
            let mut s = sub_clone.borrow_mut();

            for t in s.topics() {
                if t == topic {
                    s.new_message(payload);
                }
            }
        }
    }
}

pub trait Subscriber {
    fn new_message(&mut self, bytes: &[u8]);
    fn append_topics(&mut self, topics: &[&str]);
    fn topics(&self) -> Vec<String>;
}

#[cfg(test)]
struct TestSubscriber {
    msgs: Vec<Vec<u8>>,
    topics: Vec<String>,
}
#[cfg(test)]
impl TestSubscriber {
    fn new() -> Self {
        return TestSubscriber{msgs: vec![], topics: vec![]}
    }
}

#[cfg(test)]
impl Subscriber for TestSubscriber {

    fn new_message(&mut self, bytes: &[u8]) {
        self.msgs.push(bytes.to_vec());
    }

    fn append_topics(&mut self, topics: &[&str]) {
        let mut new_topics = topics.to_vec().into_iter().map(|t| t.to_string()).collect();
        self.topics.append(&mut new_topics);
    }

    fn topics(&self) -> Vec<String> {
        self.topics.clone()
    }
}

#[test]
fn test_subscribe() {
    let mut broker = Broker::<TestSubscriber>::new();
    let sub_ref = Rc::new(RefCell::new(TestSubscriber::new()));
    let subscriber = sub_ref.clone();
    broker.publish("topics/foo", &[0, 1, 2]);
    assert_eq!(subscriber.borrow().msgs.len(), 0);

    broker.subscribe(subscriber.clone(), &["topics/foo"]);
    broker.publish("topics/foo", &[0, 1, 9]);
    broker.publish("topics/bar", &[2, 4, 6]);
    assert_eq!(subscriber.borrow().msgs.len(), 1);
    assert_eq!(subscriber.borrow().msgs[0], &[0, 1, 9]);

    broker.subscribe(subscriber.clone(), &["topics/bar"]);
    broker.publish("topics/foo", &[1, 3, 5, 7]);
    broker.publish("topics/bar", &[2, 4]);
    assert_eq!(subscriber.borrow().msgs.len(), 3);
    assert_eq!(subscriber.borrow().msgs[0], &[0, 1, 9]);
    assert_eq!(subscriber.borrow().msgs[1], &[1, 3, 5, 7]);
    assert_eq!(subscriber.borrow().msgs[2], &[2, 4]);
}
<commit_msg>Stopped breaking the law of Demeter<commit_after>use std::rc::{Rc};
use std::cell::{RefCell};

pub struct Broker<T: Subscriber> {
    subscriptions: Subscriptions<T>,
}

impl<T: Subscriber> Broker<T> {
    pub fn new() -> Self {
        Broker{ subscriptions: Subscriptions::<T>::new() }
    }

    pub fn publish(&mut self, topic: &str, payload: &[u8]) {
        self.subscriptions.publish(topic, payload);
    }

    pub fn subscribe(&mut self, subscriber: Rc<RefCell<T>>, topics: &[&str]) {
        self.subscriptions.subscribe(subscriber, &topics)
    }
}

//is same identity
fn is_same<T>(lhs: &T, rhs: &T) -> bool {
    return lhs as *const T as i64 == rhs as *const T as i64
}

pub struct Subscriptions<T: Subscriber> {
    subscribers: Vec<Rc<RefCell<T>>>,
}

impl<T: Subscriber> Subscriptions<T> {
    fn new() -> Self {
        return Subscriptions { subscribers: vec![], }
    }

    fn subscribe(&mut self, subscriber: Rc<RefCell<T>>, topics: &[&str]) {
        let subscriber = self.add_or_find(subscriber);
        subscriber.borrow_mut().append_topics(topics);
    }

    fn add_or_find(&mut self, subscriber: Rc<RefCell<T>>) -> Rc<RefCell<T>> {
        for sub in &self.subscribers {
            let sub1 = sub.clone();
            let sub2 = subscriber.clone();
            //this is horrible
            //borrow doesn't return T, it returns Ref<T>
            //since Ref<T> itself doesn't satisfy the template contraint, we need to deref
            //it to get the to the T inside. And since is_same takes borrows, we reapply
            //the ampersand
            if is_same(&*sub1.borrow(), &*sub2.borrow()) {
                return sub2;
            }
        }

        self.subscribers.push(subscriber);
        self.subscribers[self.subscribers.len() - 1].clone()
    }

    fn publish(&mut self, topic: &str, payload: &[u8]) {
        for subscriber in &mut self.subscribers {
            let subscriber = subscriber.clone();
            let mut subscriber = subscriber.borrow_mut();

            for _ in subscriber.topics().iter().filter(|&t| topic_matches(&topic, &t)) {
                subscriber.new_message(payload);
            }
        }
    }
}

fn topic_matches(pub_topic: &str, sub_topic: &str) -> bool {
    pub_topic == sub_topic
}

pub trait Subscriber {
    fn new_message(&mut self, bytes: &[u8]);
    fn append_topics(&mut self, topics: &[&str]);
    fn topics(&self) -> Vec<String>;
}

pub trait SubscriberT {
    fn new_message(&mut self, bytes: &[u8]);
}

struct Subscription<T: SubscriberT> {
    subscriber: Rc<RefCell<T>>,
    topics: Vec<String>,
}

impl<T: SubscriberT> Subscription<T> {
    fn append_topics(&mut self, topics: &[&str]) {
        let mut new_topics = topics.to_vec().into_iter().map(|t| t.to_string()).collect();
        self.topics.append(&mut new_topics);
    }

    fn topics(&self) -> Vec<String> {
        self.topics.clone()
    }
}

#[cfg(test)]
struct TestSubscriber {
    msgs: Vec<Vec<u8>>,
    topics: Vec<String>,
}
#[cfg(test)]
impl TestSubscriber {
    fn new() -> Self {
        return TestSubscriber{msgs: vec![], topics: vec![]}
    }
}

#[cfg(test)]
impl Subscriber for TestSubscriber {

    fn new_message(&mut self, bytes: &[u8]) {
        self.msgs.push(bytes.to_vec());
    }

    fn append_topics(&mut self, topics: &[&str]) {
        let mut new_topics = topics.to_vec().into_iter().map(|t| t.to_string()).collect();
        self.topics.append(&mut new_topics);
    }

    fn topics(&self) -> Vec<String> {
        self.topics.clone()
    }
}

#[test]
fn test_subscribe() {
    let mut broker = Broker::<TestSubscriber>::new();
    let sub_ref = Rc::new(RefCell::new(TestSubscriber::new()));
    let subscriber = sub_ref.clone();
    broker.publish("topics/foo", &[0, 1, 2]);
    assert_eq!(subscriber.borrow().msgs.len(), 0);

    broker.subscribe(subscriber.clone(), &["topics/foo"]);
    broker.publish("topics/foo", &[0, 1, 9]);
    broker.publish("topics/bar", &[2, 4, 6]);
    assert_eq!(subscriber.borrow().msgs.len(), 1);
    assert_eq!(subscriber.borrow().msgs[0], &[0, 1, 9]);

    broker.subscribe(subscriber.clone(), &["topics/bar"]);
    broker.publish("topics/foo", &[1, 3, 5, 7]);
    broker.publish("topics/bar", &[2, 4]);
    assert_eq!(subscriber.borrow().msgs.len(), 3);
    assert_eq!(subscriber.borrow().msgs[0], &[0, 1, 9]);
    assert_eq!(subscriber.borrow().msgs[1], &[1, 3, 5, 7]);
    assert_eq!(subscriber.borrow().msgs[2], &[2, 4]);
}
<|endoftext|>


<commit_before>use std::cmp;
use std::iter;
use std::io::{self, Read, BufRead, Cursor};

pub struct BufReader<R> {
    buf: Cursor<Vec<u8>>,
    inner: R
}

const INIT_BUFFER_SIZE: usize = 4096;
const MAX_BUFFER_SIZE: usize = 8192 + 4096 * 100;

impl<R: Read> BufReader<R> {
    pub fn new(rdr: R) -> BufReader<R> {
        BufReader::with_capacity(rdr, INIT_BUFFER_SIZE)
    }

    pub fn with_capacity(rdr: R, cap: usize) -> BufReader<R> {
        BufReader {
            buf: Cursor::new(Vec::with_capacity(cap)),
            inner: rdr
        }
    }

    pub fn get_ref(&self) -> &R { &self.inner }

    pub fn get_mut(&mut self) -> &mut R { &mut self.inner }

    pub fn get_buf(&self) -> &[u8] {
        let pos = self.buf.position() as usize;
        if pos < self.buf.get_ref().len() {
            &self.buf.get_ref()[pos..]
        } else {
            &[]
        }
    }

    pub fn into_inner(self) -> R { self.inner }

    pub fn read_into_buf(&mut self) -> io::Result<usize> {
        let v = self.buf.get_mut();
        reserve(v);
        let inner = &mut self.inner;
        with_end_to_cap(v, |b| inner.read(b))
    }
}

impl<R: Read> Read for BufReader<R> {
    fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
        if self.buf.get_ref().len() == self.buf.position() as usize &&
            buf.len() >= self.buf.get_ref().capacity() {
            return self.inner.read(buf);
        }
        try!(self.fill_buf());
        self.buf.read(buf)
    }
}

impl<R: Read> BufRead for BufReader<R> {
    fn fill_buf(&mut self) -> io::Result<&[u8]> {
         if self.buf.position() as usize == self.buf.get_ref().len() {
            self.buf.set_position(0);
            let v = self.buf.get_mut();
            v.truncate(0);
            let inner = &mut self.inner;
            try!(with_end_to_cap(v, |b| inner.read(b)));
         }
         self.buf.fill_buf()
    }

    fn consume(&mut self, amt: usize) {
        self.buf.consume(amt)
    }
}

fn with_end_to_cap<F>(v: &mut Vec<u8>, f: F) -> io::Result<usize>
    where F: FnOnce(&mut [u8]) -> io::Result<usize>
{
    let len = v.len();
    let new_area = v.capacity() - len;
    v.extend(iter::repeat(0).take(new_area));
    match f(&mut v[len..]) {
        Ok(n) => {
            v.truncate(len + n);
            Ok(n)
        }
        Err(e) => {
            v.truncate(len);
            Err(e)
        }
    }
}

#[inline]
fn reserve(v: &mut Vec<u8>) {
    let cap = v.capacity();
    if v.len() == cap {
        v.reserve(cmp::min(cap * 4, MAX_BUFFER_SIZE) - cap);
    }
}

#[cfg(test)]
mod tests {

    use std::io::BufRead;
    use super::BufReader;

    #[test]
    fn test_consume_and_get_buf() {
        let mut rdr = BufReader::new(&b"foo bar baz"[..]);
        rdr.read_into_buf().unwrap();
        rdr.consume(8);
        assert_eq!(rdr.get_buf(), b"baz");
    }
}
<commit_msg>perf(buffer): pull in std::io::BufReader improvements<commit_after>use std::cmp;
use std::iter;
use std::io::{self, Read, BufRead};

pub struct BufReader<R> {
    inner: R,
    buf: Vec<u8>,
    pos: usize,
    cap: usize,
}

const INIT_BUFFER_SIZE: usize = 4096;
const MAX_BUFFER_SIZE: usize = 8192 + 4096 * 100;

impl<R: Read> BufReader<R> {
    #[inline]
    pub fn new(rdr: R) -> BufReader<R> {
        BufReader::with_capacity(rdr, INIT_BUFFER_SIZE)
    }

    #[inline]
    pub fn with_capacity(rdr: R, cap: usize) -> BufReader<R> {
        let mut buf = Vec::with_capacity(cap);
        buf.extend(iter::repeat(0).take(cap));
        BufReader {
            inner: rdr,
            buf: buf,
            pos: 0,
            cap: 0,
        }
    }

    #[inline]
    pub fn get_ref(&self) -> &R { &self.inner }

    #[inline]
    pub fn get_mut(&mut self) -> &mut R { &mut self.inner }

    #[inline]
    pub fn get_buf(&self) -> &[u8] {
        if self.pos < self.cap {
            debug!("slicing {:?}", (self.pos, self.cap, self.buf.len()));
            &self.buf[self.pos..self.cap]
        } else {
            &[]
        }
    }

    #[inline]
    pub fn into_inner(self) -> R { self.inner }

    #[inline]
    pub fn read_into_buf(&mut self) -> io::Result<usize> {
        self.maybe_reserve();
        let v = &mut self.buf;
        if self.cap < v.capacity() {
            let nread = try!(self.inner.read(&mut v[self.cap..]));
            self.cap += nread;
            Ok(nread)
        } else {
            Ok(0)
        }
    }

    #[inline]
    fn maybe_reserve(&mut self) {
        let cap = self.buf.capacity();
        if self.cap == cap {
            self.buf.reserve(cmp::min(cap * 4, MAX_BUFFER_SIZE) - cap);
        }
    }
}

impl<R: Read> Read for BufReader<R> {
    fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
        if self.cap == self.pos && buf.len() >= self.buf.len() {
            return self.inner.read(buf);
        }
        let nread = {
           let mut rem = try!(self.fill_buf());
           try!(rem.read(buf))
        };
        self.consume(nread);
        Ok(nread)
    }
}

impl<R: Read> BufRead for BufReader<R> {
    fn fill_buf(&mut self) -> io::Result<&[u8]> {
        if self.pos == self.cap {
            self.cap = try!(self.inner.read(&mut self.buf));
            self.pos = 0;
        }
        Ok(&self.buf[self.pos..self.cap])
    }

    #[inline]
    fn consume(&mut self, amt: usize) {
        self.pos = cmp::min(self.pos + amt, self.cap);
        if self.pos == self.cap {
            self.pos = 0;
            self.cap = 0;
        }
    }
}

#[cfg(test)]
mod tests {

    use std::io::{self, Read, BufRead};
    use super::BufReader;

    struct SlowRead(u8);

    impl Read for SlowRead {
        fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
            let state = self.0;
            self.0 += 1;
            (&match state % 3 {
                0 => b"foo",
                1 => b"bar",
                _ => b"baz",
            }[..]).read(buf)
        }
    }

    #[test]
    fn test_consume_and_get_buf() {
        let mut rdr = BufReader::new(SlowRead(0));
        rdr.read_into_buf().unwrap();
        rdr.consume(1);
        assert_eq!(rdr.get_buf(), b"oo");
        rdr.read_into_buf().unwrap();
        rdr.read_into_buf().unwrap();
        assert_eq!(rdr.get_buf(), b"oobarbaz");
        rdr.consume(5);
        assert_eq!(rdr.get_buf(), b"baz");
        rdr.consume(3);
        assert_eq!(rdr.get_buf(), b"");
        assert_eq!(rdr.pos, 0);
        assert_eq!(rdr.cap, 0);
    }
}
<|endoftext|>


<commit_before>use std::collections::BTreeMap;
use base64;
use serde_json;
use serde_json::Value as Json;

use Component;
use error::Error;

#[derive(Debug, Default, PartialEq)]
pub struct Claims {
    pub reg: Registered,
    pub private: BTreeMap<String, Json>,
}

#[derive(Debug, Default, PartialEq, Serialize, Deserialize)]
pub struct Registered {
    pub iss: Option<String>,
    pub sub: Option<String>,
    pub aud: Option<String>,
    pub exp: Option<u64>,
    pub nbf: Option<u64>,
    pub iat: Option<u64>,
    pub jti: Option<String>,
}

/// JWT Claims. Registered claims are directly accessible via the `Registered`
/// struct embedded, while private fields are a map that contains `Json`
/// values.
impl Claims {
    pub fn new(reg: Registered) -> Claims {
        Claims {
            reg: reg,
            private: BTreeMap::new(),
        }
    }
}

impl Component for Claims {
    fn from_base64(raw: &str) -> Result<Claims, Error> {
        let data = base64::decode_config(raw, base64::URL_SAFE_NO_PAD)?;
        let s = String::from_utf8(data)?;
        let tree = match serde_json::from_str(&*s)? {
            Json::Object(x) => x,
            _ => return Err(Error::Format),
        };

        const FIELDS: [&'static str; 7] = ["iss", "sub", "aud", "exp", "nbf", "iat", "jti"];

        let (_, pri): (BTreeMap<_, _>, BTreeMap<_, _>) = tree.into_iter()
            .partition(|&(ref key, _)| FIELDS.iter().any(|f| f == key));

        let reg_claims: Registered = serde_json::from_str(&*s)?;

        Ok(Claims {
            reg: reg_claims,
            private: pri,
        })
    }

    fn to_base64(&self) -> Result<String, Error> {
        let mut json_claims = match serde_json::to_value(&self.reg)? {
            Json::Object(x) => x,
            _ => return Err(Error::Format),
        };

        json_claims.extend(self.private.clone());

        let s = serde_json::to_string(&json_claims)?;
        let enc = base64::encode_config(&*s, base64::URL_SAFE_NO_PAD);
        Ok(enc)
    }
}

#[cfg(test)]
mod tests {
    use std::default::Default;
    use claims::{Claims, Registered};
    use Component;

    #[test]
    fn from_base64() {
        let enc = "ew0KICAiaXNzIjogIm1pa2t5YW5nLmNvbSIsDQogICJleHAiOiAxMzAyMzE5MTAwLA0KICAibmFtZSI6ICJNaWNoYWVsIFlhbmciLA0KICAiYWRtaW4iOiB0cnVlDQp9";
        let claims = Claims::from_base64(enc).unwrap();

        assert_eq!(claims.reg.iss.unwrap(), "mikkyang.com");
        assert_eq!(claims.reg.exp.unwrap(), 1302319100);
    }

    #[test]
    fn multiple_types() {
        let enc = "ew0KICAiaXNzIjogIm1pa2t5YW5nLmNvbSIsDQogICJleHAiOiAxMzAyMzE5MTAwLA0KICAibmFtZSI6ICJNaWNoYWVsIFlhbmciLA0KICAiYWRtaW4iOiB0cnVlDQp9";
        let claims = Registered::from_base64(enc).unwrap();

        assert_eq!(claims.iss.unwrap(), "mikkyang.com");
        assert_eq!(claims.exp.unwrap(), 1302319100);
    }

    #[test]
    fn roundtrip() {
        let mut claims: Claims = Default::default();
        claims.reg.iss = Some("mikkyang.com".into());
        claims.reg.exp = Some(1302319100);
        let enc = claims.to_base64().unwrap();
        assert_eq!(claims, Claims::from_base64(&*enc).unwrap());
    }
}
<commit_msg>Add test for custom claim<commit_after>use std::collections::BTreeMap;
use base64;
use serde_json;
use serde_json::Value as Json;

use Component;
use error::Error;

#[derive(Debug, Default, PartialEq)]
pub struct Claims {
    pub reg: Registered,
    pub private: BTreeMap<String, Json>,
}

#[derive(Debug, Default, PartialEq, Serialize, Deserialize)]
pub struct Registered {
    pub iss: Option<String>,
    pub sub: Option<String>,
    pub aud: Option<String>,
    pub exp: Option<u64>,
    pub nbf: Option<u64>,
    pub iat: Option<u64>,
    pub jti: Option<String>,
}

/// JWT Claims. Registered claims are directly accessible via the `Registered`
/// struct embedded, while private fields are a map that contains `Json`
/// values.
impl Claims {
    pub fn new(reg: Registered) -> Claims {
        Claims {
            reg: reg,
            private: BTreeMap::new(),
        }
    }
}

impl Component for Claims {
    fn from_base64(raw: &str) -> Result<Claims, Error> {
        let data = base64::decode_config(raw, base64::URL_SAFE_NO_PAD)?;
        let s = String::from_utf8(data)?;
        let tree = match serde_json::from_str(&*s)? {
            Json::Object(x) => x,
            _ => return Err(Error::Format),
        };

        const FIELDS: [&'static str; 7] = ["iss", "sub", "aud", "exp", "nbf", "iat", "jti"];

        let (_, pri): (BTreeMap<_, _>, BTreeMap<_, _>) = tree.into_iter()
            .partition(|&(ref key, _)| FIELDS.iter().any(|f| f == key));

        let reg_claims: Registered = serde_json::from_str(&*s)?;

        Ok(Claims {
            reg: reg_claims,
            private: pri,
        })
    }

    fn to_base64(&self) -> Result<String, Error> {
        let mut json_claims = match serde_json::to_value(&self.reg)? {
            Json::Object(x) => x,
            _ => return Err(Error::Format),
        };

        json_claims.extend(self.private.clone());

        let s = serde_json::to_string(&json_claims)?;
        let enc = base64::encode_config(&*s, base64::URL_SAFE_NO_PAD);
        Ok(enc)
    }
}

#[cfg(test)]
mod tests {
    use std::default::Default;
    use claims::{Claims, Registered};
    use Component;
    use serde_json::Value;

    #[test]
    fn from_base64() {
        let enc = "ew0KICAiaXNzIjogIm1pa2t5YW5nLmNvbSIsDQogICJleHAiOiAxMzAyMzE5MTAwLA0KICAibmFtZSI6ICJNaWNoYWVsIFlhbmciLA0KICAiYWRtaW4iOiB0cnVlDQp9";
        let claims = Claims::from_base64(enc).unwrap();

        assert_eq!(claims.reg.iss.unwrap(), "mikkyang.com");
        assert_eq!(claims.reg.exp.unwrap(), 1302319100);
    }

    #[test]
    fn multiple_types() {
        let enc = "ew0KICAiaXNzIjogIm1pa2t5YW5nLmNvbSIsDQogICJleHAiOiAxMzAyMzE5MTAwLA0KICAibmFtZSI6ICJNaWNoYWVsIFlhbmciLA0KICAiYWRtaW4iOiB0cnVlDQp9";
        let claims = Registered::from_base64(enc).unwrap();

        assert_eq!(claims.iss.unwrap(), "mikkyang.com");
        assert_eq!(claims.exp.unwrap(), 1302319100);
    }

    #[test]
    fn private_claim() {
        let encoded = "ewogICJpc3MiOiAibWlra3lhbmcuY29tIiwKICAiZXhwIjogMTMwMjMxOTEwMCwKICAiY3VzdG9tX2NsYWltIjogdHJ1ZQp9Cg==";
        let claims = Claims::from_base64(encoded).unwrap();

        assert_eq!(claims.private["custom_claim"], Value::Bool(true));
    }

    #[test]
    fn roundtrip() {
        let mut claims: Claims = Default::default();
        claims.reg.iss = Some("mikkyang.com".into());
        claims.reg.exp = Some(1302319100);
        let enc = claims.to_base64().unwrap();
        assert_eq!(claims, Claims::from_base64(&*enc).unwrap());
    }
}
<|endoftext|>