【Rust】フレームワーク自作: セッションが無ければリダイレクト処理

fn handle_connection(mut stream: TcpStream, get_routes: HashMap<String, String>, post_routes: HashMap<String, fn(String) -> HttpResponse>)  {
    let mut buffer = [0; 1024];
    stream.read(&mut buffer).unwrap();
 
    let request = String::from_utf8_lossy(&buffer);
    let request_line = request.lines().next().unwrap_or("");
    let mut parts = request_line.split_whitespace();
    let method = parts.next().unwrap_or("");
    let path = parts.next().unwrap_or("/");

    // println!("Received {} request for {}", method, path);

    let needs_auth = ["/index", "/hello"];
    let is_protected = needs_auth.contains(&path);

    let session_token = extract_cookie_token(&request);

    if is_protected && session_token.is_none() {
        let redirect_response = http_response_custom(redirect("/login"));
        stream.write_all(redirect_response.as_bytes()).unwrap();
        stream.flush().unwrap();
        return;
    }
 

    let response = match method {
        "GET" => {
            let body = get_routes.get(path).cloned().unwrap_or_else(|| "<h1>404 Not Found</h1>".to_string());
            http_response(200, "text/html", &body)
        }
        "POST" => {
            let body = request.split("\r\n\r\n").nth(1).unwrap_or("").to_string();
            match post_routes.get(path) {
                Some(handler) => {
                    let response = handler(body);
                    http_response_custom(response)
                },
                None => http_response(404, "text/html", "<h1>404 Not Found</h1>"),
            }
        }
        _ => http_response(405, "text/html", "<h1>405 Method Not Allowed</h1>"),
    };
 
    stream.write_all(response.as_bytes()).unwrap();
    stream.flush().unwrap();
 }

main.rs

use std::collections::HashMap;
use rand::SeedableRng; // トレイトは rand にある
use rand::rngs::OsRng;
use rand::RngCore;
use rand_chacha::ChaCha8Rng;
use pwhash::bcrypt;
mod ares;
use ares::{Router, HttpResponse, redirect, psql_connect, parse, remove_null_bytes, create_session_token};

fn main() {
    let mut router = Router::new();

    router.get("index", handle_index);
    router.get("hello", handle_hello);
    router.get("signin", handle_signin);
    router.post("signup", handle_signup);
    router.get("login", handle_login);
    router.post("logup", handle_logup);

    router.up("192.168.33.10:8000");
}


fn handle_index() -> Option<HashMap<&'static str, &'static str>> {

    let mut content1 = HashMap::new();
    content1.insert(
        "title",
        "This is index page!",
    );
    return Some(content1);
}

fn handle_hello() -> Option<HashMap<&'static str, &'static str>> {
    return None;
}

fn handle_signin() -> Option<HashMap<&'static str, &'static str>> {
    return None;
}

fn handle_signup(body: String) -> HttpResponse {
    let form =  parse(&body);
    let binding = "<unknown>".to_string();
    let name = form.get("name").unwrap_or(&binding);
    let password = form.get("password").unwrap_or(&binding);

    let pass_hash = bcrypt::hash(password).unwrap();

    let mut client = psql_connect().unwrap();
    let _ = client.execute(
        "INSERT INTO test (username, password) VALUES ($1, $2)",
        &[&remove_null_bytes(&name), &remove_null_bytes(&pass_hash)],
    ).unwrap();
    HttpResponse::new(200, "<h1>Success Signin!</h1>")
}

fn handle_login() -> Option<HashMap<&'static str, &'static str>> {
    return None;
}

fn handle_logup(body: String) -> HttpResponse {
    let form =  parse(&body);
    let binding = "<unknown>".to_string();
    let name = form.get("name").unwrap_or(&binding);
    let password = form.get("password").unwrap_or(&binding);

    let mut client = psql_connect().unwrap();
    let row = client.query(
        "SELECT * from test where username=$1",
        &[&name],
    ).unwrap();
    if row.is_empty() {
        return redirect("/login");
    }
    let value: String = row[0].get(2);

    if bcrypt::verify(password, &value) {
        let session_token = create_session_token();
        println!("{}", session_token);
        let mut response = HttpResponse::new(200, "<h1>Success Login!</h1>");
        response.headers.insert(
            "Set-Cookie".to_string(),
            format!("session_token={}; Path=/; HttpOnly", session_token),
        );
        response
    } else {
        redirect("/login")
    }
}

– どのページをログイン必須にするかは、ユーザ側で決めたい
– cookieのセットはフレームワーク側で行いたい