Yesod - Generate a form from post data stored in the session

GenerateFormFromLastPost.hs

-- Most of these pragmas probably aren't needed, just ripped from Yesod.Form.Functions
{-# LANGUAGE QuasiQuotes #-}
{-# LANGUAGE TupleSections #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE CPP #-}

import           ClassyPrelude.Yesod
import           Control.Monad.Trans.RWS (evalRWST)
import           Data.Byteable         (constEqBytes)
import qualified Data.Map              as Map
import qualified Data.Text.Encoding    as TE
import           Network.Wai           (requestMethod)

-- Create a form from last post data in the session if exists, otherwise create a blank form.
generateFormFromLastPost :: (RenderMessage (HandlerSite m) FormMessage, MonadHandler m) =>
                            (Markup -> MForm m (FormResult a, xml)) -> m (xml, Enctype)
generateFormFromLastPost form = do
    env <- getLastInvalidPost
    case env of
        Nothing -> generateFormPost form
        Just _ -> first snd <$> postHelper form env

lastInvalidPostSessionKey :: Text
lastInvalidPostSessionKey = "lastInvalidPost"

-- Sets the post data retreived from postEnv, ignoring the FileEnv.
setLastInvalidPost :: MonadHandler m => Maybe (Env, FileEnv) -> m ()
setLastInvalidPost Nothing = return ()
setLastInvalidPost (Just (env, _)) = sessionSetter lastInvalidPostSessionKey env

-- Retrieves the previous post data to be passed to postHelper.
getLastInvalidPost :: MonadHandler m => m (Maybe (Env, FileEnv))
getLastInvalidPost = do
    result <- sessionGetter lastInvalidPostSessionKey
    return $ case result of
        Nothing -> Nothing
        Just env -> Just (env, Map.fromList [])

sessionSetter :: (MonadHandler m, Show a) => Text -> a -> m ()
sessionSetter key = setSession key . pack . show

sessionGetter :: (MonadHandler m, Read b) => Text -> m (Maybe b)
sessionGetter key = do
    mBS <- lookupSession key
    return $ readMaybe . unpack =<< mBS
    
    
-- Ripped source from Yesod.Form.Functions
-- Requires cabal dependencies: wai, byteable, transformers

runFormGeneric :: Monad m
               => MForm m a
               -> HandlerSite m
               -> [Text]
               -> Maybe (Env, FileEnv)
               -> m (a, Enctype)
runFormGeneric form site langs env = evalRWST form (env, site, langs) (IntSingle 0)

postEnv :: (MonadHandler m, MonadResource m)
        => m (Maybe (Env, FileEnv))
postEnv = do
    req <- getRequest
    if requestMethod (reqWaiRequest req) == "GET"
        then return Nothing
        else do
            (p, f) <- runRequestBody
            let p' = Map.unionsWith (++) $ map (\(x, y) -> Map.singleton x [y]) p
            return $ Just (p', Map.unionsWith (++) $ map (\(k, v) -> Map.singleton k [v]) f)

postHelper  :: (MonadHandler m, RenderMessage (HandlerSite m) FormMessage)
            => (Html -> MForm m (FormResult a, xml))
            -> Maybe (Env, FileEnv)
            -> m ((FormResult a, xml), Enctype)
postHelper form env = do
    req <- getRequest
    let tokenKey = "_token" :: Text
    let token =
            case reqToken req of
                Nothing -> mempty
                Just n -> [shamlet|<input type=hidden name=#{tokenKey} value=#{n}>|]
    m <- getYesod
    langs <- languages
    ((res, xml), enctype) <- runFormGeneric (form token) m langs env
    let res' =
            case (res, env) of
                (FormSuccess{}, Just (params, _))
                    | not (Map.lookup tokenKey params === reqToken req) ->
                        FormFailure [renderMessage m langs MsgCsrfWarning]
                _ -> res
            where (Just [t1]) === (Just t2) = TE.encodeUtf8 t1 `constEqBytes` TE.encodeUtf8 t2
                  Nothing     === Nothing   = True   -- It's important to use constTimeEq
                  _           === _         = False  -- in order to avoid timing attacks.
    return ((res', xml), enctype)