From efaaf1bb5f7f108e9bd29d1f19584968a4e56341 Mon Sep 17 00:00:00 2001
From: Bryan Stitt <bryan@llamanodes.com>
Date: Tue, 16 May 2023 15:27:18 -0700
Subject: [PATCH] send the correct content type

---
 web3_proxy/src/frontend/mod.rs    |  2 +-
 web3_proxy/src/frontend/status.rs | 57 ++++++++++++++++++++++++-------
 2 files changed, 46 insertions(+), 13 deletions(-)

diff --git a/web3_proxy/src/frontend/mod.rs b/web3_proxy/src/frontend/mod.rs
index aa827909..9eb2f6d9 100644
--- a/web3_proxy/src/frontend/mod.rs
+++ b/web3_proxy/src/frontend/mod.rs
@@ -38,7 +38,7 @@ pub enum ResponseCacheKey {
 
 pub type ResponseCache = quick_cache_ttl::CacheWithTTL<
     ResponseCacheKey,
-    (StatusCode, axum::body::Bytes),
+    (StatusCode, &'static str, axum::body::Bytes),
     UnitWeighter,
     quick_cache_ttl::DefaultHashBuilder,
 >;
diff --git a/web3_proxy/src/frontend/status.rs b/web3_proxy/src/frontend/status.rs
index afa9ebb0..1f1cb94b 100644
--- a/web3_proxy/src/frontend/status.rs
+++ b/web3_proxy/src/frontend/status.rs
@@ -5,7 +5,12 @@
 
 use super::{ResponseCache, ResponseCacheKey};
 use crate::app::{Web3ProxyApp, APP_USER_AGENT};
-use axum::{body::Bytes, http::StatusCode, response::IntoResponse, Extension};
+use axum::{
+    body::{Bytes, Full},
+    http::StatusCode,
+    response::{IntoResponse, Response},
+    Extension,
+};
 use axum_macros::debug_handler;
 use once_cell::sync::Lazy;
 use serde_json::json;
@@ -17,26 +22,40 @@ static HEALTH_NOT_OK: Lazy<Bytes> = Lazy::new(|| Bytes::from(":(\n"));
 static BACKUPS_NEEDED_TRUE: Lazy<Bytes> = Lazy::new(|| Bytes::from("true\n"));
 static BACKUPS_NEEDED_FALSE: Lazy<Bytes> = Lazy::new(|| Bytes::from("false\n"));
 
+static CONTENT_TYPE_JSON: &str = "application/json";
+static CONTENT_TYPE_PLAIN: &str = "text/plain";
+
 /// Health check page for load balancers to use.
 #[debug_handler]
 pub async fn health(
     Extension(app): Extension<Arc<Web3ProxyApp>>,
     Extension(cache): Extension<Arc<ResponseCache>>,
 ) -> impl IntoResponse {
-    cache
+    let (code, content_type, body) = cache
         .get_or_insert_async::<Infallible, _>(&ResponseCacheKey::Health, async move {
             Ok(_health(app).await)
         })
         .await
+        .unwrap();
+
+    Response::builder()
+        .status(code)
+        .header("content-type", content_type)
+        .body(Full::from(body))
+        .unwrap()
 }
 
 // TODO: _health doesn't need to be async, but _quick_cache_ttl needs an async function
 #[inline]
-async fn _health(app: Arc<Web3ProxyApp>) -> (StatusCode, Bytes) {
+async fn _health(app: Arc<Web3ProxyApp>) -> (StatusCode, &'static str, Bytes) {
     if app.balanced_rpcs.synced() {
-        (StatusCode::OK, HEALTH_OK.clone())
+        (StatusCode::OK, CONTENT_TYPE_PLAIN, HEALTH_OK.clone())
     } else {
-        (StatusCode::SERVICE_UNAVAILABLE, HEALTH_NOT_OK.clone())
+        (
+            StatusCode::SERVICE_UNAVAILABLE,
+            CONTENT_TYPE_PLAIN,
+            HEALTH_NOT_OK.clone(),
+        )
     }
 }
 
@@ -46,15 +65,22 @@ pub async fn backups_needed(
     Extension(app): Extension<Arc<Web3ProxyApp>>,
     Extension(cache): Extension<Arc<ResponseCache>>,
 ) -> impl IntoResponse {
-    cache
+    let (code, content_type, body) = cache
         .get_or_insert_async::<Infallible, _>(&ResponseCacheKey::BackupsNeeded, async move {
             Ok(_backups_needed(app).await)
         })
         .await
+        .unwrap();
+
+    Response::builder()
+        .status(code)
+        .header("content-type", content_type)
+        .body(Full::from(body))
+        .unwrap()
 }
 
 #[inline]
-async fn _backups_needed(app: Arc<Web3ProxyApp>) -> (StatusCode, Bytes) {
+async fn _backups_needed(app: Arc<Web3ProxyApp>) -> (StatusCode, &'static str, Bytes) {
     let code = {
         let consensus_rpcs = app
             .balanced_rpcs
@@ -75,9 +101,9 @@ async fn _backups_needed(app: Arc<Web3ProxyApp>) -> (StatusCode, Bytes) {
     };
 
     if matches!(code, StatusCode::OK) {
-        (code, BACKUPS_NEEDED_FALSE.clone())
+        (code, CONTENT_TYPE_PLAIN, BACKUPS_NEEDED_FALSE.clone())
     } else {
-        (code, BACKUPS_NEEDED_TRUE.clone())
+        (code, CONTENT_TYPE_PLAIN, BACKUPS_NEEDED_TRUE.clone())
     }
 }
 
@@ -89,16 +115,23 @@ pub async fn status(
     Extension(app): Extension<Arc<Web3ProxyApp>>,
     Extension(cache): Extension<Arc<ResponseCache>>,
 ) -> impl IntoResponse {
-    cache
+    let (code, content_type, body) = cache
         .get_or_insert_async::<Infallible, _>(&ResponseCacheKey::Status, async move {
             Ok(_status(app).await)
         })
         .await
+        .unwrap();
+
+    Response::builder()
+        .status(code)
+        .header("content-type", content_type)
+        .body(Full::from(body))
+        .unwrap()
 }
 
 // TODO: _status doesn't need to be async, but _quick_cache_ttl needs an async function
 #[inline]
-async fn _status(app: Arc<Web3ProxyApp>) -> (StatusCode, Bytes) {
+async fn _status(app: Arc<Web3ProxyApp>) -> (StatusCode, &'static str, Bytes) {
     // TODO: what else should we include? uptime, cache hit rates, cpu load, memory used
     // TODO: the hostname is probably not going to change. only get once at the start?
     let body = json!({
@@ -120,5 +153,5 @@ async fn _status(app: Arc<Web3ProxyApp>) -> (StatusCode, Bytes) {
         StatusCode::INTERNAL_SERVER_ERROR
     };
 
-    (code, body)
+    (code, CONTENT_TYPE_JSON, body)
 }