diff --git a/src/Handler/Admin/Tokens.hs b/src/Handler/Admin/Tokens.hs index 949b0895f..698b15345 100644 --- a/src/Handler/Admin/Tokens.hs +++ b/src/Handler/Admin/Tokens.hs @@ -8,7 +8,6 @@ import Handler.Utils import qualified Data.HashSet as HashSet import qualified Data.HashMap.Strict as HashMap import qualified Data.Set as Set -import qualified Data.Map as Map import qualified Data.Aeson as Aeson import qualified Data.Aeson.Encode.Pretty as Aeson @@ -19,15 +18,12 @@ import qualified Data.Text as Text import qualified Database.Esqueleto as E hiding (random_) import qualified Database.Esqueleto.PostgreSQL as E -import qualified Database.Esqueleto.Utils as E -import qualified Control.Monad.Random.Class as Random -import Control.Monad.Random.Strict (evalRand, Rand) - -import qualified Data.Conduit.List as C (unfoldM) import qualified Data.Conduit.Combinators as C -import qualified Crypto.Random as Crypto +import Data.List (genericTake) + +import System.Random.Shuffle (shuffleM) data BTFImpersonate @@ -136,22 +132,13 @@ postAdminTokensR = do E.limit btfiCount return $ user E.^. UserId | otherwise -> do - users'' <- E.select . E.from $ \user -> - return ( user E.^. UserId - , E.maybe E.nothing (E.just . E.diffTimes (E.val now)) $ user E.^. UserLastAuthentication - ) - let users :: Map UserId (Maybe Rational) - users = Map.fromList $ users'' <&> \(E.Value uid', E.Value mDiff) -> (uid', recip . toRational <$> mDiff) - chooseUsers :: ConduitT () UserId (Rand Crypto.ChaChaDRG) () - chooseUsers = C.unfoldM chooseUsers' (users, btfiCount) - where chooseUsers' (users', n) = runMaybeT $ do - guard $ n > 0 - let getWeighted = MaybeT . Random.weightedMay . mapMaybe (\(uid', mDiff) -> (uid', ) <$> mDiff) $ Map.toList users' - getUnweighted = MaybeT . Random.uniformMay $ Map.keysSet users' - user <- getWeighted <|> getUnweighted - return (user, (Map.delete user users', pred n)) - drg <- liftIO Crypto.drgNew - return . flip evalRand drg . runConduit $ chooseUsers .| C.foldMap pure + users <- fmap (fmap E.unValue) . E.select . E.from $ \user -> do + E.orderBy [ E.asc . E.isNothing $ user E.^. UserLastAuthentication + , E.desc $ user E.^. UserLastAuthentication + ] + E.limit $ 2 * btfiCount + return $ user E.^. UserId + genericTake btfiCount <$> shuffleM users let toTokenFile :: UserId -> DB (Either Void DBFile)