From 90f101734187e159eff686dff8f89992d12b5315 Mon Sep 17 00:00:00 2001
From: Ethan Donowitz <8703826+ethowitz@users.noreply.github.com>
Date: Thu, 25 Aug 2022 09:04:16 -0400
Subject: [PATCH] feat: improve Tokenserver errors and metrics (#1385)

---
 syncstorage/src/tokenserver/db/models.rs |  3 +++
 syncstorage/src/web/extractors.rs        | 22 +++++++++++++++-------
 2 files changed, 18 insertions(+), 7 deletions(-)

diff --git a/syncstorage/src/tokenserver/db/models.rs b/syncstorage/src/tokenserver/db/models.rs
index a1e0d2a3..e48e574d 100644
--- a/syncstorage/src/tokenserver/db/models.rs
+++ b/syncstorage/src/tokenserver/db/models.rs
@@ -297,6 +297,9 @@ impl TokenserverDb {
         &self,
         params: params::AddUserToNode,
     ) -> DbResult<results::AddUserToNode> {
+        let mut metrics = self.metrics.clone();
+        metrics.start_timer("storage.add_user_to_node", None);
+
         const QUERY: &str = r#"
             UPDATE nodes
                SET current_load = current_load + 1,
diff --git a/syncstorage/src/web/extractors.rs b/syncstorage/src/web/extractors.rs
index c1326997..00aabaa3 100644
--- a/syncstorage/src/web/extractors.rs
+++ b/syncstorage/src/web/extractors.rs
@@ -42,6 +42,7 @@ use crate::tokenserver::auth::TokenserverOrigin;
 use crate::web::{
     auth::HawkPayload,
     error::{HawkErrorKind, ValidationErrorKind},
+    tags::Tags,
     DOCKER_FLOW_ENDPOINTS, X_WEAVE_RECORDS,
 };
 const BATCH_MAX_IDS: usize = 100;
@@ -1161,13 +1162,20 @@ impl FromRequest for HawkIdentifier {
             }
         };
 
-        future::ready(Self::extrude(
-            &req,
-            method.as_str(),
-            uri,
-            &connection_info,
-            secrets,
-        ))
+        let result = Self::extrude(&req, method.as_str(), uri, &connection_info, secrets);
+
+        if let Ok(ref hawk_id) = result {
+            // Store the origin of the token as an extra to be included when emitting a Sentry error
+            let mut exts = req.extensions_mut();
+            let mut tags = Tags::default();
+            tags.add_extra(
+                "tokenserver_origin",
+                &hawk_id.tokenserver_origin.to_string(),
+            );
+            tags.commit(&mut exts);
+        }
+
+        future::ready(result)
     }
 }