1 files changed, 59 insertions, 0 deletions
diff --git a/Network/Wai/Middleware/Time.hs b/Network/Wai/Middleware/Time.hs
new file mode 100644
index 0000000..2a1d034
--- /dev/null
+++ b/Network/Wai/Middleware/Time.hs
@@ -0,0 +1,59 @@
+-- {-# LANGUAGE OverloadedStrings #-}
+-- |
+-- Module     : Network.Wai.Middleware.Time
+-- License    : WTFPL2
+--
+-- Maintainer : tv@destroy.retiolum
+--
+-- WAI X-Response-Time middleware.
+
+module Network.Wai.Middleware.Time
+    ( time
+    ) where
+
+import Control.Monad.Trans    (liftIO)
+import Data.ByteString.Char8  (ByteString, pack)
+import Data.CaseInsensitive       (mk)
+import Network.HTTP.Types.Header  (HeaderName)
+import Network.HTTP.Types         (ResponseHeaders)
+import Network.Wai.Internal   (Response (..))
+import Network.Wai            (Middleware)
+import System.Clock           (getTime, Clock(Monotonic), sec, nsec, TimeSpec)
+import Text.Printf            (printf)
+
+getTimeMonotonic :: IO TimeSpec
+getTimeMonotonic = liftIO $ getTime Monotonic
+
+-- | Use time to append an X-Response-Time header to the response.
+time :: Middleware
+time app rq = do
+  t0 <- getTimeMonotonic
+  rs <- app rq
+  t1 <- getTimeMonotonic
+  return $ addResponseTime (timeSpecsToByteString t0 t1) rs
+
+addResponseTime :: ByteString -> Response -> Response
+addResponseTime xrt rs = do
+  case rs of
+    (ResponseFile st hs path part) ->
+      ResponseFile st (responseTime xrt hs) path part
+    (ResponseBuilder st hs b) ->
+      ResponseBuilder st (responseTime xrt hs) b
+    (ResponseSource st hs s) ->
+      ResponseSource st (responseTime xrt hs) s
+    (ResponseRaw rawapp r) ->
+      ResponseRaw rawapp $ addResponseTime xrt r
+
+hResponseTime :: HeaderName
+hResponseTime = mk . pack $ "X-Response-Time"
+
+responseTime :: ByteString -> ResponseHeaders -> ResponseHeaders
+responseTime xrt =
+  (:) (hResponseTime, xrt)
+
+timeSpecsToByteString :: TimeSpec -> TimeSpec -> ByteString
+timeSpecsToByteString t0 t1 =
+    pack $ printf "%d.%09ds" s ns
+  where
+    s = (sec t1) - (sec t0)
+    ns = (nsec t1) - (nsec t0)