udpt/src/webserver.rs

use std::collections::HashMap;
use std::sync::Arc;

use actix_web;
use binascii;

use config;
use tracker;

const SERVER: &str = concat!("udpt/", env!("CARGO_PKG_VERSION"));

pub struct WebServer;

mod http_responses {
    use tracker::InfoHash;

    #[derive(Serialize)]
    pub struct TorrentInfo {
        pub is_flagged: bool,
        pub leecher_count: u32,
        pub seeder_count: u32,
        pub completed: u32,
    }

    #[derive(Serialize)]
    pub struct TorrentList {
        pub offset: u32,
        pub length: u32,
        pub total: u32,
        pub torrents: Vec<InfoHash>,
    }

    #[derive(Serialize)]
    #[serde(rename_all = "snake_case")]
    pub enum APIResponse {
        Error(String),
        TorrentList(TorrentList),
        TorrentInfo(TorrentInfo),
    }
}

struct UdptState {
    // k=token, v=username.
    access_tokens: HashMap<String, String>,
    tracker: Arc<tracker::TorrentTracker>,
}

impl UdptState {
    fn new(tracker: Arc<tracker::TorrentTracker>, tokens: HashMap<String, String>) -> UdptState {
        UdptState {
            tracker,
            access_tokens: tokens,
        }
    }
}

#[derive(Debug)]
struct UdptRequestState {
    current_user: Option<String>,
}

impl Default for UdptRequestState {
    fn default() -> Self {
        UdptRequestState {
            current_user: Option::None,
        }
    }
}

impl UdptRequestState {
    fn get_user<S>(req: &actix_web::HttpRequest<S>) -> Option<String> {
        let exts = req.extensions();
        let req_state: Option<&UdptRequestState> = exts.get();
        match req_state {
            None => None,
            Option::Some(state) => match state.current_user {
                Option::Some(ref v) => Option::Some(v.clone()),
                None => {
                    error!(
                        "Invalid API token from {} @ {}",
                        req.peer_addr().unwrap(),
                        req.path()
                    );
                    return None;
                }
            },
        }
    }
}

struct UdptMiddleware;

impl actix_web::middleware::Middleware<UdptState> for UdptMiddleware {
    fn start(
        &self,
        req: &actix_web::HttpRequest<UdptState>,
    ) -> actix_web::Result<actix_web::middleware::Started> {
        let mut req_state = UdptRequestState::default();
        if let Option::Some(token) = req.query().get("token") {
            let app_state: &UdptState = req.state();
            if let Option::Some(v) = app_state.access_tokens.get(token) {
                req_state.current_user = Option::Some(v.clone());
            }
        }
        req.extensions_mut().insert(req_state);
        Ok(actix_web::middleware::Started::Done)
    }

    fn response(
        &self,
        _req: &actix_web::HttpRequest<UdptState>,
        mut resp: actix_web::HttpResponse,
    ) -> actix_web::Result<actix_web::middleware::Response> {
        resp.headers_mut().insert(
            actix_web::http::header::SERVER,
            actix_web::http::header::HeaderValue::from_static(SERVER),
        );

        Ok(actix_web::middleware::Response::Done(resp))
    }
}

impl WebServer {
    fn get_access_tokens(cfg: &config::HTTPConfig, tokens: &mut HashMap<String, String>) {
        for (user, token) in cfg.get_access_tokens().iter() {
            tokens.insert(token.clone(), user.clone());
        }
        if tokens.len() == 0 {
            warn!("No access tokens provided. HTTP API will not be useful.");
        }
    }

    pub fn new(
        tracker: Arc<tracker::TorrentTracker>,
        cfg: Arc<config::Configuration>,
    ) -> WebServer {
        let cfg_cp = cfg.clone();

        let server = actix_web::server::HttpServer::new(move || {
            let mut access_tokens = HashMap::new();

            if let Some(http_cfg) = cfg_cp.get_http_config() {
                Self::get_access_tokens(http_cfg, &mut access_tokens);
            }

            let state = UdptState::new(tracker.clone(), access_tokens);

            actix_web::App::<UdptState>::with_state(state)
                .middleware(UdptMiddleware)
                .resource("/t", |r| r.f(Self::view_torrent_list))
                .scope(r"/t/{info_hash:[\dA-Fa-f]{40,40}}", |scope| {
                    scope.resource("", |r| {
                        r.method(actix_web::http::Method::GET)
                            .f(Self::view_torrent_stats);
                        r.method(actix_web::http::Method::POST)
                            .f(Self::torrent_action);
                    })
                })
                .resource("/", |r| {
                    r.method(actix_web::http::Method::GET).f(Self::view_root)
                })
        });

        if let Some(http_cfg) = cfg.get_http_config() {
            let bind_addr = http_cfg.get_address();
            match server.bind(bind_addr) {
                Ok(v) => {
                    v.run();
                }
                Err(err) => {
                    error!("Failed to bind http server. {}", err);
                }
            }
        } else {
            unreachable!();
        }

        WebServer {}
    }

    fn view_root(_req: &actix_web::HttpRequest<UdptState>) -> actix_web::HttpResponse {
        actix_web::HttpResponse::build(actix_web::http::StatusCode::OK)
            .content_type("text/html")
            .body(r#"Powered by <a href="https://github.com/naim94a/udpt">https://github.com/naim94a/udpt</a>"#)
    }

    fn view_torrent_list(req: &actix_web::HttpRequest<UdptState>) -> impl actix_web::Responder {
        use std::str::FromStr;

        if UdptRequestState::get_user(req).is_none() {
            return actix_web::Json(http_responses::APIResponse::Error(String::from(
                "access_denied",
            )));
        }

        let req_offset = match req.query().get("offset") {
            None => 0,
            Some(v) => match u32::from_str(v.as_str()) {
                Ok(v) => v,
                Err(_) => 0,
            },
        };

        let mut req_limit = match req.query().get("limit") {
            None => 0,
            Some(v) => match u32::from_str(v.as_str()) {
                Ok(v) => v,
                Err(_) => 0,
            },
        };

        if req_limit > 4096 {
            req_limit = 4096;
        } else if req_limit == 0 {
            req_limit = 1000;
        }

        let app_state: &UdptState = req.state();
        let app_db = app_state.tracker.get_database();

        let total = app_db.len() as u32;

        let mut torrents = Vec::with_capacity(req_limit as usize);

        for (info_hash, _) in app_db
            .iter()
            .skip(req_offset as usize)
            .take(req_limit as usize)
        {
            torrents.push(info_hash.clone());
        }

        actix_web::Json(http_responses::APIResponse::TorrentList(
            http_responses::TorrentList {
                total,
                length: torrents.len() as u32,
                offset: req_offset,
                torrents,
            },
        ))
    }

    fn view_torrent_stats(req: &actix_web::HttpRequest<UdptState>) -> actix_web::HttpResponse {
        use actix_web::FromRequest;

        if UdptRequestState::get_user(req).is_none() {
            return actix_web::HttpResponse::build(actix_web::http::StatusCode::UNAUTHORIZED).json(
                http_responses::APIResponse::Error(String::from("access_denied")),
            );
        }

        let path: actix_web::Path<String> = match actix_web::Path::extract(req) {
            Ok(v) => v,
            Err(_) => {
                return actix_web::HttpResponse::build(
                    actix_web::http::StatusCode::INTERNAL_SERVER_ERROR,
                )
                .json(http_responses::APIResponse::Error(String::from(
                    "internal_error",
                )));
            }
        };

        let mut info_hash = [0u8; 20];
        if let Err(_) = binascii::hex2bin((*path).as_bytes(), &mut info_hash) {
            return actix_web::HttpResponse::build(actix_web::http::StatusCode::BAD_REQUEST).json(
                http_responses::APIResponse::Error(String::from("invalid_info_hash")),
            );
        }

        let app_state: &UdptState = req.state();

        let db = app_state.tracker.get_database();
        let entry = match db.get(&info_hash.into()) {
            Some(v) => v,
            None => {
                return actix_web::HttpResponse::build(actix_web::http::StatusCode::NOT_FOUND).json(
                    http_responses::APIResponse::Error(String::from("not_found")),
                );
            }
        };

        let is_flagged = entry.is_flagged();
        let (seeders, completed, leechers) = entry.get_stats();

        return actix_web::HttpResponse::build(actix_web::http::StatusCode::OK).json(
            http_responses::APIResponse::TorrentInfo(http_responses::TorrentInfo {
                is_flagged,
                seeder_count: seeders,
                leecher_count: leechers,
                completed,
            }),
        );
    }

    fn torrent_action(req: &actix_web::HttpRequest<UdptState>) -> actix_web::HttpResponse {
        use actix_web::FromRequest;

        if UdptRequestState::get_user(req).is_none() {
            return actix_web::HttpResponse::build(actix_web::http::StatusCode::UNAUTHORIZED).json(
                http_responses::APIResponse::Error(String::from("access_denied")),
            );
        }

        let query = req.query();
        let action_opt = query.get("action");
        let action = match action_opt {
            Some(v) => v,
            None => {
                return actix_web::HttpResponse::build(actix_web::http::StatusCode::BAD_REQUEST)
                    .json(http_responses::APIResponse::Error(String::from(
                        "action_required",
                    )));
            }
        };

        let app_state: &UdptState = req.state();

        let path: actix_web::Path<String> = match actix_web::Path::extract(req) {
            Ok(v) => v,
            Err(_err) => {
                return actix_web::HttpResponse::build(
                    actix_web::http::StatusCode::INTERNAL_SERVER_ERROR,
                )
                .json(http_responses::APIResponse::Error(String::from(
                    "internal_error",
                )));
            }
        };

        let info_hash_str = &(*path);
        let mut info_hash = [0u8; 20];
        if let Err(_) = binascii::hex2bin(info_hash_str.as_bytes(), &mut info_hash) {
            return actix_web::HttpResponse::build(actix_web::http::StatusCode::BAD_REQUEST).json(
                http_responses::APIResponse::Error(String::from("invalid_info_hash")),
            );
        }

        match action.as_str() {
            "flag" => {
                app_state.tracker.set_torrent_flag(&info_hash.into(), true);
                info!("Flagged {}", info_hash_str.as_str());
                return actix_web::HttpResponse::build(actix_web::http::StatusCode::OK).body("");
            }
            "unflag" => {
                app_state.tracker.set_torrent_flag(&info_hash.into(), false);
                info!("Unflagged {}", info_hash_str.as_str());
                return actix_web::HttpResponse::build(actix_web::http::StatusCode::OK).body("");
            }
            "add" => {
                let success = app_state.tracker.add_torrent(&info_hash.into()).is_ok();
                info!("Added {}, success={}", info_hash_str.as_str(), success);
                let code = if success {
                    actix_web::http::StatusCode::OK
                } else {
                    actix_web::http::StatusCode::INTERNAL_SERVER_ERROR
                };

                return actix_web::HttpResponse::build(code).body("");
            }
            "remove" => {
                let success = app_state
                    .tracker
                    .remove_torrent(&info_hash.into(), true)
                    .is_ok();
                info!("Removed {}, success={}", info_hash_str.as_str(), success);
                let code = if success {
                    actix_web::http::StatusCode::OK
                } else {
                    actix_web::http::StatusCode::INTERNAL_SERVER_ERROR
                };

                return actix_web::HttpResponse::build(code).body("");
            }
            _ => {
                debug!("Invalid action {}", action.as_str());
                return actix_web::HttpResponse::build(actix_web::http::StatusCode::BAD_REQUEST)
                    .json(http_responses::APIResponse::Error(String::from(
                        "invalid_action",
                    )));
            }
        }
    }
}