Rust 와 보안 2 — 웹 API 를 지키는 설정: 크기·시간 제한과 SQL 주입

use axum::extract::DefaultBodyLimit;
use axum::routing::{get, post};
use axum::Router;
use std::time::Duration;
use tokio::io::{AsyncReadExt, AsyncWriteExt};
use tower_http::timeout::TimeoutLayer;

async fn upload(body: String) -> String {
    format!("received {} bytes", body.len())
}

async fn slow() -> &'static str {
    tokio::time::sleep(Duration::from_secs(3)).await;
    "finally"
}

fn app() -> Router {
    Router::new()
        .route("/upload", post(upload))
        .route("/slow", get(slow))
        .layer(DefaultBodyLimit::max(1024))
        .layer(TimeoutLayer::with_status_code(
            axum::http::StatusCode::REQUEST_TIMEOUT,
            Duration::from_secs(1),
        ))
}

async fn send(method: &str, path: &str, body: &str) -> String {
    let mut s = tokio::net::TcpStream::connect("127.0.0.1:7878").await.unwrap();
    let req = format!(
        "{method} {path} HTTP/1.1\r\nHost: localhost\r\nContent-Length: {}\r\nConnection: close\r\n\r\n{body}",
        body.len()
    );
    s.write_all(req.as_bytes()).await.unwrap();
    let mut res = String::new();
    s.read_to_string(&mut res).await.unwrap();
    let status = res.lines().next().unwrap_or("").replace("HTTP/1.1 ", "");
    let text = res.split("\r\n\r\n").nth(1).unwrap_or("");
    format!("{method} {path} ({} bytes) -> {status} {text}", body.len())
}

#[tokio::main]
async fn main() {
    let listener = tokio::net::TcpListener::bind("127.0.0.1:7878").await.unwrap();
    tokio::spawn(async move { axum::serve(listener, app()).await.unwrap() });

    println!("{}", send("POST", "/upload", &"a".repeat(500)).await);
    println!("{}", send("POST", "/upload", &"a".repeat(5000)).await);
    println!("{}", send("GET", "/slow", "").await);
}

use rusqlite::{Connection, params};

fn setup() -> Connection {
    let db = Connection::open_in_memory().unwrap();
    db.execute_batch(
        "CREATE TABLE notes (owner TEXT, title TEXT);
         INSERT INTO notes VALUES ('ferris', 'shopping list');
         INSERT INTO notes VALUES ('corro', 'private diary');",
    )
    .unwrap();
    db
}

fn find_unsafe(db: &Connection, name: &str) -> Vec<String> {
    let sql = format!("SELECT title FROM notes WHERE owner = '{name}'");
    let mut stmt = db.prepare(&sql).unwrap();
    stmt.query_map([], |row| row.get(0)).unwrap().map(|r| r.unwrap()).collect()
}

fn find_safe(db: &Connection, name: &str) -> Vec<String> {
    let mut stmt = db.prepare("SELECT title FROM notes WHERE owner = ?1").unwrap();
    stmt.query_map(params![name], |row| row.get(0)).unwrap().map(|r| r.unwrap()).collect()
}

fn main() {
    let db = setup();
    let attack = "nobody' OR '1'='1";

    println!("unsafe, normal: {:?}", find_unsafe(&db, "ferris"));
    println!("unsafe, attack: {:?}", find_unsafe(&db, attack));
    println!("safe,   normal: {:?}", find_safe(&db, "ferris"));
    println!("safe,   attack: {:?}", find_safe(&db, attack));
}