255 lines
8.6 KiB
Haskell
255 lines
8.6 KiB
Haskell
{-# LANGUAGE DataKinds #-}
|
|
{-# LANGUAGE DeriveGeneric #-}
|
|
{-# LANGUAGE FlexibleContexts #-}
|
|
{-# LANGUAGE OverloadedStrings #-}
|
|
{-# LANGUAGE QuasiQuotes #-}
|
|
{-# LANGUAGE TemplateHaskell #-}
|
|
{-# LANGUAGE TypeApplications #-}
|
|
{-# LANGUAGE TypeFamilies #-}
|
|
|
|
|
|
module GHCup.Download.IOStreams where
|
|
|
|
|
|
import GHCup.Download.Utils
|
|
import GHCup.Errors
|
|
import GHCup.Types.Optics
|
|
import GHCup.Types.JSON ( )
|
|
import GHCup.Utils.File
|
|
import GHCup.Utils.Prelude
|
|
|
|
import Control.Applicative
|
|
import Control.Exception.Safe
|
|
import Control.Monad
|
|
import Control.Monad.Reader
|
|
import Data.ByteString ( ByteString )
|
|
import Data.ByteString.Builder
|
|
import Data.CaseInsensitive ( CI )
|
|
import Data.IORef
|
|
import Data.Maybe
|
|
import Data.Text.Read
|
|
import HPath
|
|
import HPath.IO as HIO
|
|
import Haskus.Utils.Variant.Excepts
|
|
import Network.Http.Client hiding ( URL )
|
|
import Optics
|
|
import Prelude hiding ( abs
|
|
, readFile
|
|
, writeFile
|
|
)
|
|
import "unix" System.Posix.IO.ByteString
|
|
hiding ( fdWrite )
|
|
import "unix-bytestring" System.Posix.IO.ByteString
|
|
( fdWrite )
|
|
import System.ProgressBar
|
|
import URI.ByteString
|
|
|
|
import qualified Data.ByteString as BS
|
|
import qualified Data.ByteString.Lazy as L
|
|
import qualified Data.Map.Strict as M
|
|
import qualified Data.Text.Encoding as E
|
|
import qualified System.IO.Streams as Streams
|
|
|
|
|
|
|
|
|
|
|
|
----------------------------
|
|
--[ Low-level (non-curl) ]--
|
|
----------------------------
|
|
|
|
|
|
-- | Load the result of this download into memory at once.
|
|
downloadBS' :: MonadIO m
|
|
=> Bool -- ^ https?
|
|
-> ByteString -- ^ host (e.g. "www.example.com")
|
|
-> ByteString -- ^ path (e.g. "/my/file") including query
|
|
-> Maybe Int -- ^ optional port (e.g. 3000)
|
|
-> Excepts
|
|
'[ HTTPStatusError
|
|
, URIParseError
|
|
, UnsupportedScheme
|
|
, NoLocationHeader
|
|
, TooManyRedirs
|
|
]
|
|
m
|
|
(L.ByteString)
|
|
downloadBS' https host path port = do
|
|
bref <- liftIO $ newIORef (mempty :: Builder)
|
|
let stepper bs = modifyIORef bref (<> byteString bs)
|
|
downloadInternal False https host path port stepper
|
|
liftIO (readIORef bref <&> toLazyByteString)
|
|
|
|
|
|
downloadToFile :: (MonadMask m, MonadIO m)
|
|
=> Bool -- ^ https?
|
|
-> ByteString -- ^ host (e.g. "www.example.com")
|
|
-> ByteString -- ^ path (e.g. "/my/file") including query
|
|
-> Maybe Int -- ^ optional port (e.g. 3000)
|
|
-> Path Abs -- ^ destination file to create and write to
|
|
-> Excepts '[DownloadFailed] m ()
|
|
downloadToFile https host fullPath port destFile = do
|
|
fd <- liftIO $ createRegularFileFd newFilePerms destFile
|
|
let stepper = fdWrite fd
|
|
flip finally (liftIO $ closeFd fd)
|
|
$ reThrowAll DownloadFailed $ downloadInternal True https host fullPath port stepper
|
|
|
|
|
|
downloadInternal :: MonadIO m
|
|
=> Bool -- ^ whether to show a progress bar
|
|
-> Bool -- ^ https?
|
|
-> ByteString -- ^ host
|
|
-> ByteString -- ^ path with query
|
|
-> Maybe Int -- ^ optional port
|
|
-> (ByteString -> IO a) -- ^ the consuming step function
|
|
-> Excepts
|
|
'[ HTTPStatusError
|
|
, URIParseError
|
|
, UnsupportedScheme
|
|
, NoLocationHeader
|
|
, TooManyRedirs
|
|
]
|
|
m
|
|
()
|
|
downloadInternal = go (5 :: Int)
|
|
|
|
where
|
|
go redirs progressBar https host path port consumer = do
|
|
r <- liftIO $ withConnection' https host port action
|
|
veitherToExcepts r >>= \case
|
|
Just r' ->
|
|
if redirs > 0 then followRedirectURL r' else throwE TooManyRedirs
|
|
Nothing -> pure ()
|
|
where
|
|
action c = do
|
|
let q = buildRequest1 $ http GET path
|
|
|
|
sendRequest c q emptyBody
|
|
|
|
receiveResponse
|
|
c
|
|
(\r i' -> runE $ do
|
|
let scode = getStatusCode r
|
|
if
|
|
| scode >= 200 && scode < 300 -> downloadStream r i' >> pure Nothing
|
|
| scode >= 300 && scode < 400 -> case getHeader r "Location" of
|
|
Just r' -> pure $ Just $ r'
|
|
Nothing -> throwE NoLocationHeader
|
|
| otherwise -> throwE $ HTTPStatusError scode
|
|
)
|
|
|
|
followRedirectURL bs = case parseURI strictURIParserOptions bs of
|
|
Right uri' -> do
|
|
(https', host', fullPath', port') <- liftE $ uriToQuadruple uri'
|
|
go (redirs - 1) progressBar https' host' fullPath' port' consumer
|
|
Left e -> throwE e
|
|
|
|
downloadStream r i' = do
|
|
let size = case getHeader r "Content-Length" of
|
|
Just x' -> case decimal $ E.decodeUtf8 x' of
|
|
Left _ -> 0
|
|
Right (r', _) -> r'
|
|
Nothing -> 0
|
|
|
|
mpb <- if progressBar
|
|
then Just <$> (liftIO $ newProgressBar defStyle 10 (Progress 0 size ()))
|
|
else pure Nothing
|
|
|
|
outStream <- liftIO $ Streams.makeOutputStream
|
|
(\case
|
|
Just bs -> do
|
|
forM_ mpb $ \pb -> incProgress pb (BS.length bs)
|
|
void $ consumer bs
|
|
Nothing -> pure ()
|
|
)
|
|
liftIO $ Streams.connect i' outStream
|
|
|
|
|
|
getHead :: (MonadCatch m, MonadIO m)
|
|
=> URI
|
|
-> Excepts
|
|
'[ HTTPStatusError
|
|
, URIParseError
|
|
, UnsupportedScheme
|
|
, NoLocationHeader
|
|
, TooManyRedirs
|
|
, ProcessError
|
|
]
|
|
m
|
|
(M.Map (CI ByteString) ByteString)
|
|
getHead uri' | scheme == "https" = head' True
|
|
| scheme == "http" = head' False
|
|
| otherwise = throwE UnsupportedScheme
|
|
|
|
where
|
|
scheme = view (uriSchemeL' % schemeBSL') uri'
|
|
head' https = do
|
|
(_, host', fullPath', port') <- liftE $ uriToQuadruple uri'
|
|
liftE $ headInternal https host' fullPath' port'
|
|
|
|
|
|
headInternal :: MonadIO m
|
|
=> Bool -- ^ https?
|
|
-> ByteString -- ^ host
|
|
-> ByteString -- ^ path with query
|
|
-> Maybe Int -- ^ optional port
|
|
-> Excepts
|
|
'[ HTTPStatusError
|
|
, URIParseError
|
|
, UnsupportedScheme
|
|
, TooManyRedirs
|
|
, NoLocationHeader
|
|
]
|
|
m
|
|
(M.Map (CI ByteString) ByteString)
|
|
headInternal = go (5 :: Int)
|
|
|
|
where
|
|
go redirs https host path port = do
|
|
r <- liftIO $ withConnection' https host port action
|
|
veitherToExcepts r >>= \case
|
|
Left r' ->
|
|
if redirs > 0 then followRedirectURL r' else throwE TooManyRedirs
|
|
Right hs -> pure hs
|
|
where
|
|
|
|
action c = do
|
|
let q = buildRequest1 $ http HEAD path
|
|
|
|
sendRequest c q emptyBody
|
|
|
|
unsafeReceiveResponse
|
|
c
|
|
(\r _ -> runE $ do
|
|
let scode = getStatusCode r
|
|
if
|
|
| scode >= 200 && scode < 300 -> do
|
|
let headers = getHeaderMap r
|
|
pure $ Right $ headers
|
|
| scode >= 300 && scode < 400 -> case getHeader r "Location" of
|
|
Just r' -> pure $ Left $ r'
|
|
Nothing -> throwE NoLocationHeader
|
|
| otherwise -> throwE $ HTTPStatusError scode
|
|
)
|
|
|
|
followRedirectURL bs = case parseURI strictURIParserOptions bs of
|
|
Right uri' -> do
|
|
(https', host', fullPath', port') <- liftE $ uriToQuadruple uri'
|
|
go (redirs - 1) https' host' fullPath' port'
|
|
Left e -> throwE e
|
|
|
|
|
|
withConnection' :: Bool
|
|
-> ByteString
|
|
-> Maybe Int
|
|
-> (Connection -> IO a)
|
|
-> IO a
|
|
withConnection' https host port action = bracket acquire closeConnection action
|
|
|
|
where
|
|
acquire = case https of
|
|
True -> do
|
|
ctx <- baselineContextSSL
|
|
openConnectionSSL ctx host (fromIntegral $ fromMaybe 443 port)
|
|
False -> openConnection host (fromIntegral $ fromMaybe 80 port)
|