<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 std::vec::Vec;
use std::fmt;

use rust_crypto::sha1::Sha1;
use rust_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)]
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 u8;

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

        stream.write_u8(self.fin | self.opcode as u8).unwrap();
        stream.write_u8(length).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(proc() {
                let bs = BufferedStream::new(stream);
                handle_client(bs)
            })
        }
    }

    drop(acceptor);
}
<commit_msg>send big messages<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 std::vec::Vec;
use std::fmt;

use rust_crypto::sha1::Sha1;
use rust_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)]
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(proc() {
                let bs = BufferedStream::new(stream);
                handle_client(bs)
            })
        }
    }

    drop(acceptor);
}
<|endoftext|>


<commit_before>// Copyright 2017 CoreOS, Inc.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

//! azure metadata fetcher

mod crypto;

use self::crypto::x509;
use errors::*;
use metadata::Metadata;
use retry;
use util;
use openssh_keys::PublicKey;

use std::net::{IpAddr, SocketAddr};

header! {(MSAgentName, "x-ms-agent-name") => [String]}
header! {(MSVersion, "x-ms-version") => [String]}
header! {(MSCipherName, "x-ms-cipher-name") => [String]}
header! {(MSCert, "x-ms-guest-agent-public-x509-cert") => [String]}

const OPTION_245: &str = "OPTION_245";
const MS_AGENT_NAME: &str = "com.coreos.metadata";
const MS_VERSION: &str = "2012-11-30";
const SMIME_HEADER: &str = "\
MIME-Version:1.0
Content-Disposition: attachment; filename=/home/core/encrypted-ssh-cert.pem
Content-Type: application/x-pkcs7-mime; name=/home/core/encrypted-ssh-cert.pem
Content-Transfer-Encoding: base64

";

#[derive(Debug, Deserialize)]
struct GoalState {
    #[serde(rename = "Container")]
    pub container: Container
}

#[derive(Debug, Deserialize)]
struct Container {
    #[serde(rename = "RoleInstanceList")]
    pub role_instance_list: RoleInstanceList
}

#[derive(Debug, Deserialize)]
struct RoleInstanceList {
    #[serde(rename = "RoleInstance")]
    pub role_instances: Vec<RoleInstance>
}

#[derive(Debug, Deserialize)]
struct RoleInstance {
    #[serde(rename = "Configuration")]
    pub configuration: Configuration
}

#[derive(Debug, Deserialize)]
struct Configuration {
    #[serde(rename = "Certificates", default)]
    pub certificates: String,
    #[serde(rename = "SharedConfig", default)]
    pub shared_config: String,
}

#[derive(Debug, Deserialize)]
struct CertificatesFile {
    #[serde(rename = "Data", default)]
    pub data: String
}

#[derive(Debug, Deserialize)]
struct Versions {
    #[serde(rename = "Supported")]
    pub supported: Supported
}

#[derive(Debug, Deserialize)]
struct Supported {
    #[serde(rename = "Version", default)]
    pub versions: Vec<String>
}

#[derive(Debug, Deserialize)]
struct SharedConfig {
    #[serde(rename = "Incarnation")]
    pub incarnation: Incarnation,
    #[serde(rename = "Instances")]
    pub instances: Instances,
}

#[derive(Debug, Deserialize)]
struct Incarnation {
    pub instance: String,
}

#[derive(Debug, Deserialize)]
struct Instances {
    #[serde(rename = "Instance", default)]
    pub instances: Vec<Instance>,
}

#[derive(Debug, Deserialize)]
struct Instance {
    pub id: String,
    pub address: String,
    #[serde(rename = "InputEndpoints")]
    pub input_endpoints: InputEndpoints,
}

#[derive(Debug, Deserialize)]
struct InputEndpoints {
    #[serde(rename = "Endpoint", default)]
    pub endpoints: Vec<Endpoint>,
}

#[derive(Debug, Deserialize)]
struct Endpoint {
    #[serde(rename = "loadBalancedPublicAddress", default)]
    pub load_balanced_public_address: String,
}

#[derive(Debug, Copy, Clone, Default)]
struct Attributes {
    pub virtual_ipv4: Option<IpAddr>,
    pub dynamic_ipv4: Option<IpAddr>,
}

struct Azure {
    client: retry::Client,
    endpoint: IpAddr,
}

impl Azure {
    fn new() -> Result<Self> {
        let addr = Azure::get_fabric_address()
            .chain_err(|| "failed to get fabric address")?;
        let client = retry::Client::new()?
            .header(MSAgentName(MS_AGENT_NAME.to_owned()))
            .header(MSVersion(MS_VERSION.to_owned()));
        let azure = Azure {
            client: client,
            endpoint: addr,
        };
        azure.is_fabric_compatible(MS_VERSION)
            .chain_err(|| "failed version compatibility check")?;
        Ok(azure)
    }

    fn get_fabric_address() -> Result<IpAddr> {
        let v = util::dns_lease_key_lookup(OPTION_245)?;
        // value is an 8 digit hex value. convert it to u32 and
        // then parse that into an ip. Ipv4Addr::from(u32)
        // performs conversion from big-endian
        trace!("found fabric address in hex - {:?}", v);
        let dec = u32::from_str_radix(&v, 16)
            .chain_err(|| format!("failed to convert '{}' from hex", v))?;
        Ok(IpAddr::V4(dec.into()))
    }

    fn is_fabric_compatible(&self, version: &str) -> Result<()> {
        let versions: Versions = self.client.get(retry::Xml, format!("http://{}/?comp=versions", self.endpoint)).send()
            .chain_err(|| "failed to get versions")?
            .ok_or_else(|| "failed to get versions: not found")?;

        if versions.supported.versions.iter().any(|v| v == version) {
            Ok(())
        } else {
            Err(format!("fabric version {} not compatible with fabric address {}", MS_VERSION, self.endpoint).into())
        }
    }

    fn get_goal_state(&self) -> Result<GoalState> {
        self.client.get(retry::Xml, format!("http://{}/machine/?comp=goalstate", self.endpoint)).send()
            .chain_err(|| "failed to get goal state")?
            .ok_or_else(|| "failed to get goal state: not found response".into())
    }

    fn get_certs_endpoint(&self) -> Result<String> {
        let goalstate = self.get_goal_state()?;

        // grab the certificates endpoint from the xml and return it
        let cert_endpoint: &str = &goalstate.container.role_instance_list.role_instances[0].configuration.certificates;
        Ok(String::from(cert_endpoint))
    }

    fn get_certs(&self, mangled_pem: String) -> Result<String> {
        // get the certificates
        let endpoint = self.get_certs_endpoint()
            .chain_err(|| "failed to get certs endpoint")?;

        let certs: CertificatesFile = self.client.get(retry::Xml, endpoint)
            .header(MSCipherName("DES_EDE3_CBC".to_owned()))
            .header(MSCert(mangled_pem))
            .send()
            .chain_err(|| "failed to get certificates")?
            .ok_or_else(|| "failed to get certificates: not found")?;

        // the cms decryption expects it to have MIME information on the top
        // since cms is really for email attachments...don't tell the cops.
        let mut smime = String::from(SMIME_HEADER);
        smime.push_str(&certs.data);

        Ok(smime)
    }

    // put it all together
    fn get_ssh_pubkey(&self) -> Result<PublicKey> {
        // first we have to get the certificates endoint.
        // we have to generate the rsa public/private keypair and the x509 cert
        // that we use to make the request. this is equivalent to
        // `openssl req -x509 -nodes -subj /CN=LinuxTransport -days 365 -newkey rsa:2048 -keyout private.pem -out cert.pem`
        let (x509, pkey) = x509::generate_cert(&x509::Config::new(2048, 365))
            .chain_err(|| "failed to generate keys")?;

        // mangle the pem file for the request
        let mangled_pem = crypto::mangle_pem(&x509)
            .chain_err(|| "failed to mangle pem")?;

        // fetch the encrypted cms blob from the certs endpoint
        let smime = self.get_certs(mangled_pem)
            .chain_err(|| "failed to get certs")?;

        // decrypt the cms blob
        let p12 = crypto::decrypt_cms(smime.as_bytes(), &pkey, &x509)
            .chain_err(|| "failed to decrypt cms blob")?;

        // convert that to the OpenSSH public key format
        let ssh_pubkey = crypto::p12_to_ssh_pubkey(&p12)
            .chain_err(|| "failed to convert pkcs12 blob to ssh pubkey")?;

        Ok(ssh_pubkey)
    }

    fn get_attributes(&self) -> Result<Attributes> {
        let goalstate = self.get_goal_state()?;

        let endpoint = &goalstate.container.role_instance_list.role_instances[0].configuration.shared_config;

        let shared_config: SharedConfig = self.client.get(retry::Xml, endpoint.to_string()).send()
            .chain_err(|| "failed to get shared configuration")?
            .ok_or_else(|| "failed to get shared configuration: not found")?;

        let mut attributes = Attributes::default();

        for instance in shared_config.instances.instances {
            if instance.id == shared_config.incarnation.instance {
                attributes.dynamic_ipv4 = Some(instance.address.parse()
                    .chain_err(|| format!("failed to parse instance ip address: {}", instance.address))?);
                for endpoint in instance.input_endpoints.endpoints {
                    attributes.virtual_ipv4 = match endpoint.load_balanced_public_address.parse::<SocketAddr>() {
                        Ok(lbpa) => Some(lbpa.ip()),
                        Err(_) => continue,
                    };
                }
            }
        }

        Ok(attributes)
    }
}

pub fn fetch_metadata() -> Result<Metadata> {
    let provider = Azure::new()
        .chain_err(|| "azure: failed create metadata client")?;

    let ssh_pubkey = provider.get_ssh_pubkey()
        .chain_err(|| "azure: failed to get ssh pubkey")?;

    let attributes = provider.get_attributes()
        .chain_err(|| "azure: failed to get attributes")?;

    Ok(Metadata::builder()
       .add_publickeys(vec![ssh_pubkey])
       .add_attribute_if_exists("AZURE_IPV4_DYNAMIC".to_string(), attributes.dynamic_ipv4.map(|x| x.to_string()))
       .add_attribute_if_exists("AZURE_IPV4_VIRTUAL".to_string(), attributes.virtual_ipv4.map(|x| x.to_string()))
       .build())
}
<commit_msg>providers/azure: only fetch goalstate once<commit_after>// Copyright 2017 CoreOS, Inc.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

//! azure metadata fetcher

mod crypto;

use self::crypto::x509;
use errors::*;
use metadata::Metadata;
use retry;
use util;
use openssh_keys::PublicKey;

use std::net::{IpAddr, SocketAddr};

header! {(MSAgentName, "x-ms-agent-name") => [String]}
header! {(MSVersion, "x-ms-version") => [String]}
header! {(MSCipherName, "x-ms-cipher-name") => [String]}
header! {(MSCert, "x-ms-guest-agent-public-x509-cert") => [String]}

const OPTION_245: &str = "OPTION_245";
const MS_AGENT_NAME: &str = "com.coreos.metadata";
const MS_VERSION: &str = "2012-11-30";
const SMIME_HEADER: &str = "\
MIME-Version:1.0
Content-Disposition: attachment; filename=/home/core/encrypted-ssh-cert.pem
Content-Type: application/x-pkcs7-mime; name=/home/core/encrypted-ssh-cert.pem
Content-Transfer-Encoding: base64

";

#[derive(Debug, Deserialize, Clone, Default)]
struct GoalState {
    #[serde(rename = "Container")]
    pub container: Container
}

#[derive(Debug, Deserialize, Clone, Default)]
struct Container {
    #[serde(rename = "RoleInstanceList")]
    pub role_instance_list: RoleInstanceList
}

#[derive(Debug, Deserialize, Clone, Default)]
struct RoleInstanceList {
    #[serde(rename = "RoleInstance", default)]
    pub role_instances: Vec<RoleInstance>
}

#[derive(Debug, Deserialize, Clone)]
struct RoleInstance {
    #[serde(rename = "Configuration")]
    pub configuration: Configuration
}

#[derive(Debug, Deserialize, Clone)]
struct Configuration {
    #[serde(rename = "Certificates", default)]
    pub certificates: String,
    #[serde(rename = "SharedConfig", default)]
    pub shared_config: String,
}

#[derive(Debug, Deserialize, Clone)]
struct CertificatesFile {
    #[serde(rename = "Data", default)]
    pub data: String
}

#[derive(Debug, Deserialize, Clone)]
struct Versions {
    #[serde(rename = "Supported")]
    pub supported: Supported
}

#[derive(Debug, Deserialize, Clone)]
struct Supported {
    #[serde(rename = "Version", default)]
    pub versions: Vec<String>
}

#[derive(Debug, Deserialize, Clone)]
struct SharedConfig {
    #[serde(rename = "Incarnation")]
    pub incarnation: Incarnation,
    #[serde(rename = "Instances")]
    pub instances: Instances,
}

#[derive(Debug, Deserialize, Clone)]
struct Incarnation {
    pub instance: String,
}

#[derive(Debug, Deserialize, Clone)]
struct Instances {
    #[serde(rename = "Instance", default)]
    pub instances: Vec<Instance>,
}

#[derive(Debug, Deserialize, Clone)]
struct Instance {
    pub id: String,
    pub address: String,
    #[serde(rename = "InputEndpoints")]
    pub input_endpoints: InputEndpoints,
}

#[derive(Debug, Deserialize, Clone)]
struct InputEndpoints {
    #[serde(rename = "Endpoint", default)]
    pub endpoints: Vec<Endpoint>,
}

#[derive(Debug, Deserialize, Clone)]
struct Endpoint {
    #[serde(rename = "loadBalancedPublicAddress", default)]
    pub load_balanced_public_address: String,
}

#[derive(Debug, Copy, Clone, Default)]
struct Attributes {
    pub virtual_ipv4: Option<IpAddr>,
    pub dynamic_ipv4: Option<IpAddr>,
}

#[derive(Debug, Clone)]
struct Azure {
    client: retry::Client,
    endpoint: IpAddr,
    goal_state: GoalState,
}

impl Azure {
    fn new() -> Result<Self> {
        let addr = Azure::get_fabric_address()
            .chain_err(|| "failed to get fabric address")?;
        let client = retry::Client::new()?
            .header(MSAgentName(MS_AGENT_NAME.to_owned()))
            .header(MSVersion(MS_VERSION.to_owned()));

        let mut azure = Azure {
            client: client,
            endpoint: addr,
            goal_state: GoalState::default(),
        };

        // make sure the metadata service is compatible with our version
        azure.is_fabric_compatible(MS_VERSION)
            .chain_err(|| "failed version compatibility check")?;

        // populate goalstate
        azure.goal_state = azure.get_goal_state()?;
        Ok(azure)
    }

    fn get_goal_state(&self) -> Result<GoalState> {
        self.client.get(retry::Xml, format!("http://{}/machine/?comp=goalstate", self.endpoint)).send()
            .chain_err(|| "failed to get goal state")?
        .ok_or_else(|| "failed to get goal state: not found response".into())
    }

    fn get_fabric_address() -> Result<IpAddr> {
        let v = util::dns_lease_key_lookup(OPTION_245)?;
        // value is an 8 digit hex value. convert it to u32 and
        // then parse that into an ip. Ipv4Addr::from(u32)
        // performs conversion from big-endian
        trace!("found fabric address in hex - {:?}", v);
        let dec = u32::from_str_radix(&v, 16)
            .chain_err(|| format!("failed to convert '{}' from hex", v))?;
        Ok(IpAddr::V4(dec.into()))
    }

    fn is_fabric_compatible(&self, version: &str) -> Result<()> {
        let versions: Versions = self.client.get(retry::Xml, format!("http://{}/?comp=versions", self.endpoint)).send()
            .chain_err(|| "failed to get versions")?
            .ok_or_else(|| "failed to get versions: not found")?;

        if versions.supported.versions.iter().any(|v| v == version) {
            Ok(())
        } else {
            Err(format!("fabric version {} not compatible with fabric address {}", MS_VERSION, self.endpoint).into())
        }
    }

    fn get_certs_endpoint(&self) -> Result<String> {
        // grab the certificates endpoint from the xml and return it
        let cert_endpoint: &str = &self.goal_state.container.role_instance_list.role_instances[0].configuration.certificates;
        Ok(String::from(cert_endpoint))
    }

    fn get_certs(&self, mangled_pem: String) -> Result<String> {
        // get the certificates
        let endpoint = self.get_certs_endpoint()
            .chain_err(|| "failed to get certs endpoint")?;

        let certs: CertificatesFile = self.client.get(retry::Xml, endpoint)
            .header(MSCipherName("DES_EDE3_CBC".to_owned()))
            .header(MSCert(mangled_pem))
            .send()
            .chain_err(|| "failed to get certificates")?
            .ok_or_else(|| "failed to get certificates: not found")?;

        // the cms decryption expects it to have MIME information on the top
        // since cms is really for email attachments...don't tell the cops.
        let mut smime = String::from(SMIME_HEADER);
        smime.push_str(&certs.data);

        Ok(smime)
    }

    // put it all together
    fn get_ssh_pubkey(&self) -> Result<PublicKey> {
        // first we have to get the certificates endoint.
        // we have to generate the rsa public/private keypair and the x509 cert
        // that we use to make the request. this is equivalent to
        // `openssl req -x509 -nodes -subj /CN=LinuxTransport -days 365 -newkey rsa:2048 -keyout private.pem -out cert.pem`
        let (x509, pkey) = x509::generate_cert(&x509::Config::new(2048, 365))
            .chain_err(|| "failed to generate keys")?;

        // mangle the pem file for the request
        let mangled_pem = crypto::mangle_pem(&x509)
            .chain_err(|| "failed to mangle pem")?;

        // fetch the encrypted cms blob from the certs endpoint
        let smime = self.get_certs(mangled_pem)
            .chain_err(|| "failed to get certs")?;

        // decrypt the cms blob
        let p12 = crypto::decrypt_cms(smime.as_bytes(), &pkey, &x509)
            .chain_err(|| "failed to decrypt cms blob")?;

        // convert that to the OpenSSH public key format
        let ssh_pubkey = crypto::p12_to_ssh_pubkey(&p12)
            .chain_err(|| "failed to convert pkcs12 blob to ssh pubkey")?;

        Ok(ssh_pubkey)
    }

    fn get_attributes(&self) -> Result<Attributes> {
        let endpoint = &self.goal_state.container.role_instance_list.role_instances[0].configuration.shared_config;

        let shared_config: SharedConfig = self.client.get(retry::Xml, endpoint.to_string()).send()
            .chain_err(|| "failed to get shared configuration")?
            .ok_or_else(|| "failed to get shared configuration: not found")?;

        let mut attributes = Attributes::default();

        for instance in shared_config.instances.instances {
            if instance.id == shared_config.incarnation.instance {
                attributes.dynamic_ipv4 = Some(instance.address.parse()
                    .chain_err(|| format!("failed to parse instance ip address: {}", instance.address))?);
                for endpoint in instance.input_endpoints.endpoints {
                    attributes.virtual_ipv4 = match endpoint.load_balanced_public_address.parse::<SocketAddr>() {
                        Ok(lbpa) => Some(lbpa.ip()),
                        Err(_) => continue,
                    };
                }
            }
        }

        Ok(attributes)
    }
}

pub fn fetch_metadata() -> Result<Metadata> {
    let provider = Azure::new()
        .chain_err(|| "azure: failed create metadata client")?;

    let ssh_pubkey = provider.get_ssh_pubkey()
        .chain_err(|| "azure: failed to get ssh pubkey")?;

    let attributes = provider.get_attributes()
        .chain_err(|| "azure: failed to get attributes")?;

    Ok(Metadata::builder()
       .add_publickeys(vec![ssh_pubkey])
       .add_attribute_if_exists("AZURE_IPV4_DYNAMIC".to_string(), attributes.dynamic_ipv4.map(|x| x.to_string()))
       .add_attribute_if_exists("AZURE_IPV4_VIRTUAL".to_string(), attributes.virtual_ipv4.map(|x| x.to_string()))
       .build())
}
<|endoftext|>


<commit_before>use std::fmt;
use std::str::FromStr;
use std::convert::AsRef;
use std::collections::HashMap;
use command::{Command, Query};
use util;
use selectable::fragmentable::Fragmentable;
use selectable::fragmentable::{OrderedFragment, QueryFragment};
use types::data_type::DataType;
use types::column_flag_type::ColumnFlagType;
use self::WindowFunction::{RecordNumber, ExtWindowFunction};
use self::WindowableColumn::{Window, Value};

#[derive(Clone, PartialEq, PartialOrd, Eq, Ord, Debug, Hash)]
pub enum WindowableColumn<W, V> {
    Window(W),
    Value(V),
}

#[derive(Clone, PartialEq, PartialOrd, Eq, Ord, Debug, Hash)]
pub enum WindowFunction {
    RecordNumber,
    /// For future extensibility.
    ExtWindowFunction(String),
}

impl AsRef<str> for WindowFunction {
    fn as_ref(&self) -> &str {
        match *self {
            RecordNumber => "record_number()",
            ExtWindowFunction(ref s) => s.as_ref(),
        }
    }
}

#[derive (Debug)]
pub enum WindowFunctionError {
    Empty,
}

impl FromStr for WindowFunction {
    type Err = WindowFunctionError;
    fn from_str(s: &str) -> Result<WindowFunction, WindowFunctionError> {
        if s == "" {
            Err(WindowFunctionError::Empty)
        } else {
            Ok(match s {
                "record_number" | "RecordNumber" => RecordNumber,
                /// For future extensibility.
                _ => ExtWindowFunction(s.to_owned()),
            })
        }
    }
}

impl fmt::Display for WindowFunction {
    fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
        fmt.write_str(match *self {
            RecordNumber => "record_number()",
            ExtWindowFunction(ref s) => s.as_ref(),
        })
    }
}

#[derive(Clone, PartialEq, Eq, Debug)]
pub struct StagedColumns {
    label: String,
    stage: String,
    column_type: DataType,
    value: WindowableColumn<WindowFunction, String>,
    arguments: HashMap<String, String>,
}

impl StagedColumns {
    pub fn new(label: String,
               stage: String,
               column_type: DataType,
               value: WindowableColumn<WindowFunction, String>)
               -> StagedColumns {
        StagedColumns {
            label: label,
            stage: stage,
            column_type: column_type,
            value: value,
            arguments: HashMap::new(),
        }
    }

    pub fn flags(mut self, flags: Vec<ColumnFlagType>) -> StagedColumns {
        let flags_key = util::labeled_staged_key(self.label.to_owned(), "flags".to_string());
        self.arguments.insert(flags_key,
                              format!("\'{}\'", util::split_flags_vec(flags.clone())));
        self
    }

    pub fn sort_keys(mut self, sort_keys: Vec<String>) -> StagedColumns {
        let value = util::split_values_vec(sort_keys);
        let sort_keys_key = match self.value {
            Window(_) => util::labeled_window_key(self.label.to_owned(), "sort_keys".to_string()),
            Value(_) => util::labeled_staged_key(self.label.to_owned(), "sort_keys".to_string()),
        };
        self.arguments.insert(sort_keys_key, value);
        self
    }

    pub fn construct(mut self) -> StagedColumns {
        let stage_key = util::labeled_staged_key(self.label.to_owned(), "stage".to_string());
        let column_type_key = util::labeled_staged_key(self.label.to_owned(), "type".to_string());
        let value_key = util::labeled_staged_key(self.label.to_owned(), "value".to_string());
        let value_str = match self.value.clone() {
            Window(w) => format!("{}", w.as_ref()),
            Value(v) => v,
        };
        self.arguments.insert(stage_key, self.stage.clone());
        self.arguments.insert(column_type_key, format!("{}", self.column_type));
        self.arguments.insert(value_key, value_str.clone());
        self
    }

    pub fn build(self) -> Query {
        let this = self.construct();
        let mut query: Query = vec![];
        for (key, value) in &this.arguments {
            query.push((key.to_owned(), value.to_owned()));
        }
        query
    }
}

impl Fragmentable for StagedColumns {
    // Before calling this, it must call construct().
    fn to_fragment(self) -> (Command, OrderedFragment, QueryFragment) {
        // Command::Extension is `Command` type requirement. It should be ignored.
        (Command::Extension("staged_columns".to_string()), vec![], self.arguments.to_owned())
    }
}

#[cfg(test)]
mod test {
    use super::*;
    use std::str::FromStr;
    use std::collections::HashMap;
    use types::data_type::DataType;
    use types::column_flag_type::ColumnFlagType;
    use super::WindowableColumn::Value;

    #[test]
    fn test_as_str() {
        assert_eq!(WindowFunction::RecordNumber.as_ref(), "record_number()");
        assert_eq!(WindowFunction::ExtWindowFunction("ext_func()".to_owned()).as_ref(),
                   "ext_func()");
    }

    #[test]
    fn test_from_str() {
        assert_eq!(WindowFunction::RecordNumber, FromStr::from_str("record_number").unwrap());
        assert_eq!(WindowFunction::RecordNumber, FromStr::from_str("RecordNumber").unwrap());
        assert_eq!(WindowFunction::ExtWindowFunction("ext_func()".to_owned()),
                   FromStr::from_str("ext_func()").unwrap());
        let x: Result<WindowFunction, _> = FromStr::from_str("");
        if let Err(WindowFunctionError::Empty) = x {
        } else {
            panic!("An empty window function is invalid!")
        }
    }

    #[test]
    fn test_fmt() {
        assert_eq!("record_number()".to_owned(),
                   format!("{}", WindowFunction::RecordNumber));
        assert_eq!("ext_func()".to_owned(),
                   format!("{}",
                           WindowFunction::ExtWindowFunction("ext_func()".to_owned())));
    }

    #[test]
    fn test_new() {
        let label = "label1".to_string();
        let stage = "staged".to_string();
        let value = "'_id'".to_string();
        let staged_columns = StagedColumns::new(label.clone(),
                                                stage.clone(),
                                                DataType::UInt32,
                                                Value(value.clone()))
            .flags(vec![(ColumnFlagType::Scalar)]);
        let mut arg: HashMap<String, String> = HashMap::new();
        arg.insert(format!("columns[{}].flags", label.clone()),
                   "\'COLUMN_SCALAR\'".to_string());
        let expected = StagedColumns {
            label: label.to_owned(),
            stage: stage.clone(),
            column_type: DataType::UInt32,
            value: Value(value.clone()),
            arguments: arg,
        };
        assert_eq!(expected, staged_columns);
    }

    #[test]
    fn test_construct() {
        let label = "label1".to_string();
        let stage = "staged".to_string();
        let value = "'_id'".to_string();
        let staged_columns = StagedColumns::new(label.clone(),
                                                stage.clone(),
                                                DataType::UInt32,
                                                Value(value.clone()))
            .flags(vec![(ColumnFlagType::Scalar)])
            .construct();
        let mut arg: HashMap<String, String> = HashMap::new();
        arg.insert(format!("columns[{}].flags", label.clone()),
                   "\'COLUMN_SCALAR\'".to_string());
        assert_eq!(staged_columns.arguments
                       .get(&format!("columns[{}].stage", label.to_owned()))
                       .unwrap(),
                   &"staged");
        assert_eq!(staged_columns.arguments
                       .get(&format!("columns[{}].type", label.to_owned()))
                       .unwrap(),
                   &"UInt32");
        assert_eq!(staged_columns.arguments
                       .get(&format!("columns[{}].flags", label.to_owned()))
                       .unwrap(),
                   &"\'COLUMN_SCALAR\'");
        assert_eq!(staged_columns.arguments
                       .get(&format!("columns[{}].value", label.to_owned()))
                       .unwrap(),
                   &"\'_id\'");
    }
}
<commit_msg>rustfmt<commit_after>use std::fmt;
use std::str::FromStr;
use std::convert::AsRef;
use std::collections::HashMap;
use command::{Command, Query};
use util;
use selectable::fragmentable::Fragmentable;
use selectable::fragmentable::{OrderedFragment, QueryFragment};
use types::data_type::DataType;
use types::column_flag_type::ColumnFlagType;
use self::WindowFunction::{RecordNumber, ExtWindowFunction};
use self::WindowableColumn::{Window, Value};

#[derive(Clone, PartialEq, PartialOrd, Eq, Ord, Debug, Hash)]
pub enum WindowableColumn<W, V> {
    Window(W),
    Value(V),
}

#[derive(Clone, PartialEq, PartialOrd, Eq, Ord, Debug, Hash)]
pub enum WindowFunction {
    RecordNumber,
    /// For future extensibility.
    ExtWindowFunction(String),
}

impl AsRef<str> for WindowFunction {
    fn as_ref(&self) -> &str {
        match *self {
            RecordNumber => "record_number()",
            ExtWindowFunction(ref s) => s.as_ref(),
        }
    }
}

#[derive (Debug)]
pub enum WindowFunctionError {
    Empty,
}

impl FromStr for WindowFunction {
    type Err = WindowFunctionError;
    fn from_str(s: &str) -> Result<WindowFunction, WindowFunctionError> {
        if s == "" {
            Err(WindowFunctionError::Empty)
        } else {
            Ok(match s {
                "record_number" | "RecordNumber" => RecordNumber,
                /// For future extensibility.
                _ => ExtWindowFunction(s.to_owned()),
            })
        }
    }
}

impl fmt::Display for WindowFunction {
    fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
        fmt.write_str(match *self {
            RecordNumber => "record_number()",
            ExtWindowFunction(ref s) => s.as_ref(),
        })
    }
}

#[derive(Clone, PartialEq, Eq, Debug)]
pub struct StagedColumns {
    label: String,
    stage: String,
    column_type: DataType,
    value: WindowableColumn<WindowFunction, String>,
    arguments: HashMap<String, String>,
}

impl StagedColumns {
    pub fn new(label: String,
               stage: String,
               column_type: DataType,
               value: WindowableColumn<WindowFunction, String>)
               -> StagedColumns {
        StagedColumns {
            label: label,
            stage: stage,
            column_type: column_type,
            value: value,
            arguments: HashMap::new(),
        }
    }

    pub fn flags(mut self, flags: Vec<ColumnFlagType>) -> StagedColumns {
        let flags_key = util::labeled_staged_key(self.label.to_owned(), "flags".to_string());
        self.arguments.insert(flags_key,
                              format!("\'{}\'", util::split_flags_vec(flags.clone())));
        self
    }

    pub fn sort_keys(mut self, sort_keys: Vec<String>) -> StagedColumns {
        let value = util::split_values_vec(sort_keys);
        let sort_keys_key = match self.value {
            Window(_) => util::labeled_window_key(self.label.to_owned(), "sort_keys".to_string()),
            Value(_) => util::labeled_staged_key(self.label.to_owned(), "sort_keys".to_string()),
        };
        self.arguments.insert(sort_keys_key, value);
        self
    }

    pub fn construct(mut self) -> StagedColumns {
        let stage_key = util::labeled_staged_key(self.label.to_owned(), "stage".to_string());
        let column_type_key = util::labeled_staged_key(self.label.to_owned(), "type".to_string());
        let value_key = util::labeled_staged_key(self.label.to_owned(), "value".to_string());
        let value_str = match self.value.clone() {
            Window(w) => format!("{}", w.as_ref()),
            Value(v) => v,
        };
        self.arguments.insert(stage_key, self.stage.clone());
        self.arguments.insert(column_type_key, format!("{}", self.column_type));
        self.arguments.insert(value_key, value_str.clone());
        self
    }

    pub fn build(self) -> Query {
        let this = self.construct();
        let mut query: Query = vec![];
        for (key, value) in &this.arguments {
            query.push((key.to_owned(), value.to_owned()));
        }
        query
    }
}

impl Fragmentable for StagedColumns {
    // Before calling this, it must call construct().
    fn to_fragment(self) -> (Command, OrderedFragment, QueryFragment) {
        // Command::Extension is `Command` type requirement. It should be ignored.
        (Command::Extension("staged_columns".to_string()), vec![], self.arguments.to_owned())
    }
}

#[cfg(test)]
mod test {
    use super::*;
    use std::str::FromStr;
    use std::collections::HashMap;
    use types::data_type::DataType;
    use types::column_flag_type::ColumnFlagType;
    use super::WindowableColumn::Value;

    #[test]
    fn test_as_str() {
        assert_eq!(WindowFunction::RecordNumber.as_ref(), "record_number()");
        assert_eq!(WindowFunction::ExtWindowFunction("ext_func()".to_owned()).as_ref(),
                   "ext_func()");
    }

    #[test]
    fn test_from_str() {
        assert_eq!(WindowFunction::RecordNumber,
                   FromStr::from_str("record_number").unwrap());
        assert_eq!(WindowFunction::RecordNumber,
                   FromStr::from_str("RecordNumber").unwrap());
        assert_eq!(WindowFunction::ExtWindowFunction("ext_func()".to_owned()),
                   FromStr::from_str("ext_func()").unwrap());
        let x: Result<WindowFunction, _> = FromStr::from_str("");
        if let Err(WindowFunctionError::Empty) = x {
        } else {
            panic!("An empty window function is invalid!")
        }
    }

    #[test]
    fn test_fmt() {
        assert_eq!("record_number()".to_owned(),
                   format!("{}", WindowFunction::RecordNumber));
        assert_eq!("ext_func()".to_owned(),
                   format!("{}",
                           WindowFunction::ExtWindowFunction("ext_func()".to_owned())));
    }

    #[test]
    fn test_new() {
        let label = "label1".to_string();
        let stage = "staged".to_string();
        let value = "'_id'".to_string();
        let staged_columns = StagedColumns::new(label.clone(),
                                                stage.clone(),
                                                DataType::UInt32,
                                                Value(value.clone()))
            .flags(vec![(ColumnFlagType::Scalar)]);
        let mut arg: HashMap<String, String> = HashMap::new();
        arg.insert(format!("columns[{}].flags", label.clone()),
                   "\'COLUMN_SCALAR\'".to_string());
        let expected = StagedColumns {
            label: label.to_owned(),
            stage: stage.clone(),
            column_type: DataType::UInt32,
            value: Value(value.clone()),
            arguments: arg,
        };
        assert_eq!(expected, staged_columns);
    }

    #[test]
    fn test_construct() {
        let label = "label1".to_string();
        let stage = "staged".to_string();
        let value = "'_id'".to_string();
        let staged_columns = StagedColumns::new(label.clone(),
                                                stage.clone(),
                                                DataType::UInt32,
                                                Value(value.clone()))
            .flags(vec![(ColumnFlagType::Scalar)])
            .construct();
        let mut arg: HashMap<String, String> = HashMap::new();
        arg.insert(format!("columns[{}].flags", label.clone()),
                   "\'COLUMN_SCALAR\'".to_string());
        assert_eq!(staged_columns.arguments
                       .get(&format!("columns[{}].stage", label.to_owned()))
                       .unwrap(),
                   &"staged");
        assert_eq!(staged_columns.arguments
                       .get(&format!("columns[{}].type", label.to_owned()))
                       .unwrap(),
                   &"UInt32");
        assert_eq!(staged_columns.arguments
                       .get(&format!("columns[{}].flags", label.to_owned()))
                       .unwrap(),
                   &"\'COLUMN_SCALAR\'");
        assert_eq!(staged_columns.arguments
                       .get(&format!("columns[{}].value", label.to_owned()))
                       .unwrap(),
                   &"\'_id\'");
    }
}
<|endoftext|>


<commit_before>extern crate teleborg;

#[cfg(test)]
mod tests {
    use teleborg::{Dispatcher, Updater, Bot};
    use teleborg::objects::Update;

    fn test_updater() {
        let mut dispatcher = Dispatcher::new();
        dispatcher.add_command_handler("test", test, false);
        Updater::start(None, None, None, None, dispatcher);
    }

    fn test(bot: &Bot, update: Update, args: Option<Vec<&str>>) {
        bot.reply_to_message(&update, "It works!").unwrap();
    }
}
<commit_msg>removed blank line<commit_after>extern crate teleborg;

#[cfg(test)]
mod tests {
    use teleborg::{Dispatcher, Updater, Bot};
    use teleborg::objects::Update;

    fn test_updater() {
        let mut dispatcher = Dispatcher::new();
        dispatcher.add_command_handler("test", test, false);
        Updater::start(None, None, None, None, dispatcher);
    }

    fn test(bot: &Bot, update: Update, args: Option<Vec<&str>>) {
        bot.reply_to_message(&update, "It works!").unwrap();
    }
}
<|endoftext|>


<commit_before>/*******************************************************************************
 *
 * kit/kernel/terminal/vga.rs
 *
 * vim:ft=rust:ts=4:sw=4:et:tw=80
 *
 * Copyright (C) 2015-2021, Devyn Cairns
 * Redistribution of this file is permitted under the terms of the simplified
 * BSD license. See LICENSE for more information.
 *
 ******************************************************************************/

use core::fmt;
use core::ops::Range;

use alloc::collections::BTreeMap;

use crate::framebuffer::Framebuffer;
use crate::util::align_up;

use super::{Color, CharStream, Terminal};

/// Emulates a terminal on a graphical framebuffer
#[derive(Debug)]
pub struct Graphical<T, U> {
    fb:         T,
    font:       U,
    row:        usize,
    col:        usize,
    rows:       usize,
    cols:       usize,
    font_w:     usize,
    font_h:     usize,
    fg:         Color,
    bg:         Color,
    native_fg:  u32,
    native_bg:  u32,
    char_buf:   CharStream,
}

impl<T, U> Graphical<T, U> where T: Framebuffer, U: Font {
    pub fn new(fb: T, font: U) -> Graphical<T, U> {
        let mut g = Graphical {
            row: 0,
            col: 0,
            rows: 0,
            cols: 0,
            font_w: font.char_width() + 1,
            font_h: font.char_height(),
            fg: Color::White,
            bg: Color::Black,
            native_fg: 0,
            native_bg: 0,
            char_buf: CharStream::default(),
            fb,
            font,
        };

        g.rows = g.fb.height()/g.font_h;
        g.cols = g.fb.width()/g.font_w;

        g.update_colors();

        g
    }

    fn update_colors(&mut self) {
        self.native_fg = self.fb.color_format().format(self.fg.to_rgb());
        self.native_bg = self.fb.color_format().format(self.bg.to_rgb());
    }

    fn position_xy(&self, row: usize, col: usize) -> (usize, usize) {
        (
            self.font_w * col,
            self.font_h * row
        )
    }

    fn render(&self, ch: char, row: usize, col: usize, colors: (u32, u32)) {
        let (x, y) = self.position_xy(row, col);

        let fontdata = self.font.get(ch);
        let c_width = self.font.char_width();

        let get_bit =
            |px, py, _| {
                let fg = if let Some(bits) = fontdata {
                    if px >= c_width {
                        false
                    } else {
                        let index = py * c_width + px;
                        let byte = index >> 3;
                        let bit  = index & 0b111;
                        let mask = 1u8 << (7 - bit);

                        bits[byte] & mask != 0u8
                    }
                } else {
                    false
                };

                if fg { colors.0 } else { colors.1 }
            };

        self.fb.edit(x, y, self.font_w, self.font_h, get_bit);
    }

    fn update_cursor(&mut self) {
        // TODO
    }

    fn put_here(&self, ch: char) {
        self.render(ch, self.row, self.col, (self.native_fg, self.native_bg));
    }

    fn scroll(&self) {
        let (x0, y0) = self.position_xy(1, 0);
        let (x1, y1) = self.position_xy(0, 0);

        // Move the contents of everything but the first row up one row
        self.fb.copy_within(x0, y0, x1, y1,
            self.font_w * self.cols, self.font_h * (self.rows-1));

        // Clear the last row
        let (x2, y2) = self.position_xy(self.rows - 1, 0);
        self.fb.fill(x2, y2,
            self.font_w * self.cols, self.font_h,
            self.native_bg);
    }

    fn new_line(&mut self) {
        let (x, y) = self.position_xy(self.row, self.col);

        // Fill end of line
        let remaining = self.cols - self.col;
        self.fb.fill(x, y, self.font_w * remaining, self.font_h, self.native_bg);

        self.row += 1;
        self.col = 0;

        if self.row >= self.rows {
            self.row = self.rows - 1;
            self.scroll();
        }
    }

    fn write_char(&mut self, ch: char) {
        match ch {
            '\n' => {
                self.new_line();
            },

            '\x08' /* backspace */ => {
                if self.col > 0 {
                    self.col -= 1;
                }

                self.put_here(' ');
            },

            _ => {
                self.put_here(ch);
                self.col += 1;

                if self.col >= self.cols {
                    self.new_line();
                }
            }
        }
    }
}

impl<T, U> Terminal for Graphical<T, U> where T: Framebuffer, U: Font {
    fn reset(&mut self) -> fmt::Result {
        self.fg = Color::LightGrey;
        self.bg = Color::Black;
        self.update_colors();
        self.clear()
    }

    fn clear(&mut self) -> fmt::Result {
        self.row = 0;
        self.col = 0;

        self.fb.fill(0, 0, self.fb.width(), self.fb.height(), self.native_bg);

        Ok(())
    }

    fn get_cursor(&self) -> (usize, usize) {
        (self.row, self.col)
    }

    fn set_cursor(&mut self, row: usize, col: usize) -> fmt::Result {
        self.row = row;
        self.col = col;

        self.update_cursor();
        Ok(())
    }

    fn get_color(&self) -> (Color, Color) {
        (self.fg, self.bg)
    }

    fn set_color(&mut self, fg: Color, bg: Color) -> fmt::Result {
        self.fg = fg;
        self.bg = bg;
        self.update_colors();
        Ok(())
    }

    fn put_raw_byte(&mut self,
                    byte: u8,
                    fg:   Color,
                    bg:   Color,
                    row:  usize,
                    col:  usize) -> fmt::Result {
        // I think we should get rid of this method, it doesn't handle unicode properly
        self.render(
            char::from(byte),
            row,
            col,
            (self.fb.color_format().format(fg.to_rgb()),
             self.fb.color_format().format(bg.to_rgb())));
        Ok(())
    }

    fn write_raw_byte(&mut self, byte: u8) -> fmt::Result {
        // Build unicode
        self.char_buf.push(byte);

        while let Some(ch) = self.char_buf.consume() {
            self.write_char(ch);
        }

        Ok(())
    }

    fn set_double_buffer(&mut self, enabled: bool) {
        self.fb.set_double_buffer(enabled);
    }

    fn flush(&mut self) -> fmt::Result {
        self.update_cursor();
        Ok(())
    }
}

impl<T, U> fmt::Write for Graphical<T, U> where T: Framebuffer, U: Font {
    fn write_char(&mut self, ch: char) -> fmt::Result {
        let mut buf = [0; 4];
        self.write_raw_bytes(ch.encode_utf8(&mut buf).as_bytes())?;
        self.flush()?;
        Ok(())
    }

    fn write_str(&mut self, s: &str) -> fmt::Result {
        self.write_raw_bytes(s.as_bytes())?;
        self.flush()?;
        Ok(())
    }
}

static U_VGA16: &[u8] = include_bytes!("font/u_vga16.psf");

pub fn u_vga16() -> PsfFont<&'static [u8]> {
    PsfFont::new(U_VGA16).unwrap()
}

pub trait Font {
    fn char_width(&self) -> usize;
    fn char_height(&self) -> usize;
    fn char_bytes(&self) -> usize {
        align_up(self.char_width() * self.char_height(), 8) / 8
    }

    fn get(&self, ch: char) -> Option<&[u8]>;
}

#[derive(Debug, Clone)]
pub struct PsfFont<T: AsRef<[u8]>> {
    buf: T,
    chars_range: Range<usize>,
    unicode_table: BTreeMap<char, usize>,
    char_width: usize,
    char_height: usize,
}

const PSF1_MAGIC: u16 = 0x0436; /* little endian */
const PSF1_MODE_512: u8 = 0x01;
const PSF1_MODE_HASTAB: u8 = 0x02;
const PSF1_MODE_HASSEQ: u8 = 0x04;

const PSF1_SEPARATOR: u32 = 0xffff;
const PSF1_STARTSEQ: u32 = 0xfffe;

#[derive(Debug, Clone)]
pub enum PsfError {
    /// this is not a recognized PC Screen Font file
    WrongMagic,
    /// the mode field ({:x}) has unrecognized bits
    BadMode(u8),
    /// file is smaller than it should be
    FileTooSmall,
}

fn read_u16(slice: &[u8]) -> u16 {
    let mut bytes: [u8; 2] = [0; 2];
    bytes.copy_from_slice(&slice[0..2]);

    u16::from_le_bytes(bytes)
}

impl<T> PsfFont<T> where T: AsRef<[u8]> {
    pub fn new(buf: T) -> Result<PsfFont<T>, PsfError> {
        let file: &[u8] = buf.as_ref();

        if read_u16(&file[0..2]) != PSF1_MAGIC {
            return Err(PsfError::WrongMagic);
        }

        let mode: u8 = file[2];

        if mode & 0x07 != mode {
            return Err(PsfError::BadMode(mode));
        }

        let mode_512    = mode & PSF1_MODE_512 != 0;
        let mode_hastab = mode & PSF1_MODE_HASTAB != 0;
        let mode_hasseq = mode & PSF1_MODE_HASSEQ != 0;

        let charsize: usize = file[3] as usize; // height, also number of bytes

        let char_width = 8;
        let char_height = charsize;

        let length = if mode_512 { 512 } else { 256 };

        let chars_range = 4 .. (4 + charsize * length);

        // File too small if chars_range doesn't fit
        if chars_range.end > file.len() {
            return Err(PsfError::FileTooSmall);
        }

        let mut unicode_table = BTreeMap::new();

        // Parse the unicode table, or generate it if it's not there
        if mode_hastab {
            let unicode_table_buf = &file[chars_range.end ..];

            let mut pos = 0;

            'eof: for index in 0..length {
                loop {
                    if pos >= file.len() {
                        break 'eof;
                    }

                    let codepoint = read_u16(&unicode_table_buf[pos..]) as u32;

                    pos += 2;

                    if codepoint == PSF1_SEPARATOR {
                        // End of sequences for this index
                        break;
                    } else if codepoint == PSF1_STARTSEQ && mode_hasseq {
                        // Sequences not supported yet.
                        pos += 4;
                    } else if let Some(c) = char::from_u32(codepoint) {
                        // Insert codepoint.
                        unicode_table.insert(c, index);
                    }
                }
            }
        } else {
            for index in 0..length {
                if let Some(c) = char::from_u32(index as u32) {
                    unicode_table.insert(c, index);
                }
            }
        }

        Ok(PsfFont {
            buf,
            chars_range,
            unicode_table,
            char_width,
            char_height
        })
    }
}

impl<T> Font for PsfFont<T> where T: AsRef<[u8]> {
    fn char_width(&self) -> usize {
        self.char_width
    }

    fn char_height(&self) -> usize {
        self.char_height
    }

    fn get(&self, ch: char) -> Option<&[u8]> {
        self.unicode_table.get(&ch).map(|index| {
            let start = self.char_height * index;
            &self.buf.as_ref()[self.chars_range.clone()][start .. (start + self.char_height)]
        })
    }
}
<commit_msg>partial implementation of terminal cursor<commit_after>/*******************************************************************************
 *
 * kit/kernel/terminal/vga.rs
 *
 * vim:ft=rust:ts=4:sw=4:et:tw=80
 *
 * Copyright (C) 2015-2021, Devyn Cairns
 * Redistribution of this file is permitted under the terms of the simplified
 * BSD license. See LICENSE for more information.
 *
 ******************************************************************************/

use core::fmt;
use core::ops::Range;

use alloc::collections::BTreeMap;

use crate::framebuffer::Framebuffer;
use crate::util::align_up;

use super::{Color, CharStream, Terminal};

/// Emulates a terminal on a graphical framebuffer
#[derive(Debug)]
pub struct Graphical<T, U> {
    fb:         T,
    font:       U,
    row:        usize,
    col:        usize,
    rows:       usize,
    cols:       usize,
    font_w:     usize,
    font_h:     usize,
    fg:         Color,
    bg:         Color,
    native_fg:  u32,
    native_bg:  u32,
    char_buf:   CharStream,
}

impl<T, U> Graphical<T, U> where T: Framebuffer, U: Font {
    pub fn new(fb: T, font: U) -> Graphical<T, U> {
        let mut g = Graphical {
            row: 0,
            col: 0,
            rows: 0,
            cols: 0,
            font_w: font.char_width() + 1,
            font_h: font.char_height(),
            fg: Color::White,
            bg: Color::Black,
            native_fg: 0,
            native_bg: 0,
            char_buf: CharStream::default(),
            fb,
            font,
        };

        g.rows = g.fb.height()/g.font_h;
        g.cols = g.fb.width()/g.font_w;

        g.update_colors();

        g
    }

    fn update_colors(&mut self) {
        self.native_fg = self.fb.color_format().format(self.fg.to_rgb());
        self.native_bg = self.fb.color_format().format(self.bg.to_rgb());
    }

    fn position_xy(&self, row: usize, col: usize) -> (usize, usize) {
        (
            self.font_w * col,
            self.font_h * row
        )
    }

    fn render(&self, ch: char, row: usize, col: usize, colors: (u32, u32)) {
        let (x, y) = self.position_xy(row, col);

        let fontdata = self.font.get(ch);
        let c_width = self.font.char_width();

        let get_bit =
            |px, py, _| {
                let fg = if let Some(bits) = fontdata {
                    if px >= c_width {
                        false
                    } else {
                        let index = py * c_width + px;
                        let byte = index >> 3;
                        let bit  = index & 0b111;
                        let mask = 1u8 << (7 - bit);

                        bits[byte] & mask != 0u8
                    }
                } else {
                    false
                };

                if fg { colors.0 } else { colors.1 }
            };

        self.fb.edit(x, y, self.font_w, self.font_h, get_bit);
    }

    fn update_cursor(&mut self) {
        // TODO: make this non-destructive so we can put back whatever we
        // replaced. Should merely invert the character really.
        let (x, y) = self.position_xy(self.row, self.col);

        self.fb.fill(x, y, self.font_w, self.font_h, self.native_fg);
    }

    fn put_here(&self, ch: char) {
        self.render(ch, self.row, self.col, (self.native_fg, self.native_bg));
    }

    fn scroll(&self) {
        let (x0, y0) = self.position_xy(1, 0);
        let (x1, y1) = self.position_xy(0, 0);

        // Move the contents of everything but the first row up one row
        self.fb.copy_within(x0, y0, x1, y1,
            self.font_w * self.cols, self.font_h * (self.rows-1));

        // Clear the last row
        let (x2, y2) = self.position_xy(self.rows - 1, 0);
        self.fb.fill(x2, y2,
            self.font_w * self.cols, self.font_h,
            self.native_bg);
    }

    fn new_line(&mut self) {
        let (x, y) = self.position_xy(self.row, self.col);

        // Fill end of line
        let remaining = self.cols - self.col;
        self.fb.fill(x, y, self.font_w * remaining, self.font_h, self.native_bg);

        self.row += 1;
        self.col = 0;

        if self.row >= self.rows {
            self.row = self.rows - 1;
            self.scroll();
        }

        self.update_cursor();
    }

    fn write_char(&mut self, ch: char) {
        match ch {
            '\n' => {
                self.new_line();
            },

            '\x08' /* backspace */ => {
                self.put_here(' ');

                if self.col > 0 {
                    self.col -= 1;
                }

                self.update_cursor();
            },

            _ => {
                self.put_here(ch);
                self.col += 1;

                if self.col >= self.cols {
                    self.new_line();
                } else {
                    self.update_cursor();
                }
            }
        }
    }
}

impl<T, U> Terminal for Graphical<T, U> where T: Framebuffer, U: Font {
    fn reset(&mut self) -> fmt::Result {
        self.fg = Color::LightGrey;
        self.bg = Color::Black;
        self.update_colors();
        self.clear()
    }

    fn clear(&mut self) -> fmt::Result {
        self.row = 0;
        self.col = 0;

        self.fb.fill(0, 0, self.fb.width(), self.fb.height(), self.native_bg);

        self.update_cursor();

        Ok(())
    }

    fn get_cursor(&self) -> (usize, usize) {
        (self.row, self.col)
    }

    fn set_cursor(&mut self, row: usize, col: usize) -> fmt::Result {
        self.row = row;
        self.col = col;

        self.update_cursor();
        Ok(())
    }

    fn get_color(&self) -> (Color, Color) {
        (self.fg, self.bg)
    }

    fn set_color(&mut self, fg: Color, bg: Color) -> fmt::Result {
        self.fg = fg;
        self.bg = bg;
        self.update_colors();
        Ok(())
    }

    fn put_raw_byte(&mut self,
                    byte: u8,
                    fg:   Color,
                    bg:   Color,
                    row:  usize,
                    col:  usize) -> fmt::Result {
        // I think we should get rid of this method, it doesn't handle unicode properly
        self.render(
            char::from(byte),
            row,
            col,
            (self.fb.color_format().format(fg.to_rgb()),
             self.fb.color_format().format(bg.to_rgb())));
        Ok(())
    }

    fn write_raw_byte(&mut self, byte: u8) -> fmt::Result {
        // Build unicode
        self.char_buf.push(byte);

        while let Some(ch) = self.char_buf.consume() {
            self.write_char(ch);
        }

        Ok(())
    }

    fn set_double_buffer(&mut self, enabled: bool) {
        self.fb.set_double_buffer(enabled);
    }

    fn flush(&mut self) -> fmt::Result {
        self.update_cursor();
        Ok(())
    }
}

impl<T, U> fmt::Write for Graphical<T, U> where T: Framebuffer, U: Font {
    fn write_char(&mut self, ch: char) -> fmt::Result {
        let mut buf = [0; 4];
        self.write_raw_bytes(ch.encode_utf8(&mut buf).as_bytes())?;
        self.flush()?;
        Ok(())
    }

    fn write_str(&mut self, s: &str) -> fmt::Result {
        self.write_raw_bytes(s.as_bytes())?;
        self.flush()?;
        Ok(())
    }
}

static U_VGA16: &[u8] = include_bytes!("font/u_vga16.psf");

pub fn u_vga16() -> PsfFont<&'static [u8]> {
    PsfFont::new(U_VGA16).unwrap()
}

pub trait Font {
    fn char_width(&self) -> usize;
    fn char_height(&self) -> usize;
    fn char_bytes(&self) -> usize {
        align_up(self.char_width() * self.char_height(), 8) / 8
    }

    fn get(&self, ch: char) -> Option<&[u8]>;
}

#[derive(Debug, Clone)]
pub struct PsfFont<T: AsRef<[u8]>> {
    buf: T,
    chars_range: Range<usize>,
    unicode_table: BTreeMap<char, usize>,
    char_width: usize,
    char_height: usize,
}

const PSF1_MAGIC: u16 = 0x0436; /* little endian */
const PSF1_MODE_512: u8 = 0x01;
const PSF1_MODE_HASTAB: u8 = 0x02;
const PSF1_MODE_HASSEQ: u8 = 0x04;

const PSF1_SEPARATOR: u32 = 0xffff;
const PSF1_STARTSEQ: u32 = 0xfffe;

#[derive(Debug, Clone)]
pub enum PsfError {
    /// this is not a recognized PC Screen Font file
    WrongMagic,
    /// the mode field ({:x}) has unrecognized bits
    BadMode(u8),
    /// file is smaller than it should be
    FileTooSmall,
}

fn read_u16(slice: &[u8]) -> u16 {
    let mut bytes: [u8; 2] = [0; 2];
    bytes.copy_from_slice(&slice[0..2]);

    u16::from_le_bytes(bytes)
}

impl<T> PsfFont<T> where T: AsRef<[u8]> {
    pub fn new(buf: T) -> Result<PsfFont<T>, PsfError> {
        let file: &[u8] = buf.as_ref();

        if read_u16(&file[0..2]) != PSF1_MAGIC {
            return Err(PsfError::WrongMagic);
        }

        let mode: u8 = file[2];

        if mode & 0x07 != mode {
            return Err(PsfError::BadMode(mode));
        }

        let mode_512    = mode & PSF1_MODE_512 != 0;
        let mode_hastab = mode & PSF1_MODE_HASTAB != 0;
        let mode_hasseq = mode & PSF1_MODE_HASSEQ != 0;

        let charsize: usize = file[3] as usize; // height, also number of bytes

        let char_width = 8;
        let char_height = charsize;

        let length = if mode_512 { 512 } else { 256 };

        let chars_range = 4 .. (4 + charsize * length);

        // File too small if chars_range doesn't fit
        if chars_range.end > file.len() {
            return Err(PsfError::FileTooSmall);
        }

        let mut unicode_table = BTreeMap::new();

        // Parse the unicode table, or generate it if it's not there
        if mode_hastab {
            let unicode_table_buf = &file[chars_range.end ..];

            let mut pos = 0;

            'eof: for index in 0..length {
                loop {
                    if pos >= file.len() {
                        break 'eof;
                    }

                    let codepoint = read_u16(&unicode_table_buf[pos..]) as u32;

                    pos += 2;

                    if codepoint == PSF1_SEPARATOR {
                        // End of sequences for this index
                        break;
                    } else if codepoint == PSF1_STARTSEQ && mode_hasseq {
                        // Sequences not supported yet.
                        pos += 4;
                    } else if let Some(c) = char::from_u32(codepoint) {
                        // Insert codepoint.
                        unicode_table.insert(c, index);
                    }
                }
            }
        } else {
            for index in 0..length {
                if let Some(c) = char::from_u32(index as u32) {
                    unicode_table.insert(c, index);
                }
            }
        }

        Ok(PsfFont {
            buf,
            chars_range,
            unicode_table,
            char_width,
            char_height
        })
    }
}

impl<T> Font for PsfFont<T> where T: AsRef<[u8]> {
    fn char_width(&self) -> usize {
        self.char_width
    }

    fn char_height(&self) -> usize {
        self.char_height
    }

    fn get(&self, ch: char) -> Option<&[u8]> {
        self.unicode_table.get(&ch).map(|index| {
            let start = self.char_height * index;
            &self.buf.as_ref()[self.chars_range.clone()][start .. (start + self.char_height)]
        })
    }
}
<|endoftext|>


<commit_before>#[test]
fn path_finding_test() {
  use cgmath::Point2;
  use shaders::Position;
  use terrain::path_finding::calc_route;

  let impassable_tiles = vec![[1, 1], [2, 2], [2, 3], [2, 4], [2, 5], [3, 0], [3, 5], [4, 0], [4, 2], [4, 3], [4, 4], [4, 5]];

  let expected_result: Vec<Point2<i32>> = vec![[0, 0], [1, 0], [2, 1], [3, 2], [4, 1], [5, 2], [5, 3], [5, 4], [5, 5]]
    .iter()
    .map(|e| Point2::new(e[0], e[1]))
    .collect();

  assert_eq!(expected_result,
             calc_route(Position::new([0.0, -1500.0]), Position::new([0.0, -1500.0 + 5.0 * 46.0]), &impassable_tiles)
               .map_or_else(|| vec![],
                            |(f, _)| f),
             "path_finding_test");
}
<commit_msg>Add more unit tests<commit_after>const IMPASSABLE_TILES: [[usize; 2]; 12] = [[1, 1], [2, 2], [2, 3], [2, 4], [2, 5], [3, 0], [3, 5], [4, 0], [4, 2], [4, 3], [4, 4], [4, 5]];

#[test]
fn path_finding_test_1() {
  use cgmath::Point2;
  use shaders::Position;
  use terrain::path_finding::calc_route;

  let expected_result_1: Vec<Point2<i32>> = vec![[0, 0], [1, 0], [2, 1], [3, 2], [4, 1], [5, 2], [5, 3], [5, 4], [5, 5]]
    .iter()
    .map(|e| Point2::new(e[0], e[1]))
    .collect();

  assert_eq!(expected_result_1,
             calc_route(Position::new([0.0, -1500.0]), Position::new([0.0, -1500.0 + 5.0 * 46.0]), &IMPASSABLE_TILES.to_vec())
               .map_or_else(|| vec![],
                            |(f, _)| f),
             "path_finding_test_1");
}

#[test]
fn path_finding_test_2() {
  use cgmath::Point2;
  use shaders::Position;
  use terrain::path_finding::calc_route;

  let expected_result: Vec<Point2<i32>> = vec![[3, 3], [3, 2], [4, 1], [5, 2], [5, 3], [5, 4], [5, 5]]
    .iter()
    .map(|e| Point2::new(e[0], e[1]))
    .collect();

  assert_eq!(expected_result,
             calc_route(Position::new([0.0, -1500.0 + 3.0 * 46.0]), Position::new([0.0, -1500.0 + 5.0 * 46.0]), &IMPASSABLE_TILES.to_vec())
               .map_or_else(|| vec![],
                            |(f, _)| f),
             "path_finding_test_2");
}
<|endoftext|>


<commit_before>use errors::*;
use notifications::*;
use rustup_utils::utils;
use temp;

use openssl::crypto::hash::{Type, Hasher};
use itertools::Itertools;

use std::path::Path;
use std::process::Command;

#[derive(Debug)]
pub struct DownloadCfg<'a> {
    pub temp_cfg: &'a temp::Cfg,
    pub notify_handler: NotifyHandler<'a>,
    pub gpg_key: Option<&'a str>,
}

impl<'a> DownloadCfg<'a> {
    pub fn get(&self, url: &str) -> Result<temp::File<'a>> {
        if let Some(key) = self.gpg_key {
            // Download and verify with GPG key

            let sig_url = try!(utils::parse_url(&format!("{}.asc", url)));
            let sig_file = try!(self.temp_cfg.new_file());
            try!(utils::download_file(sig_url, &sig_file, None, ntfy!(&self.notify_handler)));

            let target_url = try!(utils::parse_url(url));
            let target_file = try!(self.temp_cfg.new_file());

            {
                let target_filename: &Path = &target_file;
                try!(utils::download_file(target_url,
                                          &target_file,
                                          None,
                                          ntfy!(&self.notify_handler)));

                let key_file = try!(self.temp_cfg.new_file());
                let key_filename: &Path = &key_file;
                try!(utils::write_file("key", &key_file, key));

                let gpg = try!(utils::find_cmd(&["gpg2", "gpg"])
                               .ok_or(Error::from("could not find 'gpg' on PATH")));

                try!(utils::cmd_status("gpg",
                                       Command::new(gpg)
                                           .arg("--no-permission-warning")
                                           .arg("--dearmor")
                                           .arg(key_filename)));

                try!(utils::cmd_status("gpg",
                                       Command::new(gpg)
                                           .arg("--no-permission-warning")
                                           .arg("--keyring")
                                           .arg(&key_filename.with_extension("gpg"))
                                           .arg("--verify")
                                           .arg(target_filename)));
            }

            Ok(target_file)
        } else {
            // Download and verify with checksum

            let hash_url = try!(utils::parse_url(&format!("{}.sha256", url)));
            let hash_file = try!(self.temp_cfg.new_file());
            try!(utils::download_file(hash_url, &hash_file, None, ntfy!(&self.notify_handler)));

            let hash = try!(utils::read_file("hash", &hash_file).map(|s| s[0..64].to_owned()));
            let mut hasher = Hasher::new(Type::SHA256);

            let target_url = try!(utils::parse_url(url));
            let target_file = try!(self.temp_cfg.new_file());
            try!(utils::download_file(target_url,
                                      &target_file,
                                      Some(&mut hasher),
                                      ntfy!(&self.notify_handler)));

            let actual_hash = hasher.finish()
                                    .iter()
                                    .map(|b| format!("{:02x}", b))
                                    .join("");

            if hash != actual_hash {
                // Incorrect hash
                return Err(ErrorKind::ChecksumFailed {
                    url: url.to_owned(),
                    expected: hash,
                    calculated: actual_hash,
                }.into());
            } else {
                self.notify_handler.call(Notification::ChecksumValid(url));
            }

            Ok(target_file)
        }

    }
}
<commit_msg>Remove an Error::from<commit_after>use errors::*;
use notifications::*;
use rustup_utils::utils;
use temp;

use openssl::crypto::hash::{Type, Hasher};
use itertools::Itertools;

use std::path::Path;
use std::process::Command;

#[derive(Debug)]
pub struct DownloadCfg<'a> {
    pub temp_cfg: &'a temp::Cfg,
    pub notify_handler: NotifyHandler<'a>,
    pub gpg_key: Option<&'a str>,
}

impl<'a> DownloadCfg<'a> {
    pub fn get(&self, url: &str) -> Result<temp::File<'a>> {
        if let Some(key) = self.gpg_key {
            // Download and verify with GPG key

            let sig_url = try!(utils::parse_url(&format!("{}.asc", url)));
            let sig_file = try!(self.temp_cfg.new_file());
            try!(utils::download_file(sig_url, &sig_file, None, ntfy!(&self.notify_handler)));

            let target_url = try!(utils::parse_url(url));
            let target_file = try!(self.temp_cfg.new_file());

            {
                let target_filename: &Path = &target_file;
                try!(utils::download_file(target_url,
                                          &target_file,
                                          None,
                                          ntfy!(&self.notify_handler)));

                let key_file = try!(self.temp_cfg.new_file());
                let key_filename: &Path = &key_file;
                try!(utils::write_file("key", &key_file, key));

                let gpg = try!(utils::find_cmd(&["gpg2", "gpg"])
                               .ok_or("could not find 'gpg' on PATH"));

                try!(utils::cmd_status("gpg",
                                       Command::new(gpg)
                                           .arg("--no-permission-warning")
                                           .arg("--dearmor")
                                           .arg(key_filename)));

                try!(utils::cmd_status("gpg",
                                       Command::new(gpg)
                                           .arg("--no-permission-warning")
                                           .arg("--keyring")
                                           .arg(&key_filename.with_extension("gpg"))
                                           .arg("--verify")
                                           .arg(target_filename)));
            }

            Ok(target_file)
        } else {
            // Download and verify with checksum

            let hash_url = try!(utils::parse_url(&format!("{}.sha256", url)));
            let hash_file = try!(self.temp_cfg.new_file());
            try!(utils::download_file(hash_url, &hash_file, None, ntfy!(&self.notify_handler)));

            let hash = try!(utils::read_file("hash", &hash_file).map(|s| s[0..64].to_owned()));
            let mut hasher = Hasher::new(Type::SHA256);

            let target_url = try!(utils::parse_url(url));
            let target_file = try!(self.temp_cfg.new_file());
            try!(utils::download_file(target_url,
                                      &target_file,
                                      Some(&mut hasher),
                                      ntfy!(&self.notify_handler)));

            let actual_hash = hasher.finish()
                                    .iter()
                                    .map(|b| format!("{:02x}", b))
                                    .join("");

            if hash != actual_hash {
                // Incorrect hash
                return Err(ErrorKind::ChecksumFailed {
                    url: url.to_owned(),
                    expected: hash,
                    calculated: actual_hash,
                }.into());
            } else {
                self.notify_handler.call(Notification::ChecksumValid(url));
            }

            Ok(target_file)
        }

    }
}
<|endoftext|>