canteen 0.2.1

// Copyright (c) 2016
// Jeff Nettleton
//
// Licensed under the MIT license (http://opensource.org/licenses/MIT). This
// file may not be copied, modified, or distributed except according to those
// terms

extern crate chrono;

use std::collections::HashMap;

/// A trait that converts data from the handler function to a u8 slice.
pub trait ToOutput {
    fn to_output(&self) -> &[u8];
}

impl ToOutput for str {
    fn to_output(&self) -> &[u8] {
        self.as_bytes()
    }
}

impl ToOutput for &'static str {
    fn to_output(&self) -> &[u8] {
        self.as_bytes()
    }
}

impl ToOutput for String {
    fn to_output(&self) -> &[u8] {
        self.as_bytes()
    }
}

impl ToOutput for Vec<u8> {
    fn to_output(&self) -> &[u8] {
        self.as_slice()
    }
}

/// This struct reprsents the response to an HTTP client.
#[derive(Debug)]
pub struct Response {
    code:       u16,
    cmsg:       String,
    ctype:      String,
    headers:    HashMap<String, String>,
    payload:    Vec<u8>,
}

impl Response {
    /// Create a new, empty Response.
    pub fn new() -> Response {
        let mut res = Response {
            code:       200,
            cmsg:       String::from("OK"),
            ctype:      String::from("text/plain"),
            headers:    HashMap::new(),
            payload:    Vec::with_capacity(2048),
        };

        let now = chrono::UTC::now().format("%a, %d %b %Y, %H:%M:%S GMT").to_string();

        res.add_header("Connection", "close");
        res.add_header("Server", "canteen/0.0.1");
        res.add_header("Date", &now);

        res
    }

    /// Gets the HTTP message for a given code.
    pub fn get_http_message(code: u16) -> String {
        let msg = match code {
            100 => "Continue",
            101 => "Switching Protocols",
            200 => "OK",
            201 => "Created",
            202 => "Accepted",
            203 => "Non-Authoritative Information",
            204 => "No Content",
            205 => "Reset Content",
            206 => "Partial Content",
            300 => "Multiple Choices",
            301 => "Moved Permanently",
            302 => "Found",
            303 => "See Other",
            304 => "Not Modified",
            305 => "Use Proxy",
            307 => "Temporary Redirect",
            400 => "Bad Request",
            401 => "Unauthorized",
            402 => "Payment Required",
            403 => "Forbidden",
            404 => "Not Found",
            405 => "Method Not Allowed",
            406 => "Not Acceptable",
            407 => "Proxy Authentication Required",
            408 => "Request Time Out",
            409 => "Conflict",
            410 => "Gone",
            411 => "Length Required",
            412 => "Precondition Failed",
            413 => "Request Entity Too Large",
            414 => "Request-URI Too Large",
            415 => "Unsupported Media Type",
            416 => "Requested Range Not Satisfiable",
            417 => "Expectation Failed",
            500 => "Internal Server Error",
            501 => "Not Implemented",
            502 => "Bad Gateway",
            503 => "Service Unavailable",
            504 => "Gateway Time-out",
            505 => "HTTP Version Not Supported",
            _   => "OK",
        };

        String::from(msg)
    }

    /// Sets the response code for the HTTP response.
    ///
    /// # Examples
    ///
    /// ```
    /// res.set_code(200);
    /// ```
    pub fn set_code(&mut self, code: u16) {
        self.code = code;
        self.cmsg = Response::get_http_message(code);
    }

    /// Sets the Content-Type header for the HTTP response.
    ///
    /// # Examples
    ///
    /// ```
    /// res.set_content_type("text/html");
    /// ```
    pub fn set_content_type(&mut self, ctype: &str) {
        self.ctype = String::from(ctype);
    }

    /// Adds a header to the HTTP response.
    ///
    /// # Examples
    ///
    /// ```
    /// res.add_header("Content-Type", "text/html");
    /// ```
    pub fn add_header(&mut self, key: &str, value: &str) {
        if !self.headers.contains_key(key) {
            self.headers.insert(String::from(key), String::from(value));
        }
    }

    /// Appends data to the body of the HTTP response. The trait ToOutput must
    /// be implemented for the type passed.
    ///
    /// # Examples
    ///
    /// ```
    /// let data = "{ message: \"Hello, world!\" }";
    /// res.append(data);
    /// ```
    pub fn append<T: ToOutput>(&mut self, payload: T) {
        self.payload.extend(payload.to_output().into_iter());
    }

    /// Returns a byte array containing the full contents of the HTTP response,
    /// for use by the Canteen struct.
    pub fn gen_output(&self) -> Vec<u8> {
        let mut output: Vec<u8> = Vec::with_capacity(self.payload.len() + 500);
        let mut inter = String::new();

        inter.push_str(&format!("HTTP/1.1 {} {}\r\n", self.code, self.cmsg));

        for (key, value) in &self.headers {
            inter.push_str(&format!("{}: {}\r\n", key, value));
        }

        inter.push_str(&format!("Content-Type: {}\r\n", self.ctype));
        inter.push_str(&format!("Content-Length: {}\r\n", self.payload.len()));
        inter.push_str("\r\n");

        output.extend(inter.as_bytes());
        output.extend(self.payload.iter());

        output
    }
}