Authorization

use std::collections::HashMap;
use std::sync::{Arc, RwLock};

use axum::{
    RequestPartsExt, Router,
    extract::{FromRef, FromRequestParts, Path, State},
    http::{StatusCode, request::Parts},
    middleware::from_extractor_with_state,
    routing::{delete, get, post},
};

#[derive(Clone, Copy, Debug, PartialEq, Eq)]
enum Role {
    Member,
    Librarian,
}

#[derive(Clone, Debug)]
struct CurrentUser {
    id: u64,
    role: Role,
}

#[derive(Clone, Default, FromRef)]
struct AppState {
    tokens: Arc<HashMap<String, CurrentUser>>,       // stands in for sessions or JWT validation
    reviews: Arc<RwLock<HashMap<u64, (u64, String)>>>, // review id -> (author id, text)
}

impl<S> FromRequestParts<S> for CurrentUser
where
    S: Send + Sync,
    Arc<HashMap<String, CurrentUser>>: FromRef<S>,
{
    type Rejection = StatusCode;

    async fn from_request_parts(parts: &mut Parts, state: &S) -> Result<Self, StatusCode> {
        let token = parts
            .headers
            .get("authorization")
            .and_then(|v| v.to_str().ok())
            .and_then(|v| v.strip_prefix("Bearer "))
            .ok_or(StatusCode::UNAUTHORIZED)?;
        let tokens: Arc<HashMap<String, CurrentUser>> = FromRef::from_ref(state);
        tokens.get(token).cloned().ok_or(StatusCode::UNAUTHORIZED) // 401: we do not know who you are
    }
}

/// Only librarians. Rejects with 401 (not signed in) or 403 (signed in, wrong role).
struct Librarian(CurrentUser);

impl<S> FromRequestParts<S> for Librarian
where
    S: Send + Sync,
    Arc<HashMap<String, CurrentUser>>: FromRef<S>,
{
    type Rejection = StatusCode;

    async fn from_request_parts(parts: &mut Parts, state: &S) -> Result<Self, StatusCode> {
        let user = parts.extract_with_state::<CurrentUser, S>(state).await?;
        if user.role == Role::Librarian { Ok(Librarian(user)) } else { Err(StatusCode::FORBIDDEN) }
    }
}

async fn add_book(Librarian(user): Librarian) -> String {
    format!("book added by librarian {}\n", user.id)
}

/// Resource ownership: members may delete their own reviews, librarians any review.
async fn delete_review(
    user: CurrentUser,
    State(state): State<AppState>,
    Path(review_id): Path<u64>,
) -> Result<StatusCode, StatusCode> {
    let mut reviews = state.reviews.write().expect("lock poisoned");
    let (author, _) = reviews.get(&review_id).ok_or(StatusCode::NOT_FOUND)?;
    if *author != user.id && user.role != Role::Librarian {
        // reviews are public, so 403 leaks nothing; for private resources answer 404 instead
        return Err(StatusCode::FORBIDDEN);
    }
    reviews.remove(&review_id);
    Ok(StatusCode::NO_CONTENT)
}

async fn reports() -> &'static str {
    "loan statistics\n"
}

#[tokio::main]
async fn main() {
    let state = AppState::default();
    let mut tokens = HashMap::new();
    tokens.insert("member-demo-token".to_owned(), CurrentUser { id: 7, role: Role::Member });
    tokens.insert("librarian-demo-token".to_owned(), CurrentUser { id: 1, role: Role::Librarian });
    let state = AppState { tokens: Arc::new(tokens), ..state };
    state.reviews.write().unwrap().insert(100, (7, "A classic.".into()));
    state.reviews.write().unwrap().insert(101, (8, "Too long.".into()));

    // a whole router guarded by one extractor: every route needs a librarian
    let admin = Router::new()
        .route("/reports", get(reports))
        .route_layer(from_extractor_with_state::<Librarian, _>(state.clone()));

    let app = Router::new()
        .route("/api/v1/books", post(add_book))
        .route("/api/v1/reviews/{id}", delete(delete_review))
        .nest("/admin", admin)
        .with_state(state);
    let listener = tokio::net::TcpListener::bind("127.0.0.1:3000").await.unwrap();
    axum::serve(listener, app).await.unwrap();
}

$ curl -s -o /dev/null -w '%{http_code}\n' -X POST localhost:3000/api/v1/books
401
$ curl -s -o /dev/null -w '%{http_code}\n' -X POST -H 'authorization: Bearer member-demo-token' localhost:3000/api/v1/books
403
$ curl -s -X POST -H 'authorization: Bearer librarian-demo-token' localhost:3000/api/v1/books
book added by librarian 1
$ curl -s -o /dev/null -w '%{http_code}\n' -X DELETE -H 'authorization: Bearer member-demo-token' localhost:3000/api/v1/reviews/101
403
$ curl -s -o /dev/null -w '%{http_code}\n' -X DELETE -H 'authorization: Bearer member-demo-token' localhost:3000/api/v1/reviews/100
204
$ curl -s -o /dev/null -w '%{http_code}\n' -H 'authorization: Bearer member-demo-token' localhost:3000/admin/reports
403

use std::collections::HashSet;

use axum::{Router, routing::get};
use axum_login::{AuthUser, AuthnBackend, AuthzBackend, UserId, permission_required};

#[derive(Clone, Debug)]
struct User {
    id: i64,
    permissions: HashSet<String>,
}

impl AuthUser for User {
    type Id = i64;
    fn id(&self) -> i64 {
        self.id
    }
    fn session_auth_hash(&self) -> &[u8] {
        b"password-hash-goes-here"
    }
}

#[derive(Clone)]
struct Backend;

impl AuthnBackend for Backend {
    type User = User;
    type Credentials = ();
    type Error = std::convert::Infallible;

    async fn authenticate(&self, _: ()) -> Result<Option<User>, Self::Error> {
        Ok(None)
    }
    async fn get_user(&self, id: &UserId<Self>) -> Result<Option<User>, Self::Error> {
        Ok(Some(User { id: *id, permissions: HashSet::from(["books.write".to_owned()]) }))
    }
}

impl AuthzBackend for Backend {
    type Permission = String;

    async fn get_user_permissions(&self, user: &User) -> Result<HashSet<String>, Self::Error> {
        Ok(user.permissions.clone())
    }
}

fn librarian_routes() -> Router {
    Router::new()
        .route("/catalogue/edit", get(|| async { "edit form" }))
        .route_layer(permission_required!(Backend, login_url = "/login", "books.write"))
}

fn main() {
    let _router = librarian_routes(); // add AuthManagerLayer as on the Password Authentication page
}