Merge pull request #12 from njasm/feature_authorization

pietroalbini · web-flow · commit cc856e0fd0a4 · 2020-11-25T13:30:08.000+01:00
MONITORBOT_SECRET authorization
diff --git a/src/config.rs b/src/config.rs
@@ -6,6 +6,8 @@ const ENVIRONMENT_VARIABLE_PREFIX: &str = "MONITORBOT_";
 
 #[derive(Clone, Debug)]
 pub struct Config {
+    // authorization secret (token) to be able to scrape the metrics endpoint
+    pub secret: String,
     // http server port to bind to
     pub port: u16,
     // github api tokens to collect rate limit statistics
@@ -17,6 +19,7 @@ pub struct Config {
 impl Config {
     pub fn from_env() -> Result<Self, Error> {
         Ok(Self {
+            secret: require_env("SECRET")?,
             port: default_env("PORT", 3001)?,
             gh_rate_limit_tokens: require_env("RATE_LIMIT_TOKENS")?,
             gh_rate_limit_stats_cache_refresh: default_env("GH_RATE_LIMIT_STATS_REFRESH", 120)?,
diff --git a/src/lib.rs b/src/lib.rs
@@ -11,8 +11,10 @@ use prometheus::{Encoder, Registry};
 use anyhow::{Error, Result};
 use futures::future;
 use futures::task::{Context, Poll};
+use hyper::header::AUTHORIZATION;
+use hyper::http::HeaderValue;
 use hyper::service::Service;
-use hyper::{Body, Method, Request, Response, StatusCode};
+use hyper::{Body, HeaderMap, Method, Request, Response, StatusCode};
 use log::{debug, error};
 
 #[derive(Clone, Debug)]
@@ -59,9 +61,10 @@ impl Service<Request<Body>> for MetricProvider {
     fn call(&mut self, req: Request<Body>) -> Self::Future {
         debug!("New Request to endpoint {}", req.uri().path());
 
-        let output = match (req.method(), req.uri().path()) {
+        let authorized = is_auth_token_valid(&self.config.secret, req.headers());
+        let output = match (req.method(), req.uri().path(), authorized) {
             // Metrics handler
-            (&Method::GET, "/metrics") => {
+            (&Method::GET, "/metrics", true) => {
                 let encoder = prometheus::TextEncoder::new();
                 let mut buffer = Vec::<u8>::new();
                 match self.gather_with_encoder(encoder, &mut buffer) {
@@ -78,6 +81,11 @@ impl Service<Request<Body>> for MetricProvider {
                     }
                 }
             }
+            // Unauthorized request
+            (&Method::GET, "/metrics", false) => Response::builder()
+                .status(StatusCode::UNAUTHORIZED)
+                .body(Body::empty())
+                .unwrap(),
             // All other paths and methods
             _ => Response::builder()
                 .status(StatusCode::OK)
@@ -89,6 +97,22 @@ impl Service<Request<Body>> for MetricProvider {
     }
 }
 
+fn is_auth_token_valid(secret: &str, headers: &HeaderMap<HeaderValue>) -> bool {
+    match headers.get(AUTHORIZATION) {
+        Some(value) => value.to_str().map_or_else(
+            |_| false,
+            |t| {
+                if let Some(t) = t.strip_prefix("Bearer ") {
+                    t == secret
+                } else {
+                    false
+                }
+            },
+        ),
+        None => false,
+    }
+}
+
 pub struct MetricProviderFactory(pub MetricProvider);
 
 impl<T> Service<T> for MetricProviderFactory {
@@ -104,3 +128,40 @@ impl<T> Service<T> for MetricProviderFactory {
         future::ok(self.0.clone())
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use crate::is_auth_token_valid;
+    use hyper::http::HeaderValue;
+    use hyper::HeaderMap;
+
+    #[test]
+    fn auth_token_strip_bearer() {
+        use hyper::header::AUTHORIZATION;
+
+        let secret = "aASgwyfbFAKETOKEN44562uj36";
+        let token = "Bearer aASgwyfbFAKETOKEN44562uj36";
+        let v = HeaderValue::from_static(token);
+        let mut hv = HeaderMap::new();
+        hv.insert(AUTHORIZATION, v);
+
+        // should be true
+        let result = is_auth_token_valid(secret, &hv);
+        assert!(result);
+    }
+
+    #[test]
+    fn auth_token_strip_bearer_fail() {
+        use hyper::header::AUTHORIZATION;
+
+        let secret = "aASgwyfbFAKETOKEN44562uj36";
+        let token = "Bearer aASgwyfbFAKETOKEN44562uj36 "; // notice the whitespace in the end
+        let v = HeaderValue::from_static(token);
+        let mut hv = HeaderMap::new();
+        hv.insert(AUTHORIZATION, v);
+
+        // should be true
+        let result = is_auth_token_valid(secret, &hv);
+        assert_eq!(false, result);
+    }
+}
