<commit_before>use std::collections::HashMap;
use std::net::SocketAddr;
use std::rc::{Rc, Weak};
use std::cell::RefCell;

use mioco::udp::UdpSocket;
use mioco::sync::mpsc::{channel, Receiver, Sender};
use bincode::rustc_serialize as bcode_rcs;
use sodiumoxide::crypto::secretbox as crypto_secretbox;

use ::Result;
use identity::{Identity, Extension, RemoteServer};
use packet::*;
use keys::*;
use boxes::*;
use nonces::*;

pub struct Stream {

}

impl Stream {
    pub fn new() -> Stream {
        Stream {}
    }
}

trait PacketProcessor {
    fn process_packet(&mut self, packet: Packet, sock: &mut UdpSocket, rem_addr: SocketAddr) -> Result<()>;
}

pub struct Demultiplexor {
    sock: UdpSocket,
    packet_processors: HashMap<Extension, Weak<RefCell<PacketProcessor>>>,
}

impl Demultiplexor {
    pub fn new(sock: UdpSocket) -> Demultiplexor {
        Demultiplexor {
            sock: sock,
            packet_processors: HashMap::new(),
        }
    }

    pub fn create_listener(&mut self, listener_id: Identity) -> Rc<RefCell<Listener>> {
        let listener_extension = listener_id.extension.clone();
        let listener = Rc::new(RefCell::new(Listener::new(listener_id)));

        self.packet_processors.insert(listener_extension, Rc::downgrade(&listener) as Weak<RefCell<PacketProcessor>>);

        listener
    }

    pub fn create_client(&mut self, client_id: Identity, remote_id: RemoteServer) -> Rc<RefCell<ClientSock>> {
        let client_extension = client_id.extension.clone();
        let client = Rc::new(RefCell::new(ClientSock::new(client_id, remote_id)));

        self.packet_processors.insert(client_extension, Rc::downgrade(&client) as Weak<RefCell<PacketProcessor>>);

        client
    }

    pub fn listen(&mut self) -> Result<()> {
        let mut buf = [0u8; PACKET_MAX_SIZE];
        loop {
            let (recv_len, rem_addr) = try!(self.sock.recv(&mut buf));
            let packet: Packet = try!(bcode_rcs::decode(&buf[..recv_len]));

            if let Some(ref mut packet_processor) = {
                let dst_extension = packet.get_destination_extension();
                self.packet_processors.get_mut(dst_extension).and_then(|packet_processor_weak|
                    packet_processor_weak.upgrade()
                )
            } {
                try!(packet_processor.borrow_mut().process_packet(packet, &mut self.sock, rem_addr));
            }
        }
    }
}

pub struct ClientSock {
    // recv_rx: Receiver<Vec<u8>>,
    // recv_tx: Sender<Vec<u8>>,
}

impl ClientSock {
    pub fn new(my_id: Identity, remote_id: RemoteServer) -> ClientSock {
        // let (recv_rx, recv_tx) = channel();

        ClientSock {}
    }

    fn process_server_msg(&mut self, server_msg_packet: ServerMessagePacket) -> Result<()> {
        Err("Not implemented".into())
    }
}

impl PacketProcessor for ClientSock {
    fn process_packet(&mut self, packet: Packet, sock: &mut UdpSocket, rem_addr: SocketAddr) -> Result<()> {
        match packet {
            Packet::ServerMessage(server_msg_packet) => {
                try!(self.process_server_msg(server_msg_packet));
            },
            Packet::Cookie(cookie_packet) => {
            },
            _ => {
                debug!("Unvalid packet type");
            }
        }

        Ok(())
    }
}

pub struct ServerSocket {
    recv_rx: Receiver<Vec<u8>>,
    sock: UdpSocket,
    rem_addr: SocketAddr,
    client_extension: Extension,
    my_extension: Extension,
    precomputed_key: PrecomputedKey,
    next_send_nonce: Nonce8,
}

impl ServerSocket {
    pub fn recv(&mut self) -> Result<Vec<u8>> {
        self.recv_rx.recv().or(Err("Coudlnt read from recv channel".into()))
    }

    pub fn send(&mut self, msg: &[u8]) -> Result<usize> {
        let server_msg_packet: Packet = ServerMessagePacket {
            client_extension: self.client_extension.clone(),
            server_extension: self.my_extension.clone(),
            payload_box: ServerMessageBox::seal_precomputed(msg, &self.next_send_nonce, &self.precomputed_key),
        }.into();

        self.next_send_nonce.increment();

        let sent = try!(server_msg_packet.send(&mut self.sock, &self.rem_addr));

        if sent >= SERVER_MSG_PACKET_BASE_SIZE {
            Ok(sent - SERVER_MSG_PACKET_BASE_SIZE)
        } else {
            Err("Header trimmed".into())
        }
    }
}

pub struct Listener {
    my_extension: Extension,
    my_long_term_sk: server_long_term::SecretKey,
    minute_key: crypto_secretbox::Key,
    last_minute_key: crypto_secretbox::Key,
    conns: HashMap<(Extension, client_short_term::PublicKey), ServerConnection>,
    accept_tx: Sender<ServerSocket>,
    accept_rx: Receiver<ServerSocket>,
}

impl Listener {
    fn new(my_id: Identity) -> Listener {
        let (_, my_long_term_sk, my_extension) = my_id.as_server();
        let (accept_tx, accept_rx) = channel();

        Listener {
            my_extension: my_extension,
            my_long_term_sk: my_long_term_sk,
            minute_key: crypto_secretbox::gen_key(),
            last_minute_key: crypto_secretbox::gen_key(),
            conns: HashMap::new(),
            accept_tx: accept_tx,
            accept_rx: accept_rx,
        }
    }

    pub fn accept(&mut self) -> Result<ServerSocket> {
        self.accept_rx.recv().or(Err("Couldnt read accept channel".into()))
    }

    fn process_hello(&self, hello_packet: HelloPacket, sock: &mut UdpSocket, rem_addr: SocketAddr) -> Result<()> {
        let client_short_term_pk = hello_packet.client_short_term_pk;
        let conn_key = (hello_packet.client_extension.clone(), client_short_term_pk.clone());

        if self.conns.contains_key(&conn_key) {
            debug!("Hello packet received on open connection");
            return Ok(());
        }

        hello_packet.hello_box.open(&client_short_term_pk, &self.my_long_term_sk).unwrap();

        let (server_short_term_pk, server_short_term_sk) = server_short_term::gen_keypair();

        let cookie_packet: Packet = CookiePacket {
            client_extension: hello_packet.client_extension.clone(),
            server_extension: self.my_extension.clone(),
            cookie_box: PlainCookieBox {
                server_short_term_pk: server_short_term_pk,
                server_cookie: PlainCookie {
                    client_short_term_pk: client_short_term_pk.clone(),
                    server_short_term_sk: server_short_term_sk.clone(),
                }.seal(&CookieNonce::new_random(), &self.minute_key),
            }.seal(&Nonce16::new_random(), &client_short_term_pk, &self.my_long_term_sk),
        }.into();

        try!(cookie_packet.send(sock, &rem_addr));

        Ok(())
    }

    fn process_initiate(&mut self, initiate_packet: InitiatePacket, sock: &mut UdpSocket, rem_addr: SocketAddr) -> Result<()> {
        let conn_key = (initiate_packet.client_extension.clone(), initiate_packet.client_short_term_pk.clone());

        if self.conns.contains_key(&conn_key) {
            debug!("Recv'd INITIATE on already initiated connection");
            return Ok(());
        }

        let cookie = initiate_packet.server_cookie.open(&self.minute_key, &self.last_minute_key).unwrap();
        let precomputed_key = PrecomputedKey::precompute_at_server(&initiate_packet.client_short_term_pk, &cookie.server_short_term_sk);
        let (initiate_box, payload) = initiate_packet.initiate_box.open_precomputed_with_payload(&precomputed_key).unwrap();

        let client_long_term_pk = initiate_box.client_long_term_pk;
        let vouch = initiate_box.vouch.open(&client_long_term_pk, &self.my_long_term_sk).unwrap();

        if vouch.client_short_term_pk != initiate_packet.client_short_term_pk {
            return try!(Err("Invalid vouch"));
        }

        //TODO: Check if client_long_term_pk is accepted

        let (recv_tx, recv_rx) = channel();

        if let Some(msg) = payload {
            recv_tx.send(msg).unwrap();
        }

        let new_conn = ServerConnection {
            precomputed_key: precomputed_key.clone(),
            recv_tx: recv_tx,
            last_recv_nonce: initiate_packet.initiate_box.nonce.clone(),
        };

        self.conns.insert(conn_key, new_conn);
        info!("\"{}\" accepting connection \"{}\"@{}", &self.my_extension, &initiate_packet.client_extension, &rem_addr);

        let new_sock = ServerSocket {
            recv_rx: recv_rx,
            sock: sock.try_clone().unwrap(),
            rem_addr: rem_addr,
            client_extension: initiate_packet.client_extension.clone(),
            my_extension: self.my_extension.clone(),
            precomputed_key: precomputed_key,
            next_send_nonce: Nonce8::new_low(),
        };

        self.accept_tx.send(new_sock).or(Err("Couldnt read accept channel".into()))
    }

    fn process_client_msg(&mut self, client_msg_packet: ClientMessagePacket, _sock: &mut UdpSocket, _rem_addr: SocketAddr) -> Result<()> {
        let conn_key = (client_msg_packet.client_extension.clone(),
                        client_msg_packet.client_short_term_pk.clone());

        if let Some(server_conn) = self.conns.get_mut(&conn_key) {
            try!(server_conn.process_packet(client_msg_packet));
        }

        Ok(())
    }
}

impl PacketProcessor for Listener {
    fn process_packet(&mut self, packet: Packet, sock: &mut UdpSocket, rem_addr: SocketAddr) -> Result<()> {
        match packet {
            Packet::ClientMessage(client_msg_packet) => {
                try!(self.process_client_msg(client_msg_packet, sock, rem_addr));
            },
            Packet::Initiate(initiate_packet) => {
                try!(self.process_initiate(initiate_packet, sock, rem_addr));
            },
            Packet::Hello(hello_packet) => {
                try!(self.process_hello(hello_packet, sock, rem_addr));
            },
            _ => {
                debug!("Unvalid packet type");
            }
        }

        Ok(())
    }
}

struct ServerConnection {
    precomputed_key: PrecomputedKey,
    recv_tx: Sender<Vec<u8>>,
    last_recv_nonce: Nonce8,
}

impl ServerConnection {
    pub fn process_packet(&mut self, client_msg_packet: ClientMessagePacket) -> Result<()> {
        if client_msg_packet.payload_box.nonce <= self.last_recv_nonce {
            return Err("Bad nonce".into()); // TODO: Rework to allow some packet reordering
        }

        let msg = try!(client_msg_packet.payload_box.open_precomputed(&self.precomputed_key)
            .or(Err("Bad encrypted message")));

        self.last_recv_nonce = client_msg_packet.payload_box.nonce;

        try!(self.recv_tx.send(msg).or(Err("Couldnt write to recv channel")));

        Ok(())
    }
}
<commit_msg>ClientSock recv<commit_after>use std::collections::HashMap;
use std::net::SocketAddr;
use std::rc::{Rc, Weak};
use std::cell::RefCell;

use mioco::udp::UdpSocket;
use mioco::sync::mpsc::{channel, Receiver, Sender};
use bincode::rustc_serialize as bcode_rcs;
use sodiumoxide::crypto::secretbox as crypto_secretbox;

use ::Result;
use identity::{Identity, Extension, RemoteServer};
use packet::*;
use keys::*;
use boxes::*;
use nonces::*;

pub struct Stream {

}

impl Stream {
    pub fn new() -> Stream {
        Stream {}
    }
}

trait PacketProcessor {
    fn process_packet(&mut self, packet: Packet, sock: &mut UdpSocket, rem_addr: SocketAddr) -> Result<()>;
}

pub struct Demultiplexor {
    sock: UdpSocket,
    packet_processors: HashMap<Extension, Weak<RefCell<PacketProcessor>>>,
}

impl Demultiplexor {
    pub fn new(sock: UdpSocket) -> Demultiplexor {
        Demultiplexor {
            sock: sock,
            packet_processors: HashMap::new(),
        }
    }

    pub fn create_listener(&mut self, listener_id: Identity) -> Rc<RefCell<Listener>> {
        let listener_extension = listener_id.extension.clone();
        let listener = Rc::new(RefCell::new(Listener::new(listener_id)));

        self.packet_processors.insert(listener_extension, Rc::downgrade(&listener) as Weak<RefCell<PacketProcessor>>);

        listener
    }

    pub fn create_client(&mut self, client_id: Identity, remote_id: RemoteServer) -> Rc<RefCell<ClientSock>> {
        let client_extension = client_id.extension.clone();
        let client = Rc::new(RefCell::new(ClientSock::new(client_id, remote_id)));

        self.packet_processors.insert(client_extension, Rc::downgrade(&client) as Weak<RefCell<PacketProcessor>>);

        client
    }

    pub fn listen(&mut self) -> Result<()> {
        let mut buf = [0u8; PACKET_MAX_SIZE];
        loop {
            let (recv_len, rem_addr) = try!(self.sock.recv(&mut buf));
            let packet: Packet = try!(bcode_rcs::decode(&buf[..recv_len]));

            if let Some(ref mut packet_processor) = {
                let dst_extension = packet.get_destination_extension();
                self.packet_processors.get_mut(dst_extension).and_then(|packet_processor_weak|
                    packet_processor_weak.upgrade()
                )
            } {
                try!(packet_processor.borrow_mut().process_packet(packet, &mut self.sock, rem_addr));
            }
        }
    }
}

pub struct ClientSock {
    recv_rx: Receiver<Vec<u8>>,
    recv_tx: Sender<Vec<u8>>,
}

impl ClientSock {
    pub fn new(my_id: Identity, remote_id: RemoteServer) -> ClientSock {
        let (recv_tx, recv_rx) = channel();

        ClientSock {
            recv_rx: recv_rx,
            recv_tx: recv_tx,
        }
    }

    pub fn recv(&mut self) -> Result<Vec<u8>> {
        self.recv_rx.recv().or(Err("Coudlnt read from recv channel".into()))
    }

    fn process_server_msg(&mut self, server_msg_packet: ServerMessagePacket) -> Result<()> {
        Err("Not implemented".into())
    }
}

impl PacketProcessor for ClientSock {
    fn process_packet(&mut self, packet: Packet, sock: &mut UdpSocket, rem_addr: SocketAddr) -> Result<()> {
        match packet {
            Packet::ServerMessage(server_msg_packet) => {
                try!(self.process_server_msg(server_msg_packet));
            },
            Packet::Cookie(cookie_packet) => {
            },
            _ => {
                debug!("Unvalid packet type");
            }
        }

        Ok(())
    }
}

pub struct ServerSocket {
    recv_rx: Receiver<Vec<u8>>,
    sock: UdpSocket,
    rem_addr: SocketAddr,
    client_extension: Extension,
    my_extension: Extension,
    precomputed_key: PrecomputedKey,
    next_send_nonce: Nonce8,
}

impl ServerSocket {
    pub fn recv(&mut self) -> Result<Vec<u8>> {
        self.recv_rx.recv().or(Err("Coudlnt read from recv channel".into()))
    }

    pub fn send(&mut self, msg: &[u8]) -> Result<usize> {
        let server_msg_packet: Packet = ServerMessagePacket {
            client_extension: self.client_extension.clone(),
            server_extension: self.my_extension.clone(),
            payload_box: ServerMessageBox::seal_precomputed(msg, &self.next_send_nonce, &self.precomputed_key),
        }.into();

        self.next_send_nonce.increment();

        let sent = try!(server_msg_packet.send(&mut self.sock, &self.rem_addr));

        if sent >= SERVER_MSG_PACKET_BASE_SIZE {
            Ok(sent - SERVER_MSG_PACKET_BASE_SIZE)
        } else {
            Err("Header trimmed".into())
        }
    }
}

pub struct Listener {
    my_extension: Extension,
    my_long_term_sk: server_long_term::SecretKey,
    minute_key: crypto_secretbox::Key,
    last_minute_key: crypto_secretbox::Key,
    conns: HashMap<(Extension, client_short_term::PublicKey), ServerConnection>,
    accept_tx: Sender<ServerSocket>,
    accept_rx: Receiver<ServerSocket>,
}

impl Listener {
    fn new(my_id: Identity) -> Listener {
        let (_, my_long_term_sk, my_extension) = my_id.as_server();
        let (accept_tx, accept_rx) = channel();

        Listener {
            my_extension: my_extension,
            my_long_term_sk: my_long_term_sk,
            minute_key: crypto_secretbox::gen_key(),
            last_minute_key: crypto_secretbox::gen_key(),
            conns: HashMap::new(),
            accept_tx: accept_tx,
            accept_rx: accept_rx,
        }
    }

    pub fn accept(&mut self) -> Result<ServerSocket> {
        self.accept_rx.recv().or(Err("Couldnt read accept channel".into()))
    }

    fn process_hello(&self, hello_packet: HelloPacket, sock: &mut UdpSocket, rem_addr: SocketAddr) -> Result<()> {
        let client_short_term_pk = hello_packet.client_short_term_pk;
        let conn_key = (hello_packet.client_extension.clone(), client_short_term_pk.clone());

        if self.conns.contains_key(&conn_key) {
            debug!("Hello packet received on open connection");
            return Ok(());
        }

        hello_packet.hello_box.open(&client_short_term_pk, &self.my_long_term_sk).unwrap();

        let (server_short_term_pk, server_short_term_sk) = server_short_term::gen_keypair();

        let cookie_packet: Packet = CookiePacket {
            client_extension: hello_packet.client_extension.clone(),
            server_extension: self.my_extension.clone(),
            cookie_box: PlainCookieBox {
                server_short_term_pk: server_short_term_pk,
                server_cookie: PlainCookie {
                    client_short_term_pk: client_short_term_pk.clone(),
                    server_short_term_sk: server_short_term_sk.clone(),
                }.seal(&CookieNonce::new_random(), &self.minute_key),
            }.seal(&Nonce16::new_random(), &client_short_term_pk, &self.my_long_term_sk),
        }.into();

        try!(cookie_packet.send(sock, &rem_addr));

        Ok(())
    }

    fn process_initiate(&mut self, initiate_packet: InitiatePacket, sock: &mut UdpSocket, rem_addr: SocketAddr) -> Result<()> {
        let conn_key = (initiate_packet.client_extension.clone(), initiate_packet.client_short_term_pk.clone());

        if self.conns.contains_key(&conn_key) {
            debug!("Recv'd INITIATE on already initiated connection");
            return Ok(());
        }

        let cookie = initiate_packet.server_cookie.open(&self.minute_key, &self.last_minute_key).unwrap();
        let precomputed_key = PrecomputedKey::precompute_at_server(&initiate_packet.client_short_term_pk, &cookie.server_short_term_sk);
        let (initiate_box, payload) = initiate_packet.initiate_box.open_precomputed_with_payload(&precomputed_key).unwrap();

        let client_long_term_pk = initiate_box.client_long_term_pk;
        let vouch = initiate_box.vouch.open(&client_long_term_pk, &self.my_long_term_sk).unwrap();

        if vouch.client_short_term_pk != initiate_packet.client_short_term_pk {
            return try!(Err("Invalid vouch"));
        }

        //TODO: Check if client_long_term_pk is accepted

        let (recv_tx, recv_rx) = channel();

        if let Some(msg) = payload {
            recv_tx.send(msg).unwrap();
        }

        let new_conn = ServerConnection {
            precomputed_key: precomputed_key.clone(),
            recv_tx: recv_tx,
            last_recv_nonce: initiate_packet.initiate_box.nonce.clone(),
        };

        self.conns.insert(conn_key, new_conn);
        info!("\"{}\" accepting connection \"{}\"@{}", &self.my_extension, &initiate_packet.client_extension, &rem_addr);

        let new_sock = ServerSocket {
            recv_rx: recv_rx,
            sock: sock.try_clone().unwrap(),
            rem_addr: rem_addr,
            client_extension: initiate_packet.client_extension.clone(),
            my_extension: self.my_extension.clone(),
            precomputed_key: precomputed_key,
            next_send_nonce: Nonce8::new_low(),
        };

        self.accept_tx.send(new_sock).or(Err("Couldnt read accept channel".into()))
    }

    fn process_client_msg(&mut self, client_msg_packet: ClientMessagePacket, _sock: &mut UdpSocket, _rem_addr: SocketAddr) -> Result<()> {
        let conn_key = (client_msg_packet.client_extension.clone(),
                        client_msg_packet.client_short_term_pk.clone());

        if let Some(server_conn) = self.conns.get_mut(&conn_key) {
            try!(server_conn.process_packet(client_msg_packet));
        }

        Ok(())
    }
}

impl PacketProcessor for Listener {
    fn process_packet(&mut self, packet: Packet, sock: &mut UdpSocket, rem_addr: SocketAddr) -> Result<()> {
        match packet {
            Packet::ClientMessage(client_msg_packet) => {
                try!(self.process_client_msg(client_msg_packet, sock, rem_addr));
            },
            Packet::Initiate(initiate_packet) => {
                try!(self.process_initiate(initiate_packet, sock, rem_addr));
            },
            Packet::Hello(hello_packet) => {
                try!(self.process_hello(hello_packet, sock, rem_addr));
            },
            _ => {
                debug!("Unvalid packet type");
            }
        }

        Ok(())
    }
}

struct ServerConnection {
    precomputed_key: PrecomputedKey,
    recv_tx: Sender<Vec<u8>>,
    last_recv_nonce: Nonce8,
}

impl ServerConnection {
    pub fn process_packet(&mut self, client_msg_packet: ClientMessagePacket) -> Result<()> {
        if client_msg_packet.payload_box.nonce <= self.last_recv_nonce {
            return Err("Bad nonce".into()); // TODO: Rework to allow some packet reordering
        }

        let msg = try!(client_msg_packet.payload_box.open_precomputed(&self.precomputed_key)
            .or(Err("Bad encrypted message")));

        self.last_recv_nonce = client_msg_packet.payload_box.nonce;

        try!(self.recv_tx.send(msg).or(Err("Couldnt write to recv channel")));

        Ok(())
    }
}
<|endoftext|>


<commit_before>// Copyright 2020 Google LLC
//
// Use of this source code is governed by an MIT-style license that can be found
// in the LICENSE file or at https://opensource.org/licenses/MIT.

//! Utilities for working with the chunked format.
//!
//! The chunked file format is extremely simple serialization procedure where
//! each chunk of raw bytes is prepended with the size of the chunk (as 64-bit
//! unsigned big-endian integer).
//!
//! Its primary application is streaming encoding and decoding of blobs of data
//! in formats that do not support or are inefficient for these purposes (such
//! as serialized Protocol Buffer messages).

use std::io::Cursor;

use byteorder::BigEndian;

/// Encodes a given iterator over protobuf messages into the chunked format.
///
/// This is a streaming encoder and performs the encoding in a lazy way. It
/// should compose well with other streaming encoders (e.g. these offered by
/// the [`flate2`] crate).
///
/// [`flate2`]: https://crates.io/crates/flate2
///
/// # Examples
///
/// ```no_run
/// use std::fs::File;
///
/// let data = vec! {
///     String::from("foo"),
///     String::from("bar"),
///     String::from("baz")
/// };
///
/// let mut stream = rrg::chunked::encode(data.into_iter());
/// let mut file = File::create("output.chunked").unwrap();
/// std::io::copy(&mut stream, &mut file).unwrap();
/// ```
pub fn encode<I>(iter: I) -> Encode<I>
where
    I: Iterator,
    I::Item: prost::Message,
{
    Encode {
        iter: iter,
        cur: Cursor::new(vec!()),
    }
}

/// Decodes a buffer in the chunked format into a stream of protobuf messages.
///
/// This is a streaming decoder and performs the decoding in a lazy way. It
/// should compose well with other streaming decoders (e.g. these offered by the
/// [`flate2`] crate).
///
/// [`flate2`]: https://crates.io/crates/flate2
///
/// # Examples
///
/// ```no_run
/// use std::fs::File;
///
/// let file = File::open("input.chunked").unwrap();
/// for (idx, msg) in rrg::chunked::decode(file).enumerate() {
///     let msg: String = msg.unwrap();
///     println!("message #{}: {:?}", idx, msg);
/// }
/// ```
pub fn decode<R, M>(buf: R) -> Decode<R, M>
where
    R: std::io::Read,
    M: prost::Message,
{
    Decode {
        reader: buf,
        buf: vec!(),
        marker: std::marker::PhantomData,
    }
}

/// Streaming encoder for the chunked format.
///
/// It implements the `Read` trait, lazily polling the underlying iterator over
/// Protocol Buffers messages as needed.
///
/// Instances of this type can be constructed using the [`encode`] function.
///
/// [`encode`]: fn.encode.html
pub struct Encode<I> {
    iter: I,
    cur: Cursor<Vec<u8>>,
}

impl<I, M> Encode<I>
where
    I: Iterator<Item = M>,
    M: prost::Message,
{
    /// Checks whether all the data from the underlying cursor has been read.
    fn is_empty(&self) -> bool {
        self.cur.position() == self.cur.get_ref().len() as u64
    }

    /// Pulls another message from the underlying iterator.
    fn pull(&mut self) -> std::io::Result<()> {
        use byteorder::WriteBytesExt as _;

        let msg = match self.iter.next() {
            Some(msg) => msg,
            None => return Ok(()),
        };

        self.cur.get_mut().clear();
        self.cur.set_position(0);

        self.cur.write_u64::<BigEndian>(msg.encoded_len() as u64)?;
        msg.encode(&mut self.cur.get_mut())?;
        self.cur.set_position(0);

        Ok(())
    }
}

impl<I, M> std::io::Read for Encode<I>
where
    I: Iterator<Item = M>,
    M: prost::Message,
{
    fn read(&mut self, buf: &mut [u8]) -> std::io::Result<usize> {
        if self.is_empty() {
            self.pull()?;
        }

        self.cur.read(buf)
    }
}

/// Streaming decoder for the chunked format.
///
/// It implements the `Iterator` trait yielding Protocol Buffers messages,
/// lazily decoding data from the underlying buffer.
///
/// Instances of this type can be constructed using the [`decode`] function.
///
/// [`decode`]: fn.decode.html
pub struct Decode<R, M> {
    reader: R,
    buf: Vec<u8>,
    marker: std::marker::PhantomData<M>,
}

impl<R: std::io::Read, M> Decode<R, M> {

    /// Reads a size tag from the underlying buffer.
    ///
    /// It will return `None` if the is no more data in the buffer.
    fn read_len(&mut self) -> std::io::Result<Option<usize>> {
        use byteorder::ReadBytesExt as _;

        // `read` might not always read all 8 bytes. On the other hand, we also
        // cannot use just `read_exact` because the stream might have ended
        // already. Hence, we combine the two. First we attempt to read some
        // bytes with `read`: it should either return 0 (indicating end of the
        // stream), 8 (indicating that we have filled the whole buffer fully)
        // or something in between. In the last case, we use `read_exact to get
        // the remaining bytes (which should be non-zero now).
        let mut buf = [0; 8];
        match self.reader.read(&mut buf[..])? {
            8 => (),
            0 => return Ok(None),
            len => self.reader.read_exact(&mut buf[len..])?,
        }

        let len = (&buf[..]).read_u64::<BigEndian>()? as usize;
        Ok(Some(len))
    }
}

impl<R, M> Iterator for Decode<R, M>
where
    R: std::io::Read,
    M: prost::Message + Default,
{
    type Item = std::io::Result<M>;

    fn next(&mut self) -> Option<std::io::Result<M>> {
        let len = match self.read_len() {
            Ok(Some(len)) => len,
            Ok(None) => return None,
            Err(error) => return Some(Err(error)),
        };

        self.buf.resize(len, u8::default());
        match self.reader.read_exact(&mut self.buf[..]) {
            Ok(()) => (),
            Err(error) => return Some(Err(error)),
        }

        let msg = match M::decode(&self.buf[..]) {
            Ok(msg) => msg,
            Err(error) => return Some(Err(error.into())),
        };

        Some(Ok(msg))
    }
}

#[cfg(test)]
pub mod tests {

    use super::*;

    #[test]
    pub fn test_encode_empty_iter() {
        use std::io::Read as _;

        let mut stream = encode(std::iter::empty::<()>());

        let mut output = vec!();
        stream.read_to_end(&mut output).unwrap();

        assert!(output.is_empty());
    }

    #[test]
    pub fn test_decode_empty_buf() {
        let buf: &[u8] = b"";
        let mut iter = decode::<_, ()>(buf);

        assert!(iter.next().is_none());
    }

    #[test]
    pub fn test_decode_incorrect_size_tag() {
        let buf: &[u8] = b"\x12\x34\x56";
        let mut iter = decode::<_, ()>(buf);

        let error = iter.next().unwrap().unwrap_err();
        assert_eq!(error.kind(), std::io::ErrorKind::UnexpectedEof);
    }

    #[test]
    pub fn test_decode_zero_size_tag() {
        let buf: &[u8] = b"\x00\x00\x00\x00\x00\x00\x00\x00";

        let mut iter = decode(buf).map(Result::unwrap);

        assert_eq!(iter.next(), Some(()));
        assert_eq!(iter.next(), None);
    }

    #[test]
    pub fn test_decode_partial_size_tag() {

        // A simple reader that yields a 0-valued size tag byte by byte.
        struct Reader(u8);

        impl std::io::Read for Reader {

            fn read(&mut self, buf: &mut [u8]) -> std::io::Result<usize> {
                if self.0 == 8 {
                    Ok(0)
                } else {
                    buf[0] = 0;
                    self.0 += 1;
                    Ok(1)
                }
            }
        }

        let mut iter = decode(Reader(0)).map(Result::unwrap);

        assert_eq!(iter.next(), Some(()));
        assert_eq!(iter.next(), None);
    }

    #[test]
    pub fn test_decode_short_input() {
        let buf: &[u8] = b"\x00\x00\x00\x00\x00\x00\x00\x42foo";
        let mut iter = decode::<_, ()>(buf);

        let error = iter.next().unwrap().unwrap_err();
        assert_eq!(error.kind(), std::io::ErrorKind::UnexpectedEof);
    }

    #[test]
    pub fn test_encode_and_decode_single_message() {
        let mut iter = decode(encode(std::iter::once(String::from("foo"))))
            .map(Result::unwrap);

        assert_eq!(iter.next(), Some(String::from("foo")));
        assert_eq!(iter.next(), None);
    }

    #[test]
    pub fn test_encode_and_decode_single_unit_message() {
        let mut iter = decode(encode(std::iter::once(())))
            .map(Result::unwrap);

        assert_eq!(iter.next(), Some(()));
        assert_eq!(iter.next(), None);
    }

    #[test]
    pub fn test_encode_and_decode_multiple_messages() {
        let msgs = vec! {
            b"foo".to_vec(),
            b"bar".to_vec(),
            b"baz".to_vec(),
        };

        let mut iter = decode(encode(msgs.into_iter()))
            .map(Result::unwrap);

        assert_eq!(iter.next(), Some(b"foo".to_vec()));
        assert_eq!(iter.next(), Some(b"bar".to_vec()));
        assert_eq!(iter.next(), Some(b"baz".to_vec()));
        assert_eq!(iter.next(), None);
    }

    #[test]
    pub fn test_encode_and_decode_multiple_unit_messages() {
        let msgs = vec!((), (), ());

        let mut iter = decode(encode(msgs.into_iter()))
            .map(Result::unwrap);

        assert_eq!(iter.next(), Some(()));
        assert_eq!(iter.next(), Some(()));
        assert_eq!(iter.next(), Some(()));
        assert_eq!(iter.next(), None);
    }
}
<commit_msg>Remove an unnecessary global import.<commit_after>// Copyright 2020 Google LLC
//
// Use of this source code is governed by an MIT-style license that can be found
// in the LICENSE file or at https://opensource.org/licenses/MIT.

//! Utilities for working with the chunked format.
//!
//! The chunked file format is extremely simple serialization procedure where
//! each chunk of raw bytes is prepended with the size of the chunk (as 64-bit
//! unsigned big-endian integer).
//!
//! Its primary application is streaming encoding and decoding of blobs of data
//! in formats that do not support or are inefficient for these purposes (such
//! as serialized Protocol Buffer messages).

use byteorder::BigEndian;

/// Encodes a given iterator over protobuf messages into the chunked format.
///
/// This is a streaming encoder and performs the encoding in a lazy way. It
/// should compose well with other streaming encoders (e.g. these offered by
/// the [`flate2`] crate).
///
/// [`flate2`]: https://crates.io/crates/flate2
///
/// # Examples
///
/// ```no_run
/// use std::fs::File;
///
/// let data = vec! {
///     String::from("foo"),
///     String::from("bar"),
///     String::from("baz")
/// };
///
/// let mut stream = rrg::chunked::encode(data.into_iter());
/// let mut file = File::create("output.chunked").unwrap();
/// std::io::copy(&mut stream, &mut file).unwrap();
/// ```
pub fn encode<I>(iter: I) -> Encode<I>
where
    I: Iterator,
    I::Item: prost::Message,
{
    Encode {
        iter: iter,
        cur: std::io::Cursor::new(vec!()),
    }
}

/// Decodes a buffer in the chunked format into a stream of protobuf messages.
///
/// This is a streaming decoder and performs the decoding in a lazy way. It
/// should compose well with other streaming decoders (e.g. these offered by the
/// [`flate2`] crate).
///
/// [`flate2`]: https://crates.io/crates/flate2
///
/// # Examples
///
/// ```no_run
/// use std::fs::File;
///
/// let file = File::open("input.chunked").unwrap();
/// for (idx, msg) in rrg::chunked::decode(file).enumerate() {
///     let msg: String = msg.unwrap();
///     println!("message #{}: {:?}", idx, msg);
/// }
/// ```
pub fn decode<R, M>(buf: R) -> Decode<R, M>
where
    R: std::io::Read,
    M: prost::Message,
{
    Decode {
        reader: buf,
        buf: vec!(),
        marker: std::marker::PhantomData,
    }
}

/// Streaming encoder for the chunked format.
///
/// It implements the `Read` trait, lazily polling the underlying iterator over
/// Protocol Buffers messages as needed.
///
/// Instances of this type can be constructed using the [`encode`] function.
///
/// [`encode`]: fn.encode.html
pub struct Encode<I> {
    iter: I,
    cur: std::io::Cursor<Vec<u8>>,
}

impl<I, M> Encode<I>
where
    I: Iterator<Item = M>,
    M: prost::Message,
{
    /// Checks whether all the data from the underlying cursor has been read.
    fn is_empty(&self) -> bool {
        self.cur.position() == self.cur.get_ref().len() as u64
    }

    /// Pulls another message from the underlying iterator.
    fn pull(&mut self) -> std::io::Result<()> {
        use byteorder::WriteBytesExt as _;

        let msg = match self.iter.next() {
            Some(msg) => msg,
            None => return Ok(()),
        };

        self.cur.get_mut().clear();
        self.cur.set_position(0);

        self.cur.write_u64::<BigEndian>(msg.encoded_len() as u64)?;
        msg.encode(&mut self.cur.get_mut())?;
        self.cur.set_position(0);

        Ok(())
    }
}

impl<I, M> std::io::Read for Encode<I>
where
    I: Iterator<Item = M>,
    M: prost::Message,
{
    fn read(&mut self, buf: &mut [u8]) -> std::io::Result<usize> {
        if self.is_empty() {
            self.pull()?;
        }

        self.cur.read(buf)
    }
}

/// Streaming decoder for the chunked format.
///
/// It implements the `Iterator` trait yielding Protocol Buffers messages,
/// lazily decoding data from the underlying buffer.
///
/// Instances of this type can be constructed using the [`decode`] function.
///
/// [`decode`]: fn.decode.html
pub struct Decode<R, M> {
    reader: R,
    buf: Vec<u8>,
    marker: std::marker::PhantomData<M>,
}

impl<R: std::io::Read, M> Decode<R, M> {

    /// Reads a size tag from the underlying buffer.
    ///
    /// It will return `None` if the is no more data in the buffer.
    fn read_len(&mut self) -> std::io::Result<Option<usize>> {
        use byteorder::ReadBytesExt as _;

        // `read` might not always read all 8 bytes. On the other hand, we also
        // cannot use just `read_exact` because the stream might have ended
        // already. Hence, we combine the two. First we attempt to read some
        // bytes with `read`: it should either return 0 (indicating end of the
        // stream), 8 (indicating that we have filled the whole buffer fully)
        // or something in between. In the last case, we use `read_exact to get
        // the remaining bytes (which should be non-zero now).
        let mut buf = [0; 8];
        match self.reader.read(&mut buf[..])? {
            8 => (),
            0 => return Ok(None),
            len => self.reader.read_exact(&mut buf[len..])?,
        }

        let len = (&buf[..]).read_u64::<BigEndian>()? as usize;
        Ok(Some(len))
    }
}

impl<R, M> Iterator for Decode<R, M>
where
    R: std::io::Read,
    M: prost::Message + Default,
{
    type Item = std::io::Result<M>;

    fn next(&mut self) -> Option<std::io::Result<M>> {
        let len = match self.read_len() {
            Ok(Some(len)) => len,
            Ok(None) => return None,
            Err(error) => return Some(Err(error)),
        };

        self.buf.resize(len, u8::default());
        match self.reader.read_exact(&mut self.buf[..]) {
            Ok(()) => (),
            Err(error) => return Some(Err(error)),
        }

        let msg = match M::decode(&self.buf[..]) {
            Ok(msg) => msg,
            Err(error) => return Some(Err(error.into())),
        };

        Some(Ok(msg))
    }
}

#[cfg(test)]
pub mod tests {

    use super::*;

    #[test]
    pub fn test_encode_empty_iter() {
        use std::io::Read as _;

        let mut stream = encode(std::iter::empty::<()>());

        let mut output = vec!();
        stream.read_to_end(&mut output).unwrap();

        assert!(output.is_empty());
    }

    #[test]
    pub fn test_decode_empty_buf() {
        let buf: &[u8] = b"";
        let mut iter = decode::<_, ()>(buf);

        assert!(iter.next().is_none());
    }

    #[test]
    pub fn test_decode_incorrect_size_tag() {
        let buf: &[u8] = b"\x12\x34\x56";
        let mut iter = decode::<_, ()>(buf);

        let error = iter.next().unwrap().unwrap_err();
        assert_eq!(error.kind(), std::io::ErrorKind::UnexpectedEof);
    }

    #[test]
    pub fn test_decode_zero_size_tag() {
        let buf: &[u8] = b"\x00\x00\x00\x00\x00\x00\x00\x00";

        let mut iter = decode(buf).map(Result::unwrap);

        assert_eq!(iter.next(), Some(()));
        assert_eq!(iter.next(), None);
    }

    #[test]
    pub fn test_decode_partial_size_tag() {

        // A simple reader that yields a 0-valued size tag byte by byte.
        struct Reader(u8);

        impl std::io::Read for Reader {

            fn read(&mut self, buf: &mut [u8]) -> std::io::Result<usize> {
                if self.0 == 8 {
                    Ok(0)
                } else {
                    buf[0] = 0;
                    self.0 += 1;
                    Ok(1)
                }
            }
        }

        let mut iter = decode(Reader(0)).map(Result::unwrap);

        assert_eq!(iter.next(), Some(()));
        assert_eq!(iter.next(), None);
    }

    #[test]
    pub fn test_decode_short_input() {
        let buf: &[u8] = b"\x00\x00\x00\x00\x00\x00\x00\x42foo";
        let mut iter = decode::<_, ()>(buf);

        let error = iter.next().unwrap().unwrap_err();
        assert_eq!(error.kind(), std::io::ErrorKind::UnexpectedEof);
    }

    #[test]
    pub fn test_encode_and_decode_single_message() {
        let mut iter = decode(encode(std::iter::once(String::from("foo"))))
            .map(Result::unwrap);

        assert_eq!(iter.next(), Some(String::from("foo")));
        assert_eq!(iter.next(), None);
    }

    #[test]
    pub fn test_encode_and_decode_single_unit_message() {
        let mut iter = decode(encode(std::iter::once(())))
            .map(Result::unwrap);

        assert_eq!(iter.next(), Some(()));
        assert_eq!(iter.next(), None);
    }

    #[test]
    pub fn test_encode_and_decode_multiple_messages() {
        let msgs = vec! {
            b"foo".to_vec(),
            b"bar".to_vec(),
            b"baz".to_vec(),
        };

        let mut iter = decode(encode(msgs.into_iter()))
            .map(Result::unwrap);

        assert_eq!(iter.next(), Some(b"foo".to_vec()));
        assert_eq!(iter.next(), Some(b"bar".to_vec()));
        assert_eq!(iter.next(), Some(b"baz".to_vec()));
        assert_eq!(iter.next(), None);
    }

    #[test]
    pub fn test_encode_and_decode_multiple_unit_messages() {
        let msgs = vec!((), (), ());

        let mut iter = decode(encode(msgs.into_iter()))
            .map(Result::unwrap);

        assert_eq!(iter.next(), Some(()));
        assert_eq!(iter.next(), Some(()));
        assert_eq!(iter.next(), Some(()));
        assert_eq!(iter.next(), None);
    }
}
<|endoftext|>


<commit_before>use std::fs::{create_dir_all, File, remove_file};
use std::io::Write;

use clap::ArgMatches;
use rusoto::ChainProvider;

use aws::credentials_provider;
use encryption::Encryptor;
use error::KawsResult;
use process::execute_child_process;

pub struct Cluster<'a> {
    aws_account_id: Option<&'a str>,
    aws_credentials_provider: ChainProvider,
    ca_cert_path: String,
    ca_key_path: String,
    coreos_ami: Option<&'a str>,
    domain: Option<&'a str>,
    encrypted_ca_key_path: String,
    encrypted_master_key_path: String,
    encrypted_node_key_path: String,
    iam_users: Option<Vec<&'a str>>,
    instance_size: Option<&'a str>,
    kms_master_key_id: Option<&'a str>,
    kms_policy_path: String,
    kubernetes_version: Option<&'a str>,
    master_cert_path: String,
    master_csr_path: String,
    master_key_path: String,
    masters_max_size: Option<&'a str>,
    masters_min_size: Option<&'a str>,
    name: &'a str,
    node_cert_path: String,
    node_csr_path: String,
    node_key_path: String,
    nodes_max_size: Option<&'a str>,
    nodes_min_size: Option<&'a str>,
    openssl_config_path: String,
    region: &'a str,
    ssh_key: Option<&'a str>,
    tfvars_path: String,
    zone_id: Option<&'a str>,
}

impl<'a> Cluster<'a> {
    pub fn new(matches: &'a ArgMatches) -> Self {
        let name =  matches.value_of("cluster").expect("clap should have required cluster");

        Cluster {
            aws_account_id: matches.value_of("aws-account-id"),
            aws_credentials_provider: credentials_provider(
                matches.value_of("aws-credentials-path"),
                matches.value_of("aws-credentials-profile"),
            ),
            ca_cert_path: format!("clusters/{}/ca.pem", name),
            ca_key_path: format!("clusters/{}/ca-key.pem", name),
            coreos_ami: matches.value_of("ami"),
            domain: matches.value_of("domain"),
            encrypted_ca_key_path: format!("clusters/{}/ca-key-encrypted.base64", name),
            encrypted_master_key_path: format!("clusters/{}/master-key-encrypted.base64", name),
            encrypted_node_key_path: format!("clusters/{}/node-key-encrypted.base64", name),
            name: name,
            iam_users: matches.values_of("iam-users").map(|values| values.collect()),
            instance_size: matches.value_of("size"),
            kms_master_key_id: matches.value_of("kms-key"),
            kms_policy_path: format!("clusters/{}/kms-policy.json", name),
            kubernetes_version: matches.value_of("k8s-version"),
            master_cert_path: format!("clusters/{}/master.pem", name),
            master_csr_path: format!("clusters/{}/master.csr", name),
            master_key_path: format!("clusters/{}/master-key.pem", name),
            masters_max_size: matches.value_of("masters-max-size"),
            masters_min_size: matches.value_of("masters-min-size"),
            node_cert_path: format!("clusters/{}/node.pem", name),
            node_csr_path: format!("clusters/{}/node.csr", name),
            node_key_path: format!("clusters/{}/node-key.pem", name),
            nodes_max_size: matches.value_of("nodes-max-size"),
            nodes_min_size: matches.value_of("nodes-min-size"),
            openssl_config_path: format!("clusters/{}/openssl.cnf", name),
            region: matches.value_of("region").expect("clap should have required region"),
            ssh_key: matches.value_of("ssh-key"),
            tfvars_path: format!("clusters/{}/terraform.tfvars", name),
            zone_id: matches.value_of("zone-id"),
        }
    }

    pub fn init(&mut self) -> KawsResult {
        try!(self.create_directories());
        try!(self.create_tfvars());
        try!(self.create_openssl_config());
        try!(self.create_kms_policy());
        try!(self.create_pki_stubs());

        Ok(Some(format!(
            "Cluster \"{}\" initialized!",
            self.name,
        )))
    }

    pub fn generate_pki(&mut self) -> KawsResult {
        try!(self.create_ca());
        try!(self.create_master_credentials());
        try!(self.create_node_credentials());
        try!(self.encrypt_secrets(
            self.kms_master_key_id.expect("clap should have required kms-key"),
        ));

        Ok(None)
    }

    fn create_directories(&self) -> KawsResult {
        log_wrap!("Creating directories for the new cluster", {
            try!(create_dir_all(format!("clusters/{}", self.name)));
        });

        Ok(None)
    }

    fn create_kms_policy(&self) -> KawsResult {
        log_wrap!("Creating KMS policy file", {
            let mut file = try!(File::create(&self.kms_policy_path));

            let aws_account_id = self.aws_account_id.expect(
                "AWS account ID should have been required by clap"
            );

            let iam_user_arns = self.iam_users.as_ref().expect(
                "IAM users should have been required by clap"
            ).iter().map(|iam_user| {
                format!("\"arn:aws:iam::{}:user/{}\"", aws_account_id, iam_user)
            }).collect::<Vec<String>>().join(",");

            try!(write!(
                file,
                r#"{{
  "Version": "2012-10-17",
  "Statement": [
    {{
      "Sid": "Enable IAM User Permissions",
      "Effect": "Allow",
      "Principal": {{
        "AWS": [
          "arn:aws:iam::{aws_account_id}:root"
        ]
      }},
      "Action": "kms:*",
      "Resource": "*"
    }},
    {{
      "Sid": "Allow access for Key Administrators",
      "Effect": "Allow",
      "Principal": {{
        "AWS": [
          {iam_user_arns}
        ]
      }},
      "Action": [
        "kms:Create*",
        "kms:Describe*",
        "kms:Enable*",
        "kms:List*",
        "kms:Put*",
        "kms:Update*",
        "kms:Revoke*",
        "kms:Disable*",
        "kms:Get*",
        "kms:Delete*",
        "kms:ScheduleKeyDeletion",
        "kms:CancelKeyDeletion"
      ],
      "Resource": "*"
    }},
    {{
      "Sid": "Allow use of the key",
      "Effect": "Allow",
      "Principal": {{
        "AWS": [
          "arn:aws:iam::{aws_account_id}:role/kaws-k8s-master-{cluster}",
          "arn:aws:iam::{aws_account_id}:role/kaws-k8s-node-{cluster}",
          {iam_user_arns}
        ]
      }},
      "Action": [
        "kms:Encrypt",
        "kms:Decrypt",
        "kms:ReEncrypt*",
        "kms:GenerateDataKey*",
        "kms:DescribeKey"
      ],
      "Resource": "*"
    }},
    {{
      "Sid": "Allow attachment of persistent resources",
      "Effect": "Allow",
      "Principal": {{
        "AWS": [
          "arn:aws:iam::{aws_account_id}:role/kaws-k8s-master-{cluster}",
          "arn:aws:iam::{aws_account_id}:role/kaws-k8s-node-{cluster}",
          {iam_user_arns}
        ]
      }},
      "Action": [
        "kms:CreateGrant",
        "kms:ListGrants",
        "kms:RevokeGrant"
      ],
      "Resource": "*",
      "Condition": {{
        "Bool": {{
          "kms:GrantIsForAWSResource": true
        }}
      }}
    }}
  ]
}}
"#,
                aws_account_id = aws_account_id,
                cluster = self.name,
                iam_user_arns = iam_user_arns,
            ));
        });

        Ok(None)
    }

    fn create_master_credentials(&self) -> KawsResult {
        log_wrap!("Creating Kubernetes master private key", {
            try!(execute_child_process("openssl", &[
                "genrsa",
                "-out",
                &self.master_key_path,
                "2048",
            ]));
        });

        log_wrap!("Creating Kubernetes master certificate signing request", {
            try!(execute_child_process("openssl", &[
                "req",
                "-new",
                "-key",
                &self.master_key_path,
                "-out",
                &self.master_csr_path,
                "-subj",
                &format!("/CN=kaws-master-{}", self.name),
                "-config",
                &self.openssl_config_path,
            ]));
        });

        log_wrap!("Creating Kubernetes master certificate", {
            try!(execute_child_process("openssl", &[
                "x509",
                "-req",
                "-in",
                &self.master_csr_path,
                "-CA",
                &self.ca_cert_path,
                "-CAkey",
                &self.ca_key_path,
                "-CAcreateserial",
                "-out",
                &self.master_cert_path,
                "-days",
                "365",
                "-extensions",
                "v3_req",
                "-extfile",
                &self.openssl_config_path,
            ]));
        });

        log_wrap!("Removing Kubernetes master certificate signing request", {
            try!(remove_file(&self.master_csr_path));
        });

        Ok(None)
    }

    fn create_node_credentials(&self) -> KawsResult {
        log_wrap!("Creating Kubernetes node private key", {
            try!(execute_child_process("openssl", &[
                "genrsa",
                "-out",
                &self.node_key_path,
                "2048",
            ]));
        });

        log_wrap!("Creating Kubernetes node certificate signing request", {
            try!(execute_child_process("openssl", &[
                "req",
                "-new",
                "-key",
                &self.node_key_path,
                "-out",
                &self.node_csr_path,
                "-subj",
                &format!("/CN=kaws-node-{}", self.name),
            ]));
        });

        log_wrap!("Creating Kubernetes node certificate", {
            try!(execute_child_process("openssl", &[
                "x509",
                "-req",
                "-in",
                &self.node_csr_path,
                "-CA",
                &self.ca_cert_path,
                "-CAkey",
                &self.ca_key_path,
                "-CAcreateserial",
                "-out",
                &self.node_cert_path,
                "-days",
                "365",
            ]));
        });

        log_wrap!("Removing Kubernetes node certificate signing request", {
            try!(remove_file(&self.node_csr_path));
        });

        Ok(None)
    }

    fn create_openssl_config(&self) -> KawsResult {
        log_wrap!("Creating OpenSSL config file", {
            let mut file = try!(File::create(&self.openssl_config_path));

            try!(write!(
                file,
                "\
[req]
req_extensions = v3_req
distinguished_name = req_distinguished_name
[req_distinguished_name]
[v3_req]
basicConstraints = CA:FALSE
keyUsage = nonRepudiation, digitalSignature, keyEncipherment
subjectAltName = @alt_names
[alt_names]
DNS.1 = kubernetes
DNS.2 = kubernetes.default
DNS.3 = kubernetes.{}
IP.1 = 10.3.0.1
",
                self.domain.expect("domain should have been required by clap"),
            ));
        });

        Ok(None)
    }

    fn create_ca(&self) -> KawsResult {
        log_wrap!("Creating Kubernetes certificate authority private key", {
            try!(execute_child_process("openssl", &[
                "genrsa",
                "-out",
                &self.ca_key_path,
                "2048",
            ]));
        });

        log_wrap!("Creating Kubernetes certificate authority certificate", {
            try!(execute_child_process("openssl", &[
                "req",
                "-x509",
                "-new",
                "-nodes",
                "-key",
                &self.ca_key_path,
                "-days",
                "10000",
                "-out",
                &self.ca_cert_path,
                "-subj",
                &format!("/CN=kaws-ca-{}", self.name),
            ]));
        });

        Ok(None)
    }

    fn create_pki_stubs(&self) -> KawsResult {
        let paths = [
            &self.ca_cert_path,
            &self.master_cert_path,
            &self.encrypted_master_key_path,
            &self.node_cert_path,
            &self.encrypted_node_key_path,
        ];

        for path in paths.iter() {
            try!(File::create(path));
        }

        Ok(None)
    }

    fn create_tfvars(&self) -> KawsResult {
        log_wrap!("Creating tfvars file", {
            let mut file = try!(File::create(&self.tfvars_path));

            try!(write!(
                file,
                "\
cluster = \"{}\"
coreos_ami = \"{}\"
domain = \"{}\"
etcd_01_initial_cluster_state = \"new\"
etcd_02_initial_cluster_state = \"new\"
etcd_03_initial_cluster_state = \"new\"
instance_size = \"{}\"
masters_max_size = \"{}\"
masters_min_size = \"{}\"
nodes_max_size = \"{}\"
nodes_min_size = \"{}\"
region = \"{}\"
ssh_key = \"{}\"
version = \"{}\"
zone_id = \"{}\"
",
                self.name,
                self.coreos_ami.expect("AMI should have been required by clap"),
                self.domain.expect("domain should have been required by clap"),
                self.instance_size.expect("instance size should have been required by clap"),
                self.masters_max_size.expect("masters max size should have been required by clap"),
                self.masters_min_size.expect("masters min size should have been required by clap"),
                self.nodes_max_size.expect("nodes max size should have been required by clap"),
                self.nodes_min_size.expect("nodes min size should have been required by clap"),
                self.region,
                self.ssh_key.expect("ssh key should have been required by clap"),
                self.kubernetes_version.expect("k8s version should have been required by clap"),
                self.zone_id.expect("zone ID should have been required by clap"),
            ));
        });

        Ok(None)
    }

    fn encrypt_secrets<'b>(&self, kms_key_id: &'b str) -> KawsResult {
        let mut encryptor = Encryptor::new(
            self.aws_credentials_provider.clone(),
            try!(self.region.parse()),
            kms_key_id,
        );

        log_wrap!("Encrypting Kubernetes certificate authority private key", {
            try!(encryptor.encrypt_file(&self.ca_key_path, &self.encrypted_ca_key_path));
        });

        log_wrap!("Encrypting Kubernetes master private key", {
            try!(encryptor.encrypt_file(&self.master_key_path, &self.encrypted_master_key_path));
        });

        log_wrap!("Encrypting Kubernetes node private key", {
            try!(encryptor.encrypt_file(&self.node_key_path, &self.encrypted_node_key_path));
        });

        Ok(None)
    }
}
<commit_msg>Add message to commit the cluster dir to git after it's created.<commit_after>use std::fs::{create_dir_all, File, remove_file};
use std::io::Write;

use clap::ArgMatches;
use rusoto::ChainProvider;

use aws::credentials_provider;
use encryption::Encryptor;
use error::KawsResult;
use process::execute_child_process;

pub struct Cluster<'a> {
    aws_account_id: Option<&'a str>,
    aws_credentials_provider: ChainProvider,
    ca_cert_path: String,
    ca_key_path: String,
    coreos_ami: Option<&'a str>,
    domain: Option<&'a str>,
    encrypted_ca_key_path: String,
    encrypted_master_key_path: String,
    encrypted_node_key_path: String,
    iam_users: Option<Vec<&'a str>>,
    instance_size: Option<&'a str>,
    kms_master_key_id: Option<&'a str>,
    kms_policy_path: String,
    kubernetes_version: Option<&'a str>,
    master_cert_path: String,
    master_csr_path: String,
    master_key_path: String,
    masters_max_size: Option<&'a str>,
    masters_min_size: Option<&'a str>,
    name: &'a str,
    node_cert_path: String,
    node_csr_path: String,
    node_key_path: String,
    nodes_max_size: Option<&'a str>,
    nodes_min_size: Option<&'a str>,
    openssl_config_path: String,
    region: &'a str,
    ssh_key: Option<&'a str>,
    tfvars_path: String,
    zone_id: Option<&'a str>,
}

impl<'a> Cluster<'a> {
    pub fn new(matches: &'a ArgMatches) -> Self {
        let name =  matches.value_of("cluster").expect("clap should have required cluster");

        Cluster {
            aws_account_id: matches.value_of("aws-account-id"),
            aws_credentials_provider: credentials_provider(
                matches.value_of("aws-credentials-path"),
                matches.value_of("aws-credentials-profile"),
            ),
            ca_cert_path: format!("clusters/{}/ca.pem", name),
            ca_key_path: format!("clusters/{}/ca-key.pem", name),
            coreos_ami: matches.value_of("ami"),
            domain: matches.value_of("domain"),
            encrypted_ca_key_path: format!("clusters/{}/ca-key-encrypted.base64", name),
            encrypted_master_key_path: format!("clusters/{}/master-key-encrypted.base64", name),
            encrypted_node_key_path: format!("clusters/{}/node-key-encrypted.base64", name),
            name: name,
            iam_users: matches.values_of("iam-users").map(|values| values.collect()),
            instance_size: matches.value_of("size"),
            kms_master_key_id: matches.value_of("kms-key"),
            kms_policy_path: format!("clusters/{}/kms-policy.json", name),
            kubernetes_version: matches.value_of("k8s-version"),
            master_cert_path: format!("clusters/{}/master.pem", name),
            master_csr_path: format!("clusters/{}/master.csr", name),
            master_key_path: format!("clusters/{}/master-key.pem", name),
            masters_max_size: matches.value_of("masters-max-size"),
            masters_min_size: matches.value_of("masters-min-size"),
            node_cert_path: format!("clusters/{}/node.pem", name),
            node_csr_path: format!("clusters/{}/node.csr", name),
            node_key_path: format!("clusters/{}/node-key.pem", name),
            nodes_max_size: matches.value_of("nodes-max-size"),
            nodes_min_size: matches.value_of("nodes-min-size"),
            openssl_config_path: format!("clusters/{}/openssl.cnf", name),
            region: matches.value_of("region").expect("clap should have required region"),
            ssh_key: matches.value_of("ssh-key"),
            tfvars_path: format!("clusters/{}/terraform.tfvars", name),
            zone_id: matches.value_of("zone-id"),
        }
    }

    pub fn init(&mut self) -> KawsResult {
        try!(self.create_directories());
        try!(self.create_tfvars());
        try!(self.create_openssl_config());
        try!(self.create_kms_policy());
        try!(self.create_pki_stubs());

        Ok(Some(format!(
            "Cluster \"{name}\" initialized! Commit clusters/{name} to Git.",
            name = self.name,
        )))
    }

    pub fn generate_pki(&mut self) -> KawsResult {
        try!(self.create_ca());
        try!(self.create_master_credentials());
        try!(self.create_node_credentials());
        try!(self.encrypt_secrets(
            self.kms_master_key_id.expect("clap should have required kms-key"),
        ));

        Ok(None)
    }

    fn create_directories(&self) -> KawsResult {
        log_wrap!("Creating directories for the new cluster", {
            try!(create_dir_all(format!("clusters/{}", self.name)));
        });

        Ok(None)
    }

    fn create_kms_policy(&self) -> KawsResult {
        log_wrap!("Creating KMS policy file", {
            let mut file = try!(File::create(&self.kms_policy_path));

            let aws_account_id = self.aws_account_id.expect(
                "AWS account ID should have been required by clap"
            );

            let iam_user_arns = self.iam_users.as_ref().expect(
                "IAM users should have been required by clap"
            ).iter().map(|iam_user| {
                format!("\"arn:aws:iam::{}:user/{}\"", aws_account_id, iam_user)
            }).collect::<Vec<String>>().join(",");

            try!(write!(
                file,
                r#"{{
  "Version": "2012-10-17",
  "Statement": [
    {{
      "Sid": "Enable IAM User Permissions",
      "Effect": "Allow",
      "Principal": {{
        "AWS": [
          "arn:aws:iam::{aws_account_id}:root"
        ]
      }},
      "Action": "kms:*",
      "Resource": "*"
    }},
    {{
      "Sid": "Allow access for Key Administrators",
      "Effect": "Allow",
      "Principal": {{
        "AWS": [
          {iam_user_arns}
        ]
      }},
      "Action": [
        "kms:Create*",
        "kms:Describe*",
        "kms:Enable*",
        "kms:List*",
        "kms:Put*",
        "kms:Update*",
        "kms:Revoke*",
        "kms:Disable*",
        "kms:Get*",
        "kms:Delete*",
        "kms:ScheduleKeyDeletion",
        "kms:CancelKeyDeletion"
      ],
      "Resource": "*"
    }},
    {{
      "Sid": "Allow use of the key",
      "Effect": "Allow",
      "Principal": {{
        "AWS": [
          "arn:aws:iam::{aws_account_id}:role/kaws-k8s-master-{cluster}",
          "arn:aws:iam::{aws_account_id}:role/kaws-k8s-node-{cluster}",
          {iam_user_arns}
        ]
      }},
      "Action": [
        "kms:Encrypt",
        "kms:Decrypt",
        "kms:ReEncrypt*",
        "kms:GenerateDataKey*",
        "kms:DescribeKey"
      ],
      "Resource": "*"
    }},
    {{
      "Sid": "Allow attachment of persistent resources",
      "Effect": "Allow",
      "Principal": {{
        "AWS": [
          "arn:aws:iam::{aws_account_id}:role/kaws-k8s-master-{cluster}",
          "arn:aws:iam::{aws_account_id}:role/kaws-k8s-node-{cluster}",
          {iam_user_arns}
        ]
      }},
      "Action": [
        "kms:CreateGrant",
        "kms:ListGrants",
        "kms:RevokeGrant"
      ],
      "Resource": "*",
      "Condition": {{
        "Bool": {{
          "kms:GrantIsForAWSResource": true
        }}
      }}
    }}
  ]
}}
"#,
                aws_account_id = aws_account_id,
                cluster = self.name,
                iam_user_arns = iam_user_arns,
            ));
        });

        Ok(None)
    }

    fn create_master_credentials(&self) -> KawsResult {
        log_wrap!("Creating Kubernetes master private key", {
            try!(execute_child_process("openssl", &[
                "genrsa",
                "-out",
                &self.master_key_path,
                "2048",
            ]));
        });

        log_wrap!("Creating Kubernetes master certificate signing request", {
            try!(execute_child_process("openssl", &[
                "req",
                "-new",
                "-key",
                &self.master_key_path,
                "-out",
                &self.master_csr_path,
                "-subj",
                &format!("/CN=kaws-master-{}", self.name),
                "-config",
                &self.openssl_config_path,
            ]));
        });

        log_wrap!("Creating Kubernetes master certificate", {
            try!(execute_child_process("openssl", &[
                "x509",
                "-req",
                "-in",
                &self.master_csr_path,
                "-CA",
                &self.ca_cert_path,
                "-CAkey",
                &self.ca_key_path,
                "-CAcreateserial",
                "-out",
                &self.master_cert_path,
                "-days",
                "365",
                "-extensions",
                "v3_req",
                "-extfile",
                &self.openssl_config_path,
            ]));
        });

        log_wrap!("Removing Kubernetes master certificate signing request", {
            try!(remove_file(&self.master_csr_path));
        });

        Ok(None)
    }

    fn create_node_credentials(&self) -> KawsResult {
        log_wrap!("Creating Kubernetes node private key", {
            try!(execute_child_process("openssl", &[
                "genrsa",
                "-out",
                &self.node_key_path,
                "2048",
            ]));
        });

        log_wrap!("Creating Kubernetes node certificate signing request", {
            try!(execute_child_process("openssl", &[
                "req",
                "-new",
                "-key",
                &self.node_key_path,
                "-out",
                &self.node_csr_path,
                "-subj",
                &format!("/CN=kaws-node-{}", self.name),
            ]));
        });

        log_wrap!("Creating Kubernetes node certificate", {
            try!(execute_child_process("openssl", &[
                "x509",
                "-req",
                "-in",
                &self.node_csr_path,
                "-CA",
                &self.ca_cert_path,
                "-CAkey",
                &self.ca_key_path,
                "-CAcreateserial",
                "-out",
                &self.node_cert_path,
                "-days",
                "365",
            ]));
        });

        log_wrap!("Removing Kubernetes node certificate signing request", {
            try!(remove_file(&self.node_csr_path));
        });

        Ok(None)
    }

    fn create_openssl_config(&self) -> KawsResult {
        log_wrap!("Creating OpenSSL config file", {
            let mut file = try!(File::create(&self.openssl_config_path));

            try!(write!(
                file,
                "\
[req]
req_extensions = v3_req
distinguished_name = req_distinguished_name
[req_distinguished_name]
[v3_req]
basicConstraints = CA:FALSE
keyUsage = nonRepudiation, digitalSignature, keyEncipherment
subjectAltName = @alt_names
[alt_names]
DNS.1 = kubernetes
DNS.2 = kubernetes.default
DNS.3 = kubernetes.{}
IP.1 = 10.3.0.1
",
                self.domain.expect("domain should have been required by clap"),
            ));
        });

        Ok(None)
    }

    fn create_ca(&self) -> KawsResult {
        log_wrap!("Creating Kubernetes certificate authority private key", {
            try!(execute_child_process("openssl", &[
                "genrsa",
                "-out",
                &self.ca_key_path,
                "2048",
            ]));
        });

        log_wrap!("Creating Kubernetes certificate authority certificate", {
            try!(execute_child_process("openssl", &[
                "req",
                "-x509",
                "-new",
                "-nodes",
                "-key",
                &self.ca_key_path,
                "-days",
                "10000",
                "-out",
                &self.ca_cert_path,
                "-subj",
                &format!("/CN=kaws-ca-{}", self.name),
            ]));
        });

        Ok(None)
    }

    fn create_pki_stubs(&self) -> KawsResult {
        let paths = [
            &self.ca_cert_path,
            &self.master_cert_path,
            &self.encrypted_master_key_path,
            &self.node_cert_path,
            &self.encrypted_node_key_path,
        ];

        for path in paths.iter() {
            try!(File::create(path));
        }

        Ok(None)
    }

    fn create_tfvars(&self) -> KawsResult {
        log_wrap!("Creating tfvars file", {
            let mut file = try!(File::create(&self.tfvars_path));

            try!(write!(
                file,
                "\
cluster = \"{}\"
coreos_ami = \"{}\"
domain = \"{}\"
etcd_01_initial_cluster_state = \"new\"
etcd_02_initial_cluster_state = \"new\"
etcd_03_initial_cluster_state = \"new\"
instance_size = \"{}\"
masters_max_size = \"{}\"
masters_min_size = \"{}\"
nodes_max_size = \"{}\"
nodes_min_size = \"{}\"
region = \"{}\"
ssh_key = \"{}\"
version = \"{}\"
zone_id = \"{}\"
",
                self.name,
                self.coreos_ami.expect("AMI should have been required by clap"),
                self.domain.expect("domain should have been required by clap"),
                self.instance_size.expect("instance size should have been required by clap"),
                self.masters_max_size.expect("masters max size should have been required by clap"),
                self.masters_min_size.expect("masters min size should have been required by clap"),
                self.nodes_max_size.expect("nodes max size should have been required by clap"),
                self.nodes_min_size.expect("nodes min size should have been required by clap"),
                self.region,
                self.ssh_key.expect("ssh key should have been required by clap"),
                self.kubernetes_version.expect("k8s version should have been required by clap"),
                self.zone_id.expect("zone ID should have been required by clap"),
            ));
        });

        Ok(None)
    }

    fn encrypt_secrets<'b>(&self, kms_key_id: &'b str) -> KawsResult {
        let mut encryptor = Encryptor::new(
            self.aws_credentials_provider.clone(),
            try!(self.region.parse()),
            kms_key_id,
        );

        log_wrap!("Encrypting Kubernetes certificate authority private key", {
            try!(encryptor.encrypt_file(&self.ca_key_path, &self.encrypted_ca_key_path));
        });

        log_wrap!("Encrypting Kubernetes master private key", {
            try!(encryptor.encrypt_file(&self.master_key_path, &self.encrypted_master_key_path));
        });

        log_wrap!("Encrypting Kubernetes node private key", {
            try!(encryptor.encrypt_file(&self.node_key_path, &self.encrypted_node_key_path));
        });

        Ok(None)
    }
}
<|endoftext|>


<commit_before>//! Partially ordered elements with a least upper bound.

/// A bounded partially ordered type supporting joins and meets.
pub trait Lattice : PartialOrd {
    /// The smallest element of the type.
    fn min() -> Self;
    /// The largest element of the type.
    fn max() -> Self;
    /// The smallest element greater than both arguments.
    fn join(&self, &Self) -> Self;
    /// The largest element lesser than both arguments.
    fn meet(&self, &Self) -> Self;
}

use timely::progress::nested::product::Product;

impl<T1: Lattice, T2: Lattice> Lattice for Product<T1, T2> {
    #[inline(always)]
    fn min() -> Self { Product::new(T1::min(), T2::min()) }
    #[inline(always)]
    fn max() -> Self { Product::new(T1::max(), T2::max()) }

    fn join(&self, other: &Product<T1, T2>) -> Product<T1, T2> {
        Product {
            outer: self.outer.join(&other.outer),
            inner: self.inner.join(&other.inner),
        }
    }
    fn meet(&self, other: &Product<T1, T2>) -> Product<T1, T2> {
        Product {
            outer: self.outer.meet(&other.outer),
            inner: self.inner.meet(&other.inner),
        }
    }
}

use timely::progress::timestamp::RootTimestamp;

impl Lattice for RootTimestamp {
    fn min() -> RootTimestamp { RootTimestamp }
    fn max() -> RootTimestamp { RootTimestamp }
    fn join(&self, _: &RootTimestamp) -> RootTimestamp { RootTimestamp }
    fn meet(&self, _: &RootTimestamp) -> RootTimestamp { RootTimestamp }
}

impl Lattice for u64 {
    fn min() -> u64 { u64::min_value() }
    fn max() -> u64 { u64::max_value() }
    fn join(&self, other: &u64) -> u64 { ::std::cmp::max(*self, *other) }
    fn meet(&self, other: &u64) -> u64 { ::std::cmp::min(*self, *other) }
}

impl Lattice for u32 {
    fn min() -> u32 { u32::min_value() }
    fn max() -> u32 { u32::max_value() }
    fn join(&self, other: &u32) -> u32 { ::std::cmp::max(*self, *other) }
    fn meet(&self, other: &u32) -> u32 { ::std::cmp::min(*self, *other) }
}

impl Lattice for i32 {
    fn min() -> i32 { i32::min_value() }
    fn max() -> i32 { i32::max_value() }
    fn join(&self, other: &i32) -> i32 { ::std::cmp::max(*self, *other) }
    fn meet(&self, other: &i32) -> i32 { ::std::cmp::min(*self, *other) }
}

impl Lattice for () {
    fn min() -> () { () }
    fn max() -> () { () }
    fn join(&self, _other: &()) -> () { () }
    fn meet(&self, _other: &()) -> () { () }
}

/// Extends `vector` to contain all joins of pairs of elements.
pub fn close_under_join<T: Lattice>(vector: &mut Vec<T>) {
    // compares each element to those elements after it.
    let mut first = 0;
    while first < vector.len() {
        let mut next = first + 1;
        while next < vector.len() {
            let lub = vector[first].join(&vector[next]);
            if !vector.contains(&lub) {
                vector.push(lub);
            }

            next += 1;
        }

        first += 1;
    }
}
<commit_msg>added advance_by method to Lattice<commit_after>//! Partially ordered elements with a least upper bound.

/// A bounded partially ordered type supporting joins and meets.
pub trait Lattice : PartialOrd {
    /// The smallest element of the type.
    fn min() -> Self;
    /// The largest element of the type.
    fn max() -> Self;
    /// The smallest element greater than both arguments.
    fn join(&self, &Self) -> Self;
    /// The largest element lesser than both arguments.
    fn meet(&self, &Self) -> Self;

    /// Advances self to the largest time indistinguishable under `frontier`.
    // TODO : add an example here demonstrating the result invariant.
    fn advance_by(&self, frontier: &[Self]) -> Option<Self> where Self: Sized{
        if frontier.len() > 0 {
            let mut result = self.join(&frontier[0]);
            for f in &frontier[1..] {
                result = result.meet(&self.join(f));
            }
            Some(result)
        }
        else {
            None
        }
    }
}

use timely::progress::nested::product::Product;

impl<T1: Lattice, T2: Lattice> Lattice for Product<T1, T2> {
    #[inline(always)]
    fn min() -> Self { Product::new(T1::min(), T2::min()) }
    #[inline(always)]
    fn max() -> Self { Product::new(T1::max(), T2::max()) }

    fn join(&self, other: &Product<T1, T2>) -> Product<T1, T2> {
        Product {
            outer: self.outer.join(&other.outer),
            inner: self.inner.join(&other.inner),
        }
    }
    fn meet(&self, other: &Product<T1, T2>) -> Product<T1, T2> {
        Product {
            outer: self.outer.meet(&other.outer),
            inner: self.inner.meet(&other.inner),
        }
    }
}

use timely::progress::timestamp::RootTimestamp;

impl Lattice for RootTimestamp {
    fn min() -> RootTimestamp { RootTimestamp }
    fn max() -> RootTimestamp { RootTimestamp }
    fn join(&self, _: &RootTimestamp) -> RootTimestamp { RootTimestamp }
    fn meet(&self, _: &RootTimestamp) -> RootTimestamp { RootTimestamp }
}

impl Lattice for u64 {
    fn min() -> u64 { u64::min_value() }
    fn max() -> u64 { u64::max_value() }
    fn join(&self, other: &u64) -> u64 { ::std::cmp::max(*self, *other) }
    fn meet(&self, other: &u64) -> u64 { ::std::cmp::min(*self, *other) }
}

impl Lattice for u32 {
    fn min() -> u32 { u32::min_value() }
    fn max() -> u32 { u32::max_value() }
    fn join(&self, other: &u32) -> u32 { ::std::cmp::max(*self, *other) }
    fn meet(&self, other: &u32) -> u32 { ::std::cmp::min(*self, *other) }
}

impl Lattice for i32 {
    fn min() -> i32 { i32::min_value() }
    fn max() -> i32 { i32::max_value() }
    fn join(&self, other: &i32) -> i32 { ::std::cmp::max(*self, *other) }
    fn meet(&self, other: &i32) -> i32 { ::std::cmp::min(*self, *other) }
}

impl Lattice for () {
    fn min() -> () { () }
    fn max() -> () { () }
    fn join(&self, _other: &()) -> () { () }
    fn meet(&self, _other: &()) -> () { () }
}

/// Extends `vector` to contain all joins of pairs of elements.
pub fn close_under_join<T: Lattice>(vector: &mut Vec<T>) {
    // compares each element to those elements after it.
    let mut first = 0;
    while first < vector.len() {
        let mut next = first + 1;
        while next < vector.len() {
            let lub = vector[first].join(&vector[next]);
            if !vector.contains(&lub) {
                vector.push(lub);
            }

            next += 1;
        }

        first += 1;
    }
}
<|endoftext|>


<commit_before>// Copyright 2015 Nicholas Bishop
//
// Closest-point and point-distance methods adapted from "Real-Time
// Collision Detection" by Christer Ericson, published by Morgan
// Kaufmann Publishers, Copyright 2005 Elsevier 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.

use vector::{distance3, dot3, lerp3, Vec3f};
use range::Rangef;
use std::ops::Add;

/// Line segment between two points
#[derive(Clone, Copy, Debug, PartialEq)]
pub struct Segment3f {
    pub start: Vec3f,
    pub end: Vec3f
}

impl Segment3f {
    /// Create a segment between two points
    pub fn new(start: &Vec3f, end: &Vec3f) -> Segment3f {
        Segment3f {
            start: *start,
            end: *end
        }
    }

    /// Get vector from start to end (not normalized)
    pub fn to_vec3f(&self) -> Vec3f {
        self.end - self.start
    }

    /// Length of the line segment
    pub fn length(&self) -> f32 {
        distance3(self.start, self.end)
    }

    /// Convert a distance in coordinate space to a distance in the
    /// line segment's parametric space. The sign of the input is
    /// kept.
    pub fn distance_to_parametric_delta(&self, distance: f32) -> f32 {
        distance / self.length()
    }

    /// Convert a parametric delta to coordinate space. The sign of
    /// the input is kept.
    pub fn distance_from_parametric_delta(&self, delta: f32) -> f32 {
        delta * self.length()
    }

    /// Linearly interpolate between the segment's endpoints by the
    /// factor `t`. When `t` is zero the result is `self.start`, and
    /// when `t` is one the result is `self.end`. The range of `t` is
    /// not clamped.
    pub fn point_from_parametric(&self, t: f32) -> Vec3f {
        lerp3(&self.start, &self.end, t)
    }

    /// Treat the range's start and end as parametric coords. Use
    /// `point_from_parametric` to interpolate the range into a new
    /// segment. The output segment is not clamped.
    pub fn segment_from_parametric_range(&self, r: Rangef) -> Segment3f {
        Segment3f::new(&self.point_from_parametric(r.min),
                       &self.point_from_parametric(r.max))
    }

    /// Project another segment onto `self`. The result is a
    /// parametric range of `self` clamped to [0, 1].
    pub fn project_segment_as_range(&self, other: &Segment3f) -> Rangef {
        Rangef::from_sorting(self.closest_point_to_point(&other.start).0,
                             self.closest_point_to_point(&other.end).0)
    }

    /// Return the squared distance to the input `point`.
    ///
    /// Adapted from "Real-Time Collision Detection" by Christer
    /// Ericson, published by Morgan Kaufmann Publishers, Copyright
    /// 2005 Elsevier Inc
    pub fn point_distance_squared(&self, point: Vec3f) -> f32 {
        let ab = self.end - self.start;
        let ac = point - self.start;
        let bc = point - self.end;
        let e = dot3(ac, ab);
        // Handle cases where c projects outside ab
        if e <= 0.0 {
            dot3(ac, ac)
        }
        else {
            let f = dot3(ab, ab);
            if e >= f {
                dot3(bc, bc)
            }
            else {
                // Handle cases where c projects onto ab
                dot3(ac, ac) - e * e / f
            }
        }
    }

    /// Find the point on the segment closest to the input point. The
    /// return value contains both the parametric and actual location
    /// of the closest point.
    ///
    /// Adapted from "Real-Time Collision Detection" by Christer
    /// Ericson, published by Morgan Kaufmann Publishers, Copyright
    /// 2005 Elsevier Inc
    pub fn closest_point_to_point(&self, point: &Vec3f) -> (f32, Vec3f) {
        let a = self.start;
        let b = self.end;

        let ab = b - a;
        // Project point onto ab, but deferring divide by dot3(ab, ab)
        let t = dot3(*point - a, ab);
        if t <= 0.0f32 {
            // point projects outside the [a,b] interval, on the a
            // side; clamp to a
            (0.0f32, a)
        } else {
            // Always nonnegative since denom = ||ab|| ∧ 2
            let denom = dot3(ab, ab);
            if t >= denom {
                // point projects outside the [a,b] interval, on the b
                // side; clamp to b
                (1.0f32, b)
            } else {
                // point projects inside the [a,b] interval; must do
                // deferred divide now
                let t = t / denom;
                (t, a + ab * t)
            }
        }
    }
}

impl Add<Vec3f> for Segment3f {
    type Output = Segment3f;
    fn add(self, v: Vec3f) -> Self::Output {
        Segment3f::new(&(self.start + v), &(self.end + v))
    }
}

#[cfg(test)]
mod test {
    use super::*;
    use range::Rangef;
    use vector::vec3f;

    fn make_seg(a: i32, b: i32) -> Segment3f {
        Segment3f::new(&vec3f(a, 0, 0), &vec3f(b, 0, 0))
    }

    #[test]
    fn test_to_vec3f() {
        assert_eq!(Segment3f::new(&vec3f(0, 0, 0),
                                  &vec3f(2, 3, 4)).to_vec3f(),
                   vec3f(2, 3, 4));

    }

    #[test]
    fn test_segment_length() {
        let s = Segment3f::new(&vec3f(0, 0, 0),
                               &vec3f(0, 0, 9));
        assert_eq!(s.length(), 9.0);
    }

    #[test]
    fn test_segment_distance_conversion() {
        let s = Segment3f::new(&vec3f(0, 1, 0),
                               &vec3f(0, 7, 0));

        let inputs = [
            (0.0f32, 0.0f32),
            (6.0, 1.0),
            (12.0, 2.0),
            (-3.0, -0.5)];

        for &(distance, delta) in inputs.iter() {
            assert_eq!(s.distance_to_parametric_delta(distance), delta);
            assert_eq!(s.distance_from_parametric_delta(delta), distance);
        }
    }

    #[test]
    fn test_point_from_parametric() {
        let s = Segment3f::new(&vec3f(0, 0, -1),
                               &vec3f(0, 0, 3));
        assert_eq!(s.point_from_parametric(0.0), vec3f(0, 0, -1));
        assert_eq!(s.point_from_parametric(1.0), vec3f(0, 0, 3));
        assert_eq!(s.point_from_parametric(0.5), vec3f(0, 0, 1));
    }

    #[test]
    fn test_segment_from_parametric_range() {
        let s = make_seg(0, 4);
        assert_eq!(s.segment_from_parametric_range(Rangef::new(0.0, 1.0)),
                   make_seg(0, 4));
        assert_eq!(s.segment_from_parametric_range(Rangef::new(0.25, 0.75)),
                   make_seg(1, 3));
    }

    #[test]
    fn test_segment_closest_point_to_point() {
        let s = Segment3f::new(&vec3f(2, 0, 0),
                               &vec3f(3, 0, 0));
        assert_eq!(s.closest_point_to_point(&vec3f(1, 0, 0)),
                   (0.0, vec3f(2, 0, 0)));
        assert_eq!(s.closest_point_to_point(&vec3f(4, 0, 0)),
                   (1.0, vec3f(3, 0, 0)));
        assert_eq!(s.closest_point_to_point(&vec3f(2.5, 1, 0)),
                   (0.5, vec3f(2.5, 0, 0)));
    }

    #[test]
    fn test_project_segment_as_range() {
        let s = make_seg(0, 4);
        assert_eq!(s.project_segment_as_range(&make_seg(0, 4)),
                   Rangef::new(0.0, 1.0));
        assert_eq!(s.project_segment_as_range(&make_seg(-1, 5)),
                   Rangef::new(0.0, 1.0));
        assert_eq!(s.project_segment_as_range(&make_seg(-1, -1)),
                   Rangef::new(0.0, 0.0));
        assert_eq!(s.project_segment_as_range(&make_seg(1, 3)),
                   Rangef::new(0.25, 0.75));
        assert_eq!(s.project_segment_as_range(&make_seg(3, 1)),
                   Rangef::new(0.25, 0.75));
    }
}
<commit_msg>Add Segment3f::point_distance<commit_after>// Copyright 2015 Nicholas Bishop
//
// Closest-point and point-distance methods adapted from "Real-Time
// Collision Detection" by Christer Ericson, published by Morgan
// Kaufmann Publishers, Copyright 2005 Elsevier 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.

use vector::{distance3, dot3, lerp3, Vec3f};
use range::Rangef;
use std::ops::Add;

/// Line segment between two points
#[derive(Clone, Copy, Debug, PartialEq)]
pub struct Segment3f {
    pub start: Vec3f,
    pub end: Vec3f
}

impl Segment3f {
    /// Create a segment between two points
    pub fn new(start: &Vec3f, end: &Vec3f) -> Segment3f {
        Segment3f {
            start: *start,
            end: *end
        }
    }

    /// Get vector from start to end (not normalized)
    pub fn to_vec3f(&self) -> Vec3f {
        self.end - self.start
    }

    /// Length of the line segment
    pub fn length(&self) -> f32 {
        distance3(self.start, self.end)
    }

    /// Convert a distance in coordinate space to a distance in the
    /// line segment's parametric space. The sign of the input is
    /// kept.
    pub fn distance_to_parametric_delta(&self, distance: f32) -> f32 {
        distance / self.length()
    }

    /// Convert a parametric delta to coordinate space. The sign of
    /// the input is kept.
    pub fn distance_from_parametric_delta(&self, delta: f32) -> f32 {
        delta * self.length()
    }

    /// Linearly interpolate between the segment's endpoints by the
    /// factor `t`. When `t` is zero the result is `self.start`, and
    /// when `t` is one the result is `self.end`. The range of `t` is
    /// not clamped.
    pub fn point_from_parametric(&self, t: f32) -> Vec3f {
        lerp3(&self.start, &self.end, t)
    }

    /// Treat the range's start and end as parametric coords. Use
    /// `point_from_parametric` to interpolate the range into a new
    /// segment. The output segment is not clamped.
    pub fn segment_from_parametric_range(&self, r: Rangef) -> Segment3f {
        Segment3f::new(&self.point_from_parametric(r.min),
                       &self.point_from_parametric(r.max))
    }

    /// Project another segment onto `self`. The result is a
    /// parametric range of `self` clamped to [0, 1].
    pub fn project_segment_as_range(&self, other: &Segment3f) -> Rangef {
        Rangef::from_sorting(self.closest_point_to_point(&other.start).0,
                             self.closest_point_to_point(&other.end).0)
    }

    /// Return the squared distance from this segment to the input
    /// `point`.
    ///
    /// Adapted from "Real-Time Collision Detection" by Christer
    /// Ericson, published by Morgan Kaufmann Publishers, Copyright
    /// 2005 Elsevier Inc
    pub fn point_distance_squared(&self, point: Vec3f) -> f32 {
        let ab = self.end - self.start;
        let ac = point - self.start;
        let bc = point - self.end;
        let e = dot3(ac, ab);
        // Handle cases where c projects outside ab
        if e <= 0.0 {
            dot3(ac, ac)
        }
        else {
            let f = dot3(ab, ab);
            if e >= f {
                dot3(bc, bc)
            }
            else {
                // Handle cases where c projects onto ab
                dot3(ac, ac) - e * e / f
            }
        }
    }

    /// Return the distance from this segment to the input `point`.
    pub fn point_distance(&self, point: Vec3f) -> f32 {
        self.point_distance_squared(point).sqrt()
    }

    /// Find the point on the segment closest to the input point. The
    /// return value contains both the parametric and actual location
    /// of the closest point.
    ///
    /// Adapted from "Real-Time Collision Detection" by Christer
    /// Ericson, published by Morgan Kaufmann Publishers, Copyright
    /// 2005 Elsevier Inc
    pub fn closest_point_to_point(&self, point: &Vec3f) -> (f32, Vec3f) {
        let a = self.start;
        let b = self.end;

        let ab = b - a;
        // Project point onto ab, but deferring divide by dot3(ab, ab)
        let t = dot3(*point - a, ab);
        if t <= 0.0f32 {
            // point projects outside the [a,b] interval, on the a
            // side; clamp to a
            (0.0f32, a)
        } else {
            // Always nonnegative since denom = ||ab|| ∧ 2
            let denom = dot3(ab, ab);
            if t >= denom {
                // point projects outside the [a,b] interval, on the b
                // side; clamp to b
                (1.0f32, b)
            } else {
                // point projects inside the [a,b] interval; must do
                // deferred divide now
                let t = t / denom;
                (t, a + ab * t)
            }
        }
    }
}

impl Add<Vec3f> for Segment3f {
    type Output = Segment3f;
    fn add(self, v: Vec3f) -> Self::Output {
        Segment3f::new(&(self.start + v), &(self.end + v))
    }
}

#[cfg(test)]
mod test {
    use super::*;
    use range::Rangef;
    use vector::vec3f;

    fn make_seg(a: i32, b: i32) -> Segment3f {
        Segment3f::new(&vec3f(a, 0, 0), &vec3f(b, 0, 0))
    }

    #[test]
    fn test_to_vec3f() {
        assert_eq!(Segment3f::new(&vec3f(0, 0, 0),
                                  &vec3f(2, 3, 4)).to_vec3f(),
                   vec3f(2, 3, 4));

    }

    #[test]
    fn test_segment_length() {
        let s = Segment3f::new(&vec3f(0, 0, 0),
                               &vec3f(0, 0, 9));
        assert_eq!(s.length(), 9.0);
    }

    #[test]
    fn test_segment_distance_conversion() {
        let s = Segment3f::new(&vec3f(0, 1, 0),
                               &vec3f(0, 7, 0));

        let inputs = [
            (0.0f32, 0.0f32),
            (6.0, 1.0),
            (12.0, 2.0),
            (-3.0, -0.5)];

        for &(distance, delta) in inputs.iter() {
            assert_eq!(s.distance_to_parametric_delta(distance), delta);
            assert_eq!(s.distance_from_parametric_delta(delta), distance);
        }
    }

    #[test]
    fn test_point_from_parametric() {
        let s = Segment3f::new(&vec3f(0, 0, -1),
                               &vec3f(0, 0, 3));
        assert_eq!(s.point_from_parametric(0.0), vec3f(0, 0, -1));
        assert_eq!(s.point_from_parametric(1.0), vec3f(0, 0, 3));
        assert_eq!(s.point_from_parametric(0.5), vec3f(0, 0, 1));
    }

    #[test]
    fn test_segment_from_parametric_range() {
        let s = make_seg(0, 4);
        assert_eq!(s.segment_from_parametric_range(Rangef::new(0.0, 1.0)),
                   make_seg(0, 4));
        assert_eq!(s.segment_from_parametric_range(Rangef::new(0.25, 0.75)),
                   make_seg(1, 3));
    }

    #[test]
    fn test_segment_closest_point_to_point() {
        let s = Segment3f::new(&vec3f(2, 0, 0),
                               &vec3f(3, 0, 0));
        assert_eq!(s.closest_point_to_point(&vec3f(1, 0, 0)),
                   (0.0, vec3f(2, 0, 0)));
        assert_eq!(s.closest_point_to_point(&vec3f(4, 0, 0)),
                   (1.0, vec3f(3, 0, 0)));
        assert_eq!(s.closest_point_to_point(&vec3f(2.5, 1, 0)),
                   (0.5, vec3f(2.5, 0, 0)));
    }

    #[test]
    fn test_project_segment_as_range() {
        let s = make_seg(0, 4);
        assert_eq!(s.project_segment_as_range(&make_seg(0, 4)),
                   Rangef::new(0.0, 1.0));
        assert_eq!(s.project_segment_as_range(&make_seg(-1, 5)),
                   Rangef::new(0.0, 1.0));
        assert_eq!(s.project_segment_as_range(&make_seg(-1, -1)),
                   Rangef::new(0.0, 0.0));
        assert_eq!(s.project_segment_as_range(&make_seg(1, 3)),
                   Rangef::new(0.25, 0.75));
        assert_eq!(s.project_segment_as_range(&make_seg(3, 1)),
                   Rangef::new(0.25, 0.75));
    }
}
<|endoftext|>