Move state token handling generic oauth2Plugin

Resolves #23

Yesod/Auth/OAuth2.hs

@@ -14,15 +14,18 @@ module Yesod.Auth.OAuth2
     , module Network.OAuth.OAuth2
     ) where
 
+import Control.Applicative ((<$>))
 import Control.Exception.Lifted
 import Control.Monad.IO.Class
 import Data.ByteString (ByteString)
-import Data.Text (Text)
+import Data.Monoid ((<>))
+import Data.Text (Text, pack)
 import Data.Text.Encoding (decodeUtf8With, encodeUtf8)
 import Data.Text.Encoding.Error (lenientDecode)
 import Data.Typeable
 import Network.OAuth.OAuth2
 import Network.HTTP.Conduit(Manager)
+import System.Random
 import Yesod.Auth
 import Yesod.Core
 import Yesod.Form
@@ -52,28 +55,46 @@ authOAuth2 name oauth getCreds = AuthPlugin name dispatch login
     where
         url = PluginR name ["callback"]
 
-        withCallback = do
+        withCallback csrfToken = do
             tm <- getRouteToParent
             render <- lift $ getUrlRender
-            return $ oauth { oauthCallback = Just $ encodeUtf8 $ render $ tm url }
+            let newEndpoint = oauthOAuthorizeEndpoint oauth <> "&state=" <> encodeUtf8 csrfToken
+            return $ oauth { 
+                    oauthCallback = Just $ encodeUtf8 $ render $ tm url, 
+                    oauthOAuthorizeEndpoint = newEndpoint
+                }
 
         dispatch "GET" ["forward"] = do
-            authUrl <- fmap (bsToText . authorizationUrl) withCallback
+            csrfToken <- liftIO $ generateToken
+            setSession tokenSessionKey csrfToken
+            authUrl <- (bsToText . authorizationUrl) <$> withCallback csrfToken
             lift $ redirect authUrl
 
         dispatch "GET" ["callback"] = do
-            code <- lift $ runInputGet $ ireq textField "code"
+            newToken <- lookupGetParam "state"
-            oauth' <- withCallback
+            oldToken <- lookupSession tokenSessionKey
-            master <- lift getYesod
+            deleteSession tokenSessionKey
-            result <- liftIO $ fetchAccessToken (authHttpManager master) oauth' (encodeUtf8 code)
+            case newToken of
-            case result of
+                Just csrfToken | newToken == oldToken -> do
-                Left _ -> permissionDenied "Unable to retreive OAuth2 token"
+                    code <- lift $ runInputGet $ ireq textField "code"
-                Right token -> do
+                    oauth' <- withCallback csrfToken
-                    creds <- liftIO $ getCreds (authHttpManager master) token
+                    master <- lift getYesod
-                    lift $ setCredsRedirect creds
+                    result <- liftIO $ fetchAccessToken (authHttpManager master) oauth' (encodeUtf8 code)
+                    case result of
+                        Left _ -> permissionDenied "Unable to retreive OAuth2 token"
+                        Right token -> do
+                            creds <- liftIO $ getCreds (authHttpManager master) token
+                            lift $ setCredsRedirect creds
+                _ ->
+                    permissionDenied "Invalid OAuth2 state token"
 
         dispatch _ _ = notFound
 
+        generateToken = (pack . take 30 . randomRs ('a','z')) <$> newStdGen
+
+        tokenSessionKey :: Text
+        tokenSessionKey = "_yesod_oauth2_" <> name
+
         login tm = do
             render <- getUrlRender
             let oaUrl = render $ tm $ oauth2Url name

Yesod/Auth/OAuth2/Github.hs

@@ -18,16 +18,11 @@ import Control.Exception.Lifted
 import Control.Monad (mzero)
 import Data.Aeson
 import Data.Text (Text)
-import Data.Monoid (mappend)
+import Data.Monoid ((<>))
 import Data.Text.Encoding (encodeUtf8, decodeUtf8)
 import Yesod.Auth
 import Yesod.Auth.OAuth2
-import Yesod.Core
-import Yesod.Form
 import Network.HTTP.Conduit(Manager)
-import Data.UUID (toString)
-import Data.UUID.V4 (nextRandom)
-import qualified Data.ByteString as BS
 import qualified Data.Text as T
 
 data GithubUser = GithubUser
@@ -67,38 +62,16 @@ oauth2GithubScoped :: YesodAuth m
              -> Text -- ^ Client Secret
              -> [Text] -- ^ List of scopes to request
              -> AuthPlugin m
-oauth2GithubScoped clientId clientSecret scopes = basicPlugin {apDispatch = dispatch}
+oauth2GithubScoped clientId clientSecret scopes = authOAuth2 "github" oauth fetchGithubProfile
     where
         oauth = OAuth2
                 { oauthClientId            = encodeUtf8 clientId
                 , oauthClientSecret        = encodeUtf8 clientSecret
-                , oauthOAuthorizeEndpoint  = encodeUtf8 $ "https://github.com/login/oauth/authorize?scope=" `T.append` T.intercalate "," scopes
+                , oauthOAuthorizeEndpoint  = encodeUtf8 $ "https://github.com/login/oauth/authorize?scope=" <> T.intercalate "," scopes
                 , oauthAccessTokenEndpoint = "https://github.com/login/oauth/access_token"
                 , oauthCallback            = Nothing
                 }
 
-        withState state = authOAuth2 "github"
-            (oauth {oauthOAuthorizeEndpoint = oauthOAuthorizeEndpoint oauth `BS.append` "&state=" `BS.append` encodeUtf8 state})
-            fetchGithubProfile
-
-        basicPlugin = authOAuth2 "github" oauth fetchGithubProfile
-
-        dispatch "GET" ["forward"] = do
-            state <- liftIO $ fmap (T.pack . toString) nextRandom
-            setSession "githubState" state
-            apDispatch (withState state) "GET" ["forward"]
-
-        dispatch "GET" ["callback"] = do
-            state <- lift $ runInputGet $ ireq textField "state"
-            savedState <- lookupSession "githubState"
-            _ <- apDispatch basicPlugin "GET" ["callback"]
-            case savedState of
-                Just saved | saved == state -> apDispatch basicPlugin "GET" ["callback"]
-                Just saved -> invalidArgs ["state: " `mappend` state `mappend` ", and not: " `mappend` saved]
-                _ -> invalidArgs ["state: " `mappend` state]
-
-        dispatch method ps = apDispatch basicPlugin method ps
-
 fetchGithubProfile :: Manager -> AccessToken -> IO (Creds m)
 fetchGithubProfile manager token = do
     userResult <- authGetJSON manager token "https://api.github.com/user"

yesod-auth-oauth2.cabal

@@ -36,13 +36,13 @@ library
                    , aeson                   >= 0.6       && < 0.9
                    , yesod-core              >= 1.2       && < 1.5
                    , authenticate            >= 1.3.2.7   && < 1.4
+                   , random
                    , yesod-auth              >= 1.3       && < 1.5
                    , text                    >= 0.7       && < 2.0
                    , yesod-form              >= 1.3       && < 1.5
                    , transformers            >= 0.2.2     && < 0.5
                    , hoauth2                 >= 0.4.1     && < 0.5
                    , lifted-base             >= 0.2       && < 0.4
-                   , uuid                    >= 1.3       && < 1.4
 
     exposed-modules: Yesod.Auth.OAuth2
                      Yesod.Auth.OAuth2.Google