appview/rust/src/router/xrpc.rs

use crate::atproto::Nsid;
use std::{
  collections::HashMap,
  pin::Pin,
  future::Future,
};
use axum::{
  extract::{
    Json,
    Query,
    Request,
    FromRequest,
    FromRequestParts,
    rejection::QueryRejection,
  },
  body::Bytes,
  routing::{
    get,
    post,
    method_routing::MethodRouter,
  },
  http::{
    StatusCode,
    request::Parts,
  },
  Router as axumRouter,
};
use serde_json::{Value, json};

enum Path {
  Nsid(Nsid),
  NotImplemented,
}

pub struct XrpcEndpoint {
  path: Path,
  resolver: MethodRouter,
}

pub type Response = (StatusCode, Json<Value>);
pub fn error(code: StatusCode, error: &str, message: &str) -> Response {
  (
    code,
    Json(json!({
      "error": error,
      "message": message
    }))
  )
}
pub fn response(code: StatusCode, message: &str) -> Response {
  error(code, "", message)
}

pub struct QueryInput {
  parameters: HashMap<String, String>,
}
impl<S> FromRequestParts<S> for QueryInput
where
  S: Send + Sync,
{
  type Rejection = Response;

  async fn from_request_parts(parts: &mut Parts, _state: &S)
    -> Result<Self, Self::Rejection> {
    let query_params: Result<Query<HashMap<String, String>>, QueryRejection> = Query::try_from_uri(&parts.uri);
    match query_params {
      Ok(p) => Ok(QueryInput { parameters: p.0 }),
      Err(e) => Err(error(StatusCode::BAD_REQUEST, "Bad Parameters", &e.body_text())),
    }
  }
}
pub struct ProcedureInput {
  parameters: HashMap<String, String>,
  input: Json<Value>,
}
impl<S> FromRequest<S> for ProcedureInput 
where
  Bytes: FromRequest<S>,
  S: Send + Sync,
{
  type Rejection = Response;

  async fn from_request(req: Request, state: &S)
    -> Result<Self, Self::Rejection> {
    let query_params: Result<Query<HashMap<String, String>>, QueryRejection> = Query::try_from_uri(req.uri());
    let parameters = match query_params {
      Ok(p) => p.0,
      Err(e) => return Err(error(StatusCode::BAD_REQUEST, "Bad Parameters", &e.body_text())),
    };

    let json_value = Json::<Value>::from_request(req, state).await;
    let input: Json<Value> = match json_value {
      Ok(v) => v,
      Err(e) => return Err(error(StatusCode::BAD_REQUEST, "Bad Parameters", &e.body_text())),
    };

    Ok(ProcedureInput { parameters, input })
  }
}

pub trait XrpcHandler<Input>: Send + Sync + 'static {
  fn call(&self, input: Input) 
    -> Pin<Box<dyn Future<Output = Response> + Send>>;
}
impl<F, Fut> XrpcHandler<QueryInput> for F
where
  F: Fn(QueryInput) -> Fut + Send + Sync + 'static,
  Fut: Future<Output = Response> + Send + 'static,
{
  fn call(&self, input: QueryInput) 
    -> Pin<Box<dyn Future<Output = Response>+ Send>> {
    Box::pin((self)(input))
  }
}
impl<F, Fut> XrpcHandler<ProcedureInput> for F
where
  F: Fn(ProcedureInput) -> Fut + Send + Sync + 'static,
  Fut: Future<Output = Response> + Send + 'static,
{
  fn call(&self, input: ProcedureInput) 
    -> Pin<Box<dyn Future<Output = Response>+ Send>> {
    Box::pin((self)(input))
  }
}

impl XrpcEndpoint {
  pub fn new_query<Q>(nsid: Nsid, query: Q) -> Self
  where
    Q: XrpcHandler<QueryInput> + Clone
  {
    XrpcEndpoint {
      path: Path::Nsid(nsid),
      resolver: get(async move | mut parts: Parts | -> Response {
        match QueryInput::from_request_parts(&mut parts, &()).await {
          Ok(qi) => query.call(qi).await,
          Err(e) => e
        }
      })
    }
  }

  pub fn new_procedure<P>(nsid: Nsid, procedure: P) -> Self
  where
    P: XrpcHandler<ProcedureInput> + Clone
  {
    XrpcEndpoint {
      path: Path::Nsid(nsid),
      resolver: post(async move | req: Request | -> Response {
        match ProcedureInput::from_request(req, &()).await {
          Ok(pi) => procedure.call(pi).await,
          Err(e) => e
        }
      })
    }
  }

  pub fn add_to_router(self, router: axumRouter) -> axumRouter {
    let path = match self.path {
      Path::Nsid(nsid) => &("/xrpc/".to_owned() + nsid.as_str()),
      Path::NotImplemented => "/xrpc/{*nsid}",
    };

    router.route(path, self.resolver)
  }

  pub fn not_implemented() -> Self {
    let resolver = (
      StatusCode::NOT_IMPLEMENTED,
      Json(json!({
        "error": "MethodNotImplemented",
        "message": "Method Not Implemented"
      }))
    );

    XrpcEndpoint {
      path: Path::NotImplemented,
      resolver: get(resolver.clone()).post(resolver),
    }
  }
}