Count downloads from content delivery networks via a redirection scheme

src/Distribution/Server/Features/Core.hs

@@ -376,7 +376,7 @@ coreFeature :: ServerEnv
             -> ( CoreFeature
                , IO IndexTarballInfo )
 
-coreFeature ServerEnv{serverBlobStore = store} UserFeature{..}
+coreFeature env@ServerEnv{serverBlobStore = store} UserFeature{..}
             Store.Backend{backendStore = packagesStore, backendState} cacheIndexTarball
             packageChangeHook
             preIndexUpdateHook
@@ -706,11 +706,39 @@ coreFeature ServerEnv{serverBlobStore = store} UserFeature{..}
                      [MText "No tarball exists for this package version."]
         Just (tarball, (uploadtime, _uid), _revNo) -> do
           let blobId = blobInfoId $ pkgTarballGz tarball
-          cacheControl [Public, NoTransform, maxAgeDays 30]
-                       (BlobStorage.blobETag blobId)
-          file <- liftIO $ BlobStorage.fetch store blobId
-          runHook_ packageDownloadHook pkgid
-          return $ toResponse $ Resource.PackageTarball file blobId uploadtime
+          host' <- requestHost env
+          -- Accurately counting downloads in the presence of a content-delivery network
+          -- is tricky. We want every download to be counted, even if the /data/ itself
+          -- is served from some other cache.
+          --
+          -- A solution to this is implemented below, using two endpoints. The first
+          -- endpoint counts the download, and redirects to the /real/ download endpoint.
+          -- The key is that this first endpoint set the `Cache-Control` header to `No-Store`,
+          -- such that content-delivery networks don't cache the redirect.
+          -- Then, the second endpoint serves the actual content, and its response can be cached.
+          --
+          -- One additional wrinkle is that Hackage can respond from multiple hosts.
+          -- If using the 'MainHost'/'UserContentHost' setup, we count downloads from the
+          -- main host, and serve content -- without counting downloads -- from the user-content host.
+          -- If the host is unrecognised (see 'UnrecognisedHost'), there is no well-defined redirect scheme
+          -- we can use, so we count and serve at the same time.
+          case host' of
+            MainHost -> do
+              runHook_ packageDownloadHook pkgid
+              setCacheControl [NoStore]
+              uri <- userContentRequestURI env
+              found (show uri) $ contentLength $ toResponse ()
+            UserContentHost -> serveBlob blobId uploadtime
+            UnrecognisedHost -> do
+              runHook_ packageDownloadHook pkgid
+              serveBlob blobId uploadtime
+
+    serveBlob :: BlobStorage.BlobId -> UTCTime -> ServerPartE Response
+    serveBlob blobId uploadtime = do
+      cacheControl [Public, NoTransform, maxAgeDays 30]
+                   (BlobStorage.blobETag blobId)
+      file <- liftIO $ BlobStorage.fetch store blobId
+      return $ toResponse $ Resource.PackageTarball file blobId uploadtime
 
     -- result: cabal file or not-found error
     serveCabalFile :: DynamicPath -> ServerPartE Response

src/Distribution/Server/Framework/CacheControl.hs

@@ -4,6 +4,7 @@
 module Distribution.Server.Framework.CacheControl (
     cacheControl,
     cacheControlWithoutETag,
+    setCacheControl,
     CacheControl(..),
     ETag(..),
     etagFromHash,
@@ -21,7 +22,7 @@ import qualified Data.ByteString.Char8 as BS8
 import Data.Hashable
 import Numeric
 
-data CacheControl = MaxAge Int | Public | Private | NoCache | NoTransform
+data CacheControl = MaxAge Int | Public | Private | NoCache | NoStore | NoTransform
 
 maxAgeSeconds, maxAgeMinutes, maxAgeHours,
   maxAgeDays, maxAgeMonths :: Int -> CacheControl
@@ -36,6 +37,7 @@ formatCacheControl (MaxAge n)  = "max-age=" ++ show n
 formatCacheControl Public      = "public"
 formatCacheControl Private     = "private"
 formatCacheControl NoCache     = "no-cache"
+formatCacheControl NoStore     = "no-store"
 formatCacheControl NoTransform = "no-transform"
 
 -- | Adds a @Cache-Control@ and @ETag@ header to the response. Also handles the

src/Distribution/Server/Framework/ServerEnv.hs

@@ -89,6 +89,31 @@ getHost = do
       Just hostHeaderPair | [oneValue] <- hValue hostHeaderPair -> Just oneValue
       _ -> Nothing
 
+data RequestHost
+  = MainHost -- ^ Request host matching 'serverRequiredBaseHostHeader'.
+  | UserContentHost -- ^ Request host matching 'serverUserContentBaseURI'.
+  | UnrecognisedHost -- ^ A host that matches neither the main or user-content host.
+
+requestHost :: ServerMonad m => ServerEnv -> m RequestHost
+requestHost ServerEnv {serverUserContentBaseURI, serverRequiredBaseHostHeader} = do
+  mHost <- getHost
+  let isMain = mHost == Just (encodeUtf8 (T.pack serverRequiredBaseHostHeader))
+      isUserContent = case URI.uriAuthority serverUserContentBaseURI of
+        Just auth -> mHost == Just (encodeUtf8 (T.pack (URI.uriRegName auth ++ URI.uriPort auth)))
+        Nothing -> False
+  pure $ case (isMain, isUserContent) of
+    (True, False) -> MainHost
+    (False, True) -> UserContentHost
+    _ -> UnrecognisedHost
+
+userContentRequestURI :: ServerMonad m => ServerEnv -> m URI.URI
+userContentRequestURI ServerEnv {serverUserContentBaseURI} = do
+  rq <- askRq
+  pure serverUserContentBaseURI
+    { URI.uriPath = rqUri rq
+    , URI.uriQuery = rqQuery rq
+    }
+
 requireUserContent :: ServerEnv -> Response -> ServerPartE Response
 requireUserContent ServerEnv {serverUserContentBaseURI, serverRequiredBaseHostHeader} action = do
   Just hostHeaderValue <- getHost

tests/HackageClientUtils.hs

@@ -26,6 +26,7 @@ import Util
 import HttpUtils ( ExpectedCode
                  , isOk
                  , isAccepted
+                 , isFound
                  , isSeeOther
                  , isNotModified
                  , isUnauthorized
@@ -311,6 +312,19 @@ getUrl auth url = Http.execRequest auth (mkGetReq url)
 getUserContentUrl :: Authorization -> RelativeURL -> IO String
 getUserContentUrl auth url = Http.execRequest auth (mkGetUserContentReq url)
 
+-- | Check that the URL, requested from the main host, is a temporary redirect
+-- to the same path on the user content host, and that the redirect itself is
+-- not cacheable (so that every download reaches the server to be counted).
+checkRedirectsToUserContent :: RelativeURL -> IO ()
+checkRedirectsToUserContent url = do
+  void $ Http.execRequest' NoAuth (mkGetReq url) isFound
+  cacheControl <- Http.responseHeader HdrCacheControl (mkGetReq url)
+  unless (cacheControl == "no-store") $
+    die $ "Expected 'Cache-Control: no-store' but got " ++ show cacheControl
+  location <- Http.responseHeader HdrLocation (mkGetReq url)
+  unless (location == mkUserContentUrl url) $
+    die $ "Expected redirect to " ++ mkUserContentUrl url ++ " but got " ++ show location
+
 getETag :: RelativeURL -> IO String
 getETag url = Http.responseHeader HdrETag (mkGetReq url)
 

tests/HighLevelTest.hs

@@ -342,8 +342,10 @@ runPackageTests = do
        cabalFile <- getUrl NoAuth "/package/testpackage-1.0.0.0/testpackage.cabal"
        unless (cabalFile == testpackageCabalFile) $
            die "Bad Cabal file"
+    do info "Getting testpackage tar file from the main host redirects to the user content host"
+       checkRedirectsToUserContent "/package/testpackage/testpackage-1.0.0.0.tar.gz"
     do info "Getting testpackage tar file"
-       tarFile <- getUrl NoAuth "/package/testpackage/testpackage-1.0.0.0.tar.gz"
+       tarFile <- getUserContentUrl NoAuth "/package/testpackage/testpackage-1.0.0.0.tar.gz"
        unless (tarFile == testpackageTarFileContent) $
            die "Bad tar file"
     do info "Getting testpackage source"

tests/HttpUtils.hs

@@ -7,6 +7,7 @@ module HttpUtils (
   , isOk
   , isAccepted
   , isNoContent
+  , isFound
   , isSeeOther
   , isNotModified
   , isUnauthorized
@@ -50,12 +51,13 @@ import Util
 
 type ExpectedCode = (Int, Int, Int) -> Bool
 
-isOk, isAccepted, isNoContent, isSeeOther :: ExpectedCode
+isOk, isAccepted, isNoContent, isFound, isSeeOther :: ExpectedCode
 isNotModified, isUnauthorized, isForbidden :: ExpectedCode
 isNotFound :: ExpectedCode
 isOk           = (== (2, 0, 0))
 isAccepted     = (== (2, 0, 2))
 isNoContent    = (== (2, 0, 4))
+isFound        = (== (3, 0, 2))
 isSeeOther     = (== (3, 0, 3))
 isNotModified  = (== (3, 0, 4))
 isUnauthorized = (== (4, 0, 1))