athene 2.0.4

use athene::prelude::*;
use jsonwebtoken::{
    decode, encode, errors::ErrorKind, DecodingKey, EncodingKey, Header, Validation,
};

const SECRET_KEY: &str = "afafava-rust-lang";

#[derive(Debug, Serialize, Deserialize,Default)]
pub struct JWTClaims {
    username: String,
    exp: usize,
}

#[middleware]
impl JWTClaims {
    pub fn generate_token(&self, secret: &str) -> Result<String> {
        return match encode(
            &Header::default(),
            &self,
            &EncodingKey::from_secret(secret.as_ref()),
        ) {
            Ok(t) => Ok(t),
            Err(_) => Err(Error::Response(
                StatusCode::UNAUTHORIZED,
                json!("JWTToken encode fail!"),
            )),
        };
    }

    pub fn verify(secret: &str, token: &str) -> Result<Self, Error> {
        let validation = Validation::default();
        return match decode::<JWTClaims>(
            token,
            &DecodingKey::from_secret(secret.as_ref()),
            &validation,
        ) {
            Ok(c) => Ok(c.claims),
            Err(err) => match *err.kind() {
                ErrorKind::InvalidToken => {
                    return Err(Error::Response(
                        StatusCode::UNAUTHORIZED,
                        json!("InvalidToken"),
                    ))
                }
                ErrorKind::InvalidIssuer => {
                    return Err(Error::Response(
                        StatusCode::UNAUTHORIZED,
                        json!("InvalidIssuer"),
                    ))
                }
                _ => {
                    return Err(Error::Response(
                        StatusCode::UNAUTHORIZED,
                        json!("InvalidToken other errors"),
                    ))
                }
            },
        };
    }

    pub fn checked_token(token: &str) -> Result<JWTClaims> {
        let claims = JWTClaims::verify(SECRET_KEY, token);
        match claims {
            Ok(token) => Ok(token),
            Err(e) => Err(Error::Other(e.to_string())),
        }
    }

    // JWT 中间件
    async fn next(&self, ctx: Context, next: &dyn Next) -> Result {
        let req = ctx.state.request()?;
        let token = req
            .headers()
            .get("access_token")
            .map(|v| v.to_str().unwrap_or_default().to_string())
            .unwrap_or_default();
        match Self::checked_token(&token) {
            Ok(_) => next.next(ctx).await,
            Err(e) => Err(Error::Other(e.to_string())),
        }
    }
}

#[derive(Serialize, Deserialize, Default, Debug)]
pub struct UserController {
    pub name: String,
    pub age: u16,
}

#[derive(Serialize, Deserialize, Default)]
pub struct ApiResponse<T> {
    code: u32,
    msg: String,
    data: T,
}

// http://127.0.0.1:7878/user
#[controller(name = "user")]
impl UserController {

    // http://127.0.0.1:7878/user/register
    #[post("/register")]
    async fn register(&self, mut req: Request) -> impl Responder {
        let user = req.parse::<Self>().await?;
        let jwt = JWTClaims {
            username: user.name.clone(),
            exp: 10000000000,
        };

        let token = jwt.generate_token(SECRET_KEY)?;
        Ok::<_, Error>((
            200,
            Json(ApiResponse {
                code: 200,
                msg: String::from("successful"),
                data: token,
            }),
        ))
    }

    // http://127.0.0.1:7878/user/login
    #[post("/login")]
    async fn login(&self, mut req: Request) -> impl Responder {
        let user = req.parse::<Self>().await?;
        Ok::<_, Error>((200, Json(user)))
    }
}

#[tokio::main]
async fn main() -> Result<()> {
    let app = athene::new();
    let app = app
        .router(|r| r.controller(UserController::default()))
        .middleware(|m| m.apply(JWTClaims::default(), vec!["/user/login"], vec!["/user/register"]));
    app.listen("127.0.0.1:7878").await
}