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 as E hiding (random_) import qualified Database.Esqueleto.PostgreSQL as E data TestDownloadMode = TestDownloadDirect | TestDownloadInTransaction | TestDownloadFromDatabase deriving (Eq, Ord, Enum, Bounded, Read, Show, Generic, Typeable) 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, Typeable) 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 $!!) . fileContentContent . entityVal) .| takeLimit dlMaxSize where sourceDBFiles = E.selectSource . E.from $ \fileContent -> do E.orderBy [E.asc $ E.random_ @Int64] return fileContent 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 }