Read the dispositon header when downloading

1 files changed, 59 insertions, 31 deletions

diff --git a/lib/SlackBuilder/Download.hs b/lib/SlackBuilder/Download.hs
index 50eaff1..719efe9 100644
--- a/lib/SlackBuilder/Download.hs
+++ b/lib/SlackBuilder/Download.hs
@@ -18,6 +18,7 @@ module SlackBuilder.Download
 import Data.ByteString (ByteString)
 import qualified Data.ByteString as ByteString
 import qualified Data.ByteString.Char8 as Char8
+import qualified Data.List.NonEmpty as NonEmpty
 import Data.Foldable (find)
 import Data.Map.Strict (Map)
 import qualified Data.Map.Strict as Map
@@ -40,7 +41,8 @@ import System.Process
     )
 import System.Exit (ExitCode(..))
 import Control.Monad (unless)
-import Text.URI (URI(..), mkURI)
+import Text.URI (URI(..))
+import qualified Text.URI as URI
 import Network.HTTP.Req
     ( useHttpsURI
     , HEAD(..)
@@ -52,7 +54,7 @@ import Network.HTTP.Req
     , responseStatusCode
     , HttpConfig(..)
     , GET(..)
-    , reqBr
+    , reqBr, MonadHttp
     )
 import Data.Functor ((<&>))
 import Network.HTTP.Client (BodyReader, Response(..), brRead)
@@ -72,6 +74,7 @@ import Data.Void (Void)
 import qualified Data.Conduit.Lzma as Lzma
 import qualified Data.Conduit.Zlib as Zlib
 import Control.Monad.Catch (MonadThrow(..))
+import Data.Maybe (fromMaybe)
 
 updateSlackBuildVersion :: Text -> Text -> Map Text Text -> SlackBuilderT ()
 updateSlackBuildVersion packagePath version additionalDownloads = do
@@ -127,7 +130,7 @@ commit packagePath version = do
 
 hostedSources :: Text -> SlackBuilderT URI
 hostedSources absoluteURL = SlackBuilderT (asks downloadURL)
-    >>= liftIO . mkURI . (<> absoluteURL)
+    >>= liftIO . URI.mkURI . (<> absoluteURL)
 
 remoteFileExists :: Text -> SlackBuilderT Bool
 remoteFileExists url = hostedSources url
@@ -196,17 +199,7 @@ sinkHash = sink hashInit
     sink ctx = await
         >>= maybe (pure $ hashFinalize ctx) (sink . hashUpdate ctx)
 
-download :: URI -> FilePath -> SlackBuilderT (Maybe (Digest MD5))
-download uri target = traverse (runReq defaultHttpConfig . go . fst)
-    $ useHttpsURI uri
-  where
-    go uri' = reqBr GET uri' NoReqBody mempty readResponse
-    readResponse :: Response BodyReader -> IO (Digest MD5)
-    readResponse response = runConduitRes
-        $ responseBodySource response
-        .| getZipSink (ZipSink (sinkFile target) *> ZipSink sinkHash)
-
-cloneAndUpload :: Text -> FilePath -> Text -> SlackBuilderT (Maybe (URI, Digest MD5))
+cloneAndUpload :: Text -> FilePath -> Text -> SlackBuilderT (URI, Digest MD5)
 cloneAndUpload repo tarballPath tagPrefix = do
     localPath <- relativeToRepository $ tarballPath <.> "tar.xz"
     let packageName = takeFileName $ takeDirectory tarballPath
@@ -216,12 +209,40 @@ cloneAndUpload repo tarballPath tagPrefix = do
     remoteFileExists' <- remoteFileExists remoteArchivePath
 
     if remoteFileExists'
-    then fmap (remoteResultURI,) <$> download remoteResultURI localPath
+    then (remoteResultURI,) . snd
+        <$> download remoteResultURI (takeDirectory localPath)
     else
         let go = sourceFile localPath .| sinkHash
          in cloneAndArchive repo tarballPath tagPrefix
             >> uploadCommand localPath remoteArchivePath
-            >> liftIO (runConduitRes go) <&> Just . (remoteResultURI,)
+            >> liftIO (runConduitRes go) <&> (remoteResultURI,)
+
+-- | Downlaods a file into the directory. Returns name of the downloaded file
+-- and checksum.
+--
+-- The filename is read from the disposition header or from the URL if the
+-- Content-Disposition is missing.
+download :: URI -> FilePath -> SlackBuilderT (FilePath, Digest MD5)
+download uri packagePath = runReq defaultHttpConfig go
+  where
+    go
+        | Just uriPath <- URI.uriPath uri =
+            reqGet uri
+                $ readResponse
+                $ Text.unpack
+                $ URI.unRText
+                $ NonEmpty.last
+                $ snd uriPath
+        | otherwise = throwM $ HttpsUrlExpected uri
+    readResponse :: FilePath -> Response BodyReader -> IO (FilePath, Digest MD5)
+    readResponse downloadFileName response = do
+        let attachmentName = dispositionAttachment response
+            targetFileName = fromMaybe downloadFileName attachmentName
+            target = packagePath </> fromMaybe downloadFileName attachmentName
+        digest <- runConduitRes
+            $ responseBodySource response
+            .| getZipSink (ZipSink (sinkFile target) *> ZipSink sinkHash)
+        pure (targetFileName, digest)
 
 -- | Downloads a compressed tar archive and extracts its contents on the fly to
 -- a directory.
@@ -230,22 +251,14 @@ cloneAndUpload repo tarballPath tagPrefix = do
 -- recognized as tar archive, returns the attachment name from the
 -- disposition header without the extension. So if the disposition header
 -- is "attachment; filename=package-1.2.3.tar.gz", returns "package-1.2.3".
-extractRemote :: URI -> Text -> SlackBuilderT (Maybe Text)
-extractRemote uri' packagePath = do
-    repository' <- SlackBuilderT $ asks repository
-    let localToRepository = repository' </> Text.unpack packagePath
-    case useHttpsURI uri' of
-        Just (httpsURI, _httpsOptions) ->
-            runReq defaultHttpConfig $ go localToRepository httpsURI
-        Nothing -> throwM $ HttpsUrlExpected uri'
+extractRemote :: URI -> FilePath -> SlackBuilderT (Maybe FilePath)
+extractRemote uri' packagePath =
+    runReq defaultHttpConfig $ go packagePath
   where
-    go toTarget url' = reqBr GET url' NoReqBody mempty $ readResponse toTarget
-    readResponse :: FilePath -> Response BodyReader -> IO (Maybe Text)
+    go toTarget = reqGet uri' $ readResponse toTarget
+    readResponse :: FilePath -> Response BodyReader -> IO (Maybe FilePath)
     readResponse toTarget response = do
-        let attachmentName
-                = fmap (Char8.unpack . snd . Char8.breakEnd (== '=') . snd)
-                $ find ((== "Content-Disposition") . fst)
-                $ responseHeaders response
+        let attachmentName = dispositionAttachment response
             (decompress, attachmentDirectory) =
                 case attachmentName of
                     Just attachmentName'
@@ -257,8 +270,23 @@ extractRemote uri' packagePath = do
         runConduitRes $ responseBodySource response
             .| decompress
             .| untar (withDecompressedFile toTarget)
-        pure $ Text.pack <$> attachmentDirectory
+        pure attachmentDirectory
     withDecompressedFile toTarget FileInfo{..}
         | Char8.last filePath /= '/' =
             sinkFile (toTarget </> Char8.unpack filePath)
         | otherwise = liftIO (createDirectory (toTarget </> Char8.unpack filePath))
+
+dispositionAttachment :: Response BodyReader -> Maybe FilePath
+dispositionAttachment response
+    = fmap (Char8.unpack . snd . Char8.breakEnd (== '=') . snd)
+    $ find ((== "Content-Disposition") . fst)
+    $ responseHeaders response
+
+reqGet :: (MonadThrow m, MonadHttp m)
+    => URI
+    -> (Response BodyReader -> IO a)
+    -> m a
+reqGet uri bodyReader
+    | Just (httpsURI, httpsOptions) <- useHttpsURI uri =
+        reqBr GET httpsURI NoReqBody httpsOptions bodyReader
+    | otherwise = throwM $ HttpsUrlExpected uri