LCOV - code coverage report
Current view: top level - compute_tools/src/http - server.rs (source / functions) Coverage Total Hit
Test: 5fe7fa8d483b39476409aee736d6d5e32728bfac.info Lines: 0.0 % 142 0
Test Date: 2025-03-12 16:10:49 Functions: 0.0 % 18 0

            Line data    Source code
       1              : use std::fmt::Display;
       2              : use std::net::{IpAddr, Ipv6Addr, SocketAddr};
       3              : use std::sync::Arc;
       4              : use std::time::Duration;
       5              : 
       6              : use anyhow::Result;
       7              : use axum::Router;
       8              : use axum::extract::Request;
       9              : use axum::middleware::{self, Next};
      10              : use axum::response::{IntoResponse, Response};
      11              : use axum::routing::{get, post};
      12              : use http::StatusCode;
      13              : use jsonwebtoken::jwk::JwkSet;
      14              : use tokio::net::TcpListener;
      15              : use tower::ServiceBuilder;
      16              : use tower_http::{
      17              :     auth::AsyncRequireAuthorizationLayer, request_id::PropagateRequestIdLayer, trace::TraceLayer,
      18              : };
      19              : use tracing::{Span, error, info};
      20              : use uuid::Uuid;
      21              : 
      22              : use super::{
      23              :     headers::X_REQUEST_ID,
      24              :     middleware::authorize::Authorize,
      25              :     routes::{
      26              :         check_writability, configure, database_schema, dbs_and_roles, extension_server, extensions,
      27              :         grants, insights, metrics, metrics_json, status, terminate,
      28              :     },
      29              : };
      30              : use crate::compute::ComputeNode;
      31              : 
      32              : /// `compute_ctl` has two servers: internal and external. The internal server
      33              : /// binds to the loopback interface and handles communication from clients on
      34              : /// the compute. The external server is what receives communication from the
      35              : /// control plane, the metrics scraper, etc. We make the distinction because
      36              : /// certain routes in `compute_ctl` only need to be exposed to local processes
      37              : /// like Postgres via the neon extension and local_proxy.
      38              : #[derive(Clone, Debug)]
      39              : pub enum Server {
      40              :     Internal {
      41              :         port: u16,
      42              :     },
      43              :     External {
      44              :         port: u16,
      45              :         jwks: JwkSet,
      46              :         compute_id: String,
      47              :     },
      48              : }
      49              : 
      50              : impl Display for Server {
      51            0 :     fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
      52            0 :         match self {
      53            0 :             Server::Internal { .. } => f.write_str("internal"),
      54            0 :             Server::External { .. } => f.write_str("external"),
      55              :         }
      56            0 :     }
      57              : }
      58              : 
      59              : impl From<&Server> for Router<Arc<ComputeNode>> {
      60            0 :     fn from(server: &Server) -> Self {
      61            0 :         let mut router = Router::<Arc<ComputeNode>>::new();
      62              : 
      63            0 :         router = match server {
      64              :             Server::Internal { .. } => {
      65            0 :                 router = router
      66            0 :                     .route(
      67            0 :                         "/extension_server/{*filename}",
      68            0 :                         post(extension_server::download_extension),
      69            0 :                     )
      70            0 :                     .route("/extensions", post(extensions::install_extension))
      71            0 :                     .route("/grants", post(grants::add_grant));
      72            0 : 
      73            0 :                 // Add in any testing support
      74            0 :                 if cfg!(feature = "testing") {
      75            0 :                     use super::routes::failpoints;
      76            0 : 
      77            0 :                     router = router.route("/failpoints", post(failpoints::configure_failpoints));
      78            0 :                 }
      79              : 
      80            0 :                 router
      81              :             }
      82              :             Server::External {
      83            0 :                 jwks, compute_id, ..
      84            0 :             } => {
      85            0 :                 let unauthenticated_router =
      86            0 :                     Router::<Arc<ComputeNode>>::new().route("/metrics", get(metrics::get_metrics));
      87            0 : 
      88            0 :                 let authenticated_router = Router::<Arc<ComputeNode>>::new()
      89            0 :                     .route("/check_writability", post(check_writability::is_writable))
      90            0 :                     .route("/configure", post(configure::configure))
      91            0 :                     .route("/database_schema", get(database_schema::get_schema_dump))
      92            0 :                     .route("/dbs_and_roles", get(dbs_and_roles::get_catalog_objects))
      93            0 :                     .route("/insights", get(insights::get_insights))
      94            0 :                     .route("/metrics.json", get(metrics_json::get_metrics))
      95            0 :                     .route("/status", get(status::get_status))
      96            0 :                     .route("/terminate", post(terminate::terminate))
      97            0 :                     .layer(AsyncRequireAuthorizationLayer::new(Authorize::new(
      98            0 :                         compute_id.clone(),
      99            0 :                         jwks.clone(),
     100            0 :                     )));
     101            0 : 
     102            0 :                 router
     103            0 :                     .merge(unauthenticated_router)
     104            0 :                     .merge(authenticated_router)
     105              :             }
     106              :         };
     107              : 
     108            0 :         router
     109            0 :             .fallback(Server::handle_404)
     110            0 :             .method_not_allowed_fallback(Server::handle_405)
     111            0 :             .layer(
     112            0 :                 ServiceBuilder::new()
     113            0 :                     .layer(tower_otel::trace::HttpLayer::server(tracing::Level::INFO))
     114            0 :                     // Add this middleware since we assume the request ID exists
     115            0 :                     .layer(middleware::from_fn(maybe_add_request_id_header))
     116            0 :                     .layer(
     117            0 :                         TraceLayer::new_for_http()
     118            0 :                             .on_request(|request: &http::Request<_>, _span: &Span| {
     119            0 :                                 let request_id = request
     120            0 :                                     .headers()
     121            0 :                                     .get(X_REQUEST_ID)
     122            0 :                                     .unwrap()
     123            0 :                                     .to_str()
     124            0 :                                     .unwrap();
     125            0 : 
     126            0 :                                 info!(%request_id, "{} {}", request.method(), request.uri());
     127            0 :                             })
     128            0 :                             .on_response(
     129            0 :                                 |response: &http::Response<_>, latency: Duration, _span: &Span| {
     130            0 :                                     let request_id = response
     131            0 :                                         .headers()
     132            0 :                                         .get(X_REQUEST_ID)
     133            0 :                                         .unwrap()
     134            0 :                                         .to_str()
     135            0 :                                         .unwrap();
     136            0 : 
     137            0 :                                     info!(
     138              :                                         %request_id,
     139            0 :                                         code = response.status().as_u16(),
     140            0 :                                         latency = latency.as_millis()
     141              :                                     );
     142            0 :                                 },
     143            0 :                             ),
     144            0 :                     )
     145            0 :                     .layer(PropagateRequestIdLayer::x_request_id()),
     146            0 :             )
     147            0 :     }
     148              : }
     149              : 
     150              : impl Server {
     151            0 :     async fn handle_404() -> impl IntoResponse {
     152            0 :         StatusCode::NOT_FOUND
     153            0 :     }
     154              : 
     155            0 :     async fn handle_405() -> impl IntoResponse {
     156            0 :         StatusCode::METHOD_NOT_ALLOWED
     157            0 :     }
     158              : 
     159            0 :     async fn listener(&self) -> Result<TcpListener> {
     160            0 :         let addr = SocketAddr::new(self.ip(), self.port());
     161            0 :         let listener = TcpListener::bind(&addr).await?;
     162              : 
     163            0 :         Ok(listener)
     164            0 :     }
     165              : 
     166            0 :     fn ip(&self) -> IpAddr {
     167            0 :         match self {
     168              :             // TODO: Change this to Ipv6Addr::LOCALHOST when the GitHub runners
     169              :             // allow binding to localhost
     170            0 :             Server::Internal { .. } => IpAddr::from(Ipv6Addr::UNSPECIFIED),
     171            0 :             Server::External { .. } => IpAddr::from(Ipv6Addr::UNSPECIFIED),
     172              :         }
     173            0 :     }
     174              : 
     175            0 :     fn port(&self) -> u16 {
     176            0 :         match self {
     177            0 :             Server::Internal { port, .. } => *port,
     178            0 :             Server::External { port, .. } => *port,
     179              :         }
     180            0 :     }
     181              : 
     182            0 :     async fn serve(self, compute: Arc<ComputeNode>) {
     183            0 :         let listener = self.listener().await.unwrap_or_else(|e| {
     184            0 :             // If we can't bind, the compute cannot operate correctly
     185            0 :             panic!(
     186            0 :                 "failed to bind the compute_ctl {} HTTP server to {}: {}",
     187            0 :                 self,
     188            0 :                 SocketAddr::new(self.ip(), self.port()),
     189            0 :                 e
     190            0 :             );
     191            0 :         });
     192            0 : 
     193            0 :         if tracing::enabled!(tracing::Level::INFO) {
     194            0 :             let local_addr = match listener.local_addr() {
     195            0 :                 Ok(local_addr) => local_addr,
     196            0 :                 Err(_) => SocketAddr::new(self.ip(), self.port()),
     197              :             };
     198              : 
     199            0 :             info!(
     200            0 :                 "compute_ctl {} HTTP server listening at {}",
     201              :                 self, local_addr
     202              :             );
     203            0 :         }
     204              : 
     205            0 :         let router = Router::from(&self)
     206            0 :             .with_state(compute)
     207            0 :             .into_make_service_with_connect_info::<SocketAddr>();
     208              : 
     209            0 :         if let Err(e) = axum::serve(listener, router).await {
     210            0 :             error!("compute_ctl {} HTTP server error: {}", self, e);
     211            0 :         }
     212            0 :     }
     213              : 
     214            0 :     pub fn launch(self, compute: &Arc<ComputeNode>) {
     215            0 :         let state = Arc::clone(compute);
     216            0 : 
     217            0 :         info!("Launching the {} server", self);
     218              : 
     219            0 :         tokio::spawn(self.serve(state));
     220            0 :     }
     221              : }
     222              : 
     223              : /// This middleware function allows compute_ctl to generate its own request ID
     224              : /// if one isn't supplied. The control plane will always send one as a UUID. The
     225              : /// neon Postgres extension on the other hand does not send one.
     226            0 : async fn maybe_add_request_id_header(mut request: Request, next: Next) -> Response {
     227            0 :     let headers = request.headers_mut();
     228            0 :     if headers.get(X_REQUEST_ID).is_none() {
     229            0 :         headers.append(X_REQUEST_ID, Uuid::new_v4().to_string().parse().unwrap());
     230            0 :     }
     231              : 
     232            0 :     next.run(request).await
     233            0 : }
        

Generated by: LCOV version 2.1-beta