110 lines
3.8 KiB
Haskell
110 lines
3.8 KiB
Haskell
-- SPDX-FileCopyrightText: 2022 Gregor Kleen <gregor.kleen@ifi.lmu.de>
|
|
--
|
|
-- SPDX-License-Identifier: AGPL-3.0-or-later
|
|
|
|
module Handler.Admin.Test.Download
|
|
( testDownload
|
|
, generateDownload', Random.seedNew
|
|
) where
|
|
|
|
import Import hiding (Builder)
|
|
import Handler.Utils
|
|
|
|
import qualified Crypto.Random as Random
|
|
|
|
import qualified Data.ByteString.Base64.URL as Base64
|
|
|
|
import qualified Data.Conduit.Combinators as C
|
|
|
|
import qualified Data.Binary as Binary
|
|
|
|
import Data.Binary.Builder (Builder)
|
|
|
|
import Control.Monad.Random.Lazy (evalRandT, liftRandT)
|
|
|
|
import qualified Database.Esqueleto.Legacy as E hiding (random_)
|
|
import qualified Database.Esqueleto.PostgreSQL as E
|
|
|
|
|
|
data TestDownloadMode
|
|
= TestDownloadDirect
|
|
| TestDownloadInTransaction
|
|
| TestDownloadFromDatabase
|
|
deriving (Eq, Ord, Enum, Bounded, Read, Show, Generic)
|
|
deriving anyclass (Universe, Finite, Binary)
|
|
|
|
nullaryPathPiece ''TestDownloadMode $ camelToPathPiece' 2
|
|
embedRenderMessage ''UniWorX ''TestDownloadMode id
|
|
|
|
data TestDownloadOptions = TestDownloadOptions
|
|
{ dlSeed :: Random.Seed
|
|
, dlMaxSize, dlChunkSize :: Int
|
|
, dlMode :: TestDownloadMode
|
|
} deriving (Generic)
|
|
deriving anyclass (Binary)
|
|
|
|
testDownloadForm :: Form TestDownloadOptions
|
|
testDownloadForm = identifyForm FIDTestDownload . renderWForm FormStandard $ do
|
|
randomSeed <- liftIO Random.seedNew
|
|
|
|
maxSizeRes <- wreq intField (fslI MsgTestDownloadMaxSize) . Just $ 2 * 2^30
|
|
modeRes <- wpopt (selectField optionsFinite) (fslI MsgTestDownloadMode) $ Just TestDownloadDirect
|
|
|
|
return $ TestDownloadOptions
|
|
randomSeed
|
|
<$> maxSizeRes
|
|
<*> pure (2^20)
|
|
<*> modeRes
|
|
|
|
|
|
generateDownload :: Monad m => TestDownloadOptions -> ConduitT i ByteString m ()
|
|
generateDownload TestDownloadOptions{..}
|
|
= C.unfold genChunk dlMaxSize .| generateDownload' dlSeed
|
|
where genChunk remaining
|
|
| remaining >= dlChunkSize = Just (dlChunkSize, remaining - dlChunkSize)
|
|
| remaining <= 0 = Nothing
|
|
| otherwise = Just (remaining, 0)
|
|
|
|
generateDownload' :: Monad m => Random.Seed -> ConduitT Int ByteString m ()
|
|
generateDownload' seed = transPipe (evalRandT ?? Random.drgNewSeed seed) $
|
|
C.mapM (liftRandT . (return .) . Random.randomBytesGenerate)
|
|
|
|
|
|
testDownload :: Handler Widget
|
|
testDownload = do
|
|
((dlRes, dlForm), dlEnctype) <- runFormPost testDownloadForm
|
|
|
|
formResult dlRes $ \opts@TestDownloadOptions{..} -> do
|
|
addHeader "Test-Download-Options" . decodeUtf8 . Base64.encode . toStrict $ Binary.encode opts
|
|
setContentDisposition ContentAttachment $ Just "u2w-test-download.bin"
|
|
let
|
|
sendDownload :: forall m. Monad m => ConduitT ByteString (Flush Builder) m ()
|
|
sendDownload = awaitForever $ \bs -> sendChunkBS bs >> sendFlush
|
|
|
|
sourceDBChunks :: ConduitT () Int DB ()
|
|
sourceDBChunks = forever sourceDBFiles
|
|
.| C.mapM (\x -> x <$ $logDebugS "testDownload.sourceDBChunks" (tshow $ entityKey x))
|
|
.| C.map ((length $!!) . fileContentChunkContent . entityVal)
|
|
.| takeLimit dlMaxSize
|
|
where
|
|
sourceDBFiles = E.selectSource . E.from $ \fileContentChunk -> do
|
|
E.orderBy [E.asc $ E.random_ @Int64]
|
|
return fileContentChunk
|
|
|
|
takeLimit n | n <= 0 = return ()
|
|
takeLimit n = do
|
|
c <- await
|
|
$logDebugS "testDownload.takeLimit" $ tshow c
|
|
case c of
|
|
Nothing -> return ()
|
|
Just c' -> yield c' >> takeLimit (n - c')
|
|
(sendResponse =<<) $ case dlMode of
|
|
TestDownloadDirect -> respondSource typeOctet $ generateDownload opts .| sendDownload
|
|
TestDownloadInTransaction -> respondSourceDB typeOctet $ generateDownload opts .| sendDownload
|
|
TestDownloadFromDatabase -> respondSourceDB typeOctet $ sourceDBChunks .| generateDownload' dlSeed .| sendDownload
|
|
|
|
return $
|
|
wrapForm dlForm def
|
|
{ formEncoding = dlEnctype
|
|
}
|