Brig: Ensure servant APIs are recorded by the metrics middleware

libs/metrics-wai/metrics-wai.cabal

@@ -4,7 +4,7 @@ cabal-version: 1.12
 --
 -- see: https://github.com/sol/hpack
 --
--- hash: 1a7ee5db749a3d0036c292dd1b4e3f518dc51ae7c2c7c51ea8c0c387c09ecbe3
+-- hash: 3d924e65711daf57942ff893071402feeeaf2993f05a2302322064960c4bbedb
 
 name:           metrics-wai
 version:        0.5.7
@@ -47,3 +47,33 @@ library
     , wai-route >=0.3
     , wai-routing
   default-language: Haskell2010
+
+test-suite unit
+  type: exitcode-stdio-1.0
+  main-is: Spec.hs
+  other-modules:
+      Test.Data.Metrics.TypesSpec
+      Paths_metrics_wai
+  hs-source-dirs:
+      test
+  default-extensions: AllowAmbiguousTypes BangPatterns ConstraintKinds DataKinds DefaultSignatures DerivingStrategies DeriveFunctor DeriveGeneric DeriveLift DeriveTraversable EmptyCase FlexibleContexts FlexibleInstances FunctionalDependencies GADTs InstanceSigs KindSignatures LambdaCase MultiParamTypeClasses MultiWayIf NamedFieldPuns NoImplicitPrelude OverloadedStrings PackageImports PatternSynonyms PolyKinds QuasiQuotes RankNTypes ScopedTypeVariables StandaloneDeriving TemplateHaskell TupleSections TypeApplications TypeFamilies TypeFamilyDependencies TypeOperators UndecidableInstances ViewPatterns
+  ghc-options: -O2 -Wall -Wincomplete-uni-patterns -Wincomplete-record-updates -Wpartial-fields -fwarn-tabs -optP-Wno-nonportable-include-path -threaded -with-rtsopts=-N
+  build-tool-depends:
+      hspec-discover:hspec-discover
+  build-depends:
+      base ==4.*
+    , bytestring >=0.10
+    , containers
+    , hspec
+    , http-types >=0.8
+    , imports
+    , metrics-core >=0.3
+    , metrics-wai
+    , servant
+    , string-conversions
+    , text >=0.11
+    , wai >=3
+    , wai-middleware-prometheus
+    , wai-route >=0.3
+    , wai-routing
+  default-language: Haskell2010

libs/metrics-wai/package.yaml

@@ -25,3 +25,15 @@ dependencies:
 library:
   source-dirs: src
   ghc-prof-options: -fprof-auto
+tests:
+  unit:
+    main: Spec.hs
+    source-dirs: test
+    ghc-options:
+      - -threaded
+      - -with-rtsopts=-N
+    build-tools:
+      - hspec-discover:hspec-discover
+    dependencies:
+      - hspec
+      - metrics-wai

libs/metrics-wai/src/Data/Metrics/Middleware/Prometheus.hs

@@ -17,6 +17,7 @@
 
 module Data.Metrics.Middleware.Prometheus
   ( waiPrometheusMiddleware,
+    normalizeWaiRequestRoute,
   )
 where
 

libs/metrics-wai/src/Data/Metrics/Servant.hs

@@ -26,32 +26,51 @@
 -- | Given a servant API type, this module gives you a 'Paths' for 'withPathTemplate'.
 module Data.Metrics.Servant where
 
+import Data.Metrics.Middleware.Prometheus (normalizeWaiRequestRoute)
 import Data.Metrics.Types
 import qualified Data.Metrics.Types as Metrics
+import Data.Metrics.WaiRoute (treeToPaths)
 import Data.Proxy
 import Data.String.Conversions
 import Data.Tree
 import GHC.TypeLits
 import Imports
 import qualified Network.Wai as Wai
+import Network.Wai.Middleware.Prometheus
 import qualified Network.Wai.Middleware.Prometheus as Promth
+import Network.Wai.Routing (Routes, prepare)
 import Servant.API
 
 -- | This does not catch errors, so it must be called outside of 'WU.catchErrors'.
 servantPrometheusMiddleware :: forall proxy api. (RoutesToPaths api) => proxy api -> Wai.Middleware
 servantPrometheusMiddleware _ = Promth.prometheus conf . Promth.instrumentHandlerValue promthNormalize
   where
-    conf =
-      Promth.def
-        { Promth.prometheusEndPoint = ["i", "metrics"],
-          Promth.prometheusInstrumentApp = False
-        }
     promthNormalize :: Wai.Request -> Text
     promthNormalize req = pathInfo
       where
         mPathInfo = Metrics.treeLookup (routesToPaths @api) $ cs <$> Wai.pathInfo req
         pathInfo = cs $ fromMaybe "N/A" mPathInfo
 
+servantPlusWAIPrometheusMiddleware :: forall proxy api a m b. (RoutesToPaths api, Monad m) => Routes a m b -> proxy api -> Wai.Middleware
+servantPlusWAIPrometheusMiddleware routes _ = do
+  Promth.prometheus conf . instrument (normalizeWaiRequestRoute paths)
+  where
+    -- See Note [Raw Response]
+    instrument = Promth.instrumentHandlerValueWithFilter Promth.ignoreRawResponses
+
+    paths =
+      let Paths servantPaths = routesToPaths @api
+          Paths waiPaths = treeToPaths (prepare routes)
+       in Paths (meltTree (servantPaths <> waiPaths))
+
+conf :: PrometheusSettings
+conf =
+  Promth.def
+    { Promth.prometheusEndPoint = ["i", "metrics"],
+      -- We provide our own instrumentation so we can normalize routes
+      Promth.prometheusInstrumentApp = False
+    }
+
 routesToPaths :: forall routes. RoutesToPaths routes => Paths
 routesToPaths = Paths (meltTree (getRoutes @routes))
 
@@ -68,17 +87,17 @@ instance
   where
   getRoutes = [Node (Right . cs $ symbolVal (Proxy @seg)) (getRoutes @segs)]
 
--- <capture> <:> routes
+-- <capture> :> routes
 instance
   {-# OVERLAPPING #-}
   ( KnownSymbol capture,
     RoutesToPaths segs
   ) =>
   RoutesToPaths (Capture' mods capture a :> segs)
   where
-  getRoutes = [Node (Left ":_") (getRoutes @segs)]
+  getRoutes = [Node (Left (cs (":" <> symbolVal (Proxy @capture)))) (getRoutes @segs)]
 
--- route <:> routes
+-- route :> routes
 instance
   {-# OVERLAPPING #-}
   ( RoutesToPaths route,
@@ -104,10 +123,31 @@ instance {-# OVERLAPPING #-} RoutesToPaths (Verb 'PUT status ctypes content) whe
 instance {-# OVERLAPPING #-} RoutesToPaths (Verb 'DELETE status ctypes content) where
   getRoutes = []
 
+instance {-# OVERLAPPING #-} RoutesToPaths (Verb 'PATCH status ctypes content) where
+  getRoutes = []
+
 instance RoutesToPaths (NoContentVerb 'DELETE) where
   getRoutes = []
 
-instance {-# OVERLAPPING #-} RoutesToPaths (Verb 'PATCH status ctypes content) where
+instance {-# OVERLAPPING #-} RoutesToPaths (UVerb 'HEAD ctypes content) where
+  getRoutes = []
+
+instance {-# OVERLAPPING #-} RoutesToPaths (UVerb 'GET ctypes content) where
+  getRoutes = []
+
+instance {-# OVERLAPPING #-} RoutesToPaths (UVerb 'POST ctypes content) where
+  getRoutes = []
+
+instance {-# OVERLAPPING #-} RoutesToPaths (UVerb 'PUT ctypes content) where
+  getRoutes = []
+
+instance {-# OVERLAPPING #-} RoutesToPaths (UVerb 'DELETE ctypes content) where
+  getRoutes = []
+
+instance {-# OVERLAPPING #-} RoutesToPaths (UVerb 'PATCH ctypes content) where
+  getRoutes = []
+
+instance RoutesToPaths Raw where
   getRoutes = []
 
 instance

libs/metrics-wai/src/Data/Metrics/Types.hs

@@ -2,7 +2,6 @@
 {-# LANGUAGE OverloadedStrings #-}
 {-# LANGUAGE ScopedTypeVariables #-}
 {-# LANGUAGE TypeFamilies #-}
-{-# LANGUAGE ViewPatterns #-}
 
 -- This file is part of the Wire Server implementation.
 --
@@ -52,7 +51,7 @@ mkTree = fmap (Paths . meltTree) . mapM mkbranch . sortBy (flip compare) . fmap
   where
     mkbranch :: [PathSegment] -> Either String (Tree PathSegment)
     mkbranch (seg : segs@(_ : _)) = Node seg . (: []) <$> mkbranch segs
-    mkbranch (seg : []) = Right $ Node seg []
+    mkbranch [seg] = Right $ Node seg []
     mkbranch [] = Left "internal error: path with on segments."
     mknode :: ByteString -> PathSegment
     mknode seg = if BS.head seg /= ':' then Right seg else Left seg
@@ -69,14 +68,54 @@ meltTree = go
 -- | A variant of 'Network.Wai.Route.Tree.lookup'.  The segments contain values to be captured
 -- when running the 'App', here we simply replace them with their identifier;
 -- e.g. @/user/1234@ might become @/user/userid@
+--
+-- This lookup will do its best when it sees ambiguous paths like
+-- /users/:uid/clients and /users/:domain/:uid. In this case, for input
+-- ["users", "some-uuid", "clients"] it will see that /users/:uid/clients has
+-- more verbatim matches than /users/:domain/:uid and so, it will prefer
+-- /users/:uid/clients.
+--
+-- This is not the case all the time, for instance if there are two paths like:
+-- /users/clients/:cid and /users/:uid/clients and the input is ["users",
+-- "clients", "clients"], the lookup will see exactly the same number of
+-- verbatim matches and return any one of them. This may not really be a
+-- problem, because we may not have such a request in our real paths.
+--
+-- Note [Trees for Metrics]
+--
+-- The use of trees hides information about which of the
+-- partial paths are real paths. Consider a tree like this:
+--
+-- /users
+--   /:uid
+--   /:domain
+--     /:uid
+--   /clients
+--     /:uid
+--
+-- Here, it is impossible to tell if /users or /users/:domain is a valid
+-- endpoint or not, so this function will always return these as a match, even
+-- if in reality there is no such API. This smells like we shouldn't be using
+-- trees here and instead just try to match a given path with list of routes.
 treeLookup :: Paths -> [ByteString] -> Maybe ByteString
-treeLookup (Paths forest) = go [] forest
+treeLookup (Paths forest) = mungeSegments <=< go forest
   where
-    go :: [PathSegment] -> Forest PathSegment -> [ByteString] -> Maybe ByteString
-    go path _ [] = Just . ("/" <>) . BS.intercalate "/" . fmap (either id id) . reverse $ path
-    go _ [] _ = Nothing
-    go path trees (seg : segs) =
-      find (seg `fits`) trees >>= \(Node root trees') -> go (root : path) trees' segs
+    go :: Forest PathSegment -> [ByteString] -> Maybe [PathSegment]
+    go _trees [] = Just []
+    go [] _segs = Nothing
+    go trees (seg : segs) =
+      let allMatches = mapMaybe (matchTree seg segs) trees
+          sorted = sortOn (Down . length . filter isRight) allMatches
+       in listToMaybe sorted
+
+    matchTree :: ByteString -> [ByteString] -> Tree PathSegment -> Maybe [PathSegment]
+    matchTree seg segs tree =
+      if seg `fits` tree
+        then (rootLabel tree :) <$> go (subForest tree) segs
+        else Nothing
+
     fits :: ByteString -> Tree PathSegment -> Bool
     fits _ (Node (Left _) _) = True
     fits seg (Node (Right seg') _) = seg == seg'
+
+    mungeSegments path = Just . ("/" <>) . BS.intercalate "/" . fmap (either id id) $ path

libs/metrics-wai/test/Spec.hs

@@ -0,0 +1 @@
+{-# OPTIONS_GHC -F -pgmF hspec-discover #-}

libs/metrics-wai/test/Test/Data/Metrics/TypesSpec.hs

@@ -0,0 +1,101 @@
+module Test.Data.Metrics.TypesSpec where
+
+import Data.Metrics.Types
+import Data.Tree
+import Imports
+import Test.Hspec
+
+spec :: Spec
+spec = describe "Data.Metrics.Types" $ do
+  describe "treeLookup" $ do
+    it "should match exact string when PathSegment is Right" $ do
+      let testPaths = Paths [Node (Right "users") []]
+      treeLookup testPaths ["users"] `shouldBe` Just "/users"
+
+    it "should match anything when PathSegment is Left" $ do
+      let testPaths = Paths [Node (Right "users") [Node (Left ":uid") []]]
+      treeLookup testPaths ["users", "some-uuid"] `shouldBe` Just "/users/:uid"
+
+    it "should match the correct endpoint when there are multiple Left options" $ do
+      let testPaths =
+            Paths
+              [ Node
+                  (Right "users")
+                  [ Node (Left ":uid") [],
+                    Node (Left ":domain") [Node (Left ":uid") []]
+                  ]
+              ]
+      treeLookup testPaths ["users", "some-uuid"] `shouldBe` Just "/users/:uid"
+      treeLookup testPaths ["users", "example.com", "some-uuid"] `shouldBe` Just "/users/:domain/:uid"
+
+    -- See Note [Trees for Metrics]
+    it "should match even when prefix of a path is matched" $ do
+      let testPaths =
+            Paths
+              [ Node
+                  (Right "users")
+                  [ Node (Left ":uid") [],
+                    Node (Left ":domain") [Node (Left ":uid") []]
+                  ]
+              ]
+      treeLookup testPaths ["users"] `shouldBe` Just "/users"
+
+    it "shouldn't match when endpoint doesn't fit" $ do
+      let testPaths =
+            Paths
+              [ Node
+                  (Right "users")
+                  [ Node (Left ":uid") [],
+                    Node (Left ":domain") [Node (Left ":uid") []]
+                  ]
+              ]
+      treeLookup testPaths ["aliens"] `shouldBe` Nothing
+      treeLookup testPaths ["users", "some-domain", "some-uuid", "extra-thing"] `shouldBe` Nothing
+
+    it "should prioritize matches with most number of Rights" $ do
+      let testPaths =
+            Paths
+              [ Node
+                  (Right "users")
+                  [ Node (Left ":domain") [Node (Left ":uid") []],
+                    Node (Left ":uid") [Node (Right "clients") []]
+                  ]
+              ]
+      let testPathsReverse =
+            Paths
+              [ Node
+                  (Right "users")
+                  [ Node (Left ":uid") [Node (Right "clients") []],
+                    Node (Left ":domain") [Node (Left ":uid") []]
+                  ]
+              ]
+
+      treeLookup testPaths ["users", "example.com", "some-uuid"] `shouldBe` Just "/users/:domain/:uid"
+      treeLookup testPathsReverse ["users", "example.com", "some-uuid"] `shouldBe` Just "/users/:domain/:uid"
+
+      treeLookup testPaths ["users", "some-uuid", "clients"] `shouldBe` Just "/users/:uid/clients"
+      treeLookup testPathsReverse ["users", "some-uuid", "clients"] `shouldBe` Just "/users/:uid/clients"
+
+    it "should prioritize matches with most number of Rights" $ do
+      let testPaths =
+            Paths
+              [ Node
+                  (Right "users")
+                  [ Node (Left ":domain") [Node (Left ":uid") []],
+                    Node (Left ":uid") [Node (Right "clients") []]
+                  ]
+              ]
+      let testPathsReverse =
+            Paths
+              [ Node
+                  (Right "users")
+                  [ Node (Left ":uid") [Node (Right "clients") []],
+                    Node (Left ":domain") [Node (Left ":uid") []]
+                  ]
+              ]
+
+      treeLookup testPaths ["users", "example.com", "some-uuid"] `shouldBe` Just "/users/:domain/:uid"
+      treeLookup testPathsReverse ["users", "example.com", "some-uuid"] `shouldBe` Just "/users/:domain/:uid"
+
+      treeLookup testPaths ["users", "some-uuid", "clients"] `shouldBe` Just "/users/:uid/clients"
+      treeLookup testPathsReverse ["users", "some-uuid", "clients"] `shouldBe` Just "/users/:uid/clients"

services/brig/src/Brig/API/Public.hs

@@ -717,6 +717,9 @@ sitemap o = do
     Doc.errorResponse invalidUser
     Doc.errorResponse (noIdentity 5)
 
+  -- This endpoint is used to test /i/metrics, when this is servantified, please
+  -- make sure some other endpoint is used to test that routes defined in this
+  -- function are recorded and reported correctly in /i/metrics.
   get "/connections" (continue listConnectionsH) $
     accept "application" "json"
       .&. zauthUserId