hglib 0.1.1

//! Raw command server API.
//!
//! Creating a [`Connection`](struct.Connection.html) spawns an instance
//! of the command server and allows you to interact with it. When it
//! starts up, the command server sends a hello message on its output
//! channel, which can be read and parsed by
//! [`read_hello()`](struct.Connection.html#method.read_hello):
//!
//! ```rust
//! # use std::io;
//! # use hglib::connection::Connection;
//! # use hglib::Chunk;
//! let mut conn = Connection::new().ok().expect("failed to start command server");
//! let (capabilities, encoding) =
//!     conn.read_hello().ok().expect("failed to read server hello");
//! ```

use std::ascii::AsciiExt;
use std::io;
use std::io::prelude::*;
use std::process::{Command, Stdio, Child, ChildStdout, ExitStatus};
use std::str;

use byteorder::{BigEndian, ReadBytesExt, WriteBytesExt};

mod channel {
    pub const OUTPUT: u8 = b'o';
    pub const ERROR: u8 = b'e';
    pub const RESULT: u8 = b'r';
    #[allow(dead_code)]
    pub const DEBUG: u8 = b'd';
    #[allow(dead_code)]
    pub const INPUT: u8 = b'I';
    #[allow(dead_code)]
    pub const LINE_INPUT: u8 = b'L';
}

use Chunk;

/// An iterator over the results of a single Mercurial command.
///
/// Each iteration yields a [`Chunk`](../enum.Chunk.html) with the
/// output data or length of input the server is expecting.
pub struct CommandRun<'a> {
    connection: &'a mut Connection,
    done: bool,
}

impl<'a> Iterator for CommandRun<'a> {
    type Item = io::Result<Chunk>;

    fn next(&mut self) -> Option<Self::Item> {
        if self.done {
            return None;
        }
        let (chan, length) = match self.connection.read_header() {
            Ok(t) => t,
            Err(e) => return Some(Err(e)),
        };
        match chan {
            channel::OUTPUT => {
                Some(self.connection.read_body(length).map(Chunk::Output))
            },
            channel::ERROR => {
                Some(self.connection.read_body(length).map(Chunk::Error))
            },
            channel::RESULT => {
                self.done = true;
                Some(self.connection.read_result().map(Chunk::Result))
            },
            _ => {
                if (chan as char).is_uppercase() {
                    return Some(Err(io::Error::new(
                        io::ErrorKind::Other,
                        format!("unexpected required channel: {:?}", chan as char))));
                }
                return Some(Err(io::Error::new(
                    io::ErrorKind::Other,
                    format!("unexpected channel: {:?}", chan as char))));
            },
        }
    }
}

/// A handle to a running command server instance.
pub struct Connection {
    child: Child,
}

impl Connection {
    /// Spawns a new command server process.
    pub fn new() -> io::Result<Connection> {
        let cmdserver = try!(
            Command::new("hg")
                .args(&["serve", "--cmdserver", "pipe", "--config", "ui.interactive=True"])
                .stdin(Stdio::piped())
                .stdout(Stdio::piped())
                .spawn());

        Ok(Connection {
            child: cmdserver,
        })
    }

    fn child_stdout(&mut self) -> &mut ChildStdout {
        // We just unwrap the Option<ChildStdout> because we know that
        // we set up the pipe in Connection::new(). We have to call
        // .as_mut() because .unwrap()'s signature moves the `self`
        // value out.
        self.child.stdout.as_mut().unwrap()
    }

    /// Reads and parses the server hello message. Returns a tuple of
    /// ([capabilities], encoding).
    ///
    /// ## Errors
    ///
    /// Returns an I/O error if reading or parsing the hello message
    /// failed.
    pub fn read_hello(&mut self) -> io::Result<(Vec<String>, String)> {
        fn fetch_field(line: Option<&[u8]>, field: &[u8]) -> io::Result<Vec<u8>>
        {
            let mut label = field.to_vec();
            label.extend(b": ");

            match line {
                Some(l) if l.is_ascii() && l.starts_with(&label) => {
                    Ok(l[label.len()..].to_vec())
                },
                Some(l) => {
                    let err_data = match str::from_utf8(l) {
                        Ok(s)  => s.to_string(),
                        Err(e) => format!("{:?} (bad encoding: {})", l, e),
                    };
                    return Err(io::Error::new(
                        io::ErrorKind::InvalidData,
                        format!("expected '{}: ', got {:?}",
                                String::from_utf8_lossy(field), err_data)));
                },
                None => return Err(io::Error::new(
                    // TODO: With 1.4 we can use UnexpectedEOF
                    io::ErrorKind::Other,
                    format!("missing field '{}' in server hello",
                            String::from_utf8_lossy(field)))),
            }
        }

        fn parse_capabilities(cap_line: Vec<u8>) -> Vec<String> {
            let mut caps = vec![];
            for cap in cap_line.split(|byte| *byte == b' ') {
                let cap = str::from_utf8(cap)
                    // TODO: With 1.4 we can use Result::expect directly
                    .ok().expect("failed to decode ASCII as UTF-8?!");
                caps.push(cap.to_string());
            }
            caps
        }

        let (_, length) = try!(self.read_header());
        let hello = try!(self.read_body(length));
        let mut hello = hello.split(|byte| *byte == b'\n');

        let caps = try!(fetch_field(hello.next(), b"capabilities")
                        .map(|l| parse_capabilities(l)));
        let enc = try!(fetch_field(hello.next(), b"encoding")
                       .map(|l| String::from_utf8(l)
                            .ok().expect("failed to decode ASCII as UTF-8?!")));
        Ok((caps, enc))
    }

    fn read_header(&mut self) -> io::Result<(u8, i32)> {
        let pout = self.child_stdout();
        let chan = try!(pout.read_u8());
        let length = try!(pout.read_i32::<BigEndian>());
        Ok((chan, length))
    }

    fn read_body(&mut self, length: i32) -> io::Result<Vec<u8>> {
        let pout = self.child_stdout();
        let mut buf = Vec::with_capacity(length as usize);
        try!(pout.take(length as u64).read_to_end(&mut buf));
        Ok(buf)
    }

    fn read_result(&mut self) -> io::Result<i32> {
        let pout = self.child_stdout();
        let result = try!(pout.read_i32::<BigEndian>());
        Ok(result)
    }

    fn _raw_command(&mut self, command: Vec<&[u8]>) -> io::Result<()> {
        // sum of lengths of all arguments, plus null byte separators
        let len = command.iter().map(|item| item.len())
            .fold(command.len() - 1, |acc, l| acc + l);
        if len > i32::max_value() as usize {
            return Err(io::Error::new(io::ErrorKind::InvalidInput, "message too long"));
        }

        let pin = self.child.stdin.as_mut().unwrap();
        try!(pin.write(b"runcommand\n"));
        try!(pin.write_i32::<BigEndian>(len as i32));
        try!(pin.write(command[0]));
        for arg in &command[1..] {
            try!(pin.write(b"\0"));
            try!(pin.write(arg));
        }
        Ok(())
    }

    /// Sends the given `command` to Mercurial, returning an iterator
    /// over the results.
    pub fn raw_command(&mut self, command: Vec<&[u8]>) -> io::Result<CommandRun> {
        try!(self._raw_command(command));
        Ok(CommandRun {
            connection: self,
            done: false,
        })
    }

    /// Shuts down the command server process.
    pub fn close(&mut self) -> io::Result<ExitStatus> {
        // This will close the command server's stdin, which signals
        // that it should exit. Returns the command server's exit code.
        self.child.wait()
    }
}