<commit_before>extern crate "rust-crypto" as rust_crypto;
extern crate serialize;

use std::io::{TcpListener, TcpStream};
use std::io::{Acceptor, Listener};
use std::io::{File, BufferedStream};
use std::string::String;

use rust_crypto::sha1::Sha1;
use rust_crypto::digest::Digest;
use serialize::base64::{ToBase64, STANDARD};

pub const CR: u8 = b'\r';
pub const LF: u8 = b'\n';
pub const SP: u8 = b' ';
pub const CRLF: &'static [u8] = &[CR,LF];

struct RequestedRoute {
    method: String,
    pathname: String
}

fn parse_request_line(header: &str) -> RequestedRoute {
    let mut pathname = String::new();
    let mut method = String::new();
    let mut starting = true;
    let mut found_method = false;

    for ch in header.graphemes(true) {
        let cbyte = ch.as_bytes()[0];

        if starting && cbyte != SP {
            method.push_str(ch);
        }
        if starting && cbyte == SP {
            starting = false;
            found_method = true;
            continue;
        }
        if found_method && cbyte != SP {
            pathname.push_str(ch);
        }
        if found_method && cbyte == SP {
            break;
        }

    }
    return RequestedRoute {method: method, pathname: pathname};
}

fn get_normal_body (path_on_disk: &str) -> String {
    let path = Path::new(path_on_disk);
    let display = path.display();
    let mut file = match File::open(&path) {
        Ok(f) => f,
        Err(err) => panic!("file error: {}", err)
    };

    let content = match file.read_to_end() {
        Ok(c) => c,
        Err(err) => panic!("{}", err)
    };

    let s = String::from_utf8(content).unwrap();
    return s;
}

fn get_js_body () -> String {
    let path = Path::new("../html/ws1.js");
    let display = path.display();
    let mut file = match File::open(&path) {
        Ok(f) => f,
        Err(err) => panic!("file error: {}", err)
    };

    let content = match file.read_to_end() {
        Ok(c) => c,
        Err(err) => panic!("{}", err)
    };

    let s = String::from_utf8(content).unwrap();
    return s;
}

fn sec_handshake (from_server: &[u8]) -> String {

    let guid = b"258EAFA5-E914-47DA-95CA-C5AB0DC85B11";

    let mut sha = Sha1::new();

    sha.input(from_server);
    sha.input(guid);
    let mut out = [0u8, ..20];
    sha.result(out.as_mut_slice());

    println!("{} {}", sha.result_str(), out.to_base64(STANDARD));

    return out.to_base64(STANDARD);
}

fn main () {

    let listener = TcpListener::bind("127.0.0.1:8099").unwrap();
    let mut acceptor = listener.listen().unwrap();

    fn handle_client(mut stream: BufferedStream<TcpStream>) {

        let mut body: String = "
You sent it!
".to_string();

        let mut cur_line: String;
        let mut ii = 0u;
        let mut req: RequestedRoute = RequestedRoute{
            pathname: String::new(),
            method: String::new()
        };

        loop {
            match stream.read_line() { // XXX: strange unwrap like thing
                Ok(line) => cur_line = line,
                Err(_) => break
            }
            if ii == 0u { // the Request-Line is always the first line
                req = parse_request_line(cur_line.as_slice());
                ii += 1;
            }

            body = body + cur_line + "
";
            if cur_line.as_bytes() == CRLF {
                break;
            }
        }

        println!("pathname {} method {}", req.pathname, req.method);
        let mut secHandshake = false;

        if req.pathname.as_bytes() == b"/" {
            body = get_normal_body("../html/ws1.html");
        } else if req.pathname.as_bytes() == b"/ws1.js" {
            body = get_js_body();
        } else if req.pathname.as_bytes() == b"/ws" {
            secHandshake = true;
            // TODO get handshake
        }

        let body_length = format!("Content-length: {}", body.len());

        stream.write(b"HTTP/1.1 200 OK\r\n").unwrap(); // byte literal
        stream.write(b"Content-type: text/html\r\n").unwrap();
        stream.write(b"X-header: from bytes\r\n").unwrap();
        stream.write(body_length.into_bytes().as_slice()).unwrap();
        stream.write(b"\r\n\r\n").unwrap();
        stream.write(body.into_bytes().as_slice()).unwrap();
        stream.flush().unwrap();

        println!("Handling acceptor");
    }

    for stream in acceptor.incoming() {
        match stream {
            Err(e) => { println!("connection failed: {}", e) }
            Ok(stream) => spawn(proc() {
                let bs = BufferedStream::new(stream);
                handle_client(bs)
            })
        }
    }

    drop(acceptor);
}
<commit_msg>get headers by name, get security header<commit_after>extern crate "rust-crypto" as rust_crypto;
extern crate serialize;

use std::io::{TcpListener, TcpStream};
use std::io::{Acceptor, Listener};
use std::io::{File, BufferedStream};
use std::string::String;

use rust_crypto::sha1::Sha1;
use rust_crypto::digest::Digest;
use serialize::base64::{ToBase64, STANDARD};

pub const CR: u8 = b'\r';
pub const LF: u8 = b'\n';
pub const SP: u8 = b' ';
pub const COLON: u8 = b':';
pub const CRLF: &'static [u8] = &[CR,LF];

struct RequestedRoute {
    method: String,
    pathname: String
}

struct ClientHeader {
    key: String,
    value: String
}

fn get_header_by_name (header: &[u8], headers: Vec<ClientHeader>) -> String {
    let mut result = String::new();

    for h in headers.into_iter() {
        if (h.key.as_bytes() == header) {
            result = h.value;
            break;
        }
    }

    return result;
}

fn parse_request_line (header: &str) -> RequestedRoute {
    let mut pathname = String::new();
    let mut method = String::new();
    let mut starting = true;
    let mut found_method = false;

    for ch in header.graphemes(true) {
        let cbyte = ch.as_bytes()[0];

        if starting && cbyte != SP {
            method.push_str(ch);
        }
        if starting && cbyte == SP {
            starting = false;
            found_method = true;
            continue;
        }
        if found_method && cbyte != SP {
            pathname.push_str(ch);
        }
        if found_method && cbyte == SP {
            break;
        }

    }

    return RequestedRoute {method: method, pathname: pathname};
}

fn parse_normal_header (header: &str) -> ClientHeader {
    let header_name = b"Sec-WebSocket-Key";
    let mut lhs = String::new(); // lhs of ':' in header
    let mut rhs = String::new(); // rhs of ':' in header
    let mut found_colon = false;

    for ch in header.graphemes(true) {
        let cbyte = ch.as_bytes()[0];
        if !found_colon && cbyte == COLON {
            found_colon = true;
            continue;
        }

        if cbyte == CR || cbyte == LF {
            continue;
        }

        if !found_colon {
            lhs.push_str(ch);
        } else {
            rhs.push_str(ch);
        }
    }

    println!("key: {}, value: {}", lhs.trim(), rhs.trim());

    return ClientHeader{key:lhs, value:rhs};
}

fn get_normal_body (path_on_disk: &str) -> String {
    let path = Path::new(path_on_disk);
    let display = path.display();
    let mut file = match File::open(&path) {
        Ok(f) => f,
        Err(err) => panic!("file error: {}", err)
    };

    let content = match file.read_to_end() {
        Ok(c) => c,
        Err(err) => panic!("{}", err)
    };

    let s = String::from_utf8(content).unwrap();
    return s;
}

fn sec_handshake (from_server: &[u8]) -> String {

    let guid = b"258EAFA5-E914-47DA-95CA-C5AB0DC85B11";

    let mut sha = Sha1::new();

    sha.input(from_server);
    sha.input(guid);
    let mut out = [0u8, ..20];
    sha.result(out.as_mut_slice());

    println!("{} {}", sha.result_str(), out.to_base64(STANDARD));

    return out.to_base64(STANDARD);
}

fn ws_handshake (mut stream: BufferedStream<TcpStream>) {
    println!("running ws_handshake");
    stream.write(b"HTTP/1.1 101 Switching Protocols\r\n").unwrap();
    stream.write(b"Upgrade: websocket\r\n").unwrap();
    stream.write(b"Connection: Upgrade\r\n").unwrap();
    stream.write(
        b"Sec-WebSocket-Accept: s3pPLMBiTxaQ9kYGzzhZRbK+xOo=\r\n").unwrap();
    stream.flush().unwrap();
}

fn main () {

    let addr = "127.0.0.1:8099";

    let listener = TcpListener::bind(addr).unwrap();
    let mut acceptor = listener.listen().unwrap();
    println!("Listening on {}", addr);

    fn handle_client(mut stream: BufferedStream<TcpStream>) {
        let mut body: String = "
You sent it!
".to_string();

        let mut cur_line: String;
        let mut ii = 0u;
        let mut req: RequestedRoute = RequestedRoute {
            pathname: String::new(),
            method: String::new()
        };
        let mut header_sec_key = b"Sec-WebSocket-Key";

        let mut headers: Vec<ClientHeader> = Vec::new();
        loop {
            match stream.read_line() { // XXX: strange unwrap like thing
                Ok(line) => cur_line = line,
                Err(_) => break
            }

            println!("{}", cur_line);
            body = body + cur_line + "
";
            if cur_line.as_bytes() == CRLF {
                break;
            }

            if ii == 0u { // the Request-Line is always the first line
                req = parse_request_line(cur_line.as_slice());
                println!("{}", req.pathname);
                ii += 1;
            } else { // this should be  "normal" key-value type header
                headers.push(parse_normal_header(cur_line.as_slice()));
            }

        }

        println!("pathname {} method {}", req.pathname, req.method);

        if req.pathname.as_bytes() == b"/" {
            body = get_normal_body("../html/ws1.html");
        } else if req.pathname.as_bytes() == b"/ws1.js" {
            body = get_normal_body("../html/ws1.js");
        } else if req.pathname.as_bytes() == b"/ws" {
            println!("ws {}", get_header_by_name(header_sec_key, headers));
            ws_handshake(stream);
            return;
        }

        let body_length = format!("Content-length: {}", body.len());

        stream.write(b"HTTP/1.1 200 OK\r\n").unwrap();
        stream.write(b"Content-type: text/html\r\n").unwrap();
        stream.write(b"X-header: from bytes\r\n").unwrap();
        stream.write(body_length.into_bytes().as_slice()).unwrap();
        stream.write(b"\r\n\r\n").unwrap();
        stream.write(body.into_bytes().as_slice()).unwrap();
        stream.flush().unwrap();

        println!("Done handling acceptor");
    }

    for stream in acceptor.incoming() {
        match stream {
            Err(e) => { println!("connection failed: {}", e) }
            Ok(stream) => spawn(proc() {
                let bs = BufferedStream::new(stream);
                handle_client(bs)
            })
        }
    }

    drop(acceptor);
}
<|endoftext|>


<commit_before>extern crate crypto;
extern crate serialize;

use std::io::{TcpListener, TcpStream};
use std::io::{Acceptor, Listener};
use std::io::{File, BufferedStream};
use std::string::String;
use std::os;
use std::vec::Vec;
use std::fmt;

use crypto::sha1::Sha1;
use crypto::digest::Digest;
use serialize::base64::{ToBase64, STANDARD};

use std::io::Timer;
use std::io::timer;
use std::time::duration::Duration;

pub const CR: u8 = b'\r';
pub const LF: u8 = b'\n';
pub const SP: u8 = b' ';
pub const COLON: u8 = b':';
pub const CRLF: &'static [u8] = &[CR,LF];

struct RequestedRoute {
    method: String,
    pathname: String
}

struct ClientHeader {
    key: String,
    value: String
}

enum Payload {
    Text(String),
    Binary(Vec<u8>),
    Empty
}

#[deriving(FromPrimitive)]
#[deriving(Copy)]
enum Opcode {
    EmptyOp = 0x0,
    TextOp = 0x1,
    BinaryOp = 0x2,
    CloseOp = 0x8,
    PingOp = 0x9,
    PongOp = 0xA
}

struct Message {
    payload: Payload,
    opcode: Opcode,
    fin: u8
 }

impl fmt::Show for Payload {

     fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
         let msg = match self {
             &Payload::Text(ref s) => s.as_bytes(),
             &Payload::Binary(ref s) => s.as_slice(),
             &Payload::Empty => "".as_bytes(),
         };
         write!(f, "{}", msg)
     }
}

impl Message {

    fn from_payload (payload: Payload, fin: u8) -> Message {
        let opcode: Opcode = match payload {
            Payload::Text(_) => Opcode::TextOp,
            Payload::Binary(_) => Opcode::BinaryOp,
            Payload::Empty => Opcode::EmptyOp
        };
        return Message {payload: payload, opcode: opcode, fin: fin};
    }

    fn continue_from_payload (payload: Payload) -> Message {
        let opcode = Opcode::EmptyOp;
        let fin = 0x0;
        return Message {payload: payload, opcode: opcode, fin: fin};
    }

    fn send (&self,
            mut stream: &mut BufferedStream<TcpStream>) {
        let msg = match self.payload {
            Payload::Text(ref s) => s.as_bytes(),
            Payload::Binary(ref s) => s.as_slice(),
            Payload::Empty => "".as_bytes(),
        };

        let length = msg.len() as u64;

        println!("the whole length {}, length {}", msg.len(), length);

        println!("send fin: {}, length {}, msg: {}, opcode: {}",
                 self.fin, length, msg, self.opcode as u8);

        stream.write_u8(self.fin | self.opcode as u8).unwrap();

        if length < 126 {
            stream.write_u8(length as u8).unwrap();
        } else if (length < 65535){
            stream.write_u8(126u8).unwrap();
            stream.write_be_u16(length as u16).unwrap();
        } else {
            stream.write_u8(127u8).unwrap();
            stream.write_be_u64(length as u64).unwrap();
        }

        stream.write(msg).unwrap();
        stream.flush();
    }

    fn from_stream(mut stream: &mut BufferedStream<TcpStream>) -> Message {
        let cur_byte: u8 = stream.read_byte().unwrap();

        let fin = cur_byte & 0b1000_0000;
        let rsv = cur_byte & 0b0111_0000;
        let opc = cur_byte & 0b0000_1111;
        let msk = cur_byte & 0b0000_0001;

        let cur_byte: u8 = stream.read_byte().unwrap();
        let len_indicator = (cur_byte & 0b0111_1111) as uint;

        let len = match len_indicator {
            126 => stream.read_be_u16().unwrap() as u64,
            127 => stream.read_be_u64().unwrap() as u64,
            _ => len_indicator as u64
        };

        let mskkey = stream.read_exact(4).unwrap();

        let mut msg = Vec::new();
        for ii in range(0u64, len) {
            let cur_byte: u8 = stream.read_byte().unwrap();
            let ch = mskkey[ii as uint % 4] ^ cur_byte;
            msg.push(ch);
        }

        let opcode: Opcode = FromPrimitive::from_u8(opc).expect("unknown opcode");

        let payload = match opcode {
            Opcode::TextOp => Payload::Text(String::from_utf8(msg).unwrap()),
            Opcode::BinaryOp => Payload::Binary(msg),
            _ => panic!("unknown payload")
        };

        println!(
            "fin {}, rsv {}, msk {}, opcode {}, len {}, mskkey {}, payload {}",
            fin, rsv, msk, opc, len, mskkey, payload);

        return Message::from_payload(payload, fin);
    }
}

fn get_header_by_name (header: &[u8], headers: &Vec<ClientHeader>) -> String {
    let mut result = String::new();

    for h in headers.iter() {
        if h.key.as_bytes() == header {
            result = h.value.clone();
            break;
        }
    }

    return result;
}

fn parse_request_line (header: &str) -> RequestedRoute {
    let mut pathname = String::new();
    let mut method = String::new();
    let mut starting = true;
    let mut found_method = false;

    for ch in header.graphemes(true) {
        let cbyte = ch.as_bytes()[0];

        if starting && cbyte != SP {
            method.push_str(ch);
        }
        if starting && cbyte == SP {
            starting = false;
            found_method = true;
            continue;
        }
        if found_method && cbyte != SP {
            pathname.push_str(ch);
        }
        if found_method && cbyte == SP {
            break;
        }

    }

    return RequestedRoute {method: method, pathname: pathname};
}

fn parse_normal_header (header: &str) -> ClientHeader {
    let mut lhs = String::new(); // lhs of ':' in header
    let mut rhs = String::new(); // rhs of ':' in header
    let mut found_colon = false;

    for ch in header.graphemes(true) {
        let cbyte = ch.as_bytes()[0];
        if !found_colon && cbyte == COLON {
            found_colon = true;
            continue;
        }

        if cbyte == CR || cbyte == LF {
            continue;
        }

        if !found_colon {
            lhs.push_str(ch);
        } else {
            rhs.push_str(ch);
        }
    }

    return ClientHeader {
        key:lhs.trim().to_string(),
        value:rhs.trim().to_string()
    };
}

fn get_normal_body (path_on_disk: &str) -> String {
    let path = Path::new(path_on_disk);
    let display = path.display();
    let mut file = match File::open(&path) {
        Ok(f) => f,
        Err(err) => panic!("file error: {}", err)
    };

    let content = match file.read_to_end() {
        Ok(c) => c,
        Err(err) => panic!("{}", err)
    };

    let s = String::from_utf8(content).unwrap();
    return s;
}

fn sec_handshake (from_server: &[u8]) -> String {

    // from rfc6455 [page 6]
    let guid = b"258EAFA5-E914-47DA-95CA-C5AB0DC85B11";

    let mut sha = Sha1::new();

    sha.input(from_server);
    sha.input(guid);
    let mut out = [0u8, ..20];
    sha.result(out.as_mut_slice());

    return out.to_base64(STANDARD);
}

fn ws_handshake (mut stream: BufferedStream<TcpStream>,
                 headers: Vec<ClientHeader>) {

    let header_sec_key = b"Sec-WebSocket-Key";

    let from_server = get_header_by_name(header_sec_key, &headers);
    let accept = sec_handshake(from_server.as_bytes());

    let sec_header = format!("Sec-WebSocket-Accept: {}\r\n", accept);

    stream.write(b"HTTP/1.1 101 Switching Protocols\r\n").unwrap();
    stream.write(b"Upgrade: websocket\r\n").unwrap();
    stream.write(b"Connection: Upgrade\r\n").unwrap();
    stream.write(b"Sec-WebSocket-Version: 13\r\n").unwrap();
    stream.write(b"Sec-WebSocket-Protocol: protocolOne\r\n").unwrap();
    stream.write(sec_header.as_bytes()).unwrap();
    stream.write(b"\r\n");

    ws_listen(stream, headers);
}

fn ws_listen(mut stream: BufferedStream<TcpStream>,
             mut headers: Vec<ClientHeader>) {

    let payload = Payload::Text("text ".to_string());
    let msg = Message::from_payload(payload, 0b0000_0000);
    msg.send(&mut stream);

    let payload = Payload::Text("here".to_string());
    let msg = Message::continue_from_payload(payload);
    msg.send(&mut stream);

    let payload = Payload::Empty;
    let msg = Message::from_payload(payload, 0b1000_0000);
    msg.send(&mut stream);

    println!("done sending");

    Message::from_stream(&mut stream);

    let echo_msg = Message::from_stream(&mut stream);

    echo_msg.send(&mut stream);
    echo_msg.send(&mut stream);

    let mut timer = Timer::new().unwrap();
    let interval = Duration::milliseconds(5000);
    timer::sleep(interval);

    let echo_msg = Message::from_stream(&mut stream);
    echo_msg.send(&mut stream);
}

fn main () {

    for argument in os::args().iter() {
        println!("arg: {}", argument);
    }

    let addr = "127.0.0.1:8099";

    let listener = TcpListener::bind(addr).unwrap();
    let mut acceptor = listener.listen().unwrap();
    println!("Listening on {}", addr);

    fn handle_client(mut stream: BufferedStream<TcpStream>) {
        let mut body: String = "
You sent it!
".to_string();

        let mut cur_line: String;
        let mut ii = 0u;
        let mut req: RequestedRoute = RequestedRoute {
            pathname: String::new(),
            method: String::new()
        };

        let mut headers: Vec<ClientHeader> = Vec::new();
        loop {
            match stream.read_line() { // XXX: strange unwrap like thing
                Ok(line) => cur_line = line,
                Err(_) => break
            }

            //body = body + cur_line + "
";
            if cur_line.as_bytes() == CRLF {
                break;
            }

            if ii == 0u { // the Request-Line is always the first line
                req = parse_request_line(cur_line.as_slice());
                println!("{}", req.pathname);
                ii += 1;
            } else { // this should be  "normal" key-value type header
                headers.push(parse_normal_header(cur_line.as_slice()));
            }

        }

        println!("pathname {} method {}", req.pathname, req.method);

        if req.pathname.as_bytes() == b"/" {
            body = get_normal_body("./html/ws1.html");
        } else if req.pathname.as_bytes() == b"/ws1.js" {
            body = get_normal_body("./html/ws1.js");
        } else if req.pathname.as_bytes() == b"/ws" {
            ws_handshake(stream, headers);
            return;
        }

        let body_length = format!("Content-length: {}", body.len());

        stream.write(b"HTTP/1.1 200 OK\r\n").unwrap();
        stream.write(b"Content-type: text/html\r\n").unwrap();
        stream.write(b"X-header: from bytes\r\n").unwrap();
        stream.write(body_length.as_bytes()).unwrap();
        stream.write(b"\r\n\r\n").unwrap();
        stream.write(body.as_bytes()).unwrap();
        stream.flush().unwrap();

        println!("Done handling acceptor");
    }

    for stream in acceptor.incoming() {
        match stream {
            Err(e) => { println!("connection failed: {}", e) }
            Ok(stream) => spawn(move || {
                let bs = BufferedStream::new(stream);
                handle_client(bs)
            })
        }
    }

    drop(acceptor);
}
<commit_msg>remove un-needed perens<commit_after>extern crate crypto;
extern crate serialize;

use std::io::{TcpListener, TcpStream};
use std::io::{Acceptor, Listener};
use std::io::{File, BufferedStream};
use std::string::String;
use std::os;
use std::vec::Vec;
use std::fmt;

use crypto::sha1::Sha1;
use crypto::digest::Digest;
use serialize::base64::{ToBase64, STANDARD};

use std::io::Timer;
use std::io::timer;
use std::time::duration::Duration;

pub const CR: u8 = b'\r';
pub const LF: u8 = b'\n';
pub const SP: u8 = b' ';
pub const COLON: u8 = b':';
pub const CRLF: &'static [u8] = &[CR,LF];

struct RequestedRoute {
    method: String,
    pathname: String
}

struct ClientHeader {
    key: String,
    value: String
}

enum Payload {
    Text(String),
    Binary(Vec<u8>),
    Empty
}

#[deriving(FromPrimitive)]
#[deriving(Copy)]
enum Opcode {
    EmptyOp = 0x0,
    TextOp = 0x1,
    BinaryOp = 0x2,
    CloseOp = 0x8,
    PingOp = 0x9,
    PongOp = 0xA
}

struct Message {
    payload: Payload,
    opcode: Opcode,
    fin: u8
 }

impl fmt::Show for Payload {

     fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
         let msg = match self {
             &Payload::Text(ref s) => s.as_bytes(),
             &Payload::Binary(ref s) => s.as_slice(),
             &Payload::Empty => "".as_bytes(),
         };
         write!(f, "{}", msg)
     }
}

impl Message {

    fn from_payload (payload: Payload, fin: u8) -> Message {
        let opcode: Opcode = match payload {
            Payload::Text(_) => Opcode::TextOp,
            Payload::Binary(_) => Opcode::BinaryOp,
            Payload::Empty => Opcode::EmptyOp
        };
        return Message {payload: payload, opcode: opcode, fin: fin};
    }

    fn continue_from_payload (payload: Payload) -> Message {
        let opcode = Opcode::EmptyOp;
        let fin = 0x0;
        return Message {payload: payload, opcode: opcode, fin: fin};
    }

    fn send (&self,
            mut stream: &mut BufferedStream<TcpStream>) {
        let msg = match self.payload {
            Payload::Text(ref s) => s.as_bytes(),
            Payload::Binary(ref s) => s.as_slice(),
            Payload::Empty => "".as_bytes(),
        };

        let length = msg.len() as u64;

        println!("the whole length {}, length {}", msg.len(), length);

        println!("send fin: {}, length {}, msg: {}, opcode: {}",
                 self.fin, length, msg, self.opcode as u8);

        stream.write_u8(self.fin | self.opcode as u8).unwrap();

        if length < 126 {
            stream.write_u8(length as u8).unwrap();
        } else if length < 65535 {
            stream.write_u8(126u8).unwrap();
            stream.write_be_u16(length as u16).unwrap();
        } else {
            stream.write_u8(127u8).unwrap();
            stream.write_be_u64(length as u64).unwrap();
        }

        stream.write(msg).unwrap();
        stream.flush();
    }

    fn from_stream(mut stream: &mut BufferedStream<TcpStream>) -> Message {
        let cur_byte: u8 = stream.read_byte().unwrap();

        let fin = cur_byte & 0b1000_0000;
        let rsv = cur_byte & 0b0111_0000;
        let opc = cur_byte & 0b0000_1111;
        let msk = cur_byte & 0b0000_0001;

        let cur_byte: u8 = stream.read_byte().unwrap();
        let len_indicator = (cur_byte & 0b0111_1111) as uint;

        let len = match len_indicator {
            126 => stream.read_be_u16().unwrap() as u64,
            127 => stream.read_be_u64().unwrap() as u64,
            _ => len_indicator as u64
        };

        let mskkey = stream.read_exact(4).unwrap();

        let mut msg = Vec::new();
        for ii in range(0u64, len) {
            let cur_byte: u8 = stream.read_byte().unwrap();
            let ch = mskkey[ii as uint % 4] ^ cur_byte;
            msg.push(ch);
        }

        let opcode: Opcode = FromPrimitive::from_u8(opc).expect("unknown opcode");

        let payload = match opcode {
            Opcode::TextOp => Payload::Text(String::from_utf8(msg).unwrap()),
            Opcode::BinaryOp => Payload::Binary(msg),
            _ => panic!("unknown payload")
        };

        println!(
            "fin {}, rsv {}, msk {}, opcode {}, len {}, mskkey {}, payload {}",
            fin, rsv, msk, opc, len, mskkey, payload);

        return Message::from_payload(payload, fin);
    }
}

fn get_header_by_name (header: &[u8], headers: &Vec<ClientHeader>) -> String {
    let mut result = String::new();

    for h in headers.iter() {
        if h.key.as_bytes() == header {
            result = h.value.clone();
            break;
        }
    }

    return result;
}

fn parse_request_line (header: &str) -> RequestedRoute {
    let mut pathname = String::new();
    let mut method = String::new();
    let mut starting = true;
    let mut found_method = false;

    for ch in header.graphemes(true) {
        let cbyte = ch.as_bytes()[0];

        if starting && cbyte != SP {
            method.push_str(ch);
        }
        if starting && cbyte == SP {
            starting = false;
            found_method = true;
            continue;
        }
        if found_method && cbyte != SP {
            pathname.push_str(ch);
        }
        if found_method && cbyte == SP {
            break;
        }

    }

    return RequestedRoute {method: method, pathname: pathname};
}

fn parse_normal_header (header: &str) -> ClientHeader {
    let mut lhs = String::new(); // lhs of ':' in header
    let mut rhs = String::new(); // rhs of ':' in header
    let mut found_colon = false;

    for ch in header.graphemes(true) {
        let cbyte = ch.as_bytes()[0];
        if !found_colon && cbyte == COLON {
            found_colon = true;
            continue;
        }

        if cbyte == CR || cbyte == LF {
            continue;
        }

        if !found_colon {
            lhs.push_str(ch);
        } else {
            rhs.push_str(ch);
        }
    }

    return ClientHeader {
        key:lhs.trim().to_string(),
        value:rhs.trim().to_string()
    };
}

fn get_normal_body (path_on_disk: &str) -> String {
    let path = Path::new(path_on_disk);
    let display = path.display();
    let mut file = match File::open(&path) {
        Ok(f) => f,
        Err(err) => panic!("file error: {}", err)
    };

    let content = match file.read_to_end() {
        Ok(c) => c,
        Err(err) => panic!("{}", err)
    };

    let s = String::from_utf8(content).unwrap();
    return s;
}

fn sec_handshake (from_server: &[u8]) -> String {

    // from rfc6455 [page 6]
    let guid = b"258EAFA5-E914-47DA-95CA-C5AB0DC85B11";

    let mut sha = Sha1::new();

    sha.input(from_server);
    sha.input(guid);
    let mut out = [0u8, ..20];
    sha.result(out.as_mut_slice());

    return out.to_base64(STANDARD);
}

fn ws_handshake (mut stream: BufferedStream<TcpStream>,
                 headers: Vec<ClientHeader>) {

    let header_sec_key = b"Sec-WebSocket-Key";

    let from_server = get_header_by_name(header_sec_key, &headers);
    let accept = sec_handshake(from_server.as_bytes());

    let sec_header = format!("Sec-WebSocket-Accept: {}\r\n", accept);

    stream.write(b"HTTP/1.1 101 Switching Protocols\r\n").unwrap();
    stream.write(b"Upgrade: websocket\r\n").unwrap();
    stream.write(b"Connection: Upgrade\r\n").unwrap();
    stream.write(b"Sec-WebSocket-Version: 13\r\n").unwrap();
    stream.write(b"Sec-WebSocket-Protocol: protocolOne\r\n").unwrap();
    stream.write(sec_header.as_bytes()).unwrap();
    stream.write(b"\r\n");

    ws_listen(stream, headers);
}

fn ws_listen(mut stream: BufferedStream<TcpStream>,
             mut headers: Vec<ClientHeader>) {

    let payload = Payload::Text("text ".to_string());
    let msg = Message::from_payload(payload, 0b0000_0000);
    msg.send(&mut stream);

    let payload = Payload::Text("here".to_string());
    let msg = Message::continue_from_payload(payload);
    msg.send(&mut stream);

    let payload = Payload::Empty;
    let msg = Message::from_payload(payload, 0b1000_0000);
    msg.send(&mut stream);

    println!("done sending");

    Message::from_stream(&mut stream);

    let echo_msg = Message::from_stream(&mut stream);

    echo_msg.send(&mut stream);
    echo_msg.send(&mut stream);

    let mut timer = Timer::new().unwrap();
    let interval = Duration::milliseconds(5000);
    timer::sleep(interval);

    let echo_msg = Message::from_stream(&mut stream);
    echo_msg.send(&mut stream);
}

fn main () {

    for argument in os::args().iter() {
        println!("arg: {}", argument);
    }

    let addr = "127.0.0.1:8099";

    let listener = TcpListener::bind(addr).unwrap();
    let mut acceptor = listener.listen().unwrap();
    println!("Listening on {}", addr);

    fn handle_client(mut stream: BufferedStream<TcpStream>) {
        let mut body: String = "
You sent it!
".to_string();

        let mut cur_line: String;
        let mut ii = 0u;
        let mut req: RequestedRoute = RequestedRoute {
            pathname: String::new(),
            method: String::new()
        };

        let mut headers: Vec<ClientHeader> = Vec::new();
        loop {
            match stream.read_line() { // XXX: strange unwrap like thing
                Ok(line) => cur_line = line,
                Err(_) => break
            }

            //body = body + cur_line + "
";
            if cur_line.as_bytes() == CRLF {
                break;
            }

            if ii == 0u { // the Request-Line is always the first line
                req = parse_request_line(cur_line.as_slice());
                println!("{}", req.pathname);
                ii += 1;
            } else { // this should be  "normal" key-value type header
                headers.push(parse_normal_header(cur_line.as_slice()));
            }

        }

        println!("pathname {} method {}", req.pathname, req.method);

        if req.pathname.as_bytes() == b"/" {
            body = get_normal_body("./html/ws1.html");
        } else if req.pathname.as_bytes() == b"/ws1.js" {
            body = get_normal_body("./html/ws1.js");
        } else if req.pathname.as_bytes() == b"/ws" {
            ws_handshake(stream, headers);
            return;
        }

        let body_length = format!("Content-length: {}", body.len());

        stream.write(b"HTTP/1.1 200 OK\r\n").unwrap();
        stream.write(b"Content-type: text/html\r\n").unwrap();
        stream.write(b"X-header: from bytes\r\n").unwrap();
        stream.write(body_length.as_bytes()).unwrap();
        stream.write(b"\r\n\r\n").unwrap();
        stream.write(body.as_bytes()).unwrap();
        stream.flush().unwrap();

        println!("Done handling acceptor");
    }

    for stream in acceptor.incoming() {
        match stream {
            Err(e) => { println!("connection failed: {}", e) }
            Ok(stream) => spawn(move || {
                let bs = BufferedStream::new(stream);
                handle_client(bs)
            })
        }
    }

    drop(acceptor);
}
<|endoftext|>


<commit_before>extern crate "rust-crypto" as rust_crypto;
extern crate serialize;

use std::io::{TcpListener, TcpStream};
use std::io::{Acceptor, Listener};
use std::io::{File, BufferedStream};
use std::string::String;
use std::os;

use rust_crypto::sha1::Sha1;
use rust_crypto::digest::Digest;
use serialize::base64::{ToBase64, STANDARD};

pub const CR: u8 = b'\r';
pub const LF: u8 = b'\n';
pub const SP: u8 = b' ';
pub const COLON: u8 = b':';
pub const CRLF: &'static [u8] = &[CR,LF];

struct RequestedRoute {
    method: String,
    pathname: String
}

struct ClientHeader {
    key: String,
    value: String
}

enum Payload {
    Text(String),
    Empty
}

enum Opcode {
    TextOp = 0x1,
    BinaryOp = 0x2
}

struct Message {
    payload: Payload,
    opcode: Opcode
 }

impl Message {

    fn new (msg: &String) -> Message {

        let (payload, opcode) = match msg {
            String => (Payload::Text(msg.to_string()), Opcode::TextOp)
        };

        return Message{payload: payload, opcode: opcode};
    }

    fn send (self,
            mut stream: BufferedStream<TcpStream>,
            headers: Vec<ClientHeader>) {

        let msg = match self.payload {
            Payload::Text(ref s) => s.as_bytes(),
            Payload::Empty => unimplemented!()
        };

        let length = msg.len() as u8;

        println!("msg: {}, opcode: {}", msg, self.opcode as u8);

        stream.write_u8(0b1000_0000 | self.opcode as u8).unwrap();
        stream.write_u8(length).unwrap();
        stream.write(msg).unwrap();

    }
}

fn get_header_by_name (header: &[u8], headers: &Vec<ClientHeader>) -> String {
    let mut result = String::new();

    for h in headers.iter() {
        if h.key.as_bytes() == header {
            result = h.value.clone();
            break;
        }
    }

    return result;
}

fn parse_request_line (header: &str) -> RequestedRoute {
    let mut pathname = String::new();
    let mut method = String::new();
    let mut starting = true;
    let mut found_method = false;

    for ch in header.graphemes(true) {
        let cbyte = ch.as_bytes()[0];

        if starting && cbyte != SP {
            method.push_str(ch);
        }
        if starting && cbyte == SP {
            starting = false;
            found_method = true;
            continue;
        }
        if found_method && cbyte != SP {
            pathname.push_str(ch);
        }
        if found_method && cbyte == SP {
            break;
        }

    }

    return RequestedRoute {method: method, pathname: pathname};
}

fn parse_normal_header (header: &str) -> ClientHeader {
    let mut lhs = String::new(); // lhs of ':' in header
    let mut rhs = String::new(); // rhs of ':' in header
    let mut found_colon = false;

    for ch in header.graphemes(true) {
        let cbyte = ch.as_bytes()[0];
        if !found_colon && cbyte == COLON {
            found_colon = true;
            continue;
        }

        if cbyte == CR || cbyte == LF {
            continue;
        }

        if !found_colon {
            lhs.push_str(ch);
        } else {
            rhs.push_str(ch);
        }
    }

    return ClientHeader {
        key:lhs.trim().to_string(),
        value:rhs.trim().to_string()
    };
}

fn get_normal_body (path_on_disk: &str) -> String {
    let path = Path::new(path_on_disk);
    let display = path.display();
    let mut file = match File::open(&path) {
        Ok(f) => f,
        Err(err) => panic!("file error: {}", err)
    };

    let content = match file.read_to_end() {
        Ok(c) => c,
        Err(err) => panic!("{}", err)
    };

    let s = String::from_utf8(content).unwrap();
    return s;
}

fn sec_handshake (from_server: &[u8]) -> String {

    // from rfc6455 [page 6]
    let guid = b"258EAFA5-E914-47DA-95CA-C5AB0DC85B11";

    let mut sha = Sha1::new();

    sha.input(from_server);
    sha.input(guid);
    let mut out = [0u8, ..20];
    sha.result(out.as_mut_slice());

    return out.to_base64(STANDARD);
}

fn ws_handshake (mut stream: BufferedStream<TcpStream>,
                 headers: Vec<ClientHeader>) {

    let header_sec_key = b"Sec-WebSocket-Key";

    let from_server = get_header_by_name(header_sec_key, &headers);
    let accept = sec_handshake(from_server.as_bytes());

    let sec_header = format!("Sec-WebSocket-Accept: {}\r\n", accept);

    stream.write(b"HTTP/1.1 101 Switching Protocols\r\n").unwrap();
    stream.write(b"Upgrade: websocket\r\n").unwrap();
    stream.write(b"Connection: Upgrade\r\n").unwrap();
    stream.write(b"Sec-WebSocket-Version: 13\r\n").unwrap();
    stream.write(b"Sec-WebSocket-Protocol: protocolOne\r\n").unwrap();
    stream.write(sec_header.as_bytes()).unwrap();
    stream.write(b"\r\n");

    ws_listen(stream, headers);
}

fn ws_listen(mut stream: BufferedStream<TcpStream>,
             headers: Vec<ClientHeader>) {
    let string_msg = "this is a message".to_string();
    let msg = box Message::new(&string_msg);
    msg.send(stream, headers);
}

fn main () {

    for argument in os::args().iter() {
        println!("arg: {}", argument);
    }

    let addr = "127.0.0.1:8099";

    let listener = TcpListener::bind(addr).unwrap();
    let mut acceptor = listener.listen().unwrap();
    println!("Listening on {}", addr);

    fn handle_client(mut stream: BufferedStream<TcpStream>) {
        let mut body: String = "
You sent it!
".to_string();

        let mut cur_line: String;
        let mut ii = 0u;
        let mut req: RequestedRoute = RequestedRoute {
            pathname: String::new(),
            method: String::new()
        };

        let mut headers: Vec<ClientHeader> = Vec::new();
        loop {
            match stream.read_line() { // XXX: strange unwrap like thing
                Ok(line) => cur_line = line,
                Err(_) => break
            }

            body = body + cur_line + "
";
            if cur_line.as_bytes() == CRLF {
                break;
            }

            if ii == 0u { // the Request-Line is always the first line
                req = parse_request_line(cur_line.as_slice());
                println!("{}", req.pathname);
                ii += 1;
            } else { // this should be  "normal" key-value type header
                headers.push(parse_normal_header(cur_line.as_slice()));
            }

        }

        println!("pathname {} method {}", req.pathname, req.method);

        if req.pathname.as_bytes() == b"/" {
            body = get_normal_body("../html/ws1.html");
        } else if req.pathname.as_bytes() == b"/ws1.js" {
            body = get_normal_body("../html/ws1.js");
        } else if req.pathname.as_bytes() == b"/ws" {
            ws_handshake(stream, headers);
            return;
        }

        let body_length = format!("Content-length: {}", body.len());

        stream.write(b"HTTP/1.1 200 OK\r\n").unwrap();
        stream.write(b"Content-type: text/html\r\n").unwrap();
        stream.write(b"X-header: from bytes\r\n").unwrap();
        stream.write(body_length.as_bytes()).unwrap();
        stream.write(b"\r\n\r\n").unwrap();
        stream.write(body.as_bytes()).unwrap();
        stream.flush().unwrap();

        println!("Done handling acceptor");
    }

    for stream in acceptor.incoming() {
        match stream {
            Err(e) => { println!("connection failed: {}", e) }
            Ok(stream) => spawn(proc() {
                let bs = BufferedStream::new(stream);
                handle_client(bs)
            })
        }
    }

    drop(acceptor);
}
<commit_msg>foreshadow refactoring<commit_after>extern crate "rust-crypto" as rust_crypto;
extern crate serialize;

use std::io::{TcpListener, TcpStream};
use std::io::{Acceptor, Listener};
use std::io::{File, BufferedStream};
use std::string::String;
use std::os;

use rust_crypto::sha1::Sha1;
use rust_crypto::digest::Digest;
use serialize::base64::{ToBase64, STANDARD};

pub const CR: u8 = b'\r';
pub const LF: u8 = b'\n';
pub const SP: u8 = b' ';
pub const COLON: u8 = b':';
pub const CRLF: &'static [u8] = &[CR,LF];

struct RequestedRoute {
    method: String,
    pathname: String
}

struct ClientHeader {
    key: String,
    value: String
}

enum Payload {
    Text(String),
    Binary(Vec<u8>),
    Empty
}

enum Opcode {
    TextOp = 0x1,
    BinaryOp = 0x2
}

struct Message {
    payload: Payload,
    opcode: Opcode
 }

impl Message {

    fn new (msg: &String) -> Message {

        let (payload, opcode) = match msg {
            String => (Payload::Text(msg.to_string()), Opcode::TextOp),
        };

        return Message{payload: payload, opcode: opcode};
    }

    fn from_payload (msg: Payload) -> Message {
        return Message::new(&"silly".to_string());
    }

    fn send (self,
            mut stream: BufferedStream<TcpStream>,
            headers: Vec<ClientHeader>) {

        let msg = match self.payload {
            Payload::Text(ref s) => s.as_bytes(),
            Payload::Binary(ref s) => unimplemented!(),
            Payload::Empty => unimplemented!()
        };

        let length = msg.len() as u8;

        println!("msg: {}, opcode: {}", msg, self.opcode as u8);

        stream.write_u8(0b1000_0000 | self.opcode as u8).unwrap();
        stream.write_u8(length).unwrap();
        stream.write(msg).unwrap();

    }
}

fn get_header_by_name (header: &[u8], headers: &Vec<ClientHeader>) -> String {
    let mut result = String::new();

    for h in headers.iter() {
        if h.key.as_bytes() == header {
            result = h.value.clone();
            break;
        }
    }

    return result;
}

fn parse_request_line (header: &str) -> RequestedRoute {
    let mut pathname = String::new();
    let mut method = String::new();
    let mut starting = true;
    let mut found_method = false;

    for ch in header.graphemes(true) {
        let cbyte = ch.as_bytes()[0];

        if starting && cbyte != SP {
            method.push_str(ch);
        }
        if starting && cbyte == SP {
            starting = false;
            found_method = true;
            continue;
        }
        if found_method && cbyte != SP {
            pathname.push_str(ch);
        }
        if found_method && cbyte == SP {
            break;
        }

    }

    return RequestedRoute {method: method, pathname: pathname};
}

fn parse_normal_header (header: &str) -> ClientHeader {
    let mut lhs = String::new(); // lhs of ':' in header
    let mut rhs = String::new(); // rhs of ':' in header
    let mut found_colon = false;

    for ch in header.graphemes(true) {
        let cbyte = ch.as_bytes()[0];
        if !found_colon && cbyte == COLON {
            found_colon = true;
            continue;
        }

        if cbyte == CR || cbyte == LF {
            continue;
        }

        if !found_colon {
            lhs.push_str(ch);
        } else {
            rhs.push_str(ch);
        }
    }

    return ClientHeader {
        key:lhs.trim().to_string(),
        value:rhs.trim().to_string()
    };
}

fn get_normal_body (path_on_disk: &str) -> String {
    let path = Path::new(path_on_disk);
    let display = path.display();
    let mut file = match File::open(&path) {
        Ok(f) => f,
        Err(err) => panic!("file error: {}", err)
    };

    let content = match file.read_to_end() {
        Ok(c) => c,
        Err(err) => panic!("{}", err)
    };

    let s = String::from_utf8(content).unwrap();
    return s;
}

fn sec_handshake (from_server: &[u8]) -> String {

    // from rfc6455 [page 6]
    let guid = b"258EAFA5-E914-47DA-95CA-C5AB0DC85B11";

    let mut sha = Sha1::new();

    sha.input(from_server);
    sha.input(guid);
    let mut out = [0u8, ..20];
    sha.result(out.as_mut_slice());

    return out.to_base64(STANDARD);
}

fn ws_handshake (mut stream: BufferedStream<TcpStream>,
                 headers: Vec<ClientHeader>) {

    let header_sec_key = b"Sec-WebSocket-Key";

    let from_server = get_header_by_name(header_sec_key, &headers);
    let accept = sec_handshake(from_server.as_bytes());

    let sec_header = format!("Sec-WebSocket-Accept: {}\r\n", accept);

    stream.write(b"HTTP/1.1 101 Switching Protocols\r\n").unwrap();
    stream.write(b"Upgrade: websocket\r\n").unwrap();
    stream.write(b"Connection: Upgrade\r\n").unwrap();
    stream.write(b"Sec-WebSocket-Version: 13\r\n").unwrap();
    stream.write(b"Sec-WebSocket-Protocol: protocolOne\r\n").unwrap();
    stream.write(sec_header.as_bytes()).unwrap();
    stream.write(b"\r\n");

    ws_listen(stream, headers);
}

fn ws_listen(mut stream: BufferedStream<TcpStream>,
             headers: Vec<ClientHeader>) {
    let string_msg = "this is a message".to_string();
    let msg = box Message::new(&string_msg);
    msg.send(stream, headers);

    let mut bin = Vec::new();
    bin.push('a' as u8);
}

fn main () {

    for argument in os::args().iter() {
        println!("arg: {}", argument);
    }

    let addr = "127.0.0.1:8099";

    let listener = TcpListener::bind(addr).unwrap();
    let mut acceptor = listener.listen().unwrap();
    println!("Listening on {}", addr);

    fn handle_client(mut stream: BufferedStream<TcpStream>) {
        let mut body: String = "
You sent it!
".to_string();

        let mut cur_line: String;
        let mut ii = 0u;
        let mut req: RequestedRoute = RequestedRoute {
            pathname: String::new(),
            method: String::new()
        };

        let mut headers: Vec<ClientHeader> = Vec::new();
        loop {
            match stream.read_line() { // XXX: strange unwrap like thing
                Ok(line) => cur_line = line,
                Err(_) => break
            }

            body = body + cur_line + "
";
            if cur_line.as_bytes() == CRLF {
                break;
            }

            if ii == 0u { // the Request-Line is always the first line
                req = parse_request_line(cur_line.as_slice());
                println!("{}", req.pathname);
                ii += 1;
            } else { // this should be  "normal" key-value type header
                headers.push(parse_normal_header(cur_line.as_slice()));
            }

        }

        println!("pathname {} method {}", req.pathname, req.method);

        if req.pathname.as_bytes() == b"/" {
            body = get_normal_body("../html/ws1.html");
        } else if req.pathname.as_bytes() == b"/ws1.js" {
            body = get_normal_body("../html/ws1.js");
        } else if req.pathname.as_bytes() == b"/ws" {
            ws_handshake(stream, headers);
            return;
        }

        let body_length = format!("Content-length: {}", body.len());

        stream.write(b"HTTP/1.1 200 OK\r\n").unwrap();
        stream.write(b"Content-type: text/html\r\n").unwrap();
        stream.write(b"X-header: from bytes\r\n").unwrap();
        stream.write(body_length.as_bytes()).unwrap();
        stream.write(b"\r\n\r\n").unwrap();
        stream.write(body.as_bytes()).unwrap();
        stream.flush().unwrap();

        println!("Done handling acceptor");
    }

    for stream in acceptor.incoming() {
        match stream {
            Err(e) => { println!("connection failed: {}", e) }
            Ok(stream) => spawn(proc() {
                let bs = BufferedStream::new(stream);
                handle_client(bs)
            })
        }
    }

    drop(acceptor);
}
<|endoftext|>


<commit_before>// Copyright 2014 The Rust Project Developers. See the COPYRIGHT
// file at the top-level directory of this distribution and at
// http://rust-lang.org/COPYRIGHT.
//
// 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.

//! The ChaCha random number generator.

use core::prelude::*;
use core::num::Int;

use {Rng, SeedableRng, Rand};

const KEY_WORDS    : uint =  8; // 8 words for the 256-bit key
const STATE_WORDS  : uint = 16;
const CHACHA_ROUNDS: uint = 20; // Cryptographically secure from 8 upwards as of this writing

/// A random number generator that uses the ChaCha20 algorithm [1].
///
/// The ChaCha algorithm is widely accepted as suitable for
/// cryptographic purposes, but this implementation has not been
/// verified as such. Prefer a generator like `OsRng` that defers to
/// the operating system for cases that need high security.
///
/// [1]: D. J. Bernstein, [*ChaCha, a variant of
/// Salsa20*](http://cr.yp.to/chacha.html)

pub struct ChaChaRng {
    buffer:  [u32, ..STATE_WORDS], // Internal buffer of output
    state:   [u32, ..STATE_WORDS], // Initial state
    index:   uint,                 // Index into state
}

impl Copy for ChaChaRng {}

static EMPTY: ChaChaRng = ChaChaRng {
    buffer:  [0, ..STATE_WORDS],
    state:   [0, ..STATE_WORDS],
    index:   STATE_WORDS
};

macro_rules! quarter_round{
    ($a: expr, $b: expr, $c: expr, $d: expr) => {{
        $a += $b; $d ^= $a; $d = $d.rotate_left(16);
        $c += $d; $b ^= $c; $b = $b.rotate_left(12);
        $a += $b; $d ^= $a; $d = $d.rotate_left( 8);
        $c += $d; $b ^= $c; $b = $b.rotate_left( 7);
    }}
}

macro_rules! double_round{
    ($x: expr) => {{
        // Column round
        quarter_round!($x[ 0], $x[ 4], $x[ 8], $x[12]);
        quarter_round!($x[ 1], $x[ 5], $x[ 9], $x[13]);
        quarter_round!($x[ 2], $x[ 6], $x[10], $x[14]);
        quarter_round!($x[ 3], $x[ 7], $x[11], $x[15]);
        // Diagonal round
        quarter_round!($x[ 0], $x[ 5], $x[10], $x[15]);
        quarter_round!($x[ 1], $x[ 6], $x[11], $x[12]);
        quarter_round!($x[ 2], $x[ 7], $x[ 8], $x[13]);
        quarter_round!($x[ 3], $x[ 4], $x[ 9], $x[14]);
    }}
}

#[inline]
fn core(output: &mut [u32, ..STATE_WORDS], input: &[u32, ..STATE_WORDS]) {
    *output = *input;

    for _ in range(0, CHACHA_ROUNDS / 2) {
        double_round!(output);
    }

    for i in range(0, STATE_WORDS) {
        output[i] += input[i];
    }
}

impl ChaChaRng {

    /// Create an ChaCha random number generator using the default
    /// fixed key of 8 zero words.
    pub fn new_unseeded() -> ChaChaRng {
        let mut rng = EMPTY;
        rng.init(&[0, ..KEY_WORDS]);
        rng
    }

    /// Sets the internal 128-bit ChaCha counter to
    /// a user-provided value. This permits jumping
    /// arbitrarily ahead (or backwards) in the pseudorandom stream.
    ///
    /// Since the nonce words are used to extend the counter to 128 bits,
    /// users wishing to obtain the conventional ChaCha pseudorandom stream
    /// associated with a particular nonce can call this function with
    /// arguments `0, desired_nonce`.
    pub fn set_counter(&mut self, counter_low: u64, counter_high: u64) {
        self.state[12] = (counter_low >>  0) as u32;
        self.state[13] = (counter_low >> 32) as u32;
        self.state[14] = (counter_high >>  0) as u32;
        self.state[15] = (counter_high >> 32) as u32;
        self.index = STATE_WORDS; // force recomputation
    }

    /// Initializes `self.state` with the appropriate key and constants
    ///
    /// We deviate slightly from the ChaCha specification regarding
    /// the nonce, which is used to extend the counter to 128 bits.
    /// This is provably as strong as the original cipher, though,
    /// since any distinguishing attack on our variant also works
    /// against ChaCha with a chosen-nonce. See the XSalsa20 [1]
    /// security proof for a more involved example of this.
    ///
    /// The modified word layout is:
    /// ```ignore
    /// constant constant constant constant
    /// key      key      key      key
    /// key      key      key      key
    /// counter  counter  counter  counter
    /// ```
    /// [1]: Daniel J. Bernstein. [*Extending the Salsa20
    /// nonce.*](http://cr.yp.to/papers.html#xsalsa)
    fn init(&mut self, key: &[u32, ..KEY_WORDS]) {
        self.state[0] = 0x61707865;
        self.state[1] = 0x3320646E;
        self.state[2] = 0x79622D32;
        self.state[3] = 0x6B206574;

        for i in range(0, KEY_WORDS) {
            self.state[4+i] = key[i];
        }

        self.state[12] = 0;
        self.state[13] = 0;
        self.state[14] = 0;
        self.state[15] = 0;

        self.index = STATE_WORDS;
    }

    /// Refill the internal output buffer (`self.buffer`)
    fn update(&mut self) {
        core(&mut self.buffer, &self.state);
        self.index = 0;
        // update 128-bit counter
        self.state[12] += 1;
        if self.state[12] != 0 { return };
        self.state[13] += 1;
        if self.state[13] != 0 { return };
        self.state[14] += 1;
        if self.state[14] != 0 { return };
        self.state[15] += 1;
    }
}

impl Rng for ChaChaRng {
    #[inline]
    fn next_u32(&mut self) -> u32 {
        if self.index == STATE_WORDS {
            self.update();
        }

        let value = self.buffer[self.index % STATE_WORDS];
        self.index += 1;
        value
    }
}

impl<'a> SeedableRng<&'a [u32]> for ChaChaRng {

    fn reseed(&mut self, seed: &'a [u32]) {
        // reset state
        self.init(&[0u32, ..KEY_WORDS]);
        // set key in place
        let key = self.state.slice_mut(4, 4+KEY_WORDS);
        for (k, s) in key.iter_mut().zip(seed.iter()) {
            *k = *s;
        }
    }

    /// Create a ChaCha generator from a seed,
    /// obtained from a variable-length u32 array.
    /// Only up to 8 words are used; if less than 8
    /// words are used, the remaining are set to zero.
    fn from_seed(seed: &'a [u32]) -> ChaChaRng {
        let mut rng = EMPTY;
        rng.reseed(seed);
        rng
    }
}

impl Rand for ChaChaRng {
    fn rand<R: Rng>(other: &mut R) -> ChaChaRng {
        let mut key : [u32, ..KEY_WORDS] = [0, ..KEY_WORDS];
        for word in key.iter_mut() {
            *word = other.gen();
        }
        SeedableRng::from_seed(key.as_slice())
    }
}

#[cfg(test)]
mod test {
    use std::prelude::*;

    use core::iter::order;
    use {Rng, SeedableRng};
    use super::ChaChaRng;

    #[test]
    fn test_rng_rand_seeded() {
        let s = ::test::rng().gen_iter::<u32>().take(8).collect::<Vec<u32>>();
        let mut ra: ChaChaRng = SeedableRng::from_seed(s.as_slice());
        let mut rb: ChaChaRng = SeedableRng::from_seed(s.as_slice());
        assert!(order::equals(ra.gen_ascii_chars().take(100),
                              rb.gen_ascii_chars().take(100)));
    }

    #[test]
    fn test_rng_seeded() {
        let seed : &[_] = &[0,1,2,3,4,5,6,7];
        let mut ra: ChaChaRng = SeedableRng::from_seed(seed);
        let mut rb: ChaChaRng = SeedableRng::from_seed(seed);
        assert!(order::equals(ra.gen_ascii_chars().take(100),
                              rb.gen_ascii_chars().take(100)));
    }

    #[test]
    fn test_rng_reseed() {
        let s = ::test::rng().gen_iter::<u32>().take(8).collect::<Vec<u32>>();
        let mut r: ChaChaRng = SeedableRng::from_seed(s.as_slice());
        let string1: String = r.gen_ascii_chars().take(100).collect();

        r.reseed(s.as_slice());

        let string2: String = r.gen_ascii_chars().take(100).collect();
        assert_eq!(string1, string2);
    }

    #[test]
    fn test_rng_true_values() {
        // Test vectors 1 and 2 from
        // http://tools.ietf.org/html/draft-nir-cfrg-chacha20-poly1305-04
        let seed : &[_] = &[0u32, ..8];
        let mut ra: ChaChaRng = SeedableRng::from_seed(seed);

        let v = Vec::from_fn(16, |_| ra.next_u32());
        assert_eq!(v,
                   vec!(0xade0b876, 0x903df1a0, 0xe56a5d40, 0x28bd8653,
                        0xb819d2bd, 0x1aed8da0, 0xccef36a8, 0xc70d778b,
                        0x7c5941da, 0x8d485751, 0x3fe02477, 0x374ad8b8,
                        0xf4b8436a, 0x1ca11815, 0x69b687c3, 0x8665eeb2));

        let v = Vec::from_fn(16, |_| ra.next_u32());
        assert_eq!(v,
                   vec!(0xbee7079f, 0x7a385155, 0x7c97ba98, 0x0d082d73,
                        0xa0290fcb, 0x6965e348, 0x3e53c612, 0xed7aee32,
                        0x7621b729, 0x434ee69c, 0xb03371d5, 0xd539d874,
                        0x281fed31, 0x45fb0a51, 0x1f0ae1ac, 0x6f4d794b));

        let seed : &[_] = &[0,1,2,3,4,5,6,7];
        let mut ra: ChaChaRng = SeedableRng::from_seed(seed);

        // Store the 17*i-th 32-bit word,
        // i.e., the i-th word of the i-th 16-word block
        let mut v : Vec<u32> = Vec::new();
        for _ in range(0u, 16) {
            v.push(ra.next_u32());
            for _ in range(0u, 16) {
                ra.next_u32();
            }
        }

        assert_eq!(v,
                   vec!(0xf225c81a, 0x6ab1be57, 0x04d42951, 0x70858036,
                        0x49884684, 0x64efec72, 0x4be2d186, 0x3615b384,
                        0x11cfa18e, 0xd3c50049, 0x75c775f6, 0x434c6530,
                        0x2c5bad8f, 0x898881dc, 0x5f1c86d9, 0xc1f8e7f4));
    }
}

<commit_msg>rustc: Fix `make install`<commit_after>// Copyright 2014 The Rust Project Developers. See the COPYRIGHT
// file at the top-level directory of this distribution and at
// http://rust-lang.org/COPYRIGHT.
//
// 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.

//! The ChaCha random number generator.

use core::prelude::*;
use core::num::Int;

use {Rng, SeedableRng, Rand};

const KEY_WORDS    : uint =  8; // 8 words for the 256-bit key
const STATE_WORDS  : uint = 16;
const CHACHA_ROUNDS: uint = 20; // Cryptographically secure from 8 upwards as of this writing

/// A random number generator that uses the ChaCha20 algorithm [1].
///
/// The ChaCha algorithm is widely accepted as suitable for
/// cryptographic purposes, but this implementation has not been
/// verified as such. Prefer a generator like `OsRng` that defers to
/// the operating system for cases that need high security.
///
/// [1]: D. J. Bernstein, [*ChaCha, a variant of
/// Salsa20*](http://cr.yp.to/chacha.html)

pub struct ChaChaRng {
    buffer:  [u32, ..STATE_WORDS], // Internal buffer of output
    state:   [u32, ..STATE_WORDS], // Initial state
    index:   uint,                 // Index into state
}

impl Copy for ChaChaRng {}

static EMPTY: ChaChaRng = ChaChaRng {
    buffer:  [0, ..STATE_WORDS],
    state:   [0, ..STATE_WORDS],
    index:   STATE_WORDS
};

macro_rules! quarter_round{
    ($a: expr, $b: expr, $c: expr, $d: expr) => {{
        $a += $b; $d ^= $a; $d = $d.rotate_left(16);
        $c += $d; $b ^= $c; $b = $b.rotate_left(12);
        $a += $b; $d ^= $a; $d = $d.rotate_left( 8);
        $c += $d; $b ^= $c; $b = $b.rotate_left( 7);
    }}
}

macro_rules! double_round{
    ($x: expr) => {{
        // Column round
        quarter_round!($x[ 0], $x[ 4], $x[ 8], $x[12]);
        quarter_round!($x[ 1], $x[ 5], $x[ 9], $x[13]);
        quarter_round!($x[ 2], $x[ 6], $x[10], $x[14]);
        quarter_round!($x[ 3], $x[ 7], $x[11], $x[15]);
        // Diagonal round
        quarter_round!($x[ 0], $x[ 5], $x[10], $x[15]);
        quarter_round!($x[ 1], $x[ 6], $x[11], $x[12]);
        quarter_round!($x[ 2], $x[ 7], $x[ 8], $x[13]);
        quarter_round!($x[ 3], $x[ 4], $x[ 9], $x[14]);
    }}
}

#[inline]
fn core(output: &mut [u32, ..STATE_WORDS], input: &[u32, ..STATE_WORDS]) {
    *output = *input;

    for _ in range(0, CHACHA_ROUNDS / 2) {
        double_round!(output);
    }

    for i in range(0, STATE_WORDS) {
        output[i] += input[i];
    }
}

impl ChaChaRng {

    /// Create an ChaCha random number generator using the default
    /// fixed key of 8 zero words.
    pub fn new_unseeded() -> ChaChaRng {
        let mut rng = EMPTY;
        rng.init(&[0, ..KEY_WORDS]);
        rng
    }

    /// Sets the internal 128-bit ChaCha counter to
    /// a user-provided value. This permits jumping
    /// arbitrarily ahead (or backwards) in the pseudorandom stream.
    ///
    /// Since the nonce words are used to extend the counter to 128 bits,
    /// users wishing to obtain the conventional ChaCha pseudorandom stream
    /// associated with a particular nonce can call this function with
    /// arguments `0, desired_nonce`.
    pub fn set_counter(&mut self, counter_low: u64, counter_high: u64) {
        self.state[12] = (counter_low >>  0) as u32;
        self.state[13] = (counter_low >> 32) as u32;
        self.state[14] = (counter_high >>  0) as u32;
        self.state[15] = (counter_high >> 32) as u32;
        self.index = STATE_WORDS; // force recomputation
    }

    /// Initializes `self.state` with the appropriate key and constants
    ///
    /// We deviate slightly from the ChaCha specification regarding
    /// the nonce, which is used to extend the counter to 128 bits.
    /// This is provably as strong as the original cipher, though,
    /// since any distinguishing attack on our variant also works
    /// against ChaCha with a chosen-nonce. See the XSalsa20 [1]
    /// security proof for a more involved example of this.
    ///
    /// The modified word layout is:
    /// ```text
    /// constant constant constant constant
    /// key      key      key      key
    /// key      key      key      key
    /// counter  counter  counter  counter
    /// ```
    /// [1]: Daniel J. Bernstein. [*Extending the Salsa20
    /// nonce.*](http://cr.yp.to/papers.html#xsalsa)
    fn init(&mut self, key: &[u32, ..KEY_WORDS]) {
        self.state[0] = 0x61707865;
        self.state[1] = 0x3320646E;
        self.state[2] = 0x79622D32;
        self.state[3] = 0x6B206574;

        for i in range(0, KEY_WORDS) {
            self.state[4+i] = key[i];
        }

        self.state[12] = 0;
        self.state[13] = 0;
        self.state[14] = 0;
        self.state[15] = 0;

        self.index = STATE_WORDS;
    }

    /// Refill the internal output buffer (`self.buffer`)
    fn update(&mut self) {
        core(&mut self.buffer, &self.state);
        self.index = 0;
        // update 128-bit counter
        self.state[12] += 1;
        if self.state[12] != 0 { return };
        self.state[13] += 1;
        if self.state[13] != 0 { return };
        self.state[14] += 1;
        if self.state[14] != 0 { return };
        self.state[15] += 1;
    }
}

impl Rng for ChaChaRng {
    #[inline]
    fn next_u32(&mut self) -> u32 {
        if self.index == STATE_WORDS {
            self.update();
        }

        let value = self.buffer[self.index % STATE_WORDS];
        self.index += 1;
        value
    }
}

impl<'a> SeedableRng<&'a [u32]> for ChaChaRng {

    fn reseed(&mut self, seed: &'a [u32]) {
        // reset state
        self.init(&[0u32, ..KEY_WORDS]);
        // set key in place
        let key = self.state.slice_mut(4, 4+KEY_WORDS);
        for (k, s) in key.iter_mut().zip(seed.iter()) {
            *k = *s;
        }
    }

    /// Create a ChaCha generator from a seed,
    /// obtained from a variable-length u32 array.
    /// Only up to 8 words are used; if less than 8
    /// words are used, the remaining are set to zero.
    fn from_seed(seed: &'a [u32]) -> ChaChaRng {
        let mut rng = EMPTY;
        rng.reseed(seed);
        rng
    }
}

impl Rand for ChaChaRng {
    fn rand<R: Rng>(other: &mut R) -> ChaChaRng {
        let mut key : [u32, ..KEY_WORDS] = [0, ..KEY_WORDS];
        for word in key.iter_mut() {
            *word = other.gen();
        }
        SeedableRng::from_seed(key.as_slice())
    }
}

#[cfg(test)]
mod test {
    use std::prelude::*;

    use core::iter::order;
    use {Rng, SeedableRng};
    use super::ChaChaRng;

    #[test]
    fn test_rng_rand_seeded() {
        let s = ::test::rng().gen_iter::<u32>().take(8).collect::<Vec<u32>>();
        let mut ra: ChaChaRng = SeedableRng::from_seed(s.as_slice());
        let mut rb: ChaChaRng = SeedableRng::from_seed(s.as_slice());
        assert!(order::equals(ra.gen_ascii_chars().take(100),
                              rb.gen_ascii_chars().take(100)));
    }

    #[test]
    fn test_rng_seeded() {
        let seed : &[_] = &[0,1,2,3,4,5,6,7];
        let mut ra: ChaChaRng = SeedableRng::from_seed(seed);
        let mut rb: ChaChaRng = SeedableRng::from_seed(seed);
        assert!(order::equals(ra.gen_ascii_chars().take(100),
                              rb.gen_ascii_chars().take(100)));
    }

    #[test]
    fn test_rng_reseed() {
        let s = ::test::rng().gen_iter::<u32>().take(8).collect::<Vec<u32>>();
        let mut r: ChaChaRng = SeedableRng::from_seed(s.as_slice());
        let string1: String = r.gen_ascii_chars().take(100).collect();

        r.reseed(s.as_slice());

        let string2: String = r.gen_ascii_chars().take(100).collect();
        assert_eq!(string1, string2);
    }

    #[test]
    fn test_rng_true_values() {
        // Test vectors 1 and 2 from
        // http://tools.ietf.org/html/draft-nir-cfrg-chacha20-poly1305-04
        let seed : &[_] = &[0u32, ..8];
        let mut ra: ChaChaRng = SeedableRng::from_seed(seed);

        let v = Vec::from_fn(16, |_| ra.next_u32());
        assert_eq!(v,
                   vec!(0xade0b876, 0x903df1a0, 0xe56a5d40, 0x28bd8653,
                        0xb819d2bd, 0x1aed8da0, 0xccef36a8, 0xc70d778b,
                        0x7c5941da, 0x8d485751, 0x3fe02477, 0x374ad8b8,
                        0xf4b8436a, 0x1ca11815, 0x69b687c3, 0x8665eeb2));

        let v = Vec::from_fn(16, |_| ra.next_u32());
        assert_eq!(v,
                   vec!(0xbee7079f, 0x7a385155, 0x7c97ba98, 0x0d082d73,
                        0xa0290fcb, 0x6965e348, 0x3e53c612, 0xed7aee32,
                        0x7621b729, 0x434ee69c, 0xb03371d5, 0xd539d874,
                        0x281fed31, 0x45fb0a51, 0x1f0ae1ac, 0x6f4d794b));

        let seed : &[_] = &[0,1,2,3,4,5,6,7];
        let mut ra: ChaChaRng = SeedableRng::from_seed(seed);

        // Store the 17*i-th 32-bit word,
        // i.e., the i-th word of the i-th 16-word block
        let mut v : Vec<u32> = Vec::new();
        for _ in range(0u, 16) {
            v.push(ra.next_u32());
            for _ in range(0u, 16) {
                ra.next_u32();
            }
        }

        assert_eq!(v,
                   vec!(0xf225c81a, 0x6ab1be57, 0x04d42951, 0x70858036,
                        0x49884684, 0x64efec72, 0x4be2d186, 0x3615b384,
                        0x11cfa18e, 0xd3c50049, 0x75c775f6, 0x434c6530,
                        0x2c5bad8f, 0x898881dc, 0x5f1c86d9, 0xc1f8e7f4));
    }
}

<|endoftext|>


<commit_before>use std::cmp::Ordering;
use std::fmt::{self, Display};
use std::str;

use case::CaseWhenExpression;
use common::{Literal, SqlType};
use keywords::escape_if_keyword;

#[derive(Clone, Debug, Eq, Hash, PartialEq, Serialize, Deserialize)]
pub enum FunctionExpression {
    Avg(FunctionArgument, bool),
    Count(FunctionArgument, bool),
    CountStar,
    Sum(FunctionArgument, bool),
    Max(FunctionArgument),
    Min(FunctionArgument),
    GroupConcat(FunctionArgument, String),
}

impl Display for FunctionExpression {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        match *self {
            FunctionExpression::Avg(ref col, d) if d => write!(f, "avg(distinct {})", col),
            FunctionExpression::Count(ref col, d) if d => write!(f, "count(distinct {})", col),
            FunctionExpression::Sum(ref col, d) if d => write!(f, "sum(distinct {})", col),

            FunctionExpression::Avg(ref col, _) => write!(f, "avg({})", col),
            FunctionExpression::Count(ref col, _) => write!(f, "count({})", col),
            FunctionExpression::CountStar => write!(f, "count(*)"),
            FunctionExpression::Sum(ref col, _) => write!(f, "sum({})", col),
            FunctionExpression::Max(ref col) => write!(f, "max({})", col),
            FunctionExpression::Min(ref col) => write!(f, "min({})", col),
            FunctionExpression::GroupConcat(ref col, ref s) => {
                write!(f, "group_concat({}, {})", col, s)
            }
        }
    }
}

#[derive(Clone, Debug, Eq, Hash, PartialEq, Serialize, Deserialize)]
pub enum FunctionArgument {
    Column(Column),
    Conditional(CaseWhenExpression),
}

impl Display for FunctionArgument {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        match *self {
            FunctionArgument::Column(ref col) => write!(f, "{}", col)?,
            FunctionArgument::Conditional(ref e) => {
                write!(f, "{}", e)?;
            }
        }
        Ok(())
    }
}

#[derive(Clone, Debug, Eq, Hash, PartialEq, Serialize, Deserialize)]
pub struct Column {
    pub name: String,
    pub alias: Option<String>,
    pub table: Option<String>,
    pub function: Option<Box<FunctionExpression>>,
}

impl fmt::Display for Column {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        if let Some(ref table) = self.table {
            write!(
                f,
                "{}.{}",
                escape_if_keyword(table),
                escape_if_keyword(&self.name)
            )?;
        } else if let Some(ref function) = self.function {
            write!(f, "{}", *function)?;
        } else {
            write!(f, "{}", escape_if_keyword(&self.name))?;
        }
        if let Some(ref alias) = self.alias {
            write!(f, " AS {}", escape_if_keyword(alias))?;
        }
        Ok(())
    }
}

impl<'a> From<&'a str> for Column {
    fn from(c: &str) -> Column {
        match c.find(".") {
            None => Column {
                name: String::from(c),
                alias: None,
                table: None,
                function: None,
            },
            Some(i) => Column {
                name: String::from(&c[i + 1..]),
                alias: None,
                table: Some(String::from(&c[0..i])),
                function: None,
            },
        }
    }
}

impl Ord for Column {
    fn cmp(&self, other: &Column) -> Ordering {
        if self.table.is_some() && other.table.is_some() {
            match self.table.cmp(&other.table) {
                Ordering::Equal => self.name.cmp(&other.name),
                x => x,
            }
        } else {
            self.name.cmp(&other.name)
        }
    }
}

impl PartialOrd for Column {
    fn partial_cmp(&self, other: &Column) -> Option<Ordering> {
        if self.table.is_some() && other.table.is_some() {
            match self.table.cmp(&other.table) {
                Ordering::Equal => Some(self.name.cmp(&other.name)),
                x => Some(x),
            }
        } else if self.table.is_none() && other.table.is_none() {
            Some(self.name.cmp(&other.name))
        } else {
            None
        }
    }
}

#[derive(Clone, Debug, Eq, Hash, PartialEq, Serialize, Deserialize)]
pub enum ColumnConstraint {
    NotNull,
    CharacterSet(String),
    Collation(String),
    DefaultValue(Literal),
    AutoIncrement,
    PrimaryKey,
    Unique,
}

impl fmt::Display for ColumnConstraint {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        match *self {
            ColumnConstraint::NotNull => write!(f, "NOT NULL"),
            ColumnConstraint::CharacterSet(ref charset) => write!(f, "CHARACTER SET {}", charset),
            ColumnConstraint::Collation(ref collation) => write!(f, "COLLATE {}", collation),
            ColumnConstraint::DefaultValue(ref literal) => {
                write!(f, "DEFAULT {}", literal.to_string())
            }
            ColumnConstraint::AutoIncrement => write!(f, "AUTO_INCREMENT"),
            ColumnConstraint::PrimaryKey => write!(f, "PRIMARY KEY"),
            ColumnConstraint::Unique => write!(f, "UNIQUE"),
        }
    }
}

#[derive(Clone, Debug, Eq, Hash, PartialEq, Serialize, Deserialize)]
pub struct ColumnSpecification {
    pub column: Column,
    pub sql_type: SqlType,
    pub constraints: Vec<ColumnConstraint>,
    pub comment: Option<String>,
}

impl fmt::Display for ColumnSpecification {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(
            f,
            "{} {}",
            escape_if_keyword(&self.column.name),
            self.sql_type
        )?;
        for constraint in self.constraints.iter() {
            write!(f, " {}", constraint)?;
        }
        if let Some(ref comment) = self.comment {
            write!(f, " COMMENT '{}'", comment)?;
        }
        Ok(())
    }
}

impl ColumnSpecification {
    pub fn new(column: Column, sql_type: SqlType) -> ColumnSpecification {
        ColumnSpecification {
            column,
            sql_type,
            constraints: vec![],
            comment: None,
        }
    }

    pub fn with_constraints(
        column: Column,
        sql_type: SqlType,
        constraints: Vec<ColumnConstraint>,
    ) -> ColumnSpecification {
        ColumnSpecification {
            column,
            sql_type,
            constraints,
            comment: None,
        }
    }
}

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

    #[test]
    fn column_from_str() {
        let s = "table.col";
        let c = Column::from(s);

        assert_eq!(
            c,
            Column {
                name: String::from("col"),
                alias: None,
                table: Some(String::from("table")),
                function: None,
            }
        );
    }

    #[test]
    fn print_function_column() {
        let c1 = Column {
            name: "".into(), // must be present, but will be ignored
            alias: Some("foo".into()),
            table: None,
            function: Some(Box::new(FunctionExpression::CountStar)),
        };
        let c2 = Column {
            name: "".into(), // must be present, but will be ignored
            alias: None,
            table: None,
            function: Some(Box::new(FunctionExpression::CountStar)),
        };
        let c3 = Column {
            name: "".into(), // must be present, but will be ignored
            alias: None,
            table: None,
            function: Some(Box::new(FunctionExpression::Sum(
                FunctionArgument::Column(Column::from("mytab.foo")),
                false,
            ))),
        };

        assert_eq!(format!("{}", c1), "count(*) AS foo");
        assert_eq!(format!("{}", c2), "count(*)");
        assert_eq!(format!("{}", c3), "sum(mytab.foo)");
    }
}
<commit_msg>New generic function expression type<commit_after>use std::cmp::Ordering;
use std::fmt::{self, Display};
use std::str;

use case::CaseWhenExpression;
use common::{Literal, SqlType};
use keywords::escape_if_keyword;

#[derive(Clone, Debug, Eq, Hash, PartialEq, Serialize, Deserialize)]
pub enum FunctionExpression {
    Avg(FunctionArgument, bool),
    Count(FunctionArgument, bool),
    CountStar,
    Sum(FunctionArgument, bool),
    Max(FunctionArgument),
    Min(FunctionArgument),
    GroupConcat(FunctionArgument, String),
    Generic(String, FunctionArguments),
}

impl Display for FunctionExpression {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        match *self {
            FunctionExpression::Avg(ref col, d) if d => write!(f, "avg(distinct {})", col),
            FunctionExpression::Count(ref col, d) if d => write!(f, "count(distinct {})", col),
            FunctionExpression::Sum(ref col, d) if d => write!(f, "sum(distinct {})", col),

            FunctionExpression::Avg(ref col, _) => write!(f, "avg({})", col),
            FunctionExpression::Count(ref col, _) => write!(f, "count({})", col),
            FunctionExpression::CountStar => write!(f, "count(*)"),
            FunctionExpression::Sum(ref col, _) => write!(f, "sum({})", col),
            FunctionExpression::Max(ref col) => write!(f, "max({})", col),
            FunctionExpression::Min(ref col) => write!(f, "min({})", col),
            FunctionExpression::GroupConcat(ref col, ref s) => {
                write!(f, "group_concat({}, {})", col, s)
            },
            FunctionExpression::Generic(ref name, ref args) => {
                write!(f, "{}({})", name, args)
            }
        }
    }
}

#[derive(Clone, Debug, Eq, Hash, PartialEq, Serialize, Deserialize)]
pub struct FunctionArguments {
    arguments : Vec<FunctionArgument>
}

impl Display for FunctionArguments {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        let mut first = true;
        for argument in self.arguments.iter() {
            if !first {
                write!(f, ",")?;
            } else {
                first = false;
            }
            write!(f, "{}", argument)?;
        }
        Ok(())
    }
}

#[derive(Clone, Debug, Eq, Hash, PartialEq, Serialize, Deserialize)]
pub enum FunctionArgument {
    Column(Column),
    Conditional(CaseWhenExpression),
}

impl Display for FunctionArgument {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        match *self {
            FunctionArgument::Column(ref col) => write!(f, "{}", col)?,
            FunctionArgument::Conditional(ref e) => {
                write!(f, "{}", e)?;
            }
        }
        Ok(())
    }
}

#[derive(Clone, Debug, Eq, Hash, PartialEq, Serialize, Deserialize)]
pub struct Column {
    pub name: String,
    pub alias: Option<String>,
    pub table: Option<String>,
    pub function: Option<Box<FunctionExpression>>,
}

impl fmt::Display for Column {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        if let Some(ref table) = self.table {
            write!(
                f,
                "{}.{}",
                escape_if_keyword(table),
                escape_if_keyword(&self.name)
            )?;
        } else if let Some(ref function) = self.function {
            write!(f, "{}", *function)?;
        } else {
            write!(f, "{}", escape_if_keyword(&self.name))?;
        }
        if let Some(ref alias) = self.alias {
            write!(f, " AS {}", escape_if_keyword(alias))?;
        }
        Ok(())
    }
}

impl<'a> From<&'a str> for Column {
    fn from(c: &str) -> Column {
        match c.find(".") {
            None => Column {
                name: String::from(c),
                alias: None,
                table: None,
                function: None,
            },
            Some(i) => Column {
                name: String::from(&c[i + 1..]),
                alias: None,
                table: Some(String::from(&c[0..i])),
                function: None,
            },
        }
    }
}

impl Ord for Column {
    fn cmp(&self, other: &Column) -> Ordering {
        if self.table.is_some() && other.table.is_some() {
            match self.table.cmp(&other.table) {
                Ordering::Equal => self.name.cmp(&other.name),
                x => x,
            }
        } else {
            self.name.cmp(&other.name)
        }
    }
}

impl PartialOrd for Column {
    fn partial_cmp(&self, other: &Column) -> Option<Ordering> {
        if self.table.is_some() && other.table.is_some() {
            match self.table.cmp(&other.table) {
                Ordering::Equal => Some(self.name.cmp(&other.name)),
                x => Some(x),
            }
        } else if self.table.is_none() && other.table.is_none() {
            Some(self.name.cmp(&other.name))
        } else {
            None
        }
    }
}

#[derive(Clone, Debug, Eq, Hash, PartialEq, Serialize, Deserialize)]
pub enum ColumnConstraint {
    NotNull,
    CharacterSet(String),
    Collation(String),
    DefaultValue(Literal),
    AutoIncrement,
    PrimaryKey,
    Unique,
}

impl fmt::Display for ColumnConstraint {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        match *self {
            ColumnConstraint::NotNull => write!(f, "NOT NULL"),
            ColumnConstraint::CharacterSet(ref charset) => write!(f, "CHARACTER SET {}", charset),
            ColumnConstraint::Collation(ref collation) => write!(f, "COLLATE {}", collation),
            ColumnConstraint::DefaultValue(ref literal) => {
                write!(f, "DEFAULT {}", literal.to_string())
            }
            ColumnConstraint::AutoIncrement => write!(f, "AUTO_INCREMENT"),
            ColumnConstraint::PrimaryKey => write!(f, "PRIMARY KEY"),
            ColumnConstraint::Unique => write!(f, "UNIQUE"),
        }
    }
}

#[derive(Clone, Debug, Eq, Hash, PartialEq, Serialize, Deserialize)]
pub struct ColumnSpecification {
    pub column: Column,
    pub sql_type: SqlType,
    pub constraints: Vec<ColumnConstraint>,
    pub comment: Option<String>,
}

impl fmt::Display for ColumnSpecification {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(
            f,
            "{} {}",
            escape_if_keyword(&self.column.name),
            self.sql_type
        )?;
        for constraint in self.constraints.iter() {
            write!(f, " {}", constraint)?;
        }
        if let Some(ref comment) = self.comment {
            write!(f, " COMMENT '{}'", comment)?;
        }
        Ok(())
    }
}

impl ColumnSpecification {
    pub fn new(column: Column, sql_type: SqlType) -> ColumnSpecification {
        ColumnSpecification {
            column,
            sql_type,
            constraints: vec![],
            comment: None,
        }
    }

    pub fn with_constraints(
        column: Column,
        sql_type: SqlType,
        constraints: Vec<ColumnConstraint>,
    ) -> ColumnSpecification {
        ColumnSpecification {
            column,
            sql_type,
            constraints,
            comment: None,
        }
    }
}

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

    #[test]
    fn column_from_str() {
        let s = "table.col";
        let c = Column::from(s);

        assert_eq!(
            c,
            Column {
                name: String::from("col"),
                alias: None,
                table: Some(String::from("table")),
                function: None,
            }
        );
    }

    #[test]
    fn print_function_column() {
        let c1 = Column {
            name: "".into(), // must be present, but will be ignored
            alias: Some("foo".into()),
            table: None,
            function: Some(Box::new(FunctionExpression::CountStar)),
        };
        let c2 = Column {
            name: "".into(), // must be present, but will be ignored
            alias: None,
            table: None,
            function: Some(Box::new(FunctionExpression::CountStar)),
        };
        let c3 = Column {
            name: "".into(), // must be present, but will be ignored
            alias: None,
            table: None,
            function: Some(Box::new(FunctionExpression::Sum(
                FunctionArgument::Column(Column::from("mytab.foo")),
                false,
            ))),
        };

        assert_eq!(format!("{}", c1), "count(*) AS foo");
        assert_eq!(format!("{}", c2), "count(*)");
        assert_eq!(format!("{}", c3), "sum(mytab.foo)");
    }
}
<|endoftext|>


<commit_before>use std::str;

use crate::{Error, Result, Row, Rows, Statement};

/// Information about a column of a SQLite query.
#[derive(Debug)]
pub struct Column<'stmt> {
    name: &'stmt str,
    decl_type: Option<&'stmt str>,
}

impl Column<'_> {
    /// Returns the name of the column.
    pub fn name(&self) -> &str {
        self.name
    }

    /// Returns the type of the column (`None` for expression).
    pub fn decl_type(&self) -> Option<&str> {
        self.decl_type
    }
}

impl Statement<'_> {
    /// Get all the column names in the result set of the prepared statement.
    pub fn column_names(&self) -> Vec<&str> {
        let n = self.column_count();
        let mut cols = Vec::with_capacity(n as usize);
        for i in 0..n {
            let s = self.column_name(i);
            cols.push(s);
        }
        cols
    }

    /// Return the number of columns in the result set returned by the prepared
    /// statement.
    pub fn column_count(&self) -> usize {
        self.stmt.column_count()
    }

    pub(crate) fn column_name(&self, col: usize) -> &str {
        // Just panic if the bounds are wrong for now, we never call this
        // without checking first.
        let slice = self.stmt.column_name(col).expect("Column out of bounds");
        str::from_utf8(slice.to_bytes()).unwrap()
    }

    /// Returns the column index in the result set for a given column name.
    ///
    /// If there is no AS clause then the name of the column is unspecified and
    /// may change from one release of SQLite to the next.
    ///
    /// # Failure
    ///
    /// Will return an `Error::InvalidColumnName` when there is no column with
    /// the specified `name`.
    pub fn column_index(&self, name: &str) -> Result<usize> {
        let bytes = name.as_bytes();
        let n = self.column_count();
        for i in 0..n {
            // Note: `column_name` is only fallible if `i` is out of bounds,
            // which we've already checked.
            if bytes.eq_ignore_ascii_case(self.stmt.column_name(i).unwrap().to_bytes()) {
                return Ok(i);
            }
        }
        Err(Error::InvalidColumnName(String::from(name)))
    }

    /// Returns a slice describing the columns of the result of the query.
    pub fn columns<'stmt>(&'stmt self) -> Vec<Column<'stmt>> {
        let n = self.column_count();
        let mut cols = Vec::with_capacity(n as usize);
        for i in 0..n {
            let name = self.column_name(i);
            let slice = self.stmt.column_decltype(i);
            let decl_type = slice.map(|s| str::from_utf8(s.to_bytes()).unwrap());
            cols.push(Column { name, decl_type });
        }
        cols
    }
}

impl<'stmt> Rows<'stmt> {
    /// Get all the column names.
    pub fn column_names(&self) -> Option<Vec<&str>> {
        self.stmt.map(Statement::column_names)
    }

    /// Return the number of columns.
    pub fn column_count(&self) -> Option<usize> {
        self.stmt.map(Statement::column_count)
    }

    /// Returns a slice describing the columns of the Rows.
    pub fn columns(&self) -> Option<Vec<Column<'stmt>>> {
        self.stmt.map(Statement::columns)
    }
}

impl<'stmt> Row<'stmt> {
    /// Return the number of columns in the current row.
    pub fn column_count(&self) -> usize {
        self.stmt.column_count()
    }

    /// Returns a slice describing the columns of the Row.
    pub fn columns(&self) -> Vec<Column<'stmt>> {
        self.stmt.columns()
    }
}

#[cfg(test)]
mod test {
    use super::Column;
    use crate::Connection;

    #[test]
    fn test_columns() {
        let db = Connection::open_in_memory().unwrap();
        let query = db.prepare("SELECT * FROM sqlite_master").unwrap();
        let columns = query.columns();
        let column_names: Vec<&str> = columns.iter().map(Column::name).collect();
        assert_eq!(
            column_names.as_slice(),
            &["type", "name", "tbl_name", "rootpage", "sql"]
        );
        let column_types: Vec<Option<&str>> = columns.iter().map(Column::decl_type).collect();
        assert_eq!(
            &column_types[..3],
            &[Some("text"), Some("text"), Some("text"),]
        );
    }

    #[test]
    fn test_column_name_in_error() {
        use crate::{types::Type, Error};
        let db = Connection::open_in_memory().unwrap();
        db.execute_batch(
            "BEGIN;
             CREATE TABLE foo(x INTEGER, y TEXT);
             INSERT INTO foo VALUES(4, NULL);
             END;",
        )
        .unwrap();
        let mut stmt = db.prepare("SELECT x as renamed, y FROM foo").unwrap();
        let mut rows = stmt.query(crate::NO_PARAMS).unwrap();
        let row = rows.next().unwrap().unwrap();
        match row.get::<_, String>(0).unwrap_err() {
            Error::InvalidColumnType(idx, name, ty) => {
                assert_eq!(idx, 0);
                assert_eq!(name, "renamed");
                assert_eq!(ty, Type::Integer);
            }
            e => {
                panic!("Unexpected error type: {:?}", e);
            }
        }
        match row.get::<_, String>("y").unwrap_err() {
            Error::InvalidColumnType(idx, name, ty) => {
                assert_eq!(idx, 1);
                assert_eq!(name, "y");
                assert_eq!(ty, Type::Null);
            }
            e => {
                panic!("Unexpected error type: {:?}", e);
            }
        }
    }
}
<commit_msg>Use expect() instead of unwrap() for string conversion<commit_after>use std::str;

use crate::{Error, Result, Row, Rows, Statement};

/// Information about a column of a SQLite query.
#[derive(Debug)]
pub struct Column<'stmt> {
    name: &'stmt str,
    decl_type: Option<&'stmt str>,
}

impl Column<'_> {
    /// Returns the name of the column.
    pub fn name(&self) -> &str {
        self.name
    }

    /// Returns the type of the column (`None` for expression).
    pub fn decl_type(&self) -> Option<&str> {
        self.decl_type
    }
}

impl Statement<'_> {
    /// Get all the column names in the result set of the prepared statement.
    pub fn column_names(&self) -> Vec<&str> {
        let n = self.column_count();
        let mut cols = Vec::with_capacity(n as usize);
        for i in 0..n {
            let s = self.column_name(i);
            cols.push(s);
        }
        cols
    }

    /// Return the number of columns in the result set returned by the prepared
    /// statement.
    pub fn column_count(&self) -> usize {
        self.stmt.column_count()
    }

    pub(crate) fn column_name(&self, col: usize) -> &str {
        // Just panic if the bounds are wrong for now, we never call this
        // without checking first.
        let slice = self.stmt.column_name(col).expect("Column out of bounds");
        str::from_utf8(slice.to_bytes()).unwrap()
    }

    /// Returns the column index in the result set for a given column name.
    ///
    /// If there is no AS clause then the name of the column is unspecified and
    /// may change from one release of SQLite to the next.
    ///
    /// # Failure
    ///
    /// Will return an `Error::InvalidColumnName` when there is no column with
    /// the specified `name`.
    pub fn column_index(&self, name: &str) -> Result<usize> {
        let bytes = name.as_bytes();
        let n = self.column_count();
        for i in 0..n {
            // Note: `column_name` is only fallible if `i` is out of bounds,
            // which we've already checked.
            if bytes.eq_ignore_ascii_case(self.stmt.column_name(i).unwrap().to_bytes()) {
                return Ok(i);
            }
        }
        Err(Error::InvalidColumnName(String::from(name)))
    }

    /// Returns a slice describing the columns of the result of the query.
    pub fn columns<'stmt>(&'stmt self) -> Vec<Column<'stmt>> {
        let n = self.column_count();
        let mut cols = Vec::with_capacity(n as usize);
        for i in 0..n {
            let name = self.column_name(i);
            let slice = self.stmt.column_decltype(i);
            let decl_type = slice.map(|s| str::from_utf8(s.to_bytes()).expect("Invalid UTF-8 sequence in column declaration"));
            cols.push(Column { name, decl_type });
        }
        cols
    }
}

impl<'stmt> Rows<'stmt> {
    /// Get all the column names.
    pub fn column_names(&self) -> Option<Vec<&str>> {
        self.stmt.map(Statement::column_names)
    }

    /// Return the number of columns.
    pub fn column_count(&self) -> Option<usize> {
        self.stmt.map(Statement::column_count)
    }

    /// Returns a slice describing the columns of the Rows.
    pub fn columns(&self) -> Option<Vec<Column<'stmt>>> {
        self.stmt.map(Statement::columns)
    }
}

impl<'stmt> Row<'stmt> {
    /// Return the number of columns in the current row.
    pub fn column_count(&self) -> usize {
        self.stmt.column_count()
    }

    /// Returns a slice describing the columns of the Row.
    pub fn columns(&self) -> Vec<Column<'stmt>> {
        self.stmt.columns()
    }
}

#[cfg(test)]
mod test {
    use super::Column;
    use crate::Connection;

    #[test]
    fn test_columns() {
        let db = Connection::open_in_memory().unwrap();
        let query = db.prepare("SELECT * FROM sqlite_master").unwrap();
        let columns = query.columns();
        let column_names: Vec<&str> = columns.iter().map(Column::name).collect();
        assert_eq!(
            column_names.as_slice(),
            &["type", "name", "tbl_name", "rootpage", "sql"]
        );
        let column_types: Vec<Option<&str>> = columns.iter().map(Column::decl_type).collect();
        assert_eq!(
            &column_types[..3],
            &[Some("text"), Some("text"), Some("text"),]
        );
    }

    #[test]
    fn test_column_name_in_error() {
        use crate::{types::Type, Error};
        let db = Connection::open_in_memory().unwrap();
        db.execute_batch(
            "BEGIN;
             CREATE TABLE foo(x INTEGER, y TEXT);
             INSERT INTO foo VALUES(4, NULL);
             END;",
        )
        .unwrap();
        let mut stmt = db.prepare("SELECT x as renamed, y FROM foo").unwrap();
        let mut rows = stmt.query(crate::NO_PARAMS).unwrap();
        let row = rows.next().unwrap().unwrap();
        match row.get::<_, String>(0).unwrap_err() {
            Error::InvalidColumnType(idx, name, ty) => {
                assert_eq!(idx, 0);
                assert_eq!(name, "renamed");
                assert_eq!(ty, Type::Integer);
            }
            e => {
                panic!("Unexpected error type: {:?}", e);
            }
        }
        match row.get::<_, String>("y").unwrap_err() {
            Error::InvalidColumnType(idx, name, ty) => {
                assert_eq!(idx, 1);
                assert_eq!(name, "y");
                assert_eq!(ty, Type::Null);
            }
            e => {
                panic!("Unexpected error type: {:?}", e);
            }
        }
    }
}
<|endoftext|>