From 52151f8b2239c186f1d62c1a3f4d913a1774a266 Mon Sep 17 00:00:00 2001 From: Bryan Stitt Date: Fri, 12 Aug 2022 21:41:11 +0000 Subject: [PATCH] set content-type header --- Cargo.lock | 26 +++++++++ web3_proxy/Cargo.toml | 2 +- web3_proxy/src/stats.rs | 113 ++++++++++++++++++++++++++++++++++++++++ 3 files changed, 140 insertions(+), 1 deletion(-) create mode 100644 web3_proxy/src/stats.rs diff --git a/Cargo.lock b/Cargo.lock index 76430a26..053ad76f 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -345,6 +345,7 @@ dependencies = [ "bitflags", "bytes", "futures-util", + "headers", "http", "http-body", "hyper", @@ -2188,6 +2189,31 @@ dependencies = [ "hashbrown", ] +[[package]] +name = "headers" +version = "0.3.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4cff78e5788be1e0ab65b04d306b2ed5092c815ec97ec70f4ebd5aee158aa55d" +dependencies = [ + "base64 0.13.0", + "bitflags", + "bytes", + "headers-core", + "http", + "httpdate", + "mime", + "sha-1", +] + +[[package]] +name = "headers-core" +version = "0.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e7f66481bfee273957b1f20485a4ff3362987f85b2c236580d81b4eb7a326429" +dependencies = [ + "http", +] + [[package]] name = "heck" version = "0.3.3" diff --git a/web3_proxy/Cargo.toml b/web3_proxy/Cargo.toml index 49745cf9..880c66b3 100644 --- a/web3_proxy/Cargo.toml +++ b/web3_proxy/Cargo.toml @@ -19,7 +19,7 @@ migration = { path = "../migration" } anyhow = { version = "1.0.61", features = ["backtrace"] } arc-swap = "1.5.1" argh = "0.1.8" -axum = { version = "0.5.15", features = ["serde_json", "tokio-tungstenite", "ws"] } +axum = { version = "0.5.15", features = ["headers", "serde_json", "tokio-tungstenite", "ws"] } axum-client-ip = "0.2.0" counter = "0.5.6" dashmap = "5.3.4" diff --git a/web3_proxy/src/stats.rs b/web3_proxy/src/stats.rs new file mode 100644 index 00000000..23a24b8f --- /dev/null +++ b/web3_proxy/src/stats.rs @@ -0,0 +1,113 @@ +use axum::headers::{ContentType, HeaderName}; +use axum::http::HeaderValue; +use axum::response::{IntoResponse, Response}; +use axum::{routing::get, Extension, Router, TypedHeader}; +use prometheus_client::encoding::text::encode; +use prometheus_client::encoding::text::Encode; +use prometheus_client::metrics::counter::Counter; +use prometheus_client::metrics::family::Family; +use prometheus_client::registry::Registry; +use std::net::SocketAddr; +use std::sync::Arc; +use tracing::info; + +#[derive(Clone, Hash, PartialEq, Eq, Encode)] +pub struct ProxyRequestLabels { + protocol: Protocol, + rpc_method: String, + /// anonymous is user 0 + user_id: u64, +} + +#[derive(Clone, Hash, PartialEq, Eq, Encode)] +pub enum Protocol { + HTTP, + Websocket, +} + +pub struct AppStatsRegistry { + pub registry: Registry, + pub stats: AppStats, +} + +#[derive(Clone)] +pub struct AppStats { + pub proxy_requests: Family, +} + +impl AppStatsRegistry { + pub fn new() -> Arc { + // Note the angle brackets to make sure to use the default (dynamic + // dispatched boxed metric) for the generic type parameter. + let mut registry = ::default(); + + // stats for GET and POST + let proxy_requests = Family::::default(); + registry.register( + // With the metric name. + "http_requests", + // And the metric help text. + "Number of HTTP requests received", + Box::new(proxy_requests.clone()), + ); + + let new = Self { + registry, + stats: AppStats { proxy_requests }, + }; + + Arc::new(new) + } + + pub async fn serve(self: Arc, port: u16) -> anyhow::Result<()> { + // build our application with a route + // order most to least common + // TODO: 404 any unhandled routes? + let app = Router::new() + .route("/", get(root)) + .layer(Extension(self.clone())); + + // run our app with hyper + // TODO: allow only listening on localhost? + let addr = SocketAddr::from(([0, 0, 0, 0], port)); + info!("prometheus listening on port {}", port); + // TODO: into_make_service is enough if we always run behind a proxy. make into_make_service_with_connect_info optional? + + /* + It sequentially looks for an IP in: + - x-forwarded-for header (de-facto standard) + - x-real-ip header + - forwarded header (new standard) + - axum::extract::ConnectInfo (if not behind proxy) + + So we probably won't need into_make_service_with_connect_info, but it shouldn't hurt + */ + let service = app.into_make_service_with_connect_info::(); + // let service = app.into_make_service(); + + // `axum::Server` is a re-export of `hyper::Server` + axum::Server::bind(&addr) + // TODO: option to use with_connect_info. we want it in dev, but not when running behind a proxy, but not + .serve(service) + .await + .map_err(Into::into) + } +} + +async fn root(Extension(stats_registry): Extension>) -> Response { + let mut buffer = vec![]; + + encode(&mut buffer, &stats_registry.registry).unwrap(); + + let s = String::from_utf8(buffer).unwrap(); + + let mut r = s.into_response(); + + // // TODO: is there an easier way to do this? + r.headers_mut().insert( + HeaderName::from_static("content-type"), + HeaderValue::from_static("application/openmetrics-text; version=1.0.0; charset=utf-8"), + ); + + r +}