From e3c61789bafbf7278f0be92ccbbea081ebbae070 Mon Sep 17 00:00:00 2001
From: patrick brisbin <pbrisbin@gmail.com>
Date: Tue, 18 Sep 2018 10:41:17 -0400
Subject: [PATCH] Ensure we rescue our exceptions too

For some reason, I thought tryIO would catch our own exception is we
threw them via throwIO, but that's incorrect. Our own exceptions are not
IOExceptions, so they squeak by. This fixes that.
---
 src/Yesod/Auth/OAuth2/Dispatch.hs | 12 ++++++++++--
 1 file changed, 10 insertions(+), 2 deletions(-)

diff --git a/src/Yesod/Auth/OAuth2/Dispatch.hs b/src/Yesod/Auth/OAuth2/Dispatch.hs
index 580f7f9..4e5ccab 100644
--- a/src/Yesod/Auth/OAuth2/Dispatch.hs
+++ b/src/Yesod/Auth/OAuth2/Dispatch.hs
@@ -3,6 +3,7 @@
 {-# LANGUAGE OverloadedStrings #-}
 {-# LANGUAGE RankNTypes #-}
 {-# LANGUAGE RecordWildCards #-}
+{-# LANGUAGE ScopedTypeVariables #-}
 {-# LANGUAGE TemplateHaskell #-}
 {-# LANGUAGE TypeFamilies #-}
 module Yesod.Auth.OAuth2.Dispatch
@@ -10,7 +11,7 @@ module Yesod.Auth.OAuth2.Dispatch
     , dispatchAuthRequest
     ) where
 
-import Control.Exception.Safe (throwString, tryIO)
+import Control.Exception.Safe
 import Control.Monad (unless, (<=<))
 import Data.Monoid ((<>))
 import Data.Text (Text)
@@ -22,6 +23,7 @@ import System.Random (newStdGen, randomRs)
 import URI.ByteString.Extension
 import Yesod.Auth hiding (ServerError)
 import Yesod.Auth.OAuth2.ErrorResponse
+import Yesod.Auth.OAuth2.Exception
 import Yesod.Core hiding (ErrorResponse)
 
 -- | How to take an @'OAuth2Token'@ and retrieve user credentials
@@ -66,7 +68,7 @@ dispatchCallback name oauth2 getCreds = do
     manager <- authHttpManager
     oauth2' <- withCallbackAndState name oauth2 csrf
     token <- errLeft $ fetchAccessToken manager oauth2' $ ExchangeToken code
-    creds <- errLeft $ tryIO $ getCreds manager token
+    creds <- errLeft $ tryFetchCreds $ getCreds manager token
     setCredsRedirect creds
   where
     errLeft :: Show e => IO (Either e a) -> AuthHandler m a
@@ -83,6 +85,12 @@ redirectMessage msg = do
     setMessage $ toHtml msg
     redirect $ toParent LoginR
 
+tryFetchCreds :: IO a -> IO (Either SomeException a)
+tryFetchCreds f =
+    (Right <$> f)
+        `catch` (\(ex :: IOException) -> pure $ Left $ toException ex)
+        `catch` (\(ex :: YesodOAuth2Exception) -> pure $ Left $ toException ex)
+
 withCallbackAndState :: Text -> OAuth2 -> Text -> AuthHandler m OAuth2
 withCallbackAndState name oauth2 csrf = do
     let url = PluginR name ["callback"]