[WPB-1220] servantify proxy internal

changelog.d/5-internal/WPB-1220-servantify-proxy-internal

@@ -0,0 +1 @@
+Servantify internal routing table for proxy.

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

@@ -17,6 +17,7 @@
 
 module Data.Metrics.Middleware.Prometheus
   ( waiPrometheusMiddleware,
+    waiPrometheusMiddlewarePaths,
     normalizeWaiRequestRoute,
   )
 where
@@ -33,12 +34,17 @@ import Network.Wai.Routing.Route (Routes, prepare)
 -- This middleware requires your servers 'Routes' because it does some normalization
 -- (e.g. removing params from calls)
 waiPrometheusMiddleware :: (Monad m) => Routes a m b -> Wai.Middleware
-waiPrometheusMiddleware routes =
+waiPrometheusMiddleware routes = waiPrometheusMiddlewarePaths $ treeToPaths $ prepare routes
+
+-- | Helper function that should only be needed as long as we have wai-routing code left in
+-- proxy: run `treeToPaths` on old routing tables and `routeToPaths` on the servant ones, and
+-- feed both to this function.
+waiPrometheusMiddlewarePaths :: Paths -> Wai.Middleware
+waiPrometheusMiddlewarePaths paths =
   Promth.prometheus conf . instrument (normalizeWaiRequestRoute paths)
   where
     -- See Note [Raw Response]
     instrument = Promth.instrumentHandlerValueWithFilter Promth.ignoreRawResponses
-    paths = treeToPaths $ prepare routes
     conf =
       Promth.def
         { Promth.prometheusEndPoint = ["i", "metrics"],

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

@@ -41,6 +41,7 @@ newtype PathTemplate = PathTemplate Text
 -- (e.g. user id).
 newtype Paths = Paths (Forest PathSegment)
   deriving (Eq, Show)
+  deriving newtype (Semigroup)
 
 type PathSegment = Either ByteString ByteString
 

libs/types-common/src/Data/Id.hs

@@ -48,6 +48,7 @@ module Data.Id
     -- * Other IDs
     ConnId (..),
     RequestId (..),
+    defRequestId,
     BotId (..),
     NoId,
     OAuthClientId,
@@ -418,6 +419,9 @@ newtype RequestId = RequestId
       ToBytes
     )
 
+defRequestId :: RequestId
+defRequestId = RequestId "N/A"
+
 instance ToSchema RequestId where
   schema =
     RequestId . encodeUtf8

services/proxy/default.nix

@@ -9,6 +9,7 @@
 , bytestring
 , case-insensitive
 , configurator
+, errors
 , exceptions
 , extended
 , gitignoreSource
@@ -21,6 +22,7 @@
 , lib
 , metrics-wai
 , retry
+, servant-server
 , text
 , tinylog
 , types-common
@@ -46,6 +48,7 @@ mkDerivation {
     bytestring
     case-insensitive
     configurator
+    errors
     exceptions
     extended
     http-client
@@ -56,6 +59,7 @@ mkDerivation {
     lens
     metrics-wai
     retry
+    servant-server
     text
     tinylog
     types-common

services/proxy/proxy.cabal

@@ -17,7 +17,7 @@ flag static
 
 library
   exposed-modules:
-    Proxy.API
+    Proxy.API.Internal
     Proxy.API.Public
     Proxy.Env
     Proxy.Options
@@ -80,6 +80,7 @@ library
     , bytestring             >=0.10
     , case-insensitive       >=1.2
     , configurator           >=0.3
+    , errors
     , exceptions             >=0.8
     , extended
     , http-client            >=0.7
@@ -90,6 +91,7 @@ library
     , lens                   >=4.11
     , metrics-wai            >=0.5
     , retry                  >=0.7
+    , servant-server
     , text                   >=1.2
     , tinylog                >=0.12
     , types-common           >=0.8

services/proxy/src/Proxy/API.hs → services/proxy/src/Proxy/API/Internal.hs

@@ -15,29 +15,19 @@
 -- You should have received a copy of the GNU Affero General Public License along
 -- with this program. If not, see <https://www.gnu.org/licenses/>.
 
-module Proxy.API
-  ( sitemap,
+module Proxy.API.Internal
+  ( InternalAPI,
+    servantSitemap,
   )
 where
 
 import Imports hiding (head)
-import Network.Wai.Predicate (true)
-import Network.Wai.Routing (Routes, continue, get, head)
-import Network.Wai.Utilities (empty)
-import Proxy.API.Public qualified as Public
-import Proxy.Env (Env)
-import Proxy.Proxy (Proxy)
+import Proxy.Proxy qualified
+import Servant
+import Wire.API.Routes.MultiVerb
+import Wire.API.Routes.Named (Named (Named))
 
-sitemap :: Env -> Routes a Proxy ()
-sitemap e = do
-  Public.sitemap e
-  routesInternal
+type InternalAPI = Named "status" ("i" :> "status" :> MultiVerb 'GET '[Servant.JSON] '[RespondEmpty 200 "OK"] ())
 
--- | IF YOU MODIFY THIS, BE AWARE OF:
---
--- >>> /libs/wire-api/src/Wire/API/Routes/Public/Proxy.hs
--- >>> https://wearezeta.atlassian.net/browse/SQSERVICES-1647
-routesInternal :: Routes a Proxy ()
-routesInternal = do
-  head "/i/status" (continue $ const (pure empty)) true
-  get "/i/status" (continue $ const (pure empty)) true
+servantSitemap :: ServerT InternalAPI Proxy.Proxy.Proxy
+servantSitemap = Named @"status" (pure ())

services/proxy/src/Proxy/API/Public.hs

@@ -16,7 +16,9 @@
 -- with this program. If not, see <https://www.gnu.org/licenses/>.
 
 module Proxy.API.Public
-  ( sitemap,
+  ( PublicAPI,
+    servantSitemap,
+    waiRoutingSitemap,
   )
 where
 
@@ -41,18 +43,35 @@ import Network.Wai.Internal qualified as I
 import Network.Wai.Predicate hiding (Error, err, setStatus)
 import Network.Wai.Predicate.Request (getRequest)
 import Network.Wai.Routing hiding (path, route)
+import Network.Wai.Routing qualified as Routing
 import Network.Wai.Utilities
+import Network.Wai.Utilities.Server (compile)
 import Proxy.Env
 import Proxy.Proxy
+import Servant qualified
 import System.Logger.Class hiding (Error, info, render)
 import System.Logger.Class qualified as Logger
 
+type PublicAPI = Servant.Raw -- see https://wearezeta.atlassian.net/browse/WPB-1216
+
+servantSitemap :: Env -> Servant.ServerT PublicAPI Proxy.Proxy.Proxy
+servantSitemap e = Servant.Tagged app
+  where
+    app :: Application
+    app r k = appInProxy e r (Routing.route tree r k')
+      where
+        tree :: Tree (App Proxy)
+        tree = compile (waiRoutingSitemap e)
+
+        k' :: Response -> Proxy.Proxy.Proxy ResponseReceived
+        k' = liftIO . k
+
 -- | IF YOU MODIFY THIS, BE AWARE OF:
 --
 -- >>> /libs/wire-api/src/Wire/API/Routes/Public/Proxy.hs
 -- >>> https://wearezeta.atlassian.net/browse/SQSERVICES-1647
-sitemap :: Env -> Routes a Proxy ()
-sitemap e = do
+waiRoutingSitemap :: Env -> Routes a Proxy ()
+waiRoutingSitemap e = do
   get
     "/proxy/youtube/v3/:path"
     (proxy e "key" "secrets.youtube" Prefix "/youtube/v3" youtube)
@@ -107,7 +126,7 @@ proxy e qparam keyname reroute path phost rq k = do
           then do
             threadDelay 5000
             loop runInIO (n - 1) waiReq req
-          else runProxy e waiReq (k res)
+          else appInProxy e waiReq (k res)
     onUpstreamError runInIO x _ next = do
       void . runInIO $ Logger.warn (msg (val "gateway error") ~~ field "error" (show x))
       next (errorRs error502)

services/proxy/src/Proxy/Proxy.hs

@@ -17,7 +17,7 @@
 -- You should have received a copy of the GNU Affero General Public License along
 -- with this program. If not, see <https://www.gnu.org/licenses/>.
 
-module Proxy.Proxy (Proxy, runProxy) where
+module Proxy.Proxy (Proxy, appInProxy, runProxy) where
 
 import Bilge.Request (requestIdName)
 import Control.Lens hiding ((.=))
@@ -51,11 +51,14 @@ newtype Proxy a = Proxy
 instance MonadLogger Proxy where
   log l m = ask >>= \e -> Logger.log (e ^. applog) l (reqIdMsg (e ^. reqId) . m)
 
-runProxy :: Env -> Request -> Proxy ResponseReceived -> IO ResponseReceived
-runProxy e r m = do
+appInProxy :: Env -> Request -> Proxy ResponseReceived -> IO ResponseReceived
+appInProxy e r m = do
   rid <- lookupReqId (e ^. applog) r
   runReaderT (unProxy m) (reqId .~ rid $ e)
 
+runProxy :: Env -> Proxy a -> IO a
+runProxy e m = runReaderT (unProxy m) e
+
 reqIdMsg :: RequestId -> Msg -> Msg
 reqIdMsg = ("request" .=) . unRequestId
 {-# INLINE reqIdMsg #-}

services/proxy/src/Proxy/Run.hs

@@ -20,29 +20,69 @@ module Proxy.Run
   )
 where
 
+import Bilge.Request (requestIdName)
+import Control.Error
 import Control.Lens hiding ((.=))
 import Control.Monad.Catch
-import Data.Metrics.Middleware.Prometheus (waiPrometheusMiddleware)
+import Data.Id (RequestId (RequestId), defRequestId)
+import Data.Metrics.Middleware.Prometheus (waiPrometheusMiddlewarePaths)
+import Data.Metrics.Servant
+import Data.Metrics.Types
+import Data.Metrics.WaiRoute
 import Imports hiding (head)
+import Network.Wai (Middleware, Request, requestHeaders)
 import Network.Wai.Middleware.Gunzip qualified as GZip
+import Network.Wai.Routing.Route
 import Network.Wai.Utilities.Server hiding (serverPort)
-import Proxy.API (sitemap)
+import Proxy.API.Internal as I
+import Proxy.API.Public as P
 import Proxy.Env
 import Proxy.Options
 import Proxy.Proxy
+import Servant qualified
 import Wire.API.Routes.Version
 import Wire.API.Routes.Version.Wai
 
+type CombinedAPI = PublicAPI Servant.:<|> InternalAPI
+
+combinedSitemap :: Env -> Servant.ServerT CombinedAPI Proxy
+combinedSitemap env = P.servantSitemap env Servant.:<|> I.servantSitemap
+
 run :: Opts -> IO ()
 run o = do
   e <- createEnv o
   s <- newSettings $ defaultServer (o ^. host) (o ^. port) (e ^. applog)
-  let rtree = compile (sitemap e)
-  let app r k = runProxy e r (route rtree r k)
-  let middleware =
+
+  let metricsMW :: Middleware
+      metricsMW =
+        -- once wai-routing has been removed from proxy: use `servantPrometheusMiddleware
+        -- (Servant.Proxy @CombinedAPI)` here (and probably inline the whole thing).
+        waiPrometheusMiddlewarePaths (pub <> int)
+        where
+          pub, int :: Paths
+          pub = treeToPaths $ prepare (P.waiRoutingSitemap e)
+          int = routesToPaths @InternalAPI
+
+      middleware :: Middleware
+      middleware =
         versionMiddleware (foldMap expandVersionExp (o ^. disabledAPIVersions))
           . requestIdMiddleware (e ^. applog) defaultRequestIdHeaderName
-          . waiPrometheusMiddleware (sitemap e)
+          . metricsMW
           . GZip.gunzip
           . catchErrors (e ^. applog) defaultRequestIdHeaderName
-  runSettingsWithShutdown s (middleware app) Nothing `finally` destroyEnv e
+
+  runSettingsWithShutdown s (middleware (mkApp e)) Nothing `finally` destroyEnv e
+
+mkApp :: Env -> Servant.Application
+mkApp env req = Servant.serve (Servant.Proxy @CombinedAPI) toServantSitemap req
+  where
+    toServantSitemap :: Servant.Server CombinedAPI
+    toServantSitemap = Servant.hoistServer (Servant.Proxy @CombinedAPI) toServantHandler (combinedSitemap env)
+
+    toServantHandler :: Proxy a -> Servant.Handler a
+    toServantHandler p = Servant.Handler . ExceptT $ Right <$> runProxy (injectReqId req env) p
+
+    injectReqId :: Request -> Env -> Env
+    injectReqId r = reqId .~ lookupReqId r
+      where
+        lookupReqId = maybe defRequestId RequestId . lookup requestIdName . requestHeaders