simplify api router, add transport limit 2MB

src/routing/error_cause.rs

@@ -36,6 +36,7 @@ pub enum ErrorCause {
     UnableToParseHTTPArg(String),
     LoadShed,
     RequestTooLarge,
+    IncorrectPrincipal,
     MalformedRequest(String),
     MalformedResponse(String),
     NoAuthority,
@@ -68,6 +69,7 @@ impl ErrorCause {
             Self::UnableToParseHTTPArg(_) => StatusCode::BAD_REQUEST,
             Self::LoadShed => StatusCode::TOO_MANY_REQUESTS,
             Self::RequestTooLarge => StatusCode::PAYLOAD_TOO_LARGE,
+            Self::IncorrectPrincipal => StatusCode::BAD_REQUEST,
             Self::MalformedRequest(_) => StatusCode::BAD_REQUEST,
             Self::MalformedResponse(_) => StatusCode::INTERNAL_SERVER_ERROR,
             Self::NoAuthority => StatusCode::BAD_REQUEST,
@@ -138,6 +140,7 @@ impl fmt::Display for ErrorCause {
             Self::UnableToParseHTTPArg(_) => write!(f, "unable_to_parse_http_arg"),
             Self::LoadShed => write!(f, "load_shed"),
             Self::RequestTooLarge => write!(f, "request_too_large"),
+            Self::IncorrectPrincipal => write!(f, "incorrect_principal"),
             Self::MalformedRequest(_) => write!(f, "malformed_request"),
             Self::MalformedResponse(_) => write!(f, "malformed_response"),
             Self::CanisterIdNotFound => write!(f, "canister_id_not_found"),

src/routing/ic/transport.rs

@@ -25,6 +25,7 @@ use crate::http::Client as HttpClient;
 type AgentFuture<'a, V> = Pin<Box<dyn Future<Output = Result<V, AgentError>> + Send + 'a>>;
 
 const CONTENT_TYPE_CBOR: HeaderValue = HeaderValue::from_static("application/cbor");
+const MAX_RESPONSE_SIZE: usize = 2 * 1_048_576;
 
 pub struct PassHeaders {
     pub headers_in: HeaderMap<HeaderValue>,
@@ -70,7 +71,7 @@ impl ReqwestTransport {
         Ok(Self {
             route_provider,
             client,
-            max_response_body_size: None,
+            max_response_body_size: Some(MAX_RESPONSE_SIZE),
         })
     }
 

src/routing/mod.rs

@@ -128,38 +128,29 @@ pub fn setup_router(
     // Prepare the HTTP-IC library
     let route_provider = Arc::new(RoundRobinRouteProvider::new(cli.ic.url.clone())?);
     let client = ic::setup(cli, http_client.clone(), route_provider.clone())?;
-    let handler_state = Arc::new(handler::HandlerState::new(client));
 
-    let api_state = Arc::new(proxy::ApiProxyState::new(
+    // Prepare the states
+    let state_handler = Arc::new(handler::HandlerState::new(client));
+    let state_api = Arc::new(proxy::ApiProxyState::new(
         http_client.clone(),
         route_provider,
     ));
 
+    // IC API proxy router
+    let router_api = Router::new()
+        .route("/canister/:principal/query", post(proxy::api_proxy))
+        .route("/canister/:principal/call", post(proxy::api_proxy))
+        .route("/canister/:principal/read_state", post(proxy::api_proxy))
+        .route("/subnet/:principal/read_state", post(proxy::api_proxy))
+        .route("/status", get(proxy::api_proxy))
+        .with_state(state_api);
+
     let router = Router::new()
-        .route(
-            "/api/v2/canister/:id/query",
-            post(proxy::api_proxy).with_state(api_state.clone()),
-        )
-        .route(
-            "/api/v2/canister/:id/call",
-            post(proxy::api_proxy).with_state(api_state.clone()),
-        )
-        .route(
-            "/api/v2/canister/:id/read_state",
-            post(proxy::api_proxy).with_state(api_state.clone()),
-        )
-        .route(
-            "/api/v2/subnet/:id/read_state",
-            post(proxy::api_proxy).with_state(api_state.clone()),
-        )
-        .route(
-            "/api/v2/status",
-            get(proxy::api_proxy).with_state(api_state),
-        )
+        .nest("/api/v2", router_api)
         .fallback(
             get(handler::handler)
                 .post(handler::handler)
-                .with_state(handler_state),
+                .with_state(state_handler),
         )
         .layer(common_layers);
 
@@ -173,19 +164,17 @@ pub fn setup_router(
         urls.iter_mut().for_each(|x| x.set_path(""));
 
         let state = Arc::new(proxy::IssuerProxyState::new(http_client, urls));
-
-        router
+        let router_issuer = Router::new()
             .route(
-                "/registrations/:id",
+                "/:id",
                 get(proxy::issuer_proxy)
                     .put(proxy::issuer_proxy)
-                    .delete(proxy::issuer_proxy)
-                    .with_state(state.clone()),
-            )
-            .route(
-                "/registrations",
-                post(proxy::issuer_proxy).with_state(state),
+                    .delete(proxy::issuer_proxy),
             )
+            .route("/", post(proxy::issuer_proxy))
+            .with_state(state);
+
+        router.nest("/registrations", router_issuer)
     } else {
         router
     };

src/routing/proxy.rs

@@ -15,9 +15,8 @@ use ic_agent::agent::http_transport::route_provider::RouteProvider;
 use regex::Regex;
 use url::Url;
 
-use crate::http::Client;
-
 use super::{body, error_cause::ErrorCause};
+use crate::http::Client;
 
 lazy_static::lazy_static! {
     pub static ref REGEX_REG_ID: Regex = Regex::new(r"^[a-zA-Z0-9]+$").unwrap();
@@ -60,13 +59,12 @@ pub struct ApiProxyState {
 // Proxies /api/v2/... endpoints to the IC
 pub async fn api_proxy(
     State(state): State<Arc<ApiProxyState>>,
-    id: Option<Path<String>>,
+    principal: Option<Path<String>>,
     request: Request,
 ) -> Result<impl IntoResponse, ErrorCause> {
     // Check principal for correctness
-    if let Some(v) = id {
-        Principal::from_text(v.0)
-            .map_err(|e| ErrorCause::MalformedRequest(format!("incorrect principal: {e}")))?;
+    if let Some(v) = principal {
+        Principal::from_text(v.0).map_err(|_| ErrorCause::IncorrectPrincipal)?;
     }
 
     // Obtain the next IC URL from the provider