<commit_before>/*
 * Copyright (c) 2012, Ben Noordhuis <info@bnoordhuis.nl>
 *
 * Permission to use, copy, modify, and/or distribute this software for any
 * purpose with or without fee is hereby granted, provided that the above
 * copyright notice and this permission notice appear in all copies.
 *
 * THE SOFTWARE IS PROVIDED "AS IS" AND THE AUTHOR DISCLAIMS ALL WARRANTIES
 * WITH REGARD TO THIS SOFTWARE INCLUDING ALL IMPLIED WARRANTIES OF
 * MERCHANTABILITY AND FITNESS. IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR
 * ANY SPECIAL, DIRECT, INDIRECT, OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES
 * WHATSOEVER RESULTING FROM LOSS OF USE, DATA OR PROFITS, WHETHER IN AN
 * ACTION OF CONTRACT, NEGLIGENCE OR OTHER TORTIOUS ACTION, ARISING OUT OF
 * OR IN CONNECTION WITH THE USE OR PERFORMANCE OF THIS SOFTWARE.
 */

#[license = "ISC"];

#[link(name = "errno",
       vers = "1.0",
       author = "Ben Noordhuis <info@bnoordhuis.nl>")];

import size_t = ctypes::size_t;
import c_int = ctypes::c_int;
import sbuf = str::sbuf;

#[nolink]
native mod __glibc {
  fn __xpg_strerror_r(errnum: c_int, buf: *mutable u8, len: size_t) -> c_int;
  fn __errno_location() -> *c_int;
}

fn errno() -> int {
  unsafe { *__glibc::__errno_location() as int }
}

fn strerror(errnum: int) -> str {
  let bufv = vec::init_elt_mut(1024u, 0u8);
  let buf = ptr::mut_addr_of(bufv[0]);
  let len = vec::len(bufv);

  unsafe {
    let r = __glibc::__xpg_strerror_r(errnum as c_int, buf, len as size_t);

    if (r as bool) {
      fail #fmt("__glibc::__xpg_strerror_r() failed [errno=%d]", errno());
    }

    ret str::from_cstr(buf as sbuf);
  }
}

#[test]
fn test() {
  // FIXME chokes on localized messages
  assert errno::strerror(22) == "Invalid argument";
  assert errno::strerror(0) == "Success";
}

#[should_fail]
#[test]
fn test2() {
  errno::strerror(-1);
}
<commit_msg>Fix compilation with --test<commit_after>/*
 * Copyright (c) 2012, Ben Noordhuis <info@bnoordhuis.nl>
 *
 * Permission to use, copy, modify, and/or distribute this software for any
 * purpose with or without fee is hereby granted, provided that the above
 * copyright notice and this permission notice appear in all copies.
 *
 * THE SOFTWARE IS PROVIDED "AS IS" AND THE AUTHOR DISCLAIMS ALL WARRANTIES
 * WITH REGARD TO THIS SOFTWARE INCLUDING ALL IMPLIED WARRANTIES OF
 * MERCHANTABILITY AND FITNESS. IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR
 * ANY SPECIAL, DIRECT, INDIRECT, OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES
 * WHATSOEVER RESULTING FROM LOSS OF USE, DATA OR PROFITS, WHETHER IN AN
 * ACTION OF CONTRACT, NEGLIGENCE OR OTHER TORTIOUS ACTION, ARISING OUT OF
 * OR IN CONNECTION WITH THE USE OR PERFORMANCE OF THIS SOFTWARE.
 */

#[license = "ISC"];

#[link(name = "errno",
       vers = "1.0",
       author = "Ben Noordhuis <info@bnoordhuis.nl>")];

use std; // required by the tests

import size_t = ctypes::size_t;
import c_int = ctypes::c_int;
import sbuf = str::sbuf;

#[nolink]
native mod __glibc {
  fn __xpg_strerror_r(errnum: c_int, buf: *mutable u8, len: size_t) -> c_int;
  fn __errno_location() -> *c_int;
}

fn errno() -> int {
  unsafe { *__glibc::__errno_location() as int }
}

fn strerror(errnum: int) -> str {
  let bufv = vec::init_elt_mut(1024u, 0u8);
  let buf = ptr::mut_addr_of(bufv[0]);
  let len = vec::len(bufv);

  unsafe {
    let r = __glibc::__xpg_strerror_r(errnum as c_int, buf, len as size_t);

    if (r as bool) {
      fail #fmt("__glibc::__xpg_strerror_r() failed [errno=%d]", errno());
    }

    ret str::from_cstr(buf as sbuf);
  }
}

#[test]
fn test() {
  // FIXME chokes on localized messages
  assert strerror(22) == "Invalid argument";
  assert strerror(0) == "Success";
}

#[should_fail]
#[test]
fn test2() {
  strerror(-1);
}
<|endoftext|>


<commit_before>//!Neural Network Layers

use linalg::{Matrix, MatrixSlice, BaseMatrix, BaseMatrixMut};

use learning::toolkit::activ_fn::ActivationFunc;

use rand::thread_rng;
use rand::distributions::Sample;
use rand::distributions::normal::Normal;

use std::fmt::Debug;

/// Trait for neural net layers
pub trait NetLayer : Debug {
	/// The result of propogating data forward through this layer
	fn forward(&self, input: &Matrix<f64>, params: MatrixSlice<f64>) -> Matrix<f64>;

	/// The gradient of the output of this layer with respect to its input
	fn back_input(&self, out_grad: &Matrix<f64>, input: &Matrix<f64>, params: MatrixSlice<f64>) -> Matrix<f64>;

	/// The gradient of the output of this layer with respect to its parameters
	fn back_params(&self, out_grad: &Matrix<f64>, input: &Matrix<f64>, params: MatrixSlice<f64>) -> Matrix<f64>;

	/// The default value of the parameters of this layer before training
	fn default_params(&self) -> Vec<f64>;

	/// The shape of the parameters used by this layer
	fn param_shape(&self) -> (usize, usize);

	/// The number of parameters used by this layer
	fn num_params(&self) -> usize {
		let shape = self.param_shape();
		shape.0 * shape.1
	}
}

/// Linear network layer
///
/// Represents a fully connected layer with optional bias term
///
/// The parameters are a matrix of weights of size I x O
/// where O is the dimensionality of the output and I the dimensionality of the input
#[derive(Debug, Clone, Copy)]
pub struct Linear {
	/// The number of dimensions of the input
	input_size: usize,
	/// The number of dimensions of the output
	output_size: usize,
	/// Whether or not to include a bias term
	has_bias: bool,
}

impl Linear {
	/// Construct a new Linear layer
	pub fn new(input_size: usize, output_size: usize) -> Linear {
		Linear {
			input_size: input_size + 1,
			output_size: output_size,
			has_bias: true
		}
	}

	/// Construct a Linear layer with a bias term
	pub fn without_bias(input_size: usize, output_size: usize) -> Linear {
		Linear {
			input_size: input_size,
			output_size: output_size,
			has_bias: false
		}
	}
}

impl NetLayer for Linear {
	/// Computes a matrix product
	///
	/// input should have dimensions N x I
	/// where N is the number of samples and I is the dimensionality of the input
	fn forward(&self, input: &Matrix<f64>, params: MatrixSlice<f64>) -> Matrix<f64> {
		if self.has_bias {
			assert_eq!(input.cols()+1, params.rows());
			input.hcat(&Matrix::<f64>::ones(input.rows(), 1)) * &params
		} else {
			assert_eq!(input.cols(), params.rows());
			input * &params
		}
	}

	fn back_input(&self, out_grad: &Matrix<f64>, _: &Matrix<f64>, params: MatrixSlice<f64>) -> Matrix<f64> {
		assert_eq!(out_grad.cols(), params.cols());
		//let gradient = out_grad * &params.into_matrix().transpose();
		if self.has_bias {
			//let columns: Vec<_> = (0..gradient.cols()-1).collect();
			//gradient.select_cols(&columns)
			let rows: Vec<_> = (0..params.cols()-1).collect();
			out_grad * &params.into_matrix().select_rows(&rows).transpose()
		} else {
			//gradient
			out_grad * &params.into_matrix().transpose()
		}
	}

	fn back_params(&self, out_grad: &Matrix<f64>, input: &Matrix<f64>, _: MatrixSlice<f64>) -> Matrix<f64> {
		assert_eq!(input.rows(), out_grad.rows());
		if self.has_bias {
			//input.transpose().vcat(&Matrix::<f64>::ones(1, input.rows())) * out_grad
			input.hcat(&Matrix::<f64>::ones(input.rows(), 1)).transpose() * out_grad
		} else {
			input.transpose() * out_grad
		}
	}

	/// Initializes weights using Xavier initialization
	///
	/// weights drawn from gaussian distribution with 0 mean and variance 2/(input_size+output_size)
	fn default_params(&self) -> Vec<f64> {
		let mut distro = Normal::new(0.0, (2.0/(self.input_size+self.output_size) as f64).sqrt());
		let mut rng = thread_rng();

		(0..self.input_size*self.output_size).map(|_| distro.sample(&mut rng))
											 .collect()
	}

	fn param_shape(&self) -> (usize, usize) {
		(self.input_size, self.output_size)
	}
}

impl<T: ActivationFunc + Debug> NetLayer for T {
	/// Applys the activation function to each element of the input
	fn forward(&self, input: &Matrix<f64>, _: MatrixSlice<f64>) -> Matrix<f64> {
		//Matrix::new(input.rows(), input.cols(),
		//	input.iter().map(|&x| T::func(x)).collect::<Vec<_>>());
		input.clone().apply(&T::func)
	}

	fn back_input(&self, out_grad: &Matrix<f64>, input: &Matrix<f64>, _: MatrixSlice<f64>) -> Matrix<f64> {
		let in_grad = Matrix::new(input.rows(), input.cols(),
	        			input.iter().map(|&x| T::func_grad(x)).collect::<Vec<_>>());
		out_grad.elemul(&in_grad)
	}

	fn back_params(&self, _: &Matrix<f64>, _: &Matrix<f64>, _: MatrixSlice<f64>) -> Matrix<f64> {
		Matrix::new(0, 0, Vec::new())
	}

	fn default_params(&self) -> Vec<f64> {
		Vec::new()
	}

	fn param_shape(&self) -> (usize, usize) {
		(0, 0)
	}
}<commit_msg>Fixed doc test error<commit_after>//!Neural Network Layers

use linalg::{Matrix, MatrixSlice, BaseMatrix, BaseMatrixMut};

use learning::toolkit::activ_fn::ActivationFunc;

use rand::thread_rng;
use rand::distributions::Sample;
use rand::distributions::normal::Normal;

use std::fmt::Debug;

/// Trait for neural net layers
pub trait NetLayer : Debug {
	/// The result of propogating data forward through this layer
	fn forward(&self, input: &Matrix<f64>, params: MatrixSlice<f64>) -> Matrix<f64>;

	/// The gradient of the output of this layer with respect to its input
	fn back_input(&self, out_grad: &Matrix<f64>, input: &Matrix<f64>, params: MatrixSlice<f64>) -> Matrix<f64>;

	/// The gradient of the output of this layer with respect to its parameters
	fn back_params(&self, out_grad: &Matrix<f64>, input: &Matrix<f64>, params: MatrixSlice<f64>) -> Matrix<f64>;

	/// The default value of the parameters of this layer before training
	fn default_params(&self) -> Vec<f64>;

	/// The shape of the parameters used by this layer
	fn param_shape(&self) -> (usize, usize);

	/// The number of parameters used by this layer
	fn num_params(&self) -> usize {
		let shape = self.param_shape();
		shape.0 * shape.1
	}
}

/// Linear network layer
///
/// Represents a fully connected layer with optional bias term
///
/// The parameters are a matrix of weights of size I x O
/// where O is the dimensionality of the output and I the dimensionality of the input
#[derive(Debug, Clone, Copy)]
pub struct Linear {
	/// The number of dimensions of the input
	input_size: usize,
	/// The number of dimensions of the output
	output_size: usize,
	/// Whether or not to include a bias term
	has_bias: bool,
}

impl Linear {
	/// Construct a new Linear layer
	pub fn new(input_size: usize, output_size: usize) -> Linear {
		Linear {
			input_size: input_size + 1,
			output_size: output_size,
			has_bias: true
		}
	}

	/// Construct a Linear layer with a bias term
	pub fn without_bias(input_size: usize, output_size: usize) -> Linear {
		Linear {
			input_size: input_size,
			output_size: output_size,
			has_bias: false
		}
	}
}

impl NetLayer for Linear {
	/// Computes a matrix product
	///
	/// input should have dimensions N x I
	/// where N is the number of samples and I is the dimensionality of the input
	fn forward(&self, input: &Matrix<f64>, params: MatrixSlice<f64>) -> Matrix<f64> {
		if self.has_bias {
			assert_eq!(input.cols()+1, params.rows());
			input.hcat(&Matrix::<f64>::ones(input.rows(), 1)) * &params
		} else {
			assert_eq!(input.cols(), params.rows());
			input * &params
		}
	}

	fn back_input(&self, out_grad: &Matrix<f64>, _: &Matrix<f64>, params: MatrixSlice<f64>) -> Matrix<f64> {
		assert_eq!(out_grad.cols(), params.cols());
		//let gradient = out_grad * &params.into_matrix().transpose();
		if self.has_bias {
			//let columns: Vec<_> = (0..gradient.cols()-1).collect();
			//gradient.select_cols(&columns)
			let rows: Vec<_> = (0..params.rows()-1).collect();
			out_grad * &params.into_matrix().select_rows(&rows).transpose()
		} else {
			//gradient
			out_grad * &params.into_matrix().transpose()
		}
	}

	fn back_params(&self, out_grad: &Matrix<f64>, input: &Matrix<f64>, _: MatrixSlice<f64>) -> Matrix<f64> {
		assert_eq!(input.rows(), out_grad.rows());
		if self.has_bias {
			//input.transpose().vcat(&Matrix::<f64>::ones(1, input.rows())) * out_grad
			input.hcat(&Matrix::<f64>::ones(input.rows(), 1)).transpose() * out_grad
		} else {
			input.transpose() * out_grad
		}
	}

	/// Initializes weights using Xavier initialization
	///
	/// weights drawn from gaussian distribution with 0 mean and variance 2/(input_size+output_size)
	fn default_params(&self) -> Vec<f64> {
		let mut distro = Normal::new(0.0, (2.0/(self.input_size+self.output_size) as f64).sqrt());
		let mut rng = thread_rng();

		(0..self.input_size*self.output_size).map(|_| distro.sample(&mut rng))
											 .collect()
	}

	fn param_shape(&self) -> (usize, usize) {
		(self.input_size, self.output_size)
	}
}

impl<T: ActivationFunc + Debug> NetLayer for T {
	/// Applys the activation function to each element of the input
	fn forward(&self, input: &Matrix<f64>, _: MatrixSlice<f64>) -> Matrix<f64> {
		//Matrix::new(input.rows(), input.cols(),
		//	input.iter().map(|&x| T::func(x)).collect::<Vec<_>>());
		input.clone().apply(&T::func)
	}

	fn back_input(&self, out_grad: &Matrix<f64>, input: &Matrix<f64>, _: MatrixSlice<f64>) -> Matrix<f64> {
		let in_grad = Matrix::new(input.rows(), input.cols(),
	        			input.iter().map(|&x| T::func_grad(x)).collect::<Vec<_>>());
		out_grad.elemul(&in_grad)
	}

	fn back_params(&self, _: &Matrix<f64>, _: &Matrix<f64>, _: MatrixSlice<f64>) -> Matrix<f64> {
		Matrix::new(0, 0, Vec::new())
	}

	fn default_params(&self) -> Vec<f64> {
		Vec::new()
	}

	fn param_shape(&self) -> (usize, usize) {
		(0, 0)
	}
}<|endoftext|>


<commit_before>// zhttpto.rs
//
// Starting code for PS1
// Running on Rust 0.9
//
// Note that this code has serious security risks!  You should not run it
// on any system with access to sensitive files.
//
// University of Virginia - cs4414 Spring 2014
// Weilin Xu and David Evans
// Version 0.3

#[feature(globs)];
use std::io::*;
use std::io::net::ip::{SocketAddr};
use std::{str};
use std::io::buffered::BufferedReader;
use std::io::File;

static IP: &'static str = "127.0.0.1";
static PORT:        int = 4414;
static mut count: int = 0;

fn main() {

    let addr = from_str::<SocketAddr>(format!("{:s}:{:d}", IP, PORT)).unwrap();
    let mut acceptor = net::tcp::TcpListener::bind(addr).listen();

    println(format!("Listening on [{:s}] ...", addr.to_str()));

    for stream in acceptor.incoming() {

        // Spawn a task to handle the connection
        do spawn {
            let mut stream = stream;

            match stream {
                Some(ref mut s) => {
                             match s.peer_name() {
                                Some(pn) => {println(format!("Received connection from: [{:s}]", pn.to_str()));},
                                None => ()
                             }
                           },
                None => ()
            }

            let mut buf = [0, ..500];
            stream.read(buf);
            let request_str = str::from_utf8(buf);

            let mut line1 = "";

            for str in request_str.split('\n'){ // get the first line
                line1 = str.clone();
                break;
            }

            let mut path = "";

            let temp: ~[&str] = line1.split(' ').collect();
            path = temp[1].clone(); // this gets the path

            let file_content = readFile(path);

            unsafe {    // so we can increment count 'illegally'

            count += 1;

            if file_content == ~"index" {

                let index: ~str = format!(
                    "HTTP/1.1 200 OK\r\nContent-Type: text/html; charset=UTF-8\r\n\r\n
                     <doctype !html><title>Hello, Rust!</title>

Greetings, Krusty! Total hits = {:d}

                     \r\n", count);

                stream.write(index.as_bytes());
                println!("Connection terminates");
            }

            else {
                let response: ~str = format!("{:s}", file_content);

                stream.write(response.as_bytes());
                println!("Connection terminates");
        }

            }
        }
    }
}

fn readFile(x: &str) -> ~str {

    let mut y = x.slice_from(1).clone();
    let z = x.slice(0,1).clone();
    println(z);
    if z != "/" {
        y = x.clone();
    }

    let x = y.trim().clone();

    if x == "" {
        return ~"index";
    }

    let mut html = false;
    let mut in_dir = false;

    //testing to see if looking in allowed directory/if its an HTML file
    let mut temp: ~[&str] = y.split('/').collect();

    if temp[0] != "" && temp[0] != ".." {
        in_dir = true;
    }

    temp = y.split('.').collect();
    if temp[temp.len()-1] == "html"{
        html = true;
    }

    let path = Path::new(y);

    match (path.exists(), path.is_file(), html, in_dir){

        (true,true, true, true) => {

            match File::open(&Path::new(path)) {

                Some(file) => {
                    let mut reader = BufferedReader::new(file);
                    let input = reader.read_to_str();
                    return input;

                }

                None =>{~"HTTP/1.1 500 OK\r\nContent-Type: text/html; charset=UTF-8\r\n\r\n

 500 ERROR

 Something really bad happened, lo siento
"}
            }

        }
        (true,true,_,false) => {~"HTTP/1.1 403 OK\r\nContent-Type: text/html; charset=UTF-8\r\n\r\n

 403 ERROR

 Unauthorized access!
"}

        (true,true,false,_) => {~"HTTP/1.1 403 OK\r\nContent-Type: text/html; charset=UTF-8\r\n\r\n

 403 ERROR

 Unauthorized access!
"}

        (_,_,_,_) => {~"HTTP/1.1 404 OK\r\nContent-Type: text/html; charset=UTF-8\r\n\r\n

 404 ERROR

 Page not found
"}
    }
}
<commit_msg>final version of zehttpto.rs<commit_after>// zhttpto.rs
//
// Starting code for PS1
// Running on Rust 0.9
//
// Note that this code has serious security risks!  You should not run it
// on any system with access to sensitive files.
//
// University of Virginia - cs4414 Spring 2014
// Weilin Xu and David Evans
// Version 0.3

#[feature(globs)];
use std::io::*;
use std::io::net::ip::{SocketAddr};
use std::{str};
use std::io::buffered::BufferedReader;
use std::io::File;

static IP: &'static str = "127.0.0.1";
static PORT:        int = 4414;
static mut count: int = 0;

fn main() {

    let addr = from_str::<SocketAddr>(format!("{:s}:{:d}", IP, PORT)).unwrap();
    let mut acceptor = net::tcp::TcpListener::bind(addr).listen();

    println(format!("Listening on [{:s}] ...", addr.to_str()));

    for stream in acceptor.incoming() {

        // Spawn a task to handle the connection
        do spawn {
            let mut stream = stream;

            match stream {
                Some(ref mut s) => {
                             match s.peer_name() {
                                Some(pn) => {println(format!("Received connection from: [{:s}]", pn.to_str()));},
                                None => ()
                             }
                           },
                None => ()
            }

            let mut buf = [0, ..500];
            stream.read(buf);
            let request_str = str::from_utf8(buf);

            let mut line1 = "";

            for str in request_str.split('\n'){ // get the first line
                line1 = str.clone();
                break;
            }

            let mut path = "";

            let temp: ~[&str] = line1.split(' ').collect();
            path = temp[1].clone(); // this gets the path

            let file_content = readFile(path);

            unsafe {    // so we can increment count 'illegally'

            count += 1;

            if file_content == ~"index" {

                let index: ~str = format!(
                    "HTTP/1.1 200 OK\r\nContent-Type: text/html; charset=UTF-8\r\n\r\n
                     <doctype !html><title>Hello, Rust!</title>

Greetings, Krusty! Total hits = {:d}

                     \r\n", count);

                stream.write(index.as_bytes());
                println!("Connection terminates");
            }

            else {
                let response: ~str = format!("{:s}", file_content);

                stream.write(response.as_bytes());
                println!("Connection terminates");
        }

            }
        }
    }
}

fn readFile(x: &str) -> ~str {

    let mut y = x.slice_from(1).clone();
    let z = x.slice(0,1).clone();
    //println(z);
    if z != "/" {
        y = x.clone();
    }

    let x = y.trim().clone();

    if x == "" {
        return ~"index";
    }

    let mut html = false;
    let mut in_dir = false;

    //testing to see if looking in allowed directory/if its an HTML file
    let mut temp: ~[&str] = y.split('/').collect();

    if temp[0] != "" && temp[0] != ".." {
        in_dir = true;
    }

    temp = y.split('.').collect();
    if temp[temp.len()-1] == "html"{
        html = true;
    }

    let path = Path::new(y);

    match (path.exists(), path.is_file(), html, in_dir){

        (true,true, true, true) => {

            match File::open(&Path::new(path)) {

                Some(file) => {
                    let mut reader = BufferedReader::new(file);
                    let input = reader.read_to_str();
                    return input;

                }

                None =>{~"HTTP/1.1 500 OK\r\nContent-Type: text/html; charset=UTF-8\r\n\r\n

 500 ERROR

 Something really bad happened, lo siento
"}
            }

        }
        (true,true,_,false) => {~"HTTP/1.1 403 OK\r\nContent-Type: text/html; charset=UTF-8\r\n\r\n

 403 ERROR

 Unauthorized access!
"}

        (true,true,false,_) => {~"HTTP/1.1 403 OK\r\nContent-Type: text/html; charset=UTF-8\r\n\r\n

 403 ERROR

 Unauthorized access!
"}

        (_,_,_,_) => {~"HTTP/1.1 404 OK\r\nContent-Type: text/html; charset=UTF-8\r\n\r\n

 404 ERROR

 Page not found
"}
    }
}
<|endoftext|>


<commit_before>use std::fmt;
use std::collections::RingBuf;

/// Decodes an integer encoded with a given prefix size (in bits).
/// Assumes that the buffer `buf` contains the integer to be decoded,
/// with the first byte representing the octet that contains the
/// prefix.
///
/// Returns a tuple representing the decoded integer and the number
/// of bytes from the buffer that were used.
fn decode_integer(buf: &[u8], prefix_size: u8) -> (usize, usize) {
    assert!(prefix_size >= 1 && prefix_size <= 8);
    assert!(buf.len() >= 1);
    let mask = (1u8 << prefix_size) - 1;
    let mut value: usize = (buf[0] & mask) as usize;
    if value < (mask as usize) {
        return (value, 1);
    }

    // Already one byte used (the prefix)
    let mut total = 1;
    let mut m = 0;
    for &b in buf[1..].iter() {
        total += 1;
        value += ((b & 127) as usize) * (1us << m);
        m += 7;
        if b & 128 != 128 {
            break;
        }
    }

    // TODO Error situation -- what happens if we reach here w/o having
    //      `b` byte-sized...
    (value, total)
}

/// Encode an integer to the representation defined by HPACK.
///
/// Returns a newly allocated `Vec` containing the encoded bytes.
/// Only `prefix_size` lowest-order bits of the first byte in the
/// array are guaranteed to be used.
fn encode_integer(mut value: usize, prefix_size: u8) -> Vec<u8> {
    let mask: usize = ((1u8 << prefix_size) - 1) as usize;
    if value < mask {
        // Right now, the caller would need to be the one to combine
        // the other part of the prefix byte (but would know that it's
        // safe to do so using the bit-wise or).
        return vec![value as u8];
    }

    let mut res: Vec<u8> = Vec::new();
    res.push(mask as u8);
    value -= mask;
    while value >= 128 {
        res.push(((value % 128) + 128) as u8);
        value = value / 128;
    }
    res.push(value as u8);

    res
}

/// The table represents the static header table defined by the HPACK spec.
/// (HPACK, Appendix A)
static STATIC_TABLE: &'static [(&'static [u8], &'static [u8])] = &[
  (b":authority", b""),
  (b":method", b"GET"),
  (b":method", b"POST"),
  (b":path", b"/"),
  (b":path", b"/index.html"),
  (b":scheme", b"http"),
  (b":scheme", b"https"),
  (b":status", b"200"),
  (b":status", b"204"),
  (b":status", b"206"),
  (b":status", b"304"),
  (b":status", b"400"),
  (b":status", b"404"),
  (b":status", b"500"),
  (b"accept-", b""),
  (b"accept-encoding", b"gzip, deflate"),
  (b"accept-language", b""),
  (b"accept-ranges", b""),
  (b"accept", b""),
  (b"access-control-allow-origin", b""),
  (b"age", b""),
  (b"allow", b""),
  (b"authorization", b""),
  (b"cache-control", b""),
  (b"content-disposition", b""),
  (b"content-encoding", b""),
  (b"content-language", b""),
  (b"content-length", b""),
  (b"content-location", b""),
  (b"content-range", b""),
  (b"content-type", b""),
  (b"cookie", b""),
  (b"date", b""),
  (b"etag", b""),
  (b"expect", b""),
  (b"expires", b""),
  (b"from", b""),
  (b"host", b""),
  (b"if-match", b""),
  (b"if-modified-since", b""),
  (b"if-none-match", b""),
  (b"if-range", b""),
  (b"if-unmodified-since", b""),
  (b"last-modified", b""),
  (b"link", b""),
  (b"location", b""),
  (b"max-forwards", b""),
  (b"proxy-authenticate", b""),
  (b"proxy-authorization", b""),
  (b"range", b""),
  (b"referer", b""),
  (b"refresh", b""),
  (b"retry-after", b""),
  (b"server", b""),
  (b"set-cookie", b""),
  (b"strict-transport-security", b""),
  (b"transfer-encoding", b""),
  (b"user-agent", b""),
  (b"vary", b""),
  (b"via", b""),
  (b"www-authenticate", b""),
];

/// A struct representing the dynamic table that needs to be maintained by the
/// coder.
struct DynamicTable {
    table: RingBuf<(Vec<u8>, Vec<u8>)>,
    size: usize,
    max_size: usize,
}

impl DynamicTable {
    /// Creates a new empty dynamic table with default size.
    fn new() -> DynamicTable {
        DynamicTable {
            table: RingBuf::new(),
            size: 0,
            // The default maximum size corresponds to the default HTTP/2
            // setting
            max_size: 4096,
        }
    }

    fn get_size(&self) -> usize {
        self.size
    }

    /// Converts the current state of the table to a `Vec`
    fn get_table_as_list(&self) -> Vec<(Vec<u8>, Vec<u8>)> {
        let mut ret: Vec<(Vec<u8>, Vec<u8>)> = Vec::new();
        for elem in self.table.iter() {
            ret.push(elem.clone());
        }

        ret
    }

}

impl fmt::Show for DynamicTable {
    fn fmt(&self, formatter: &mut fmt::Formatter) -> fmt::Result {
        write!(formatter, "{:?}", self.table)
    }
}

mod tests {
    use super::{decode_integer};
    use super::{encode_integer};

    #[test]
    fn test_decode_integer() {
        assert_eq!((10us, 1), decode_integer(&[10], 5));
        assert_eq!((1337us, 3), decode_integer(&[31, 154, 10], 5));
        assert_eq!((1337us, 3), decode_integer(&[31 + 32, 154, 10], 5));
        assert_eq!((1337us, 3), decode_integer(&[31 + 64, 154, 10], 5));
        assert_eq!((1337us, 3), decode_integer(&[31, 154, 10, 342, 22], 5));

        assert_eq!((127us, 2), decode_integer(&[255, 0], 7));
        assert_eq!((127us, 2), decode_integer(&[127, 0], 7));
    }

    #[test]
    fn test_encode_integer() {
        assert_eq!(encode_integer(10, 5), [10]);
        assert_eq!(encode_integer(1337, 5), [31, 154, 10]);
        assert_eq!(encode_integer(127, 7), [127, 0]);
    }
}
<commit_msg>DynamicTable: implement adding and retrieving headers<commit_after>use std::fmt;
use std::collections::RingBuf;

/// Decodes an integer encoded with a given prefix size (in bits).
/// Assumes that the buffer `buf` contains the integer to be decoded,
/// with the first byte representing the octet that contains the
/// prefix.
///
/// Returns a tuple representing the decoded integer and the number
/// of bytes from the buffer that were used.
fn decode_integer(buf: &[u8], prefix_size: u8) -> (usize, usize) {
    assert!(prefix_size >= 1 && prefix_size <= 8);
    assert!(buf.len() >= 1);
    let mask = (1u8 << prefix_size) - 1;
    let mut value: usize = (buf[0] & mask) as usize;
    if value < (mask as usize) {
        return (value, 1);
    }

    // Already one byte used (the prefix)
    let mut total = 1;
    let mut m = 0;
    for &b in buf[1..].iter() {
        total += 1;
        value += ((b & 127) as usize) * (1us << m);
        m += 7;
        if b & 128 != 128 {
            break;
        }
    }

    // TODO Error situation -- what happens if we reach here w/o having
    //      `b` byte-sized...
    (value, total)
}

/// Encode an integer to the representation defined by HPACK.
///
/// Returns a newly allocated `Vec` containing the encoded bytes.
/// Only `prefix_size` lowest-order bits of the first byte in the
/// array are guaranteed to be used.
fn encode_integer(mut value: usize, prefix_size: u8) -> Vec<u8> {
    let mask: usize = ((1u8 << prefix_size) - 1) as usize;
    if value < mask {
        // Right now, the caller would need to be the one to combine
        // the other part of the prefix byte (but would know that it's
        // safe to do so using the bit-wise or).
        return vec![value as u8];
    }

    let mut res: Vec<u8> = Vec::new();
    res.push(mask as u8);
    value -= mask;
    while value >= 128 {
        res.push(((value % 128) + 128) as u8);
        value = value / 128;
    }
    res.push(value as u8);

    res
}

/// The table represents the static header table defined by the HPACK spec.
/// (HPACK, Appendix A)
static STATIC_TABLE: &'static [(&'static [u8], &'static [u8])] = &[
  (b":authority", b""),
  (b":method", b"GET"),
  (b":method", b"POST"),
  (b":path", b"/"),
  (b":path", b"/index.html"),
  (b":scheme", b"http"),
  (b":scheme", b"https"),
  (b":status", b"200"),
  (b":status", b"204"),
  (b":status", b"206"),
  (b":status", b"304"),
  (b":status", b"400"),
  (b":status", b"404"),
  (b":status", b"500"),
  (b"accept-", b""),
  (b"accept-encoding", b"gzip, deflate"),
  (b"accept-language", b""),
  (b"accept-ranges", b""),
  (b"accept", b""),
  (b"access-control-allow-origin", b""),
  (b"age", b""),
  (b"allow", b""),
  (b"authorization", b""),
  (b"cache-control", b""),
  (b"content-disposition", b""),
  (b"content-encoding", b""),
  (b"content-language", b""),
  (b"content-length", b""),
  (b"content-location", b""),
  (b"content-range", b""),
  (b"content-type", b""),
  (b"cookie", b""),
  (b"date", b""),
  (b"etag", b""),
  (b"expect", b""),
  (b"expires", b""),
  (b"from", b""),
  (b"host", b""),
  (b"if-match", b""),
  (b"if-modified-since", b""),
  (b"if-none-match", b""),
  (b"if-range", b""),
  (b"if-unmodified-since", b""),
  (b"last-modified", b""),
  (b"link", b""),
  (b"location", b""),
  (b"max-forwards", b""),
  (b"proxy-authenticate", b""),
  (b"proxy-authorization", b""),
  (b"range", b""),
  (b"referer", b""),
  (b"refresh", b""),
  (b"retry-after", b""),
  (b"server", b""),
  (b"set-cookie", b""),
  (b"strict-transport-security", b""),
  (b"transfer-encoding", b""),
  (b"user-agent", b""),
  (b"vary", b""),
  (b"via", b""),
  (b"www-authenticate", b""),
];

/// A struct representing the dynamic table that needs to be maintained by the
/// coder.
struct DynamicTable {
    table: RingBuf<(Vec<u8>, Vec<u8>)>,
    size: usize,
    max_size: usize,
}

impl DynamicTable {
    /// Creates a new empty dynamic table with default size.
    fn new() -> DynamicTable {
        DynamicTable {
            table: RingBuf::new(),
            size: 0,
            // The default maximum size corresponds to the default HTTP/2
            // setting
            max_size: 4096,
        }
    }

    fn get_size(&self) -> usize {
        self.size
    }

    /// Add a new header to the dynamic table.
    ///
    /// The table automatically gets resized, if necessary.
    ///
    /// Do note that, under the HPACK rules, it is possible the given header
    /// is not found in the dynamic table after this operation finishes, in
    /// case the total size of the given header exceeds the maximum size of the
    /// dynamic table.
    fn add_header(&mut self, name: Vec<u8>, value: Vec<u8>) {
        // This is how the HPACK spec makes us calculate the size.  The 32 is
        // a magic number determined by them (under reasonable assumptions of
        // how the table is stored).
        self.size += name.len() + value.len() + 32;
        debug!("New dynamic table size {}", self.size);
        // Now add it to the internal buffer
        self.table.push_front((name, value));
    }

    /// Returns the number of headers in the dynamic table.
    ///
    /// This is different than the size of the dynamic table.
    fn len(&self) -> usize {
        self.table.len()
    }

    /// Converts the current state of the table to a `Vec`
    fn get_table_as_list(&self) -> Vec<(Vec<u8>, Vec<u8>)> {
        let mut ret: Vec<(Vec<u8>, Vec<u8>)> = Vec::new();
        for elem in self.table.iter() {
            ret.push(elem.clone());
        }

        ret
    }

    /// Returns a reference to the header at the given index, if found in the
    /// dynamic table.
    fn get(&self, index: usize) -> Option<&(Vec<u8>, Vec<u8>)> {
        self.table.get(index)
    }
}

impl fmt::Show for DynamicTable {
    fn fmt(&self, formatter: &mut fmt::Formatter) -> fmt::Result {
        write!(formatter, "{:?}", self.table)
    }
}

mod tests {
    use super::{decode_integer};
    use super::{encode_integer};
    use super::DynamicTable;

    #[test]
    fn test_dynamic_table_size_calculation_simple() {
        let mut table = DynamicTable::new();
        // Sanity check
        assert_eq!(0, table.get_size());

        table.add_header(b"a".to_vec(), b"b".to_vec());

        assert_eq!(32 + 2, table.get_size());
    }

    #[test]
    fn test_dynamic_table_size_calculation() {
        let mut table = DynamicTable::new();

        table.add_header(b"a".to_vec(), b"b".to_vec());
        table.add_header(b"123".to_vec(), b"456".to_vec());
        table.add_header(b"a".to_vec(), b"b".to_vec());

        assert_eq!(3 * 32 + 2 + 6 + 2, table.get_size());
    }

    #[test]
    fn test_decode_integer() {
        assert_eq!((10us, 1), decode_integer(&[10], 5));
        assert_eq!((1337us, 3), decode_integer(&[31, 154, 10], 5));
        assert_eq!((1337us, 3), decode_integer(&[31 + 32, 154, 10], 5));
        assert_eq!((1337us, 3), decode_integer(&[31 + 64, 154, 10], 5));
        assert_eq!((1337us, 3), decode_integer(&[31, 154, 10, 342, 22], 5));

        assert_eq!((127us, 2), decode_integer(&[255, 0], 7));
        assert_eq!((127us, 2), decode_integer(&[127, 0], 7));
    }

    #[test]
    fn test_encode_integer() {
        assert_eq!(encode_integer(10, 5), [10]);
        assert_eq!(encode_integer(1337, 5), [31, 154, 10]);
        assert_eq!(encode_integer(127, 7), [127, 0]);
    }
}
<|endoftext|>


<commit_before>use super::STATIC_TABLE;
use super::HeaderTable;

/// Encode an integer to the representation defined by HPACK.
///
/// Returns a newly allocated `Vec` containing the encoded bytes.
/// Only `prefix_size` lowest-order bits of the first byte in the
/// array are guaranteed to be used.
pub fn encode_integer(mut value: usize, prefix_size: u8) -> Vec<u8> {
    let mask: usize = ((1u8 << prefix_size) - 1) as usize;
    if value < mask {
        // Right now, the caller would need to be the one to combine
        // the other part of the prefix byte (but would know that it's
        // safe to do so using the bit-wise or).
        return vec![value as u8];
    }

    let mut res: Vec<u8> = Vec::new();
    res.push(mask as u8);
    value -= mask;
    while value >= 128 {
        res.push(((value % 128) + 128) as u8);
        value = value / 128;
    }
    res.push(value as u8);

    res
}

/// Represents an HPACK encoder. Allows clients to encode arbitrary header sets
/// and tracks the encoding context. That is, encoding subsequent header sets
/// will use the context built by previous encode calls.
///
/// This is the main API for performing HPACK encoding of headers.
pub struct Encoder<'a> {
    /// The header table represents the encoder's context
    header_table: HeaderTable<'a>,
}

impl<'a> Encoder<'a> {
    /// Creates a new `Encoder` with a default static table, as defined by the
    /// HPACK spec (Appendix A).
    pub fn new() -> Encoder<'a> {
        Encoder {
            header_table: HeaderTable::with_static_table(STATIC_TABLE),
        }
    }
}

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

    #[test]
    fn test_encode_integer() {
        assert_eq!(encode_integer(10, 5), [10]);
        assert_eq!(encode_integer(1337, 5), [31, 154, 10]);
        assert_eq!(encode_integer(127, 7), [127, 0]);
    }
}
<commit_msg>Encoder: Implement initial HPACK encoding<commit_after>/// Implements all functionality related to encoding header blocks using
/// HPACK.
///
/// Clients should use the `Encoder` struct as the API for performing HPACK
/// encoding.
///
/// # Examples
///
/// Encodes a header using a literal encoding.
///
/// ```rust
/// use solicit::hpack::Encoder;
///
/// let mut encoder = Encoder::new();
///
/// let headers = vec![
///     (b"custom-key".to_vec(), b"custom-value".to_vec()),
/// ];
/// // First encoding...
/// let result = encoder.encode(&headers);
/// // The result is a literal encoding of the header name and value, with an
/// // initial byte representing the type of the encoding
/// // (incremental indexing).
/// assert_eq!(
///     vec![0x40,
///          10, b'c', b'u', b's', b't', b'o', b'm', b'-', b'k', b'e', b'y',
///          12, b'c', b'u', b's', b't', b'o', b'm', b'-', b'v', b'a', b'l',
///          b'u', b'e'],
///     result);
/// ```
use super::STATIC_TABLE;
use super::HeaderTable;

/// Encode an integer to the representation defined by HPACK.
///
/// Returns a newly allocated `Vec` containing the encoded bytes.
/// Only `prefix_size` lowest-order bits of the first byte in the
/// array are guaranteed to be used.
pub fn encode_integer(mut value: usize, prefix_size: u8) -> Vec<u8> {
    let mask: usize = ((1u8 << prefix_size) - 1) as usize;
    if value < mask {
        // Right now, the caller would need to be the one to combine
        // the other part of the prefix byte (but would know that it's
        // safe to do so using the bit-wise or).
        return vec![value as u8];
    }

    let mut res: Vec<u8> = Vec::new();
    res.push(mask as u8);
    value -= mask;
    while value >= 128 {
        res.push(((value % 128) + 128) as u8);
        value = value / 128;
    }
    res.push(value as u8);

    res
}

/// Represents an HPACK encoder. Allows clients to encode arbitrary header sets
/// and tracks the encoding context. That is, encoding subsequent header sets
/// will use the context built by previous encode calls.
///
/// This is the main API for performing HPACK encoding of headers.
///
/// # Examples
///
/// Encoding a header two times in a row produces two different
/// representations, due to the utilization of HPACK compression.
///
/// ```rust
/// use solicit::hpack::Encoder;
///
/// let mut encoder = Encoder::new();
///
/// let headers = vec![
///     (b"custom-key".to_vec(), b"custom-value".to_vec()),
/// ];
/// // First encoding...
/// let result = encoder.encode(&headers);
/// // The result is a literal encoding of the header name and value, with an
/// // initial byte representing the type of the encoding
/// // (incremental indexing).
/// assert_eq!(
///     vec![0x40,
///          10, b'c', b'u', b's', b't', b'o', b'm', b'-', b'k', b'e', b'y',
///          12, b'c', b'u', b's', b't', b'o', b'm', b'-', b'v', b'a', b'l',
///          b'u', b'e'],
///     result);
///
/// // Encode the same headers again!
/// let result = encoder.encode(&headers);
/// // The result is simply the index of the header in the header table (62),
/// // with a flag representing that the decoder should use the index.
/// assert_eq!(vec![0x80 | 62], result);
/// ```
pub struct Encoder<'a> {
    /// The header table represents the encoder's context
    header_table: HeaderTable<'a>,
}

impl<'a> Encoder<'a> {
    /// Creates a new `Encoder` with a default static table, as defined by the
    /// HPACK spec (Appendix A).
    pub fn new() -> Encoder<'a> {
        Encoder {
            header_table: HeaderTable::with_static_table(STATIC_TABLE),
        }
    }

    /// Encodes the given headers using the HPACK rules and returns a newly
    /// allocated `Vec` containing the bytes representing the encoded header
    /// set.
    ///
    /// The encoder so far supports only a single, extremely simple encoding
    /// strategy, whereby each header is represented as an indexed header if
    /// already found in the header table and a literal otherwise. When a
    /// header isn't found in the table, it is added if the header name wasn't
    /// found either (i.e. there are never two header names with different
    /// values in the produced header table). Strings are always encoded as
    /// literals (Huffman encoding is not used).
    pub fn encode(&mut self, headers: &Vec<(Vec<u8>, Vec<u8>)>) -> Vec<u8> {
        let mut encoded: Vec<u8> = Vec::new();

        for header in headers.iter() {
            match self.header_table.find_header((&header.0[], &header.1[])) {
                None => {
                    // The name of the header is in no tables: need to encode
                    // it with both a literal name and value.
                    self.encode_literal(header, true, &mut encoded);
                    self.header_table.add_header(header.0.clone(), header.1.clone());
                },
                Some((index, false)) => {
                    // The name of the header is at the given index, but the
                    // value does not match the current one: need to encode
                    // only the value as a literal.
                    self.encode_literal(header, false, &mut encoded);
                },
                Some((index, true)) => {
                    // The full header was found in one of the tables, so we
                    // just encode the index.
                    self.encode_indexed(index, &mut encoded);
                }
            };
        }

        encoded
    }

    /// Encodes a header as a literal (i.e. both the name and the value are
    /// encoded as a string literal) and places the result in the given buffer
    /// `buf`.
    ///
    /// # Parameters
    ///
    /// - `header` - the header to be encoded
    /// - `should_index` - indicates whether the given header should be indexed, i.e.
    ///                    inserted into the dynamic table
    /// - `buf` - The buffer into which the result is placed
    ///
    fn encode_literal(&mut self,
                      header: &(Vec<u8>, Vec<u8>),
                      should_index: bool,
                      buf: &mut Vec<u8>) {
        let mask = if should_index {
            0x40
        } else {
            0x0
        };

        buf.push(mask);
        self.encode_string_literal(&header.0[0..], buf);
        self.encode_string_literal(&header.1[0..], buf);
    }

    /// Encodes a string literal and places the result in the given buffer
    /// `buf`.
    ///
    /// The function does not consider Huffman encoding for now, but always
    /// produces a string literal representations, according to the HPACK spec
    /// section 5.2.
    fn encode_string_literal(&mut self, octet_str: &[u8], buf: &mut Vec<u8>) {
        buf.push_all(&encode_integer(octet_str.len(), 7)[0..]);
        buf.push_all(octet_str);
    }

    /// Encodes a header whose name is indexed and places the result in the
    /// given buffer `buf`.
    fn encode_indexed_name(&mut self, header: (usize, &Vec<u8>), should_index: bool, buf: &mut Vec<u8>) {
        let (mask, prefix) = if should_index {
            (0x40, 6)
        } else {
            (0x0, 4)
        };

        let mut encoded_index = encode_integer(header.0, prefix);
        encoded_index[0] |= mask;
        buf.push_all(&encoded_index[0..]);
        // So far, we rely on just one strategy for encoding string literals.
        self.encode_string_literal(&header.1[0..], buf);
    }

    /// Encodes an indexed header (a header that is fully in the header table)
    /// and places the result in the given buffer `buf`.
    ///
    /// The encoding is according to the rules of the HPACK spec, section 6.1.
    fn encode_indexed(&self, index: usize, buf: &mut Vec<u8>) {
        let mut encoded = encode_integer(index, 7);
        // We need to set the most significant bit, since the bit-pattern is
        // `1xxxxxxx` for indexed headers.
        encoded[0] |= 0x80;

        buf.push_all(&encoded[0..]);
    }
}

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

    use super::super::Decoder;

    #[test]
    fn test_encode_integer() {
        assert_eq!(encode_integer(10, 5), [10]);
        assert_eq!(encode_integer(1337, 5), [31, 154, 10]);
        assert_eq!(encode_integer(127, 7), [127, 0]);
    }

    /// A helper function that checks whether the given buffer can be decoded
    /// into a set of headers that corresponds to the given `headers` list.
    /// Relies on using the `solicit::hpack::decoder::Decoder`` struct for
    /// performing the decoding.
    ///
    /// # Returns
    ///
    /// A `bool` indicating whether such a decoding can be performed.
    fn is_decodable(buf: &Vec<u8>, headers: &Vec<(Vec<u8>, Vec<u8>)>) -> bool {
        let mut decoder = Decoder::new();
        match decoder.decode(&buf[0..]).ok() {
            Some(h) => h == *headers,
            None => false,
        }
    }

    /// Tests that encoding only the `:method` header works.
    #[test]
    fn test_encode_only_method() {
        let mut encoder: Encoder = Encoder::new();
        let headers = vec![
            (b":method".to_vec(), b"GET".to_vec()),
        ];

        let result = encoder.encode(&headers);

        debug!("{:?}", result);
        assert!(is_decodable(&result, &headers));
    }

    /// Tests that when a single custom header is sent it gets indexed by the
    /// coder.
    #[test]
    fn test_custom_header_gets_indexed() {
        let mut encoder: Encoder = Encoder::new();
        let headers = vec![
            (b"custom-key".to_vec(), b"custom-value".to_vec()),
        ];

        let result = encoder.encode(&headers);
        assert!(is_decodable(&result, &headers));
        // The header is in the encoder's dynamic table.
        assert_eq!(encoder.header_table.dynamic_table.to_vec(), headers);
        // ...but also indicated as such in the output.
        assert!(0x40 == (0x40 & result[0]));
        debug!("{:?}", result);
    }

    /// Tests that when a header gets added to the dynamic table, the encoder
    /// will use the index, instead of the literal representation on the next
    /// encoding of the same header.
    #[test]
    fn test_uses_index_on_second_iteration() {
        let mut encoder: Encoder = Encoder::new();
        let headers = vec![
            (b"custom-key".to_vec(), b"custom-value".to_vec()),
        ];
        // First encoding...
        let _ = encoder.encode(&headers);

        // Encode the same headers again!
        let result = encoder.encode(&headers);

        // The header is in the encoder's dynamic table.
        assert_eq!(encoder.header_table.dynamic_table.to_vec(), headers);
        // The output is a single index byte?
        assert_eq!(result.len(), 1);
        // The index is correctly encoded:
        // - The most significant bit is set
        assert_eq!(0x80 & result[0], 0x80);
        // - The other 7 bits decode to an integer giving the index in the full
        //   header address space.
        assert_eq!(result[0] ^ 0x80, 62);
        // assert_eq!(encoder.header_table.get_from_table(62).unwrap(), headers[0]);
    }

    /// Tests that multiple headers are correctly encoded (i.e. can be decoded
    /// back to their original representation).
    #[test]
    fn test_multiple_headers_encoded() {
        let mut encoder = Encoder::new();
        let headers = vec![
            (b"custom-key".to_vec(), b"custom-value".to_vec()),
            (b":method".to_vec(), b"GET".to_vec()),
            (b":path".to_vec(), b"/some/path".to_vec()),
        ];

        let result = encoder.encode(&headers);

        assert!(is_decodable(&result, &headers));
    }
}
<|endoftext|>


<commit_before>use mio::Token;
use mio::net::TcpStream;

use byteorder::{ByteOrder, BigEndian};
use poll;

use std::sync::{Arc, RwLock};
use std::io;
use std::io::{Write, Read, Error, ErrorKind};

pub struct Connection {
    _token: Token,
    _stream: TcpStream,

    //cache the send message between events
    _send_queue: RwLock<Vec<Arc<Vec<u8>>>>,
    _read_next: usize,
    _write_next: bool,
}

impl Connection {
    pub fn new(stream: TcpStream, token: Token) -> Connection {
        Connection {
            _token: token,
            _stream: stream,
            _send_queue: RwLock::new(vec![]),
            _read_next: 0,
            _write_next: false,
        }
    }

    pub fn get_token(&self) -> Token {
        self._token
    }

    //result means, if read success, if fail, should cloase this
    //option means, got data?
    pub fn on_read(&mut self) -> io::Result<Option<Vec<u8>>> {
        let len_result = self.read_message_len();
        if let Ok(len) = len_result {
            if len <= 0 {
                return Ok(None);
            } else {
                let mut vec: Vec<u8> = Vec::with_capacity(len);
                unsafe {
                    vec.set_len(len);
                }
                let read_result = self._stream.read(&mut vec);
                match read_result {
                    Err(e) => {
                        if e.kind() == ErrorKind::WouldBlock {
                            self._read_next = len;
                            return Ok(None);
                        } else {
                            println!("70 {:?}", e);
                            return Err(e);
                        }
                    }
                    Ok(n) => {
                        if n < len {
                            println!("error message number");
                            return Err(Error::new(ErrorKind::InvalidData, "error message number"));
                        }
                        self._read_next = 0;
                        return Ok(Some(vec));
                    }
                } //end match?
            }
        } else {
            return Err(Error::new(ErrorKind::WriteZero, "msg len write failed"));
        }
    }

    fn read_message_len(&mut self) -> io::Result<usize> {
        if self._read_next > 0 {
            return Ok(self._read_next);
        }
        let mut buf = [0u8; 8];
        let bytes = match self._stream.read(&mut buf) {
            Ok(n) => n,
            Err(e) => {
                if e.kind() == ErrorKind::WouldBlock {
                    return Ok(0);
                } else {
                    return Err(e);
                }
            }
        };

        if bytes < 8 {
            return Err(Error::new(ErrorKind::InvalidData, "Invalid message length"));
        }
        let msg_len = BigEndian::read_u64(buf.as_ref());
        Ok(msg_len as usize)
    }

    pub fn on_write(&mut self) -> io::Result<()> {
        let msg_op = {
            let mut queue = self._send_queue.write().unwrap();
            queue.pop()
        };

        if let Some(msg) = msg_op {
            let write_len_result = self.write_message_len(msg.clone());
            match write_len_result {
                Ok(b) => {
                    if !b {
                        let mut queue = self._send_queue.write().unwrap();
                        queue.push(msg);
                        return Ok(());
                    }
                }
                Err(e) => {
                    //error happend
                    println!("write len failed: {:?}", e);
                    return Err(e);
                }
            } //end write len

            let write_result = self._stream.write(&*msg);
            match write_result {
                Ok(_) => {
                    self._write_next = false;
                    self._stream.flush()?;
                    return Ok(());
                }
                Err(e) => {
                    //error happened
                    if e.kind() == ErrorKind::WouldBlock {
                        //message len has been sent, only send
                        //message body next time.
                        self._write_next = true;
                        let mut queue = self._send_queue.write().unwrap();
                        queue.push(msg);
                        println!("on write , would block");
                        return Ok(());
                    }
                }
            }
        } else {
            //println!("all message has bee sended");
        }

        Ok(())
    }

    fn write_message_len(&mut self, msg: Arc<Vec<u8>>) -> io::Result<bool> {
        //means that, the remained message
        //must be sent in next round
        if self._write_next {
            return Ok(false);
        }

        let mut len_buf: [u8; 8] = [0u8; 8]; //a u64 buf
        println!("buf.len {}", len_buf.len());

        BigEndian::write_u64(len_buf.as_mut(), msg.len() as u64);
        let write_result = self._stream.write(&len_buf);
        match write_result {
            Err(e) => {
                if e.kind() == ErrorKind::WouldBlock {
                    return Ok(false);
                } else {
                    return Err(e);
                }
            }
            Ok(_) => {
                //means, write success
                return Ok(true);
            }
        }
    }

    pub fn register(&self, poll: &mut poll::Poller) -> io::Result<()> {
        poll.register_both(&self._stream, self._token)
    }

    pub fn send_message(&self, msg: Arc<Vec<u8>>) {
        let mut queue = self._send_queue.write().unwrap();
        queue.push(msg);
    }
}
<commit_msg>Update connection.rs<commit_after>use mio::Token;
use mio::net::TcpStream;

use byteorder::{ByteOrder, BigEndian};
use poll;

use std::sync::{Arc, RwLock};
use std::io;
use std::io::{Write, Read, Error, ErrorKind};

/* a client with an cnn*/
pub struct Connection {
    _token: Token,
    pub _stream: TcpStream,

    //cache the send message between events
    _send_queue: RwLock<Vec<Arc<Vec<u8>>>>,
    _read_next: usize,
    _write_next: bool,
}

impl Connection {
    pub fn new(stream: TcpStream, token: Token) -> Connection {
        Connection {
            _token: token,
            _stream: stream,
            _send_queue: RwLock::new(vec![]),
            _read_next: 0,
            _write_next: false,
        }
    }

    pub fn get_token(&self) -> Token {
        self._token
    }

    //result means, if read success, if fail, should cloase this
    //option means, got data?
    pub fn on_read(&mut self) -> io::Result<Option<Vec<u8>>> {
        let len_result = self.read_message_len();
        if let Ok(len) = len_result {
            if len <= 0 {
                return Ok(None);
            } else {
                let mut vec: Vec<u8> = Vec::with_capacity(len);
                unsafe {
                    vec.set_len(len);
                }
                let read_result = self._stream.read(&mut vec);
                match read_result {
                    Err(e) => {
                        if e.kind() == ErrorKind::WouldBlock {
                            self._read_next = len;
                            return Ok(None);
                        } else {
                            println!("70 {:?}", e);
                            return Err(e);
                        }
                    }
                    Ok(n) => {
                        if n < len {
                            println!("error message number");
                            return Err(Error::new(ErrorKind::InvalidData, "error message number"));
                        }
                        self._read_next = 0;
                        return Ok(Some(vec));
                    }
                } //end match?
            }
        } else {
            return Err(Error::new(ErrorKind::WriteZero, "msg len write failed"));
        }
    }

    fn read_message_len(&mut self) -> io::Result<usize> {
        if self._read_next > 0 {
            return Ok(self._read_next);
        }
        let mut buf = [0u8; 8];
        let bytes = match self._stream.read(&mut buf) {
            Ok(n) => n,
            Err(e) => {
                if e.kind() == ErrorKind::WouldBlock {
                    return Ok(0);
                } else {
                    return Err(e);
                }
            }
        };

        if bytes < 8 {
            return Err(Error::new(ErrorKind::InvalidData, "Invalid message length"));
        }
        let msg_len = BigEndian::read_u64(buf.as_ref());
        Ok(msg_len as usize)
    }

    pub fn on_write(&mut self) -> io::Result<()> {
        let msg_op = {
            let mut queue = self._send_queue.write().unwrap();
            queue.pop()
        };

        if let Some(msg) = msg_op {
            let write_len_result = self.write_message_len(msg.clone());
            match write_len_result {
                Ok(b) => {
                    if !b {
                        let mut queue = self._send_queue.write().unwrap();
                        queue.push(msg);
                        return Ok(());
                    }
                }
                Err(e) => {
                    //error happend
                    println!("write len failed: {:?}", e);
                    return Err(e);
                }
            } //end write len

            let write_result = self._stream.write(&*msg);
            match write_result {
                Ok(_) => {
                    self._write_next = false;
                    self._stream.flush()?;
                    return Ok(());
                }
                Err(e) => {
                    //error happened
                    if e.kind() == ErrorKind::WouldBlock {
                        //message len has been sent, only send
                        //message body next time.
                        self._write_next = true;
                        let mut queue = self._send_queue.write().unwrap();
                        queue.push(msg);
                        println!("on write , would block");
                        return Ok(());
                    }
                }
            }
        } else {
            //println!("all message has bee sended");
        }

        Ok(())
    }

    fn write_message_len(&mut self, msg: Arc<Vec<u8>>) -> io::Result<bool> {
        //means that, the remained message
        //must be sent in next round
        if self._write_next {
            return Ok(false);
        }

        let mut len_buf: [u8; 8] = [0u8; 8]; //a u64 buf
        println!("buf.len {}", len_buf.len());

        BigEndian::write_u64(len_buf.as_mut(), msg.len() as u64);
        let write_result = self._stream.write(&len_buf);
        match write_result {
            Err(e) => {
                if e.kind() == ErrorKind::WouldBlock {
                    return Ok(false);
                } else {
                    return Err(e);
                }
            }
            Ok(_) => {
                //means, write success
                return Ok(true);
            }
        }
    }

    pub fn register(&self, poll: &mut poll::Poller) -> io::Result<()> {
        poll.register_both(&self._stream, self._token)
    }

    pub fn send_message(&self, msg: Arc<Vec<u8>>) {
        let mut queue = self._send_queue.write().unwrap();
        queue.push(msg);
    }
}
<|endoftext|>


<commit_before>extern crate ring;
extern crate udt;

use std::net::{UdpSocket, SocketAddr, IpAddr};
use std::str;
use std::fmt;
use udt::{UdtSocket, UdtError, UdtOpts, SocketType, SocketFamily};
use ring::aead;
use ring::rand;

// TODO config
const UDT_BUF_SIZE: i32 = 4096000;
pub const MAX_MESSAGE_SIZE: usize = 1024000;

pub mod crypto {
    use ring::aead;
    use ring::aead::{SealingKey, OpeningKey, Algorithm};
    use ring::rand::{SystemRandom, SecureRandom};
    static ALGORITHM: &'static Algorithm = &aead::CHACHA20_POLY1305;

    pub struct Handler {
        _working_buf: [u8; super::MAX_MESSAGE_SIZE],
        rand: SystemRandom,
        opening_key: OpeningKey,
        sealing_key: SealingKey,
    }

    pub fn gen_key() -> Vec<u8> {
        let rand = SystemRandom::new();
        let mut keybytes = vec![0u8; ALGORITHM.key_len()];
        rand.fill(&mut keybytes);
        keybytes
    }

    impl Handler {
        pub fn new(key: &[u8]) -> Handler {
            Handler {
                _working_buf: [0u8; super::MAX_MESSAGE_SIZE],
                rand: SystemRandom::new(),
                opening_key: aead::OpeningKey::new(ALGORITHM, key).unwrap(),
                sealing_key: aead::SealingKey::new(ALGORITHM, key).unwrap(),
            }
        }

        pub fn seal(&mut self, buf: &mut [u8], len: usize) -> Result<usize, ()> {
            let nonce_len = ALGORITHM.nonce_len();
            let max_suffix_len = ALGORITHM.max_overhead_len();

            assert!(nonce_len < u8::max_value() as usize,
                    "Uh, why is the nonce size this big?");

            assert!(len <= buf.len() - max_suffix_len,
                    "Buffer doesn't have enough suffix padding.");

            let mut nonce = vec![0u8; nonce_len];
            self.rand.fill(&mut nonce).unwrap();

            let mut sealed = vec![0u8; len + max_suffix_len];
            sealed[0..len].copy_from_slice(&buf[..len]);
            match aead::seal_in_place(&self.sealing_key,
                                      &nonce,
                                      &mut sealed,
                                      max_suffix_len,
                                      &[]) {
                Ok(seal_len) => {
                    buf[..nonce_len].copy_from_slice(&nonce[..]);
                    buf[nonce_len..nonce_len+seal_len].copy_from_slice(&sealed[..seal_len]);
                    Ok(nonce_len + seal_len)
                }
                Err(e) => {
                    Err(())
                }
            }
        }

        pub fn open(&mut self, buf: &mut [u8]) -> Result<usize, String> {
            let nonce_len = ALGORITHM.nonce_len();

            if buf.len() < nonce_len {
                return Err("msg not long enough to contain nonce".into());
            } else if buf.len() > super::MAX_MESSAGE_SIZE {
                return Err("max message size exceeded".into());
            }

            let nonce = buf[..nonce_len].to_owned();

            aead::open_in_place(&self.opening_key, &nonce, nonce_len, buf, &[])
                .map_err(|_| String::from("decrypt failed"))
        }
    }

    // Tests for the crypto module
    #[cfg(test)]
    mod test {

        #[test]
        fn raw_roundtrip() {
            use ring::aead;
            use ring::aead::{SealingKey, OpeningKey, Algorithm};
            use ring::rand::{SystemRandom, SecureRandom};

            let rng = SystemRandom::new();
            let mut key_bytes = vec![0u8; super::ALGORITHM.key_len()];
            let mut nonce_bytes = vec![0u8; super::ALGORITHM.nonce_len()];
            rng.fill(&mut key_bytes).unwrap();
            rng.fill(&mut nonce_bytes).unwrap();
            let key = SealingKey::new(super::ALGORITHM, &key_bytes).unwrap();

            let data = [1u8; 1350];
            let out_suffix_capacity = super::ALGORITHM.max_overhead_len();
            let mut in_out = vec![1u8; data.len() + out_suffix_capacity];
            aead::seal_in_place(&key, &nonce_bytes,
                                &mut in_out, out_suffix_capacity,
                                &[]).unwrap();

            let opening_key = OpeningKey::new(super::ALGORITHM, &key_bytes).unwrap();
            let len = aead::open_in_place(&opening_key, &nonce_bytes,
                                          0, &mut in_out, &[]).unwrap();

            assert_eq!(len, 1350);
            assert_eq!(&in_out[..len], &data[..]);
        }

        #[test]
        fn roundtrip() {
            use ::rand;
            use ::rand::distributions::{IndependentSample, Range};
            // generate some data, seal it, and then make sure it unseals to the same thing
            let mut rng = rand::thread_rng();
            let between = Range::new(10, 10000);

            let key = super::gen_key();
            let mut handler = super::Handler::new(&key);
            let data_size: usize = between.ind_sample(&mut rng);
            let mut data = vec![0u8; super::super::MAX_MESSAGE_SIZE];
            for i in 0..data_size {
                data[i] = rand::random();
            }

            let orig = data[..data_size].to_owned();

            let cipher_len = handler.seal(&mut data, data_size).unwrap();
            let decrypted_len = handler.open(&mut data[..cipher_len]).unwrap();
            assert_eq!(decrypted_len, data_size);
            assert_eq!(orig, &data[..decrypted_len], "original and decrypted don't match!");
        }

        //#[test]
        //fn key_sanity() {
        //    use std::collections::HashSet;
        //    use sodiumoxide::crypto::secretbox;
        //    use sodiumoxide::crypto::secretbox::xsalsa20poly1305::Key;

        //    let mut set: HashSet<[u8;32]> = HashSet::with_capacity(10000);

        //    for _ in 0..10000 {
        //        let key = secretbox::gen_key();
        //        let Key(keybytes) = key;
        //        assert!(set.insert(keybytes));
        //    }
        //}

        //#[test]
        //fn nonce_sanity() {
        //    use std::collections::HashSet;
        //    use sodiumoxide::crypto::secretbox;
        //    use sodiumoxide::crypto::secretbox::xsalsa20poly1305::Nonce;

        //    let mut set: HashSet<[u8;24]> = HashSet::with_capacity(10000);

        //    for _ in 0..10000 {
        //        let nonce = secretbox::gen_nonce();
        //        let Nonce(noncebytes) = nonce;
        //        assert!(set.insert(noncebytes));
        //    }
        //}
    }

    #[cfg(all(feature = "nightly", test))]
    mod bench {
       extern crate test;

       // #[bench]
       // fn bench_seal(b: &mut test::Bencher) {
       //     let key = super::gen_key();
       //     let mut handler = super::Handler::new(&key);
       //     let mut buf = [0; super::super::MAX_MESSAGE_SIZE];
       //     b.bytes = 1300;
       //     b.iter(move || handler.seal(&mut buf, 1300))
       // }
    }
}

fn new_udt_socket() -> UdtSocket {
    udt::init();
    let sock = UdtSocket::new(SocketFamily::AFInet, SocketType::Datagram).unwrap();
    sock.setsockopt(UdtOpts::UDP_RCVBUF, UDT_BUF_SIZE).unwrap();
    sock.setsockopt(UdtOpts::UDP_SNDBUF, UDT_BUF_SIZE).unwrap();
    sock
}

fn send(sock: &UdtSocket, key: &aead::SealingKey, buf: &mut [u8], len: usize) -> Result<(), UdtError> {
    unimplemented!();
    // FIXME don't unwrap, create an Error struct that can handle everything
    // if let Ok(sealed_len) = crypto::seal(buf, len, key) {
    //     sock.sendmsg(&buf[..sealed_len]).map(|_| ())
    // } else {
    //     Err(UdtError {
    //         err_code: -1,
    //         err_msg: "encryption failure".into(),
    //     })
    // }
}

fn recv(sock: &UdtSocket, key: &aead::OpeningKey, buf: &mut [u8]) -> Result<usize, UdtError> {
    unimplemented!();
    // let size = try!(sock.recvmsg(buf));
    // crypto::open(&mut buf[..size], key).map_err(|_| {
    //     UdtError {
    //         err_code: -1,
    //         err_msg: String::from("decryption failure"),
    //     }
    // })
}

pub struct PortRange {
    start: u16,
    end: u16,
}

pub trait Transceiver {
    fn send(&self, buf: &[u8]) -> Result<(), UdtError>;
    fn recv(&self, buf: &mut [u8]) -> Result<usize, UdtError>;
    fn close(&self) -> Result<(), UdtError>;
}

pub struct Server {
    pub ip_addr: IpAddr,
    pub port: u16,
    crypto: crypto::Handler,
    sock: UdtSocket,
}

pub struct Client {
    addr: SocketAddr,
    sock: UdtSocket,
    crypto: crypto::Handler,
}

pub struct ServerConnection<'a> {
    crypto: &'a crypto::Handler,
    sock: UdtSocket,
}

impl Client {
    pub fn new(addr: SocketAddr, key: &[u8]) -> Client {
        let sock = new_udt_socket();
        Client {
            addr: addr,
            sock: sock,
            crypto: crypto::Handler::new(key),
        }
    }

    pub fn connect(&self) -> Result<(), UdtError> {
        self.sock.connect(self.addr)
    }
}

impl Transceiver for Client {
    fn send(&self, buf: &[u8]) -> Result<(), UdtError> {
        unimplemented!();
        // send(&self.sock, &self.sealing_key, buf)
    }

    fn recv(&self, buf: &mut [u8]) -> Result<usize, UdtError> {
        unimplemented!();
        // recv(&self.sock, &self.opening_key, buf)
    }

    fn close(&self) -> Result<(), UdtError> {
        self.sock.close()
    }
}

impl Server {
    pub fn get_open_port(range: &PortRange) -> Result<u16, ()> {
        for p in range.start..range.end {
            if let Ok(_) = UdpSocket::bind(&format!("0.0.0.0:{}", p)[..]) {
                return Ok(p);
            }
        }
        Err(())
    }

    pub fn new(ip_addr: IpAddr, port: u16, key: &[u8]) -> Server {
        let sock = new_udt_socket();
        sock.bind(SocketAddr::new(ip_addr, port)).unwrap();
        Server {
            sock: sock,
            ip_addr: ip_addr,
            port: port,
            crypto: crypto::Handler::new(key),
        }
    }

    pub fn listen(&self) -> Result<(), UdtError> {
        self.sock.listen(2)
    }

    pub fn accept(&self) -> Result<ServerConnection, UdtError> {
        self.sock.accept().map(|(sock, _)| {
            ServerConnection {
                crypto: &self.crypto,
                sock: sock,
            }
        })
    }
}

impl<'a> ServerConnection<'a> {
    pub fn getpeer(&self) -> Result<SocketAddr, UdtError> {
        self.sock.getpeername()
    }
}

impl<'a> Transceiver for ServerConnection<'a> {
    fn send(&self, buf: &[u8]) -> Result<(), UdtError> {
        unimplemented!();
        // send(&self.sock, self.key, buf)
    }

    fn recv(&self, buf: &mut [u8]) -> Result<usize, UdtError> {
        unimplemented!();
        // recv(&self.sock, self.key, buf)
    }

    fn close(&self) -> Result<(), UdtError> {
        self.sock.close()
    }
}

impl<'a> PortRange {
    fn new(start: u16, end: u16) -> Result<PortRange, &'a str> {
        if start > end {
            Err("range end must be greater than or equal to start")
        } else {
            Ok(PortRange {
                start: start,
                end: end,
            })
        }
    }

    pub fn from(s: &str) -> Result<PortRange, &'a str> {
        let sections: Vec<&str> = s.split('-').collect();
        if sections.len() != 2 {
            return Err("Range must be specified in the form of \"<start>-<end>\"");
        }
        let (start, end) = (sections[0].parse::<u16>(), sections[1].parse::<u16>());
        if start.is_err() || end.is_err() {
            return Err("improperly formatted port range");
        }
        PortRange::new(start.unwrap(), end.unwrap())
    }
}

impl fmt::Display for PortRange {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "{}-{}", self.start, self.end)
    }
}
<commit_msg>more complete crypto benchmarks<commit_after>extern crate ring;
extern crate udt;

use std::net::{UdpSocket, SocketAddr, IpAddr};
use std::str;
use std::fmt;
use udt::{UdtSocket, UdtError, UdtOpts, SocketType, SocketFamily};
use ring::aead;
use ring::rand;

// TODO config
const UDT_BUF_SIZE: i32 = 4096000;
pub const MAX_MESSAGE_SIZE: usize = 1024000;

pub mod crypto {
    use ring::aead;
    use ring::aead::{SealingKey, OpeningKey, Algorithm};
    use ring::rand::{SystemRandom, SecureRandom};
    static ALGORITHM: &'static Algorithm = &aead::AES_256_GCM;

    pub struct Handler {
        _working_buf: [u8; super::MAX_MESSAGE_SIZE],
        rand: SystemRandom,
        opening_key: OpeningKey,
        sealing_key: SealingKey,
    }

    pub fn gen_key() -> Vec<u8> {
        let rand = SystemRandom::new();
        let mut keybytes = vec![0u8; ALGORITHM.key_len()];
        rand.fill(&mut keybytes);
        keybytes
    }

    impl Handler {
        pub fn new(key: &[u8]) -> Handler {
            Handler {
                _working_buf: [0u8; super::MAX_MESSAGE_SIZE],
                rand: SystemRandom::new(),
                opening_key: aead::OpeningKey::new(ALGORITHM, key).unwrap(),
                sealing_key: aead::SealingKey::new(ALGORITHM, key).unwrap(),
            }
        }

        pub fn seal(&mut self, buf: &mut [u8], len: usize) -> Result<usize, ()> {
            let nonce_len = ALGORITHM.nonce_len();
            let max_suffix_len = ALGORITHM.max_overhead_len();

            assert!(nonce_len < u8::max_value() as usize,
                    "Uh, why is the nonce size this big?");

            assert!(len <= buf.len() - max_suffix_len,
                    "Buffer doesn't have enough suffix padding.");

            let mut nonce = vec![0u8; nonce_len];
            self.rand.fill(&mut nonce).unwrap();

            let mut sealed = vec![0u8; len + max_suffix_len];
            sealed[0..len].copy_from_slice(&buf[..len]);
            match aead::seal_in_place(&self.sealing_key,
                                      &nonce,
                                      &mut sealed,
                                      max_suffix_len,
                                      &[]) {
                Ok(seal_len) => {
                    buf[..nonce_len].copy_from_slice(&nonce[..]);
                    buf[nonce_len..nonce_len+seal_len].copy_from_slice(&sealed[..seal_len]);
                    Ok(nonce_len + seal_len)
                }
                Err(e) => {
                    Err(())
                }
            }
        }

        pub fn open(&mut self, buf: &mut [u8]) -> Result<usize, String> {
            let nonce_len = ALGORITHM.nonce_len();

            if buf.len() < nonce_len {
                return Err("msg not long enough to contain nonce".into());
            } else if buf.len() > super::MAX_MESSAGE_SIZE {
                return Err("max message size exceeded".into());
            }

            let nonce = buf[..nonce_len].to_owned();

            aead::open_in_place(&self.opening_key, &nonce, nonce_len, buf, &[])
                .map_err(|_| String::from("decrypt failed"))
        }
    }

    // Tests for the crypto module
    #[cfg(test)]
    mod test {

        #[test]
        fn raw_roundtrip() {
            use ring::aead;
            use ring::aead::{SealingKey, OpeningKey, Algorithm};
            use ring::rand::{SystemRandom, SecureRandom};

            let rng = SystemRandom::new();
            let mut key_bytes = vec![0u8; super::ALGORITHM.key_len()];
            let mut nonce_bytes = vec![0u8; super::ALGORITHM.nonce_len()];
            rng.fill(&mut key_bytes).unwrap();
            rng.fill(&mut nonce_bytes).unwrap();
            let key = SealingKey::new(super::ALGORITHM, &key_bytes).unwrap();

            let data = [1u8; 1350];
            let out_suffix_capacity = super::ALGORITHM.max_overhead_len();
            let mut in_out = vec![1u8; data.len() + out_suffix_capacity];
            aead::seal_in_place(&key, &nonce_bytes,
                                &mut in_out, out_suffix_capacity,
                                &[]).unwrap();

            let opening_key = OpeningKey::new(super::ALGORITHM, &key_bytes).unwrap();
            let len = aead::open_in_place(&opening_key, &nonce_bytes,
                                          0, &mut in_out, &[]).unwrap();

            assert_eq!(len, 1350);
            assert_eq!(&in_out[..len], &data[..]);
        }

        #[test]
        fn roundtrip() {
            use ::rand;
            use ::rand::distributions::{IndependentSample, Range};
            // generate some data, seal it, and then make sure it unseals to the same thing
            let mut rng = rand::thread_rng();
            let between = Range::new(10, 10000);

            let key = super::gen_key();
            let mut handler = super::Handler::new(&key);
            let data_size: usize = between.ind_sample(&mut rng);
            let mut data = vec![0u8; super::super::MAX_MESSAGE_SIZE];
            for i in 0..data_size {
                data[i] = rand::random();
            }

            let orig = data[..data_size].to_owned();

            let cipher_len = handler.seal(&mut data, data_size).unwrap();
            let decrypted_len = handler.open(&mut data[..cipher_len]).unwrap();
            assert_eq!(decrypted_len, data_size);
            assert_eq!(orig, &data[..decrypted_len], "original and decrypted don't match!");
        }

        #[test]
        fn key_sanity() {
            use std::collections::HashSet;

            let mut set: HashSet<Vec<u8>> = HashSet::with_capacity(10000);

            for _ in 0..10000 {
                let key = super::gen_key();
                assert!(set.insert(key));
            }
        }
    }

    #[cfg(all(feature = "nightly", test))]
    mod bench {
        extern crate test;
        const DATA_SIZE: usize = 1300;

        #[bench]
        fn bench_raw_seal(b: &mut test::Bencher) {
            use ring::aead;
            use ring::aead::{SealingKey, OpeningKey, Algorithm};
            use ring::rand::{SystemRandom, SecureRandom};

            let rng = SystemRandom::new();
            let mut key_bytes = vec![0u8; super::ALGORITHM.key_len()];
            let mut nonce_bytes = vec![0u8; super::ALGORITHM.nonce_len()];
            rng.fill(&mut key_bytes).unwrap();
            rng.fill(&mut nonce_bytes).unwrap();
            let key = SealingKey::new(super::ALGORITHM, &key_bytes).unwrap();

            let data = [1u8; DATA_SIZE];
            let out_suffix_capacity = super::ALGORITHM.max_overhead_len();
            let mut in_out = vec![1u8; data.len() + out_suffix_capacity];

            b.bytes = DATA_SIZE as u64;
            b.iter(move || aead::seal_in_place(&key, &nonce_bytes, &mut in_out,
                                               out_suffix_capacity, &[]).unwrap())
        }

        #[bench]
        fn bench_raw_open(b: &mut test::Bencher) {
            use ring::aead;
            use ring::aead::{SealingKey, OpeningKey, Algorithm};
            use ring::rand::{SystemRandom, SecureRandom};

            let rng = SystemRandom::new();
            let mut key_bytes = vec![0u8; super::ALGORITHM.key_len()];
            let mut nonce_bytes = vec![0u8; super::ALGORITHM.nonce_len()];
            rng.fill(&mut key_bytes).unwrap();
            rng.fill(&mut nonce_bytes).unwrap();
            let key = SealingKey::new(super::ALGORITHM, &key_bytes).unwrap();
            let opening_key = OpeningKey::new(super::ALGORITHM, &key_bytes).unwrap();

            let data = [1u8; DATA_SIZE];
            let out_suffix_capacity = super::ALGORITHM.max_overhead_len();
            let mut in_out = vec![1u8; data.len() + out_suffix_capacity];

            b.bytes = DATA_SIZE as u64;

            let sealed_len = aead::seal_in_place(&key, &nonce_bytes, &mut in_out,
                                                 out_suffix_capacity, &[]).unwrap();
            b.iter(move || aead::open_in_place(&opening_key, &nonce_bytes,
                                               0, &mut in_out, &[]))
        }

        #[bench]
        fn bench_seal(b: &mut test::Bencher) {
            let key = super::gen_key();
            let mut handler = super::Handler::new(&key);
            let mut buf = vec![0u8; super::super::MAX_MESSAGE_SIZE];
            b.bytes = DATA_SIZE as u64;
            b.iter(move || handler.seal(&mut buf, DATA_SIZE))
        }

        #[bench]
        fn bench_open(b: &mut test::Bencher) {
            let key = super::gen_key();
            let mut handler = super::Handler::new(&key);
            let mut buf = vec![0u8; super::super::MAX_MESSAGE_SIZE];
            let sealed_len = handler.seal(&mut buf, DATA_SIZE).unwrap();
            b.bytes = DATA_SIZE as u64;
            b.iter(move || handler.open(&mut buf[..sealed_len]))
        }
    }
}

fn new_udt_socket() -> UdtSocket {
    udt::init();
    let sock = UdtSocket::new(SocketFamily::AFInet, SocketType::Datagram).unwrap();
    sock.setsockopt(UdtOpts::UDP_RCVBUF, UDT_BUF_SIZE).unwrap();
    sock.setsockopt(UdtOpts::UDP_SNDBUF, UDT_BUF_SIZE).unwrap();
    sock
}

fn send(sock: &UdtSocket, key: &aead::SealingKey, buf: &mut [u8], len: usize) -> Result<(), UdtError> {
    unimplemented!();
    // FIXME don't unwrap, create an Error struct that can handle everything
    // if let Ok(sealed_len) = crypto::seal(buf, len, key) {
    //     sock.sendmsg(&buf[..sealed_len]).map(|_| ())
    // } else {
    //     Err(UdtError {
    //         err_code: -1,
    //         err_msg: "encryption failure".into(),
    //     })
    // }
}

fn recv(sock: &UdtSocket, key: &aead::OpeningKey, buf: &mut [u8]) -> Result<usize, UdtError> {
    unimplemented!();
    // let size = try!(sock.recvmsg(buf));
    // crypto::open(&mut buf[..size], key).map_err(|_| {
    //     UdtError {
    //         err_code: -1,
    //         err_msg: String::from("decryption failure"),
    //     }
    // })
}

pub struct PortRange {
    start: u16,
    end: u16,
}

pub trait Transceiver {
    fn send(&self, buf: &[u8]) -> Result<(), UdtError>;
    fn recv(&self, buf: &mut [u8]) -> Result<usize, UdtError>;
    fn close(&self) -> Result<(), UdtError>;
}

pub struct Server {
    pub ip_addr: IpAddr,
    pub port: u16,
    crypto: crypto::Handler,
    sock: UdtSocket,
}

pub struct Client {
    addr: SocketAddr,
    sock: UdtSocket,
    crypto: crypto::Handler,
}

pub struct ServerConnection<'a> {
    crypto: &'a crypto::Handler,
    sock: UdtSocket,
}

impl Client {
    pub fn new(addr: SocketAddr, key: &[u8]) -> Client {
        let sock = new_udt_socket();
        Client {
            addr: addr,
            sock: sock,
            crypto: crypto::Handler::new(key),
        }
    }

    pub fn connect(&self) -> Result<(), UdtError> {
        self.sock.connect(self.addr)
    }
}

impl Transceiver for Client {
    fn send(&self, buf: &[u8]) -> Result<(), UdtError> {
        unimplemented!();
        // send(&self.sock, &self.sealing_key, buf)
    }

    fn recv(&self, buf: &mut [u8]) -> Result<usize, UdtError> {
        unimplemented!();
        // recv(&self.sock, &self.opening_key, buf)
    }

    fn close(&self) -> Result<(), UdtError> {
        self.sock.close()
    }
}

impl Server {
    pub fn get_open_port(range: &PortRange) -> Result<u16, ()> {
        for p in range.start..range.end {
            if let Ok(_) = UdpSocket::bind(&format!("0.0.0.0:{}", p)[..]) {
                return Ok(p);
            }
        }
        Err(())
    }

    pub fn new(ip_addr: IpAddr, port: u16, key: &[u8]) -> Server {
        let sock = new_udt_socket();
        sock.bind(SocketAddr::new(ip_addr, port)).unwrap();
        Server {
            sock: sock,
            ip_addr: ip_addr,
            port: port,
            crypto: crypto::Handler::new(key),
        }
    }

    pub fn listen(&self) -> Result<(), UdtError> {
        self.sock.listen(2)
    }

    pub fn accept(&self) -> Result<ServerConnection, UdtError> {
        self.sock.accept().map(|(sock, _)| {
            ServerConnection {
                crypto: &self.crypto,
                sock: sock,
            }
        })
    }
}

impl<'a> ServerConnection<'a> {
    pub fn getpeer(&self) -> Result<SocketAddr, UdtError> {
        self.sock.getpeername()
    }
}

impl<'a> Transceiver for ServerConnection<'a> {
    fn send(&self, buf: &[u8]) -> Result<(), UdtError> {
        unimplemented!();
        // send(&self.sock, self.key, buf)
    }

    fn recv(&self, buf: &mut [u8]) -> Result<usize, UdtError> {
        unimplemented!();
        // recv(&self.sock, self.key, buf)
    }

    fn close(&self) -> Result<(), UdtError> {
        self.sock.close()
    }
}

impl<'a> PortRange {
    fn new(start: u16, end: u16) -> Result<PortRange, &'a str> {
        if start > end {
            Err("range end must be greater than or equal to start")
        } else {
            Ok(PortRange {
                start: start,
                end: end,
            })
        }
    }

    pub fn from(s: &str) -> Result<PortRange, &'a str> {
        let sections: Vec<&str> = s.split('-').collect();
        if sections.len() != 2 {
            return Err("Range must be specified in the form of \"<start>-<end>\"");
        }
        let (start, end) = (sections[0].parse::<u16>(), sections[1].parse::<u16>());
        if start.is_err() || end.is_err() {
            return Err("improperly formatted port range");
        }
        PortRange::new(start.unwrap(), end.unwrap())
    }
}

impl fmt::Display for PortRange {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "{}-{}", self.start, self.end)
    }
}
<|endoftext|>