ruver/

lib.rs

use std::error::Error;
use std::fs::File;
use std::io::prelude::*;
use std::net::TcpStream;
use std::path::Path;
use std::sync::mpsc;
use std::sync::Arc;
use std::sync::Mutex;
use std::thread;
use std::thread::sleep;
use std::time::Duration;

// # Thread Pool

trait FnBox {
    fn call_box(self: Box<Self>);
}

impl<F: FnOnce()> FnBox for F {
    fn call_box(self: Box<F>) {
        (*self)()
    }
}
// Rust has a problem itself, that can not automatically use `self: Box<Self>` to take control of
// the content of box.

type Job = Box<FnBox + Send + 'static>;

struct Worker {
    id: usize,
    thread: Option<thread::JoinHandle<()>>,
}

impl Worker {
    fn new(id: usize, receiver: Arc<Mutex<mpsc::Receiver<Message>>>) -> Worker {
        let thread = thread::spawn(move || loop {
            let message = receiver.lock().unwrap().recv().unwrap();

            match message {
                Message::NewJob(job) => job.call_box(),
                Message::Terminate => break,
            }
        });

        Worker {
            id,
            thread: Some(thread),
        }
    }
}

enum Message {
    NewJob(Job),
    Terminate,
}

/// A thread pool, containing a vector of workers and an actually "spmc" sender.
pub struct ThreadPool {
    workers: Vec<Worker>,
    sender: mpsc::Sender<Message>,
}

impl Drop for ThreadPool {
    fn drop(&mut self) {
        println!("Sending terminate message to all workers.");

        for _ in &mut self.workers {
            self.sender.send(Message::Terminate).unwrap();
        }

        println!("Shutting down all workers.");

        for worker in &mut self.workers {
            println!("Shutting down worker {}", worker.id);

            if let Some(thread) = worker.thread.take() {
                thread.join().unwrap();
            }
        }
        // Why divide loop?
        // We don't exactly know if it is the worker in current loop that receive our signal.
    }
}

impl ThreadPool {
    /// Create a new thread pool.
    ///
    /// # Examples
    ///
    /// ```
    /// use ruver::ThreadPool;
    ///
    /// let pool = ThreadPool::new(256);
    ///
    /// // These codes created a thread pool who has 256 threads.
    /// ```
    pub fn new(size: usize) -> ThreadPool {
        assert!(size > 0);

        let (sender, receiver) = mpsc::channel();

        let receiver = Arc::new(Mutex::new(receiver));

        let mut workers = Vec::with_capacity(size);

        for id in 0..size {
            workers.push(Worker::new(id, Arc::clone(&receiver)));
        }

        ThreadPool { workers, sender }
    }

    /// Throw a task into thread pool.
    ///
    /// # Examples
    ///
    /// ```
    /// use ruver::ThreadPool;
    ///
    /// let pool = ThreadPool::new(256);
    ///
    /// let add_one = |x| {
    ///     println!("{}", x + 1);
    /// };
    ///
    /// pool.execute(add_one(1))
    ///
    /// // These codes print "2".
    /// ```
    pub fn execute<F>(&self, f: F)
    where
        F: FnOnce() + Send + 'static,
    {
        let job = Box::new(f);
        // Why box?
        // Box provides a type whose size is already known, while closure do not.

        self.sender.send(Message::NewJob(job)).unwrap();
    }
}

// # Handler

struct Request {
    method: String,
    resource: String,
    protocol: String,
}

enum StatusCode {
    OK,
    NotFound,
    InternalServerError,
}

struct Responder {
    stream: TcpStream,
}

impl Responder {
    fn new(stream: TcpStream) -> Responder {
        Responder { stream }
    }

    fn write(&mut self, response: &[u8]) {
        match self.stream.write(response) {
            Ok(_) => (),
            Err(error) => {
                eprintln!(
                    "Caught an error when writing to TCP stream:\n{}",
                    error.description()
                );
            }
        };

        self.stream.flush().unwrap();
    }
}

fn parser(message: String) -> Request {
    let (mut method, mut resource, mut protocol) = (String::new(), String::new(), String::new());

    for (count, content) in message.split_whitespace().enumerate() {
        match count {
            0 => method.push_str(content),
            1 => resource.push_str(content),
            2 => protocol.push_str(content),
            _ => break,
        }
    }

    Request {
        method,
        resource,
        protocol,
    }
}

fn checker(request: &Request) -> StatusCode {
    let raw_resource = format!(".{}", &request.resource);

    let resource = Path::new(&raw_resource);

    if let Err(error) = File::open(resource) {
        eprintln!("Failed to open file:\n{}", error);
        return StatusCode::InternalServerError;
    };

    if request.method.is_empty()
        || raw_resource.is_empty()
        || request.protocol.is_empty()
        || request.method != "GET"
        || resource.is_dir()
        || !(request.protocol == "HTTP/1.0" || request.protocol == "HTTP/1.1")
        || raw_resource.contains("/../")
    {
        println!("Internal server error.");
        return StatusCode::InternalServerError;
    }

    if resource.exists() {
        println!("Request {}.", raw_resource);
        StatusCode::OK
    } else {
        println!("Can not find {}.", raw_resource);
        StatusCode::NotFound
    }
}

/// Receive and handle a TCP stream.
///
/// # Example
///
/// ```
/// use std::net::TcpListener;
///
/// let listener = match TcpListener::bind("0.0.0.0:8000").unwrap();
///
/// for stream in listener.incoming() {
///     let stream = stream.unwrap();
///
///     ruver::handle_connection(stream);
/// }
/// ```
pub fn handle_connection(mut stream: TcpStream) {
    let mut buffer = [0; 1024];

    let mut raw_request = String::new();

    stream
        .set_read_timeout(Some(Duration::from_secs(5)))
        .unwrap_or_else(|error| {
            eprintln!("set_read_timeout call failed:\n{}", error.description())
        });

    stream
        .set_write_timeout(Some(Duration::from_secs(5)))
        .unwrap_or_else(|error| {
            eprintln!("set_write_timeout call failed:\n{}", error.description())
        });

    loop {
        let buffer_num = match stream.read(&mut buffer) {
            Ok(t) => (t),
            Err(error) => {
                eprintln!(
                    "Caught an error when reading from TCP stream:\n{}",
                    error.description()
                );
                return;
            }
        };

        if buffer_num > 0 {
            raw_request.push_str(&String::from_utf8(buffer[..buffer_num].to_vec()).unwrap());
        }

        if raw_request.ends_with("\r\n\r\n") {
            break;
        }

        sleep(Duration::from_millis(1));
    }

    let request = parser(raw_request);

    let mut responder = Responder::new(stream);

    match checker(&request) {
        StatusCode::OK => {
            let resource = format!(".{}", &request.resource);

            let mut file = File::open(resource).unwrap();

            let content_length = match file.metadata() {
                Ok(t) => t.len(),
                Err(error) => {
                    eprintln!("Fail to read file:\n{}", error);
                    return;
                }
            };

            let head = format!(
                "HTTP/1.0 200 OK\r\nContent-Length: {}\r\n\r\n",
                content_length,
            );

            let head = head.as_bytes();

            responder.write(head);

            let mut buffer = [0; 1024];

            loop {
                let size = match file.read(&mut buffer) {
                    Ok(t) => t,
                    Err(error) => {
                        eprintln!("Fail to read file:\n{}", error);
                        return;
                    }
                };
                if size > 0 {
                    responder.write(&buffer);
                } else {
                    break;
                }
            }
        }
        StatusCode::NotFound => {
            responder.write("HTTP/1.0 404 Not Found\r\n\r\n".as_bytes());
        }
        StatusCode::InternalServerError => {
            responder.write("HTTP/1.0 500 Internal Server Error\r\n\r\n".as_bytes());
        }
    };
}