From 52151f8b2239c186f1d62c1a3f4d913a1774a266 Mon Sep 17 00:00:00 2001
From: Bryan Stitt <bryan@satoshiandkin.com>
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<ProxyRequestLabels, Counter>,
+}
+
+impl AppStatsRegistry {
+    pub fn new() -> Arc<Self> {
+        // Note the angle brackets to make sure to use the default (dynamic
+        // dispatched boxed metric) for the generic type parameter.
+        let mut registry = <Registry>::default();
+
+        // stats for GET and POST
+        let proxy_requests = Family::<ProxyRequestLabels, Counter>::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<Self>, 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::<SocketAddr>();
+        // 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<Arc<AppStatsRegistry>>) -> 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
+}