<commit_before>use std::sync::Arc;

use smallvec::SmallVec;

use crate::crlf;
use crate::rope::Rope;
use crate::tree::{Node, NodeChildren, NodeText, MAX_BYTES, MAX_CHILDREN, MIN_BYTES};

/// An efficient incremental `Rope` builder.
///
/// This is used to efficiently build ropes from sequences of text
/// chunks.  It is useful for creating ropes from:
///
/// - ...large text files, without pre-loading their entire contents into
///   memory (but see [`from_reader()`](Rope::from_reader) for a convenience
///   function that does this for casual use-cases).
/// - ...streaming data sources.
/// - ...non-utf8 text data, doing the encoding conversion incrementally
///   as you go.
///
/// Unlike repeatedly calling `Rope::insert()` on the end of a rope,
/// this API runs in time linear to the amount of data fed to it, and
/// is overall much faster.
///
/// # Example
/// ```
/// # use ropey::RopeBuilder;
/// #
/// let mut builder = RopeBuilder::new();
///
/// builder.append("Hello ");
/// builder.append("world!\n");
/// builder.append("How's ");
/// builder.append("it goin");
/// builder.append("g?");
///
/// let rope = builder.finish();
///
/// assert_eq!(rope, "Hello world!\nHow's it going?");
/// ```
#[derive(Debug, Clone, Default)]
pub struct RopeBuilder {
    stack: SmallVec<[Arc<Node>; 4]>,
    buffer: String,
    last_chunk_len_bytes: usize,
}

impl RopeBuilder {
    /// Creates a new RopeBuilder, ready for input.
    pub fn new() -> Self {
        RopeBuilder {
            stack: {
                let mut stack = SmallVec::new();
                stack.push(Arc::new(Node::new()));
                stack
            },
            buffer: String::new(),
            last_chunk_len_bytes: 0,
        }
    }

    /// Appends `chunk` to the end of the in-progress `Rope`.
    ///
    /// Call this method repeatedly to incrementally build up a
    /// `Rope`.  The passed text chunk can be as large or small as
    /// desired, but larger chunks are more efficient.
    ///
    /// `chunk` must be valid utf8 text.
    pub fn append(&mut self, chunk: &str) {
        self.append_internal(chunk, false);
    }

    /// NOT PART OF THE PUBLIC API (hidden from docs for a reason!).
    ///
    /// Appends `contents` to the in-progress rope as a single leaf
    /// node (chunk).  This is useful for building ropes with specific
    /// chunk configurations for testing purposes.  It will happily append
    /// both empty and more-than-max-size chunks.
    ///
    /// This makes no attempt to be consistent with the standard `append()`
    /// method, and should not be used in conjunction with it.
    #[doc(hidden)]
    pub fn _append_chunk(&mut self, contents: &str) {
        self.append_leaf_node(Arc::new(Node::Leaf(NodeText::from_str(contents))));
    }

    /// Finishes the build, and returns the `Rope`.
    ///
    /// Note: this method consumes the builder.  If you want to continue
    /// building other ropes with the same prefix, you can clone the builder
    /// before calling `finish()`.
    pub fn finish(mut self) -> Rope {
        // Append the last leaf
        self.append_internal("", true);
        self.finish_internal()
    }

    /// Builds a rope all at once from a single string slice.
    ///
    /// This avoids the creation and use of the internal buffer.  This is
    /// for internal use only, because the public-facing API has
    /// Rope::from_str(), which actually uses this for its implementation.
    pub(crate) fn build_at_once(mut self, chunk: &str) -> Rope {
        self.append_internal(chunk, true);
        self.finish_internal()
    }

    //-----------------------------------------------------------------

    // Internal workings of `append()`.
    fn append_internal(&mut self, chunk: &str, is_last_chunk: bool) {
        let mut chunk = chunk;

        // Repeatedly chop text off the end of the input, creating
        // leaf nodes out of them and appending them to the tree.
        while !chunk.is_empty() || (!self.buffer.is_empty() && is_last_chunk) {
            // Get the text for the next leaf
            let (leaf_text, remainder) = self.get_next_leaf_text(chunk, is_last_chunk);
            chunk = remainder;

            self.last_chunk_len_bytes = chunk.len();

            // Append the leaf to the rope
            match leaf_text {
                NextText::None => break,
                NextText::UseBuffer => {
                    let leaf_text = NodeText::from_str(&self.buffer);
                    self.append_leaf_node(Arc::new(Node::Leaf(leaf_text)));
                    self.buffer.clear();
                }
                NextText::String(s) => {
                    self.append_leaf_node(Arc::new(Node::Leaf(NodeText::from_str(s))));
                }
            }
        }
    }

    // Internal workings of `finish()`.
    fn finish_internal(mut self) -> Rope {
        // Zip up all the remaining nodes on the stack
        let mut stack_idx = self.stack.len() - 1;
        while stack_idx >= 1 {
            let node = self.stack.pop().unwrap();
            if let Node::Internal(ref mut children) = *Arc::make_mut(&mut self.stack[stack_idx - 1])
            {
                children.push((node.text_info(), node));
            } else {
                unreachable!();
            }
            stack_idx -= 1;
        }

        // Get root and fix any right-side nodes with too few children or.
        let mut root = self.stack.pop().unwrap();
        Arc::make_mut(&mut root).zip_fix_right();

        // Create the rope, make sure it's well-formed, and return it.
        let mut rope = Rope { root: root };
        if self.last_chunk_len_bytes < MIN_BYTES && self.last_chunk_len_bytes != rope.len_bytes() {
            // Merge the last chunk if it was too small.
            let idx =
                rope.len_chars() - rope.byte_to_char(rope.len_bytes() - self.last_chunk_len_bytes);
            Arc::make_mut(&mut rope.root).fix_tree_seam(idx);
        }
        rope.pull_up_singular_nodes();
        return rope;
    }

    // Returns (next_leaf_text, remaining_text)
    #[inline(always)]
    fn get_next_leaf_text<'a>(
        &mut self,
        text: &'a str,
        is_last_chunk: bool,
    ) -> (NextText<'a>, &'a str) {
        assert!(
            self.buffer.len() < MAX_BYTES,
            "RopeBuilder: buffer is already full when receiving a chunk! \
             This should never happen!",
        );

        // Simplest case: empty buffer and enough in `text` for a full
        // chunk, so just chop a chunk off from `text` and use that.
        if self.buffer.is_empty() && text.len() >= MAX_BYTES {
            let split_idx = crlf::find_good_split(
                MAX_BYTES.min(text.len() - 1), // - 1 to avoid CRLF split.
                text.as_bytes(),
                true,
            );
            return (NextText::String(&text[..split_idx]), &text[split_idx..]);
        }
        // If the buffer + `text` is enough for a full chunk, push enough
        // of `text` onto the buffer to fill it and use that.
        else if (text.len() + self.buffer.len()) >= MAX_BYTES {
            let mut split_idx =
                crlf::find_good_split(MAX_BYTES - self.buffer.len(), text.as_bytes(), true);
            if split_idx == text.len() && text.as_bytes()[text.len() - 1] == 0x0D {
                // Avoid CRLF split.
                split_idx -= 1;
            };
            self.buffer.push_str(&text[..split_idx]);
            return (NextText::UseBuffer, &text[split_idx..]);
        }
        // If we don't have enough text for a full chunk.
        else {
            // If it's our last chunk, wrap it all up!
            if is_last_chunk {
                if self.buffer.is_empty() {
                    return if text.is_empty() {
                        (NextText::None, "")
                    } else {
                        (NextText::String(text), "")
                    };
                } else {
                    self.buffer.push_str(text);
                    return (NextText::UseBuffer, "");
                }
            }
            // Otherwise, just push to the buffer.
            else {
                self.buffer.push_str(text);
                return (NextText::None, "");
            }
        }
    }

    fn append_leaf_node(&mut self, leaf: Arc<Node>) {
        let last = self.stack.pop().unwrap();
        match *last {
            Node::Leaf(_) => {
                if last.leaf_text().is_empty() {
                    self.stack.push(leaf);
                } else {
                    let mut children = NodeChildren::new();
                    children.push((last.text_info(), last));
                    children.push((leaf.text_info(), leaf));
                    self.stack.push(Arc::new(Node::Internal(children)));
                }
            }

            Node::Internal(_) => {
                self.stack.push(last);
                let mut left = leaf;
                let mut stack_idx = (self.stack.len() - 1) as isize;
                loop {
                    if stack_idx < 0 {
                        // We're above the root, so do a root split.
                        let mut children = NodeChildren::new();
                        children.push((left.text_info(), left));
                        self.stack.insert(0, Arc::new(Node::Internal(children)));
                        break;
                    } else if self.stack[stack_idx as usize].child_count() < (MAX_CHILDREN - 1) {
                        // There's room to add a child, so do that.
                        Arc::make_mut(&mut self.stack[stack_idx as usize])
                            .children_mut()
                            .push((left.text_info(), left));
                        break;
                    } else {
                        // Not enough room to fit a child, so split.
                        left = Arc::new(Node::Internal(
                            Arc::make_mut(&mut self.stack[stack_idx as usize])
                                .children_mut()
                                .push_split((left.text_info(), left)),
                        ));
                        std::mem::swap(&mut left, &mut self.stack[stack_idx as usize]);
                        stack_idx -= 1;
                    }
                }
            }
        }
    }
}

enum NextText<'a> {
    None,
    UseBuffer,
    String(&'a str),
}

//===========================================================================

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

    // 127 bytes, 103 chars, 4 lines
    const TEXT: &str = "Hello there!  How're you doing?\r\nIt's \
                        a fine day, isn't it?\r\nAren't you glad \
                        we're alive?\r\nこんにちは、みんなさん！";

    #[test]
    fn rope_builder_01() {
        let mut b = RopeBuilder::new();

        b.append("Hello there!  How're you doing?\r");
        b.append("\nIt's a fine ");
        b.append("d");
        b.append("a");
        b.append("y,");
        b.append(" ");
        b.append("isn't it?");
        b.append("\r");
        b.append("\nAren't you ");
        b.append("glad we're alive?\r");
        b.append("\n");
        b.append("こんにち");
        b.append("は、みんなさ");
        b.append("ん！");

        let r = b.finish();

        assert_eq!(r, TEXT);

        r.assert_integrity();
        r.assert_invariants();
    }
}
<commit_msg>Add `_finish_no_fix()` to RopeBuilder, for use in internal testing.<commit_after>use std::sync::Arc;

use smallvec::SmallVec;

use crate::crlf;
use crate::rope::Rope;
use crate::tree::{Node, NodeChildren, NodeText, MAX_BYTES, MAX_CHILDREN, MIN_BYTES};

/// An efficient incremental `Rope` builder.
///
/// This is used to efficiently build ropes from sequences of text
/// chunks.  It is useful for creating ropes from:
///
/// - ...large text files, without pre-loading their entire contents into
///   memory (but see [`from_reader()`](Rope::from_reader) for a convenience
///   function that does this for casual use-cases).
/// - ...streaming data sources.
/// - ...non-utf8 text data, doing the encoding conversion incrementally
///   as you go.
///
/// Unlike repeatedly calling `Rope::insert()` on the end of a rope,
/// this API runs in time linear to the amount of data fed to it, and
/// is overall much faster.
///
/// # Example
/// ```
/// # use ropey::RopeBuilder;
/// #
/// let mut builder = RopeBuilder::new();
///
/// builder.append("Hello ");
/// builder.append("world!\n");
/// builder.append("How's ");
/// builder.append("it goin");
/// builder.append("g?");
///
/// let rope = builder.finish();
///
/// assert_eq!(rope, "Hello world!\nHow's it going?");
/// ```
#[derive(Debug, Clone, Default)]
pub struct RopeBuilder {
    stack: SmallVec<[Arc<Node>; 4]>,
    buffer: String,
    last_chunk_len_bytes: usize,
}

impl RopeBuilder {
    /// Creates a new RopeBuilder, ready for input.
    pub fn new() -> Self {
        RopeBuilder {
            stack: {
                let mut stack = SmallVec::new();
                stack.push(Arc::new(Node::new()));
                stack
            },
            buffer: String::new(),
            last_chunk_len_bytes: 0,
        }
    }

    /// Appends `chunk` to the end of the in-progress `Rope`.
    ///
    /// Call this method repeatedly to incrementally build up a
    /// `Rope`.  The passed text chunk can be as large or small as
    /// desired, but larger chunks are more efficient.
    ///
    /// `chunk` must be valid utf8 text.
    pub fn append(&mut self, chunk: &str) {
        self.append_internal(chunk, false);
    }

    /// Finishes the build, and returns the `Rope`.
    ///
    /// Note: this method consumes the builder.  If you want to continue
    /// building other ropes with the same prefix, you can clone the builder
    /// before calling `finish()`.
    pub fn finish(mut self) -> Rope {
        // Append the last leaf
        self.append_internal("", true);
        self.finish_internal(true)
    }

    /// Builds a rope all at once from a single string slice.
    ///
    /// This avoids the creation and use of the internal buffer.  This is
    /// for internal use only, because the public-facing API has
    /// Rope::from_str(), which actually uses this for its implementation.
    pub(crate) fn build_at_once(mut self, chunk: &str) -> Rope {
        self.append_internal(chunk, true);
        self.finish_internal(true)
    }

    /// NOT PART OF THE PUBLIC API (hidden from docs for a reason!).
    ///
    /// Appends `contents` to the in-progress rope as a single leaf
    /// node (chunk).  This is useful for building ropes with specific
    /// chunk configurations for testing purposes.  It will happily append
    /// both empty and more-than-max-size chunks.
    ///
    /// This makes no attempt to be consistent with the standard `append()`
    /// method, and should not be used in conjunction with it.
    #[doc(hidden)]
    pub fn _append_chunk(&mut self, contents: &str) {
        self.append_leaf_node(Arc::new(Node::Leaf(NodeText::from_str(contents))));
    }

    /// NOT PART OF THE PUBLIC API (hidden from docs for a reason!).
    ///
    /// Finishes the build without doing any tree fixing to adhere
    /// to the btree invariants. To be used with `_append_chunk()` to
    /// construct ropes with specific chunk boundaries for testing.
    #[doc(hidden)]
    pub fn _finish_no_fix(self) -> Rope {
        self.finish_internal(false)
    }

    //-----------------------------------------------------------------

    // Internal workings of `append()`.
    fn append_internal(&mut self, chunk: &str, is_last_chunk: bool) {
        let mut chunk = chunk;

        // Repeatedly chop text off the end of the input, creating
        // leaf nodes out of them and appending them to the tree.
        while !chunk.is_empty() || (!self.buffer.is_empty() && is_last_chunk) {
            // Get the text for the next leaf
            let (leaf_text, remainder) = self.get_next_leaf_text(chunk, is_last_chunk);
            chunk = remainder;

            self.last_chunk_len_bytes = chunk.len();

            // Append the leaf to the rope
            match leaf_text {
                NextText::None => break,
                NextText::UseBuffer => {
                    let leaf_text = NodeText::from_str(&self.buffer);
                    self.append_leaf_node(Arc::new(Node::Leaf(leaf_text)));
                    self.buffer.clear();
                }
                NextText::String(s) => {
                    self.append_leaf_node(Arc::new(Node::Leaf(NodeText::from_str(s))));
                }
            }
        }
    }

    // Internal workings of `finish()`.
    //
    // When `fix_tree` is false, the resulting node tree is NOT fixed up
    // to adhere to the btree invariants.  This is useful for some testing
    // code.  But generally, `fix_tree` should be set to true.
    fn finish_internal(mut self, fix_tree: bool) -> Rope {
        // Zip up all the remaining nodes on the stack
        let mut stack_idx = self.stack.len() - 1;
        while stack_idx >= 1 {
            let node = self.stack.pop().unwrap();
            if let Node::Internal(ref mut children) = *Arc::make_mut(&mut self.stack[stack_idx - 1])
            {
                children.push((node.text_info(), node));
            } else {
                unreachable!();
            }
            stack_idx -= 1;
        }

        // Create the rope.
        let mut rope = Rope {
            root: self.stack.pop().unwrap(),
        };

        // Fix up the tree to be well-formed.
        if fix_tree {
            Arc::make_mut(&mut rope.root).zip_fix_right();
            if self.last_chunk_len_bytes < MIN_BYTES
                && self.last_chunk_len_bytes != rope.len_bytes()
            {
                // Merge the last chunk if it was too small.
                let idx = rope.len_chars()
                    - rope.byte_to_char(rope.len_bytes() - self.last_chunk_len_bytes);
                Arc::make_mut(&mut rope.root).fix_tree_seam(idx);
            }
            rope.pull_up_singular_nodes();
        }

        return rope;
    }

    // Returns (next_leaf_text, remaining_text)
    #[inline(always)]
    fn get_next_leaf_text<'a>(
        &mut self,
        text: &'a str,
        is_last_chunk: bool,
    ) -> (NextText<'a>, &'a str) {
        assert!(
            self.buffer.len() < MAX_BYTES,
            "RopeBuilder: buffer is already full when receiving a chunk! \
             This should never happen!",
        );

        // Simplest case: empty buffer and enough in `text` for a full
        // chunk, so just chop a chunk off from `text` and use that.
        if self.buffer.is_empty() && text.len() >= MAX_BYTES {
            let split_idx = crlf::find_good_split(
                MAX_BYTES.min(text.len() - 1), // - 1 to avoid CRLF split.
                text.as_bytes(),
                true,
            );
            return (NextText::String(&text[..split_idx]), &text[split_idx..]);
        }
        // If the buffer + `text` is enough for a full chunk, push enough
        // of `text` onto the buffer to fill it and use that.
        else if (text.len() + self.buffer.len()) >= MAX_BYTES {
            let mut split_idx =
                crlf::find_good_split(MAX_BYTES - self.buffer.len(), text.as_bytes(), true);
            if split_idx == text.len() && text.as_bytes()[text.len() - 1] == 0x0D {
                // Avoid CRLF split.
                split_idx -= 1;
            };
            self.buffer.push_str(&text[..split_idx]);
            return (NextText::UseBuffer, &text[split_idx..]);
        }
        // If we don't have enough text for a full chunk.
        else {
            // If it's our last chunk, wrap it all up!
            if is_last_chunk {
                if self.buffer.is_empty() {
                    return if text.is_empty() {
                        (NextText::None, "")
                    } else {
                        (NextText::String(text), "")
                    };
                } else {
                    self.buffer.push_str(text);
                    return (NextText::UseBuffer, "");
                }
            }
            // Otherwise, just push to the buffer.
            else {
                self.buffer.push_str(text);
                return (NextText::None, "");
            }
        }
    }

    fn append_leaf_node(&mut self, leaf: Arc<Node>) {
        let last = self.stack.pop().unwrap();
        match *last {
            Node::Leaf(_) => {
                if last.leaf_text().is_empty() {
                    self.stack.push(leaf);
                } else {
                    let mut children = NodeChildren::new();
                    children.push((last.text_info(), last));
                    children.push((leaf.text_info(), leaf));
                    self.stack.push(Arc::new(Node::Internal(children)));
                }
            }

            Node::Internal(_) => {
                self.stack.push(last);
                let mut left = leaf;
                let mut stack_idx = (self.stack.len() - 1) as isize;
                loop {
                    if stack_idx < 0 {
                        // We're above the root, so do a root split.
                        let mut children = NodeChildren::new();
                        children.push((left.text_info(), left));
                        self.stack.insert(0, Arc::new(Node::Internal(children)));
                        break;
                    } else if self.stack[stack_idx as usize].child_count() < (MAX_CHILDREN - 1) {
                        // There's room to add a child, so do that.
                        Arc::make_mut(&mut self.stack[stack_idx as usize])
                            .children_mut()
                            .push((left.text_info(), left));
                        break;
                    } else {
                        // Not enough room to fit a child, so split.
                        left = Arc::new(Node::Internal(
                            Arc::make_mut(&mut self.stack[stack_idx as usize])
                                .children_mut()
                                .push_split((left.text_info(), left)),
                        ));
                        std::mem::swap(&mut left, &mut self.stack[stack_idx as usize]);
                        stack_idx -= 1;
                    }
                }
            }
        }
    }
}

enum NextText<'a> {
    None,
    UseBuffer,
    String(&'a str),
}

//===========================================================================

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

    // 127 bytes, 103 chars, 4 lines
    const TEXT: &str = "Hello there!  How're you doing?\r\nIt's \
                        a fine day, isn't it?\r\nAren't you glad \
                        we're alive?\r\nこんにちは、みんなさん！";

    #[test]
    fn rope_builder_01() {
        let mut b = RopeBuilder::new();

        b.append("Hello there!  How're you doing?\r");
        b.append("\nIt's a fine ");
        b.append("d");
        b.append("a");
        b.append("y,");
        b.append(" ");
        b.append("isn't it?");
        b.append("\r");
        b.append("\nAren't you ");
        b.append("glad we're alive?\r");
        b.append("\n");
        b.append("こんにち");
        b.append("は、みんなさ");
        b.append("ん！");

        let r = b.finish();

        assert_eq!(r, TEXT);

        r.assert_integrity();
        r.assert_invariants();
    }
}
<|endoftext|>


<commit_before>use bencode::BEncode;
use std::path::PathBuf;
use std::collections::{HashMap, BTreeMap};
use std::{fmt, cmp};
use url::Url;
use ring::digest;
use util::hash_to_id;
use disk;

#[derive(Serialize, Deserialize, Clone)]
pub struct Info {
    pub name: String,
    pub announce: String,
    pub piece_len: u32,
    pub total_len: u64,
    pub hashes: Vec<Vec<u8>>,
    pub hash: [u8; 20],
    pub files: Vec<File>,
    pub file_idx: HashMap<PathBuf, usize>,
    pub private: bool,
    pub be_name: Option<Vec<u8>>,
}

impl fmt::Debug for Info {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(
            f,
            "Torrent Info {{
                name: {:?},
                announce: {:?},
                piece_len: {:?},
                total_len: {:?},
                hash: {:?},
                files: {:?}
            }}",
            self.name,
            self.announce,
            self.piece_len,
            self.total_len,
            hash_to_id(&self.hash),
            self.files
        )
    }
}

#[derive(Serialize, Deserialize, Clone, Debug)]
pub struct File {
    pub path: PathBuf,
    pub length: u64,
}

impl File {
    fn from_bencode(data: BEncode) -> Result<File, &'static str> {
        let mut d = data.into_dict().ok_or("File must be a dictionary type!")?;
        match (d.remove("name"), d.remove("path"), d.remove("length")) {
            (Some(v), None, Some(l)) => {
                let f = File {
                    path: PathBuf::from(v.into_string().ok_or("Path must be a valid string.")?),
                    length: l.into_int().ok_or("File length must be a valid int")? as u64,
                };
                Ok(f)
            }
            (None, Some(path), Some(l)) => {
                let mut p = PathBuf::new();
                for dir in path.into_list().ok_or("File path should be a list")? {
                    p.push(
                        dir.into_string().ok_or("File path parts should be strings")?,
                    );
                }
                let f = File {
                    path: p,
                    length: l.into_int().ok_or("File length must be a valid int")? as u64,
                };
                Ok(f)
            }
            _ => Err("File dict must contain length and name or path"),
        }
    }
}

impl Info {
    pub fn from_magnet(data: &str) -> Result<Info, &'static str> {
        let url = match Url::parse(data) {
            Ok(u) => u,
            Err(_) => return Err("Failed to parse magnet URL!"),
        };

        if url.scheme() != "magnet" {
            return Err("magnet URL must use magnet scheme");
        };
        // TODO: Actually implmeent this
        unimplemented!();
    }

    pub fn to_bencode(&self) -> BEncode {
        let mut torrent = BTreeMap::new();
        let mut info = BTreeMap::new();
        torrent.insert(
            "announce",
            BEncode::String(self.announce.clone().into_bytes()),
        );
        torrent.insert("info", BEncode::Dict(info));
        unimplemented!();
    }

    pub fn from_bencode(data: BEncode) -> Result<Info, &'static str> {
        data.into_dict()
            .and_then(|mut d| {
                d.remove("info").and_then(|i| i.into_dict()).map(|i| (d, i))
            })
            .ok_or("invalid info field")
            .and_then(|(mut d, mut i)| {
                let mut info_bytes = Vec::new();
                BEncode::Dict(i.clone()).encode(&mut info_bytes).unwrap();
                let mut ctx = digest::Context::new(&digest::SHA1);
                ctx.update(&info_bytes[..]);
                let digest = ctx.finish();
                let mut hash = [0u8; 20];
                hash.copy_from_slice(digest.as_ref());

                let a = d.remove("announce").and_then(|a| a.into_string()).ok_or(
                    "Info must have announce url",
                )?;
                let pl = i.remove("piece length").and_then(|i| i.into_int()).ok_or(
                    "Info must specify piece length",
                )?;
                let hashes = i.remove("pieces")
                    .and_then(|p| p.into_bytes())
                    .map(|mut p| {
                        let mut v = Vec::new();
                        while !p.is_empty() {
                            let remaining = p.split_off(20);
                            v.push(p);
                            p = remaining;
                        }
                        v
                    })
                    .ok_or("Info must provide valid hashes")?;

                let private = if let Some(v) = i.remove("private") {
                    v.into_int()
                        .and_then(|p| if p == 1 { Some(true) } else { None })
                        .ok_or("private key must be an integer equal to 1 if present!")?
                } else {
                    false
                };

                let be_name = if let Some(v) = i.remove("name") {
                    Some(v.into_bytes().ok_or("name field must be a bitstring!")?)
                } else {
                    None
                };

                let files = parse_bencode_files(i)?;
                let name = if files.is_empty() {
                    files[0]
                        .path
                        .clone()
                        .into_os_string()
                        .into_string()
                        .map_err(|_| "Only UTF8 paths are accepted")?
                } else if !files[0].path.has_root() {
                    let mut piter = files[0].path.components();
                    piter
                        .next()
                        .unwrap()
                        .as_os_str()
                        .to_os_string()
                        .into_string()
                        .map_err(|_| "Only UTF8 paths are accepted")?
                } else {
                    unreachable!();
                };

                let mut file_idx = HashMap::new();
                for (i, file) in files.iter().enumerate() {
                    file_idx.insert(file.path.clone(), i);
                }
                let total_len = files.iter().map(|f| f.length).sum();
                Ok(Info {
                    name,
                    announce: a,
                    piece_len: pl as u32,
                    hashes,
                    hash,
                    files,
                    file_idx,
                    total_len,
                    private,
                    be_name,
                })
            })
    }

    #[cfg(test)]
    pub fn with_pieces(pieces: usize) -> Info {
        Info {
            name: String::from(""),
            announce: String::from(""),
            piece_len: 16_384,
            total_len: 16_384 * pieces as u64,
            hashes: vec![vec![0u8]; pieces],
            hash: [0u8; 20],
            files: vec![],
            file_idx: HashMap::new(),
            private: false,
        }
    }

    pub fn block_len(&self, idx: u32, offset: u32) -> u32 {
        if idx != self.pieces() - 1 {
            16_384
        } else {
            let last_piece_len =
                (self.total_len - self.piece_len as u64 * (self.pieces() as u64 - 1)) as u32;
            // Note this is not the real last block len, just what it will be IF the offset really
            // is for the last block
            let last_block_len = last_piece_len - offset;
            if offset < last_piece_len && last_block_len <= 16_384 {
                last_block_len
            } else {
                16_384
            }
        }
    }

    pub fn piece_len(&self, idx: u32) -> u32 {
        if idx != self.pieces() - 1 {
            self.piece_len
        } else {
            (self.total_len - self.piece_len as u64 * (self.pieces() as u64 - 1)) as u32
        }
    }

    pub fn pieces(&self) -> u32 {
        self.hashes.len() as u32
    }

    /// Calculates the file offsets for a given block at index/begin
    pub fn block_disk_locs(&self, index: u32, begin: u32) -> Vec<disk::Location> {
        let len = self.block_len(index, begin);
        self.calc_disk_locs(index, begin, len)
    }

    /// Calculates the file offsets for a given piece at index
    pub fn piece_disk_locs(&self, index: u32) -> Vec<disk::Location> {
        let len = self.piece_len(index);
        self.calc_disk_locs(index, 0, len)
    }

    /// Calculates the file offsets for a given index, begin, and block length.
    fn calc_disk_locs(&self, index: u32, begin: u32, len: u32) -> Vec<disk::Location> {
        let mut len = len as u64;
        // The absolute byte offset where we start processing data.
        let mut cur_start = index as u64 * self.piece_len as u64 + begin as u64;
        // Current index of the data block we're writing
        let mut data_start = 0;
        // The current file end length.
        let mut fidx = 0;
        // Iterate over all file lengths, if we find any which end a bigger
        // idx than cur_start, write from cur_start..cur_start + file_write_len for that file
        // and continue if we're now at the end of the file.
        let mut locs = Vec::new();
        for f in &self.files {
            fidx += f.length;
            if cur_start < fidx {
                let file_write_len = cmp::min(fidx - cur_start, len);
                let offset = cur_start - (fidx - f.length);
                if file_write_len == len {
                    // The file is longer than our len, just write to it,
                    // exit loop
                    locs.push(disk::Location::new(
                        f.path.clone(),
                        offset,
                        data_start,
                        data_start + file_write_len,
                    ));
                    break;
                } else {
                    // Write to the end of file, continue
                    locs.push(disk::Location::new(
                        f.path.clone(),
                        offset,
                        data_start,
                        data_start + file_write_len,
                    ));
                    len -= file_write_len;
                    cur_start += file_write_len;
                    data_start += file_write_len;
                }
            }
        }
        locs
    }
}

fn parse_bencode_files(mut data: BTreeMap<String, BEncode>) -> Result<Vec<File>, &'static str> {
    match data.remove("files").and_then(|l| l.into_list()) {
        Some(fs) => {
            let mut path = PathBuf::new();
            path.push(data.remove("name").and_then(|v| v.into_string()).ok_or(
                "Multifile mode must have a name field",
            )?);
            let mut files = Vec::new();
            for f in fs {
                let mut file = File::from_bencode(f)?;
                file.path = path.join(file.path);
                files.push(file);
            }
            Ok(files)
        }
        None => File::from_bencode(BEncode::Dict(data)).map(|f| vec![f]),
    }
}
<commit_msg>Add bencode serializing to info<commit_after>use bencode::BEncode;
use std::path::PathBuf;
use std::collections::{HashMap, BTreeMap};
use std::{fmt, cmp};
use url::Url;
use ring::digest;
use util::hash_to_id;
use disk;

#[derive(Serialize, Deserialize, Clone)]
pub struct Info {
    pub name: String,
    pub announce: String,
    pub piece_len: u32,
    pub total_len: u64,
    pub hashes: Vec<Vec<u8>>,
    pub hash: [u8; 20],
    pub files: Vec<File>,
    pub file_idx: HashMap<PathBuf, usize>,
    pub private: bool,
    pub be_name: Option<Vec<u8>>,
}

impl fmt::Debug for Info {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(
            f,
            "Torrent Info {{
                name: {:?},
                announce: {:?},
                piece_len: {:?},
                total_len: {:?},
                hash: {:?},
                files: {:?}
            }}",
            self.name,
            self.announce,
            self.piece_len,
            self.total_len,
            hash_to_id(&self.hash),
            self.files
        )
    }
}

#[derive(Serialize, Deserialize, Clone, Debug)]
pub struct File {
    pub path: PathBuf,
    pub length: u64,
}

impl File {
    fn from_bencode(data: BEncode) -> Result<File, &'static str> {
        let mut d = data.into_dict().ok_or("File must be a dictionary type!")?;
        match (d.remove("name"), d.remove("path"), d.remove("length")) {
            (Some(v), None, Some(l)) => {
                let f = File {
                    path: PathBuf::from(v.into_string().ok_or("Path must be a valid string.")?),
                    length: l.into_int().ok_or("File length must be a valid int")? as u64,
                };
                Ok(f)
            }
            (None, Some(path), Some(l)) => {
                let mut p = PathBuf::new();
                for dir in path.into_list().ok_or("File path should be a list")? {
                    p.push(
                        dir.into_string().ok_or("File path parts should be strings")?,
                    );
                }
                let f = File {
                    path: p,
                    length: l.into_int().ok_or("File length must be a valid int")? as u64,
                };
                Ok(f)
            }
            _ => Err("File dict must contain length and name or path"),
        }
    }
}

impl Info {
    pub fn from_magnet(data: &str) -> Result<Info, &'static str> {
        let url = match Url::parse(data) {
            Ok(u) => u,
            Err(_) => return Err("Failed to parse magnet URL!"),
        };

        if url.scheme() != "magnet" {
            return Err("magnet URL must use magnet scheme");
        };
        // TODO: Actually implmeent this
        unimplemented!();
    }

    pub fn to_bencode(&self) -> BEncode {
        let mut torrent = BTreeMap::new();
        let mut info = BTreeMap::new();
        torrent.insert(
            "announce".to_owned(),
            BEncode::String(self.announce.clone().into_bytes()),
        );
        if let Some(ref n) = self.be_name {
            info.insert("name".to_owned(), BEncode::String(n.clone()));
        }
        info.insert(
            "piece length".to_owned(),
            BEncode::Int(self.piece_len as i64),
        );
        let mut pieces = Vec::with_capacity(self.hashes.len() * 20);
        for h in &self.hashes {
            pieces.extend_from_slice(h);
        }
        info.insert("piece".to_owned(), BEncode::String(pieces));
        if self.files.len() == 1 {
            info.insert(
                "length".to_owned(),
                BEncode::Int(self.files[0].length as i64),
            );
            // This is pretty awful, but we guarantee utf8 on creation, so should be fine.
            info.insert(
                "path".to_owned(),
                BEncode::String(
                    self.files[0]
                        .path
                        .clone()
                        .into_os_string()
                        .into_string()
                        .unwrap()
                        .into_bytes(),
                ),
            );
        } else {
            let files = self.files
                .iter()
                .map(|f| {
                    let mut fb = BTreeMap::new();
                    fb.insert("length".to_owned(), BEncode::Int(f.length as i64));
                    fb.insert(
                        "path".to_owned(),
                        BEncode::String(
                            f.path
                                .clone()
                                .into_os_string()
                                .into_string()
                                .unwrap()
                                .into_bytes(),
                        ),
                    );
                    BEncode::Dict(fb)
                })
                .collect();
            info.insert("files".to_owned(), BEncode::List(files));
        }
        torrent.insert("info".to_owned(), BEncode::Dict(info));
        BEncode::Dict(torrent)
    }

    pub fn from_bencode(data: BEncode) -> Result<Info, &'static str> {
        data.into_dict()
            .and_then(|mut d| {
                d.remove("info").and_then(|i| i.into_dict()).map(|i| (d, i))
            })
            .ok_or("invalid info field")
            .and_then(|(mut d, mut i)| {
                let mut info_bytes = Vec::new();
                BEncode::Dict(i.clone()).encode(&mut info_bytes).unwrap();
                let mut ctx = digest::Context::new(&digest::SHA1);
                ctx.update(&info_bytes[..]);
                let digest = ctx.finish();
                let mut hash = [0u8; 20];
                hash.copy_from_slice(digest.as_ref());

                let a = d.remove("announce").and_then(|a| a.into_string()).ok_or(
                    "Info must have announce url",
                )?;
                let pl = i.remove("piece length").and_then(|i| i.into_int()).ok_or(
                    "Info must specify piece length",
                )?;
                let hashes = i.remove("pieces")
                    .and_then(|p| p.into_bytes())
                    .map(|mut p| {
                        let mut v = Vec::new();
                        while !p.is_empty() {
                            let remaining = p.split_off(20);
                            v.push(p);
                            p = remaining;
                        }
                        v
                    })
                    .ok_or("Info must provide valid hashes")?;

                let private = if let Some(v) = i.remove("private") {
                    v.into_int()
                        .and_then(|p| if p == 1 { Some(true) } else { None })
                        .ok_or("private key must be an integer equal to 1 if present!")?
                } else {
                    false
                };

                let be_name = if let Some(v) = i.remove("name") {
                    Some(v.into_bytes().ok_or("name field must be a bitstring!")?)
                } else {
                    None
                };

                let files = parse_bencode_files(i)?;
                let name = if files.is_empty() {
                    files[0]
                        .path
                        .clone()
                        .into_os_string()
                        .into_string()
                        .map_err(|_| "Only UTF8 paths are accepted")?
                } else if !files[0].path.has_root() {
                    let mut piter = files[0].path.components();
                    piter
                        .next()
                        .unwrap()
                        .as_os_str()
                        .to_os_string()
                        .into_string()
                        .map_err(|_| "Only UTF8 paths are accepted")?
                } else {
                    unreachable!();
                };

                let mut file_idx = HashMap::new();
                for (i, file) in files.iter().enumerate() {
                    file_idx.insert(file.path.clone(), i);
                }
                let total_len = files.iter().map(|f| f.length).sum();
                Ok(Info {
                    name,
                    announce: a,
                    piece_len: pl as u32,
                    hashes,
                    hash,
                    files,
                    file_idx,
                    total_len,
                    private,
                    be_name,
                })
            })
    }

    #[cfg(test)]
    pub fn with_pieces(pieces: usize) -> Info {
        Info {
            name: String::from(""),
            announce: String::from(""),
            piece_len: 16_384,
            total_len: 16_384 * pieces as u64,
            hashes: vec![vec![0u8]; pieces],
            hash: [0u8; 20],
            files: vec![],
            file_idx: HashMap::new(),
            private: false,
            be_name: None,
        }
    }

    pub fn block_len(&self, idx: u32, offset: u32) -> u32 {
        if idx != self.pieces() - 1 {
            16_384
        } else {
            let last_piece_len =
                (self.total_len - self.piece_len as u64 * (self.pieces() as u64 - 1)) as u32;
            // Note this is not the real last block len, just what it will be IF the offset really
            // is for the last block
            let last_block_len = last_piece_len - offset;
            if offset < last_piece_len && last_block_len <= 16_384 {
                last_block_len
            } else {
                16_384
            }
        }
    }

    pub fn piece_len(&self, idx: u32) -> u32 {
        if idx != self.pieces() - 1 {
            self.piece_len
        } else {
            (self.total_len - self.piece_len as u64 * (self.pieces() as u64 - 1)) as u32
        }
    }

    pub fn pieces(&self) -> u32 {
        self.hashes.len() as u32
    }

    /// Calculates the file offsets for a given block at index/begin
    pub fn block_disk_locs(&self, index: u32, begin: u32) -> Vec<disk::Location> {
        let len = self.block_len(index, begin);
        self.calc_disk_locs(index, begin, len)
    }

    /// Calculates the file offsets for a given piece at index
    pub fn piece_disk_locs(&self, index: u32) -> Vec<disk::Location> {
        let len = self.piece_len(index);
        self.calc_disk_locs(index, 0, len)
    }

    /// Calculates the file offsets for a given index, begin, and block length.
    fn calc_disk_locs(&self, index: u32, begin: u32, len: u32) -> Vec<disk::Location> {
        let mut len = len as u64;
        // The absolute byte offset where we start processing data.
        let mut cur_start = index as u64 * self.piece_len as u64 + begin as u64;
        // Current index of the data block we're writing
        let mut data_start = 0;
        // The current file end length.
        let mut fidx = 0;
        // Iterate over all file lengths, if we find any which end a bigger
        // idx than cur_start, write from cur_start..cur_start + file_write_len for that file
        // and continue if we're now at the end of the file.
        let mut locs = Vec::new();
        for f in &self.files {
            fidx += f.length;
            if cur_start < fidx {
                let file_write_len = cmp::min(fidx - cur_start, len);
                let offset = cur_start - (fidx - f.length);
                if file_write_len == len {
                    // The file is longer than our len, just write to it,
                    // exit loop
                    locs.push(disk::Location::new(
                        f.path.clone(),
                        offset,
                        data_start,
                        data_start + file_write_len,
                    ));
                    break;
                } else {
                    // Write to the end of file, continue
                    locs.push(disk::Location::new(
                        f.path.clone(),
                        offset,
                        data_start,
                        data_start + file_write_len,
                    ));
                    len -= file_write_len;
                    cur_start += file_write_len;
                    data_start += file_write_len;
                }
            }
        }
        locs
    }
}

fn parse_bencode_files(mut data: BTreeMap<String, BEncode>) -> Result<Vec<File>, &'static str> {
    match data.remove("files").and_then(|l| l.into_list()) {
        Some(fs) => {
            let mut path = PathBuf::new();
            path.push(data.remove("name").and_then(|v| v.into_string()).ok_or(
                "Multifile mode must have a name field",
            )?);
            let mut files = Vec::new();
            for f in fs {
                let mut file = File::from_bencode(f)?;
                file.path = path.join(file.path);
                files.push(file);
            }
            Ok(files)
        }
        None => File::from_bencode(BEncode::Dict(data)).map(|f| vec![f]),
    }
}
<|endoftext|>


<commit_before>include!(concat!(env!("OUT_DIR"), "/colors.rs"));

#[cfg(not(feature = "no_std"))]
use std::fmt;

/// A r g b a color.
#[derive(Copy, Clone, PartialOrd, Default)]
#[repr(packed)]
pub struct Color {
    pub data: u32,
}

impl Color {
    /// Create a new color from RGB
    pub const fn rgb(r: u8, g: u8, b: u8) -> Self {
        Color {
            data: 0xFF00_0000 | ((r as u32) << 16) | ((g as u32) << 8) | (b as u32),
        }
    }

    /// Create a new color from HSV(0.0-360.0, 0.0-1.0, 0.0-1.0)
    pub fn hsv(h: f64, s: f64, v: f64) -> Self {
        Self::hsva(h, s, v, 1.0)
    }

    /// Create a new color from HSL(0.0-360.0, 0.0-1.0, 0.0-1.0)
    pub fn hsl(h: f64, s: f64, l: f64) -> Self {
        Self::hsla(h, s, l, 1.0)
    }

    /// Create a new color from RGB and alpha values
    pub const fn rgba(r: u8, g: u8, b: u8, a: u8) -> Self {
        Color {
            data: ((a as u32) << 24) | ((r as u32) << 16) | ((g as u32) << 8) | (b as u32),
        }
    }

    /// Create a new color from HSV(0.0-360.0, 0.0-1.0, 0.0-1.0) and alpha values(0.0-1.0)
    pub fn hsva(mut hue: f64, mut saturation: f64, mut value: f64, alpha: f64) -> Self {
        hue %= 360.0;
        saturation = saturation.max(0.0).min(1.0);
        value = value.max(0.0).min(1.0);
        let hh = hue / 60.0;
        let idx = hh.floor() as i32;
        let ff = hh.fract();
        let chroma = value * (1.0 - saturation);
        let second_component = value * (1.0 - (saturation * ff));
        let t = value * (1.0 - (saturation * (1.0 - ff)));
        let (r, g, b) = match idx {
            0 => (value, t, chroma),
            1 => (second_component, value, chroma),
            2 => (chroma, value, t),
            3 => (chroma, second_component, value),
            4 => (t, chroma, value),
            5 => (value, chroma, second_component),
            _ => unreachable!(),
        };
        Self::rgba(
            (r * 255.0) as u8,
            (g * 255.0) as u8,
            (b * 255.0) as u8,
            (alpha * 255.0) as u8,
        )
    }

    /// Create a new color from HSL(0.0-360.0, 0.0-1.0, 0.0-1.0) and alpha values(0.0-1.0)
    pub fn hsla(mut hue: f64, mut saturation: f64, mut lightness: f64, alpha: f64) -> Self {
        hue %= 360.0;
        saturation = saturation.max(0.0).min(1.0);
        lightness = lightness.max(0.0).min(1.0);
        let hh = hue / 60.0;
        let idx = hh.floor() as i32;
        let chroma = (1.0 - ((2.0 * lightness) - 1.0).abs()) * saturation;
        let second_component = chroma * (1.0 - (((idx % 2) as f64) - 1.0).abs());
        let (mut r, mut g, mut b) = match idx {
            0 => (chroma, second_component, 0.0),
            1 => (second_component, chroma, 0.0),
            2 => (0.0, chroma, second_component),
            3 => (0.0, second_component, chroma),
            4 => (second_component, 0.0, chroma),
            5 => (chroma, 0.0, second_component),
            _ => unreachable!(),
        };
        let adjustment = lightness - chroma / 2.0;
        r += adjustment;
        g += adjustment;
        b += adjustment;
        Self::rgba(
            (r.min(1.0) * 255.0) as u8,
            (g.min(1.0) * 255.0) as u8,
            (b.min(1.0) * 255.0) as u8,
            (alpha * 255.0) as u8,
        )
    }

    /// Get the r value
    pub fn r(self) -> u8 {
        ((self.data & 0x00FF_0000) >> 16) as u8
    }

    /// Get the g value
    pub fn g(self) -> u8 {
        ((self.data & 0x0000_FF00) >> 8) as u8
    }

    /// Get the b value
    pub fn b(self) -> u8 {
        (self.data & 0x0000_00FF) as u8
    }

    /// Get the alpha value
    pub fn a(self) -> u8 {
        ((self.data & 0xFF00_0000) >> 24) as u8
    }

    /// Attempts to get a color from its name, all the CSS colors are avaible and some other ones also.
    pub fn from_name(name: &str) -> Option<Color> {
        COLORS.get(name).cloned()
    }

    /// Interpolate between two colors
    pub fn interpolate(start_color: Color, end_color: Color, scale: f64) -> Color {
        let r = Color::interp(start_color.r(), end_color.r(), scale);
        let g = Color::interp(start_color.g(), end_color.g(), scale);
        let b = Color::interp(start_color.b(), end_color.b(), scale);
        let a = Color::interp(start_color.a(), end_color.a(), scale);
        Color::rgba(r, g, b, a)
    }

    fn interp(start_color: u8, end_color: u8, scale: f64) -> u8 {
        (end_color as f64 - start_color as f64).mul_add(scale, start_color as f64) as u8
    }
}

impl ToString for Color {
    fn to_string(&self) -> String {
        if self.a() == 0 {
            return String::from("transparent");
        }

        let data = self.data;

        let mut color = format!("#{:x}", data);
        color.remove(1);
        color.remove(1);
        color
    }
}

impl From<&str> for Color {
    fn from(s: &str) -> Color {
        if !s.starts_with('#') {
            if let Some(color) = Color::from_name(s) {
                return color;
            }
        }
        let clean_hex = s.trim_start_matches('#');
        match clean_hex.len() {
            3 | 4 => {
                let num = u32::from_str_radix(&clean_hex, 16).unwrap_or(0);

                let mut blue = (num & 0xF) << 4;
                let mut green = ((num >> 4) & 0xF) << 4;
                let mut red = ((num >> 8) & 0xF) << 4;
                let mut alpha = match clean_hex.len() == 4 {
                    true => ((num >> 12) & 0xF) << 4,
                    false => 0xF,
                };
                red |= red >> 4;
                green |= green >> 4;
                blue |= blue >> 4;
                alpha |= alpha >> 4;

                Color::rgba(red as u8, green as u8, blue as u8, alpha as u8)
            }
            6 | 8 => {
                let mut x = u32::from_str_radix(&clean_hex, 16).unwrap_or(0);

                if clean_hex.len() == 6 {
                    x |= 0xFF00_0000;
                }

                Color { data: x }
            }
            _ => Color { data: 0 },
        }
    }
}

impl From<String> for Color {
    fn from(s: String) -> Color {
        Color::from(s.as_str())
    }
}

/// Compares two colors (the alpha value is ignored)
impl PartialEq for Color {
    fn eq(&self, other: &Color) -> bool {
        self.r() == other.r() && self.g() == other.g() && self.b() == other.b()
    }
}

#[cfg(not(feature = "no_std"))]
impl fmt::Debug for Color {
    fn fmt(&self, f: &mut fmt::Formatter) -> Result<(), fmt::Error> {
        write!(f, "{:#010X}", { self.data })
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    #[test]
    fn partial_eq() {
        let color_1 = Color::rgb(1, 2, 3);
        let color_2 = Color::rgba(1, 2, 3, 200);
        let color_3 = Color::rgba(11, 2, 3, 200);

        assert!(
            color_1.eq(&color_2),
            "color_1={:?}(, color_2={:?}",
            &color_1,
            &color_2
        );
        assert!(
            color_1.eq(&color_2),
            "color_1={:?}(, color_2={:?}",
            &color_1,
            &color_2
        );
        assert!(
            color_2.ne(&color_3),
            "color_1={:?}(, color_2={:?}",
            &color_2,
            &color_3
        );
    }
}
<commit_msg>utils: update test `color`<commit_after>include!(concat!(env!("OUT_DIR"), "/colors.rs"));

#[cfg(not(feature = "no_std"))]
use std::fmt;

/// A r g b a color.
#[derive(Copy, Clone, PartialOrd, Default)]
#[repr(packed)]
pub struct Color {
    pub data: u32,
}

impl Color {
    /// Create a new color from RGB
    pub const fn rgb(r: u8, g: u8, b: u8) -> Self {
	Color {
	    data: 0xFF00_0000 | ((r as u32) << 16) | ((g as u32) << 8) | (b as u32),
	}
    }

    /// Create a new color from HSV(0.0-360.0, 0.0-1.0, 0.0-1.0)
    pub fn hsv(h: f64, s: f64, v: f64) -> Self {
	Self::hsva(h, s, v, 1.0)
    }

    /// Create a new color from HSL(0.0-360.0, 0.0-1.0, 0.0-1.0)
    pub fn hsl(h: f64, s: f64, l: f64) -> Self {
	Self::hsla(h, s, l, 1.0)
    }

    /// Create a new color from RGB and alpha values
    pub const fn rgba(r: u8, g: u8, b: u8, a: u8) -> Self {
	Color {
	    data: ((a as u32) << 24) | ((r as u32) << 16) | ((g as u32) << 8) | (b as u32),
	}
    }

    /// Create a new color from HSV(0.0-360.0, 0.0-1.0, 0.0-1.0) and alpha values(0.0-1.0)
    pub fn hsva(mut hue: f64, mut saturation: f64, mut value: f64, alpha: f64) -> Self {
	hue %= 360.0;
	saturation = saturation.max(0.0).min(1.0);
	value = value.max(0.0).min(1.0);
	let hh = hue / 60.0;
	let idx = hh.floor() as i32;
	let ff = hh.fract();
	let chroma = value * (1.0 - saturation);
	let second_component = value * (1.0 - (saturation * ff));
	let t = value * (1.0 - (saturation * (1.0 - ff)));
	let (r, g, b) = match idx {
	    0 => (value, t, chroma),
	    1 => (second_component, value, chroma),
	    2 => (chroma, value, t),
	    3 => (chroma, second_component, value),
	    4 => (t, chroma, value),
	    5 => (value, chroma, second_component),
	    _ => unreachable!(),
	};
	Self::rgba(
	    (r * 255.0) as u8,
	    (g * 255.0) as u8,
	    (b * 255.0) as u8,
	    (alpha * 255.0) as u8,
	)
    }

    /// Create a new color from HSL(0.0-360.0, 0.0-1.0, 0.0-1.0) and alpha values(0.0-1.0)
    pub fn hsla(mut hue: f64, mut saturation: f64, mut lightness: f64, alpha: f64) -> Self {
	hue %= 360.0;
	saturation = saturation.max(0.0).min(1.0);
	lightness = lightness.max(0.0).min(1.0);
	let hh = hue / 60.0;
	let idx = hh.floor() as i32;
	let chroma = (1.0 - ((2.0 * lightness) - 1.0).abs()) * saturation;
	let second_component = chroma * (1.0 - (((idx % 2) as f64) - 1.0).abs());
	let (mut r, mut g, mut b) = match idx {
	    0 => (chroma, second_component, 0.0),
	    1 => (second_component, chroma, 0.0),
	    2 => (0.0, chroma, second_component),
	    3 => (0.0, second_component, chroma),
	    4 => (second_component, 0.0, chroma),
	    5 => (chroma, 0.0, second_component),
	    _ => unreachable!(),
	};
	let adjustment = lightness - chroma / 2.0;
	r += adjustment;
	g += adjustment;
	b += adjustment;
	Self::rgba(
	    (r.min(1.0) * 255.0) as u8,
	    (g.min(1.0) * 255.0) as u8,
	    (b.min(1.0) * 255.0) as u8,
	    (alpha * 255.0) as u8,
	)
    }

    /// Get the r value
    pub fn r(self) -> u8 {
	((self.data & 0x00FF_0000) >> 16) as u8
    }

    /// Get the g value
    pub fn g(self) -> u8 {
	((self.data & 0x0000_FF00) >> 8) as u8
    }

    /// Get the b value
    pub fn b(self) -> u8 {
	(self.data & 0x0000_00FF) as u8
    }

    /// Get the alpha value
    pub fn a(self) -> u8 {
	((self.data & 0xFF00_0000) >> 24) as u8
    }

    /// Attempts to get a color from its name, all the CSS colors are avaible and some other ones also.
    pub fn from_name(name: &str) -> Option<Color> {
	COLORS.get(name).cloned()
    }

    /// Interpolate between two colors
    pub fn interpolate(start_color: Color, end_color: Color, scale: f64) -> Color {
	let r = Color::interp(start_color.r(), end_color.r(), scale);
	let g = Color::interp(start_color.g(), end_color.g(), scale);
	let b = Color::interp(start_color.b(), end_color.b(), scale);
	let a = Color::interp(start_color.a(), end_color.a(), scale);
	Color::rgba(r, g, b, a)
    }

    fn interp(start_color: u8, end_color: u8, scale: f64) -> u8 {
	(end_color as f64 - start_color as f64).mul_add(scale, start_color as f64) as u8
    }
}

impl ToString for Color {
    fn to_string(&self) -> String {
	if self.a() == 0 {
	    return String::from("transparent");
	}

	let data = self.data;

	let mut color = format!("#{:x}", data);
	color.remove(1);
	color.remove(1);
	color
    }
}

impl From<&str> for Color {
    fn from(s: &str) -> Color {
	if !s.starts_with('#') {
	    if let Some(color) = Color::from_name(s) {
		return color;
	    }
	}
	let clean_hex = s.trim_start_matches('#');
	match clean_hex.len() {
	    3 | 4 => {
		let num = u32::from_str_radix(&clean_hex, 16).unwrap_or(0);

		let mut blue = (num & 0xF) << 4;
		let mut green = ((num >> 4) & 0xF) << 4;
		let mut red = ((num >> 8) & 0xF) << 4;
		let mut alpha = match clean_hex.len() == 4 {
		    true => ((num >> 12) & 0xF) << 4,
		    false => 0xF,
		};
		red |= red >> 4;
		green |= green >> 4;
		blue |= blue >> 4;
		alpha |= alpha >> 4;

		Color::rgba(red as u8, green as u8, blue as u8, alpha as u8)
	    }
	    6 | 8 => {
		let mut x = u32::from_str_radix(&clean_hex, 16).unwrap_or(0);

		if clean_hex.len() == 6 {
		    x |= 0xFF00_0000;
		}

		Color { data: x }
	    }
	    _ => Color { data: 0 },
	}
    }
}

impl From<String> for Color {
    fn from(s: String) -> Color {
	Color::from(s.as_str())
    }
}

/// Compares two colors (the alpha value is ignored)
impl PartialEq for Color {
    fn eq(&self, other: &Color) -> bool {
	self.r() == other.r() && self.g() == other.g() && self.b() == other.b()
    }
}

#[cfg(not(feature = "no_std"))]
impl fmt::Debug for Color {
    fn fmt(&self, f: &mut fmt::Formatter) -> Result<(), fmt::Error> {
	write!(f, "{:#010X}", { self.data })
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    #[test]
    fn partial_eq() {
	let color_1 = Color::rgb(1, 2, 3);
	let color_2 = Color::rgba(1, 2, 3, 200);
	let color_3 = Color::rgba(11, 2, 3, 200);

	assert!(
	    color_1.eq(&color_2),
	    "color_1={:?}(, color_2={:?}",
	    &color_1,
	    &color_2
	);
	assert!(
	    color_1.eq(&color_2),
	    "color_1={:?}(, color_2={:?}",
	    &color_1,
	    &color_2
	);
	assert!(
	    color_2.ne(&color_3),
	    "color_1={:?}(, color_2={:?}",
	    &color_2,
	    &color_3
	);
    }
}
<|endoftext|>


<commit_before>use std::fs::File;
use std::io::Read;
use std::path::Path;
use std::process::Child;

use anyhow::{anyhow, bail, Context, Result};
use procfs::process::Process;
use pueue_lib::network::message::Message;

use super::sleep_ms;

pub fn assert_success(message: Message) {
    assert!(matches!(message, Message::Success(_)));
}

/// Get a daemon pid from a specific pueue directory.
/// This function gives the daemon a little time to boot up, but ultimately crashes if it takes too
/// long.
pub fn get_pid(pueue_dir: &Path) -> Result<i32> {
    let pid_file = pueue_dir.join("pueue.pid");

    // Give the daemon about 1 sec to boot and create the pid file.
    let tries = 10;
    let mut current_try = 0;

    while current_try < tries {
        // The daemon didn't create the pid file yet. Wait for 100ms and try again.
        if !pid_file.exists() {
            sleep_ms(50);
            current_try += 1;
            continue;
        }

        let mut file = File::open(&pid_file).context("Couldn't open pid file")?;
        let mut content = String::new();
        file.read_to_string(&mut content)
            .context("Couldn't write to file")?;

        // The file has been created but not yet been written to.
        if content.is_empty() {
            sleep_ms(50);
            current_try += 1;
            continue;
        }

        let pid = content
            .parse::<i32>()
            .map_err(|_| anyhow!("Couldn't parse value: {}", content))?;
        return Ok(pid);
    }

    bail!("Couldn't find pid file after about 1 sec.");
}

/// Waits for a daemon to shut down.
/// This is done by waiting for the pid to disappear.
pub fn wait_for_shutdown(pid: i32) -> Result<()> {
    // Try to read the process. If this fails, the daemon already exited.
    let process = match Process::new(pid) {
        Ok(process) => process,
        Err(_) => return Ok(()),
    };

    // Give the daemon about 1 sec to shutdown.
    let tries = 10;
    let mut current_try = 0;

    while current_try < tries {
        // Process is still alive, wait a little longer
        if process.is_alive() {
            sleep_ms(50);
            current_try += 1;
            continue;
        }

        return Ok(());
    }

    bail!("Couldn't find pid file after about 1 sec.");
}

pub fn kill_and_print_output(mut child: Child) -> Result<()> {
    let _ = child.kill();
    let output = child.wait_with_output()?;
    println!("Stdout: \n{:?}", String::from_utf8_lossy(&output.stdout));

    println!("Stderr: \n{:?}", String::from_utf8_lossy(&output.stderr));

    Ok(())
}
<commit_msg>fix: Increase shutdown wait timeout in tests<commit_after>use std::fs::File;
use std::io::Read;
use std::path::Path;
use std::process::Child;

use anyhow::{anyhow, bail, Context, Result};
use procfs::process::Process;
use pueue_lib::network::message::Message;

use super::sleep_ms;

pub fn assert_success(message: Message) {
    assert!(matches!(message, Message::Success(_)));
}

/// Get a daemon pid from a specific pueue directory.
/// This function gives the daemon a little time to boot up, but ultimately crashes if it takes too
/// long.
pub fn get_pid(pueue_dir: &Path) -> Result<i32> {
    let pid_file = pueue_dir.join("pueue.pid");

    // Give the daemon about 1 sec to boot and create the pid file.
    let tries = 20;
    let mut current_try = 0;

    while current_try < tries {
        // The daemon didn't create the pid file yet. Wait for 100ms and try again.
        if !pid_file.exists() {
            sleep_ms(50);
            current_try += 1;
            continue;
        }

        let mut file = File::open(&pid_file).context("Couldn't open pid file")?;
        let mut content = String::new();
        file.read_to_string(&mut content)
            .context("Couldn't write to file")?;

        // The file has been created but not yet been written to.
        if content.is_empty() {
            sleep_ms(50);
            current_try += 1;
            continue;
        }

        let pid = content
            .parse::<i32>()
            .map_err(|_| anyhow!("Couldn't parse value: {}", content))?;
        return Ok(pid);
    }

    bail!("Couldn't find pid file after about 1 sec.");
}

/// Waits for a daemon to shut down.
/// This is done by waiting for the pid to disappear.
pub fn wait_for_shutdown(pid: i32) -> Result<()> {
    // Try to read the process. If this fails, the daemon already exited.
    let process = match Process::new(pid) {
        Ok(process) => process,
        Err(_) => return Ok(()),
    };

    // Give the daemon about 1 sec to shutdown.
    let tries = 40;
    let mut current_try = 0;

    while current_try < tries {
        // Process is still alive, wait a little longer
        if process.is_alive() {
            sleep_ms(50);
            current_try += 1;
            continue;
        }

        return Ok(());
    }

    bail!("Couldn't find pid file after about 2 sec.");
}

pub fn kill_and_print_output(mut child: Child) -> Result<()> {
    let _ = child.kill();
    let output = child.wait_with_output()?;
    println!("Stdout: \n{:?}", String::from_utf8_lossy(&output.stdout));

    println!("Stderr: \n{:?}", String::from_utf8_lossy(&output.stderr));

    Ok(())
}
<|endoftext|>


<commit_before>// Copyright (c) 2018 Chef Software Inc. and/or applicable contributors
//
// 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.

//! CltGateway, short for Control Gateway, is a TCP based client and server connection for
//! sending and receiving command and control requests to a running Supervisor. Operational tasks
//! such as starting, stopping, loading, and unloading services are exposed through the
//! CtlGateway.
//!
//! The [`ctl_gateway.client`] and [`ctl_gateway.server`] speak a streaming, multiplexed, binary
//! protocol defined in [`protocol.codec`].

pub mod server;

use std::borrow::Cow;
use std::fmt;
use std::fs::{self, File};
use std::io::{self, Write};
use std::path::Path;

use regex::Regex;

use common::ui::UIWriter;
use depot_client::DisplayProgress;
use futures::prelude::*;
use hcore::output;
use hcore::util::perm;
use protocol;

use error::{Error, Result};

lazy_static! {
    /// Regular expression to, um, strip ANSI color codes out of a
    /// string.
    ///
    /// Shamelessly stolen from https://github.com/chalk/ansi-regex/blob/master/index.js.
    static ref STRIP_ANSI_CODES: Regex = Regex::new(
        r"[\x1b\x9b][\[()#;?]*(?:[0-9]{1,4}(?:;[0-9]{0,4})*)?[0-9A-PRZcf-nqry=><]")
        .unwrap();
}

/// Time to wait in milliseconds for a client connection to timeout.
pub const REQ_TIMEOUT: u64 = 10_000;
static LOGKEY: &'static str = "AG";

/// Used by modules outside of the CtlGateway for seamlessly replying to transactional messages.
/// This type is used in functions which can be called by the CtlGateway such as
/// [`Manager::service_load`] and [`Manager::service_unload`].
///
/// A _bare request_ can be instantiated with `CtlRequest::default()` and used as a no-op stub when
/// calling functions which typically respond to the client. This is useful when the caller isn't
/// a client but needs to perform the same operation. For example, loading an initial service for
/// the Supervisor has no need to reply back to a networked client, so passing a bare request will
/// no-op any calls to send messages back to the client.
#[derive(Clone, Default)]
pub struct CtlRequest {
    /// The sending side of the CtlGateway's server. Replies are sent across this channel and then
    /// eventually over the network back to the client.
    tx: Option<server::CtlSender>,
    /// Transaction for the given request.
    transaction: Option<protocol::codec::SrvTxn>,
}

impl CtlRequest {
    /// Create a new CtlRequest from an optional [`server.CtlSender`] and
    /// [`protocol.codec.SrvTxn`].
    pub fn new(
        tx: Option<server::CtlSender>,
        transaction: Option<protocol::codec::SrvTxn>,
    ) -> Self {
        CtlRequest {
            tx: tx,
            transaction: transaction,
        }
    }

    /// Reply to the transaction with the given message but indicate to the receiver that this is
    /// not the final message for the transaction.
    pub fn reply_partial<T>(&mut self, msg: T)
    where
        T: Into<protocol::codec::SrvMessage> + fmt::Debug,
    {
        self.send_msg(msg, false);
    }

    /// Reply to the transaction with the given message and indicate to the receiver that this is
    /// the final message for the transaction.
    pub fn reply_complete<T>(&mut self, msg: T)
    where
        T: Into<protocol::codec::SrvMessage> + fmt::Debug,
    {
        self.send_msg(msg, true);
    }

    /// Returns true if the request is transactional and false if not.
    pub fn transactional(&self) -> bool {
        self.transaction.is_some() && self.tx.is_some()
    }

    fn send_msg<T>(&mut self, msg: T, complete: bool)
    where
        T: Into<protocol::codec::SrvMessage> + fmt::Debug,
    {
        if !self.transactional() {
            warn!(
                "Attempted to reply to a non-transactional message with {:?}",
                msg
            );
            return;
        }
        let mut wire: protocol::codec::SrvMessage = msg.into();
        wire.reply_for(self.transaction.unwrap(), complete);
        self.tx.as_ref().unwrap().start_send(wire).unwrap();
    }
}

impl UIWriter for CtlRequest {
    type ProgressBar = NetProgressBar;

    fn out(&mut self) -> &mut io::Write {
        self
    }

    fn err(&mut self) -> &mut io::Write {
        self
    }

    // Whether output is colored, or the output is a terminal is,
    // technically, a bit more complicated than simply `true`, since
    // output from the CtlRequest ends up being streamed back to the
    // client (which may or may not be an interactive terminal, by the
    // way), as well as being rendered into the Supervisor's output
    // streams, which have their own, possibly opposing, formatting
    // constraints.
    //
    // For the time being, this is mostly addressed manually down in
    // the `io::Write` implementation. Furthermore, this is currently
    // the only implementation of UIWriter subject to these
    // constraints, so it's not clear that adding additional
    // complexity at the type / trait modeling level is worthwhile.

    fn is_out_colored(&self) -> bool {
        true
    }

    fn is_err_colored(&self) -> bool {
        true
    }

    fn is_out_a_terminal(&self) -> bool {
        true
    }

    fn is_err_a_terminal(&self) -> bool {
        true
    }

    fn progress(&self) -> Option<Self::ProgressBar> {
        if self.is_out_a_terminal() {
            Some(Self::ProgressBar::new(self.clone()))
        } else {
            None
        }
    }
}

impl io::Write for CtlRequest {
    fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
        let line = String::from_utf8_lossy(buf).into_owned();

        // The protocol reply is destined for the client, so (for now,
        // at least), we'll retain any colored output.
        //
        // `line` will also have a newline character at the end, FYI.
        let mut msg = protocol::ctl::ConsoleLine::new();
        msg.set_line(line.clone());
        self.reply_partial(msg);

        // Down here, however, we're doing double-duty by *also*
        // sending the output to the Supervisor's output stream. In
        // _this_ case, we want to honor whatever global logging flags
        // have been set; in particular, if we're emitting
        // JSON-formatted logs or if we've been instructed to not
        // output ANSI color codes, we need to strip those codes out.
        //
        // TODO (CM): Obviously, this feels hacky to the extreme. It'd
        // be nice to find a cleaner way to model this, but the fact
        // that CtlRequest is sending output to two destinations with
        // different formatting requirements complicates things a bit.
        let maybe_stripped = if output::is_json() || !output::is_color() {
            STRIP_ANSI_CODES.replace_all(&line, "")
        } else {
            Cow::Owned(line)
        };
        // TODO (CM): Consider pulling this newline trimming up into
        // the macro
        outputln!("{}", maybe_stripped.trim_right_matches('\n'));

        Ok(buf.len())
    }

    fn flush(&mut self) -> io::Result<()> {
        Ok(())
    }
}

/// A wrapper around a [`protocol.ctl.NetProgress`] and [`CtlRequest`]. This type implements
/// traits for writing it's progress to the console.
pub struct NetProgressBar {
    inner: protocol::ctl::NetProgress,
    req: CtlRequest,
}

impl NetProgressBar {
    /// Create a new progress bar.
    pub fn new(req: CtlRequest) -> Self {
        NetProgressBar {
            inner: protocol::ctl::NetProgress::new(),
            req: req,
        }
    }
}

impl DisplayProgress for NetProgressBar {
    fn size(&mut self, size: u64) {
        self.inner.set_total(size);
    }

    fn finish(&mut self) {
        ()
    }
}

impl io::Write for NetProgressBar {
    fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
        let position = self.inner.get_position() + buf.len() as u64;
        self.inner.set_position(position);
        self.req.reply_partial(self.inner.clone());
        Ok(buf.len())
    }

    fn flush(&mut self) -> io::Result<()> {
        Ok(())
    }
}

/// First attempts to read the secret key used to authenticate with the `CtlGateway` from disk
/// and, if not found, will generate a new key and write it to disk.
pub fn readgen_secret_key<T>(sup_root: T) -> Result<String>
where
    T: AsRef<Path>,
{
    let mut out = String::new();
    fs::create_dir_all(&sup_root)
        .map_err(|e| sup_error!(Error::CtlSecretIo(sup_root.as_ref().to_path_buf(), e)))?;
    if protocol::read_secret_key(&sup_root, &mut out)
        .ok()
        .unwrap_or(false)
    {
        Ok(out)
    } else {
        let secret_key_path = protocol::secret_key_path(sup_root);
        {
            let mut f = File::create(&secret_key_path)?;
            protocol::generate_secret_key(&mut out);
            f.write_all(out.as_bytes())?;
            f.sync_all()?;
        }
        perm::set_permissions(&secret_key_path, 0600)?;
        Ok(out)
    }
}
<commit_msg>Use octal for permissions of CTL_SECRET!<commit_after>// Copyright (c) 2018 Chef Software Inc. and/or applicable contributors
//
// 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.

//! CltGateway, short for Control Gateway, is a TCP based client and server connection for
//! sending and receiving command and control requests to a running Supervisor. Operational tasks
//! such as starting, stopping, loading, and unloading services are exposed through the
//! CtlGateway.
//!
//! The [`ctl_gateway.client`] and [`ctl_gateway.server`] speak a streaming, multiplexed, binary
//! protocol defined in [`protocol.codec`].

pub mod server;

use std::borrow::Cow;
use std::fmt;
use std::fs::{self, File};
use std::io::{self, Write};
use std::path::Path;

use regex::Regex;

use common::ui::UIWriter;
use depot_client::DisplayProgress;
use futures::prelude::*;
use hcore::output;
use hcore::util::perm;
use protocol;

use error::{Error, Result};

lazy_static! {
    /// Regular expression to, um, strip ANSI color codes out of a
    /// string.
    ///
    /// Shamelessly stolen from https://github.com/chalk/ansi-regex/blob/master/index.js.
    static ref STRIP_ANSI_CODES: Regex = Regex::new(
        r"[\x1b\x9b][\[()#;?]*(?:[0-9]{1,4}(?:;[0-9]{0,4})*)?[0-9A-PRZcf-nqry=><]")
        .unwrap();
}

/// Time to wait in milliseconds for a client connection to timeout.
pub const REQ_TIMEOUT: u64 = 10_000;
static LOGKEY: &'static str = "AG";

/// The control gateway secret should only be readable by the
/// Supervisor process
pub const CTL_SECRET_PERMISSIONS: u32 = 0o600;

/// Used by modules outside of the CtlGateway for seamlessly replying to transactional messages.
/// This type is used in functions which can be called by the CtlGateway such as
/// [`Manager::service_load`] and [`Manager::service_unload`].
///
/// A _bare request_ can be instantiated with `CtlRequest::default()` and used as a no-op stub when
/// calling functions which typically respond to the client. This is useful when the caller isn't
/// a client but needs to perform the same operation. For example, loading an initial service for
/// the Supervisor has no need to reply back to a networked client, so passing a bare request will
/// no-op any calls to send messages back to the client.
#[derive(Clone, Default)]
pub struct CtlRequest {
    /// The sending side of the CtlGateway's server. Replies are sent across this channel and then
    /// eventually over the network back to the client.
    tx: Option<server::CtlSender>,
    /// Transaction for the given request.
    transaction: Option<protocol::codec::SrvTxn>,
}

impl CtlRequest {
    /// Create a new CtlRequest from an optional [`server.CtlSender`] and
    /// [`protocol.codec.SrvTxn`].
    pub fn new(
        tx: Option<server::CtlSender>,
        transaction: Option<protocol::codec::SrvTxn>,
    ) -> Self {
        CtlRequest {
            tx: tx,
            transaction: transaction,
        }
    }

    /// Reply to the transaction with the given message but indicate to the receiver that this is
    /// not the final message for the transaction.
    pub fn reply_partial<T>(&mut self, msg: T)
    where
        T: Into<protocol::codec::SrvMessage> + fmt::Debug,
    {
        self.send_msg(msg, false);
    }

    /// Reply to the transaction with the given message and indicate to the receiver that this is
    /// the final message for the transaction.
    pub fn reply_complete<T>(&mut self, msg: T)
    where
        T: Into<protocol::codec::SrvMessage> + fmt::Debug,
    {
        self.send_msg(msg, true);
    }

    /// Returns true if the request is transactional and false if not.
    pub fn transactional(&self) -> bool {
        self.transaction.is_some() && self.tx.is_some()
    }

    fn send_msg<T>(&mut self, msg: T, complete: bool)
    where
        T: Into<protocol::codec::SrvMessage> + fmt::Debug,
    {
        if !self.transactional() {
            warn!(
                "Attempted to reply to a non-transactional message with {:?}",
                msg
            );
            return;
        }
        let mut wire: protocol::codec::SrvMessage = msg.into();
        wire.reply_for(self.transaction.unwrap(), complete);
        self.tx.as_ref().unwrap().start_send(wire).unwrap();
    }
}

impl UIWriter for CtlRequest {
    type ProgressBar = NetProgressBar;

    fn out(&mut self) -> &mut io::Write {
        self
    }

    fn err(&mut self) -> &mut io::Write {
        self
    }

    // Whether output is colored, or the output is a terminal is,
    // technically, a bit more complicated than simply `true`, since
    // output from the CtlRequest ends up being streamed back to the
    // client (which may or may not be an interactive terminal, by the
    // way), as well as being rendered into the Supervisor's output
    // streams, which have their own, possibly opposing, formatting
    // constraints.
    //
    // For the time being, this is mostly addressed manually down in
    // the `io::Write` implementation. Furthermore, this is currently
    // the only implementation of UIWriter subject to these
    // constraints, so it's not clear that adding additional
    // complexity at the type / trait modeling level is worthwhile.

    fn is_out_colored(&self) -> bool {
        true
    }

    fn is_err_colored(&self) -> bool {
        true
    }

    fn is_out_a_terminal(&self) -> bool {
        true
    }

    fn is_err_a_terminal(&self) -> bool {
        true
    }

    fn progress(&self) -> Option<Self::ProgressBar> {
        if self.is_out_a_terminal() {
            Some(Self::ProgressBar::new(self.clone()))
        } else {
            None
        }
    }
}

impl io::Write for CtlRequest {
    fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
        let line = String::from_utf8_lossy(buf).into_owned();

        // The protocol reply is destined for the client, so (for now,
        // at least), we'll retain any colored output.
        //
        // `line` will also have a newline character at the end, FYI.
        let mut msg = protocol::ctl::ConsoleLine::new();
        msg.set_line(line.clone());
        self.reply_partial(msg);

        // Down here, however, we're doing double-duty by *also*
        // sending the output to the Supervisor's output stream. In
        // _this_ case, we want to honor whatever global logging flags
        // have been set; in particular, if we're emitting
        // JSON-formatted logs or if we've been instructed to not
        // output ANSI color codes, we need to strip those codes out.
        //
        // TODO (CM): Obviously, this feels hacky to the extreme. It'd
        // be nice to find a cleaner way to model this, but the fact
        // that CtlRequest is sending output to two destinations with
        // different formatting requirements complicates things a bit.
        let maybe_stripped = if output::is_json() || !output::is_color() {
            STRIP_ANSI_CODES.replace_all(&line, "")
        } else {
            Cow::Owned(line)
        };
        // TODO (CM): Consider pulling this newline trimming up into
        // the macro
        outputln!("{}", maybe_stripped.trim_right_matches('\n'));

        Ok(buf.len())
    }

    fn flush(&mut self) -> io::Result<()> {
        Ok(())
    }
}

/// A wrapper around a [`protocol.ctl.NetProgress`] and [`CtlRequest`]. This type implements
/// traits for writing it's progress to the console.
pub struct NetProgressBar {
    inner: protocol::ctl::NetProgress,
    req: CtlRequest,
}

impl NetProgressBar {
    /// Create a new progress bar.
    pub fn new(req: CtlRequest) -> Self {
        NetProgressBar {
            inner: protocol::ctl::NetProgress::new(),
            req: req,
        }
    }
}

impl DisplayProgress for NetProgressBar {
    fn size(&mut self, size: u64) {
        self.inner.set_total(size);
    }

    fn finish(&mut self) {
        ()
    }
}

impl io::Write for NetProgressBar {
    fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
        let position = self.inner.get_position() + buf.len() as u64;
        self.inner.set_position(position);
        self.req.reply_partial(self.inner.clone());
        Ok(buf.len())
    }

    fn flush(&mut self) -> io::Result<()> {
        Ok(())
    }
}

/// First attempts to read the secret key used to authenticate with the `CtlGateway` from disk
/// and, if not found, will generate a new key and write it to disk.
pub fn readgen_secret_key<T>(sup_root: T) -> Result<String>
where
    T: AsRef<Path>,
{
    let mut out = String::new();
    fs::create_dir_all(&sup_root)
        .map_err(|e| sup_error!(Error::CtlSecretIo(sup_root.as_ref().to_path_buf(), e)))?;
    if protocol::read_secret_key(&sup_root, &mut out)
        .ok()
        .unwrap_or(false)
    {
        Ok(out)
    } else {
        let secret_key_path = protocol::secret_key_path(sup_root);
        {
            let mut f = File::create(&secret_key_path)?;
            protocol::generate_secret_key(&mut out);
            f.write_all(out.as_bytes())?;
            f.sync_all()?;
        }
        perm::set_permissions(&secret_key_path, CTL_SECRET_PERMISSIONS)?;
        Ok(out)
    }
}
<|endoftext|>