6장: Axum 고급 패턴

src/
  main.rs              # 진입점, 서버 시작
  lib.rs               # 공개 모듈 선언
  config.rs            # 설정 로드
  routes/
    mod.rs             # 라우터 조립
    users.rs           # /users 라우트
    posts.rs           # /posts 라우트
    auth.rs            # /auth 라우트
  handlers/
    mod.rs
    users.rs           # 사용자 핸들러
    posts.rs           # 게시물 핸들러
  models/
    mod.rs
    user.rs            # User 도메인 모델
    post.rs            # Post 도메인 모델
  middleware/
    mod.rs
    auth.rs            # 인증 미들웨어
    logging.rs         # 로깅 미들웨어
  error.rs             # 에러 타입 정의
  state.rs             # AppState 정의

use axum::Router;
use std::sync::Arc;
use crate::state::AppState;
 
pub mod auth;
pub mod posts;
pub mod users;
 
pub fn create_router(state: Arc<AppState>) -> Router {
    Router::new()
        .nest("/api/v1", api_v1(state.clone()))
        .nest("/auth", auth::routes(state))
}
 
fn api_v1(state: Arc<AppState>) -> Router {
    Router::new()
        .nest("/users", users::routes(state.clone()))
        .nest("/posts", posts::routes(state))
}

use axum::{Router, routing::{get, post, put, delete}};
use std::sync::Arc;
use crate::handlers::users;
use crate::state::AppState;
 
pub fn routes(state: Arc<AppState>) -> Router {
    Router::new()
        .route("/", get(users::list).post(users::create))
        .route("/{id}", get(users::get_one).put(users::update).delete(users::remove))
        .with_state(state)
}

use axum::{
    middleware::{self, Next},
    extract::Request,
    response::Response,
};
use std::time::Instant;
 
async fn timing_middleware(
    request: Request,
    next: Next,
) -> Response {
    let start = Instant::now();
    let method = request.method().clone();
    let uri = request.uri().clone();
 
    let response = next.run(request).await;
 
    let duration = start.elapsed();
    tracing::info!(
        method = %method,
        uri = %uri,
        duration_ms = duration.as_millis(),
        status = response.status().as_u16(),
        "요청 처리 완료"
    );
 
    response
}
 
// 라우터에 적용
let app = Router::new()
    .route("/users", get(list_users))
    .layer(middleware::from_fn(timing_middleware));

use axum::{
    extract::Request,
    http::{header, StatusCode},
    middleware::Next,
    response::Response,
    Json,
};
use jsonwebtoken::{decode, DecodingKey, Validation};
use serde::{Deserialize, Serialize};
 
#[derive(Debug, Serialize, Deserialize, Clone)]
pub struct Claims {
    pub sub: String,    // 사용자 ID
    pub role: String,   // 역할
    pub exp: usize,     // 만료 시간
}
 
pub async fn auth_middleware(
    mut request: Request,
    next: Next,
) -> Result<Response, StatusCode> {
    // Authorization 헤더에서 토큰 추출
    let token = request
        .headers()
        .get(header::AUTHORIZATION)
        .and_then(|value| value.to_str().ok())
        .and_then(|value| value.strip_prefix("Bearer "))
        .ok_or(StatusCode::UNAUTHORIZED)?;
 
    // JWT 검증
    let secret = std::env::var("JWT_SECRET").unwrap_or_else(|_| "secret".to_string());
    let token_data = decode::<Claims>(
        token,
        &DecodingKey::from_secret(secret.as_bytes()),
        &Validation::default(),
    )
    .map_err(|_| StatusCode::UNAUTHORIZED)?;
 
    // 검증된 클레임을 요청 확장에 삽입
    request.extensions_mut().insert(token_data.claims);
 
    Ok(next.run(request).await)
}

use axum::Extension;
 
async fn get_profile(
    Extension(claims): Extension<Claims>,
) -> Json<serde_json::Value> {
    Json(serde_json::json!({
        "user_id": claims.sub,
        "role": claims.role,
    }))
}
 
// 인증이 필요한 라우트에만 미들웨어 적용
let protected_routes = Router::new()
    .route("/profile", get(get_profile))
    .route("/settings", get(get_settings).put(update_settings))
    .layer(middleware::from_fn(auth_middleware));
 
let public_routes = Router::new()
    .route("/health", get(health_check))
    .route("/auth/login", post(login));
 
let app = Router::new()
    .merge(public_routes)
    .merge(protected_routes);

use axum::{
    async_trait,
    extract::{FromRequest, Request, rejection::JsonRejection},
    Json,
    http::StatusCode,
    response::{IntoResponse, Response},
};
use serde::de::DeserializeOwned;
use validator::Validate;
 
// 검증 가능한 Json 추출자
pub struct ValidatedJson<T>(pub T);
 
#[async_trait]
impl<S, T> FromRequest<S> for ValidatedJson<T>
where
    T: DeserializeOwned + Validate,
    S: Send + Sync,
    Json<T>: FromRequest<S, Rejection = JsonRejection>,
{
    type Rejection = ValidationError;
 
    async fn from_request(req: Request, state: &S) -> Result<Self, Self::Rejection> {
        let Json(value) = Json::<T>::from_request(req, state)
            .await
            .map_err(|e| ValidationError::InvalidJson(e.to_string()))?;
 
        value.validate().map_err(|e| ValidationError::Validation(e.to_string()))?;
 
        Ok(ValidatedJson(value))
    }
}
 
#[derive(Debug)]
pub enum ValidationError {
    InvalidJson(String),
    Validation(String),
}
 
impl IntoResponse for ValidationError {
    fn into_response(self) -> Response {
        let (status, message) = match self {
            ValidationError::InvalidJson(msg) => (StatusCode::BAD_REQUEST, msg),
            ValidationError::Validation(msg) => (StatusCode::UNPROCESSABLE_ENTITY, msg),
        };
 
        (status, Json(serde_json::json!({ "error": message }))).into_response()
    }
}

use validator::Validate;
 
#[derive(Debug, Deserialize, Validate)]
struct CreateUserRequest {
    #[validate(length(min = 1, max = 100, message = "이름은 1-100자여야 합니다"))]
    name: String,
 
    #[validate(email(message = "유효한 이메일 주소를 입력하세요"))]
    email: String,
 
    #[validate(length(min = 8, message = "비밀번호는 8자 이상이어야 합니다"))]
    password: String,
}
 
// 핸들러에서 ValidatedJson 사용
async fn create_user(
    ValidatedJson(payload): ValidatedJson<CreateUserRequest>,
) -> (StatusCode, Json<serde_json::Value>) {
    // payload는 이미 검증 완료
    (StatusCode::CREATED, Json(serde_json::json!({
        "message": "사용자 생성 완료",
        "name": payload.name,
    })))
}

use tower::limit::RateLimitLayer;
use std::time::Duration;
 
let app = Router::new()
    .route("/api/search", get(search))
    // 1초에 최대 10개 요청
    .layer(RateLimitLayer::new(10, Duration::from_secs(1)));

use axum::{
    extract::ws::{Message, WebSocket, WebSocketUpgrade},
    response::Response,
};
 
async fn ws_handler(ws: WebSocketUpgrade) -> Response {
    ws.on_upgrade(handle_socket)
}
 
async fn handle_socket(mut socket: WebSocket) {
    // 메시지 수신 및 에코
    while let Some(Ok(msg)) = socket.recv().await {
        match msg {
            Message::Text(text) => {
                let reply = format!("에코: {}", text);
                if socket.send(Message::Text(reply.into())).await.is_err() {
                    break; // 클라이언트 연결 끊김
                }
            }
            Message::Close(_) => break,
            _ => {}
        }
    }
}
 
// 라우트 등록
let app = Router::new()
    .route("/ws", get(ws_handler));

use axum::response::sse::{Event, Sse};
use tokio_stream::StreamExt;
use std::convert::Infallible;
 
async fn sse_handler() -> Sse<impl tokio_stream::Stream<Item = Result<Event, Infallible>>> {
    let stream = tokio_stream::wrappers::IntervalStream::new(
        tokio::time::interval(std::time::Duration::from_secs(1)),
    )
    .map(|_| {
        let now = chrono::Utc::now().to_rfc3339();
        Ok(Event::default().data(format!("서버 시간: {}", now)))
    });
 
    Sse::new(stream)
}

#[cfg(test)]
mod tests {
    use super::*;
    use axum::{
        body::Body,
        http::{Request, StatusCode},
    };
    use http_body_util::BodyExt;
    use tower::ServiceExt;
 
    fn create_test_app() -> Router {
        let state = Arc::new(AppState {
            users: Arc::new(RwLock::new(HashMap::new())),
            next_id: Arc::new(tokio::sync::Mutex::new(1)),
        });
 
        Router::new()
            .route("/users", get(list_users).post(create_user))
            .route("/users/{id}", get(get_user))
            .with_state(state)
    }
 
    #[tokio::test]
    async fn test_create_user() {
        let app = create_test_app();
 
        let response = app
            .oneshot(
                Request::builder()
                    .method("POST")
                    .uri("/users")
                    .header("content-type", "application/json")
                    .body(Body::from(
                        serde_json::to_string(&serde_json::json!({
                            "name": "Alice",
                            "email": "alice@example.com"
                        }))
                        .unwrap(),
                    ))
                    .unwrap(),
            )
            .await
            .unwrap();
 
        assert_eq!(response.status(), StatusCode::CREATED);
 
        let body = response.into_body().collect().await.unwrap().to_bytes();
        let user: serde_json::Value = serde_json::from_slice(&body).unwrap();
        assert_eq!(user["name"], "Alice");
        assert_eq!(user["id"], 1);
    }
 
    #[tokio::test]
    async fn test_get_nonexistent_user() {
        let app = create_test_app();
 
        let response = app
            .oneshot(
                Request::builder()
                    .uri("/users/999")
                    .body(Body::empty())
                    .unwrap(),
            )
            .await
            .unwrap();
 
        assert_eq!(response.status(), StatusCode::NOT_FOUND);
    }
 
    #[tokio::test]
    async fn test_list_users_empty() {
        let app = create_test_app();
 
        let response = app
            .oneshot(
                Request::builder()
                    .uri("/users")
                    .body(Body::empty())
                    .unwrap(),
            )
            .await
            .unwrap();
 
        assert_eq!(response.status(), StatusCode::OK);
 
        let body = response.into_body().collect().await.unwrap().to_bytes();
        let users: Vec<serde_json::Value> = serde_json::from_slice(&body).unwrap();
        assert!(users.is_empty());
    }
}