module Unison.Auth.HTTPClient (newAuthenticatedHTTPClient, AuthenticatedHttpClient (..)) where

import Data.Text.Encoding qualified as Text
import Network.HTTP.Client (Request)
import Network.HTTP.Client qualified as HTTP
import Network.HTTP.Client.TLS qualified as HTTP
import Unison.Auth.Tokens (TokenProvider)
import Unison.Codebase.Editor.UCMVersion (UCMVersion)
import Unison.Prelude
import Unison.Share.Types (codeserverIdFromURI)
import Unison.Util.HTTP qualified as HTTP

-- | Newtype to delineate HTTP Managers with access-token logic.
newtype AuthenticatedHttpClient = AuthenticatedHttpClient HTTP.Manager

-- | Returns a new http manager which applies the appropriate Authorization header to
-- any hosts our UCM is authenticated with.
newAuthenticatedHTTPClient :: (MonadIO m) => TokenProvider -> UCMVersion -> m AuthenticatedHttpClient
newAuthenticatedHTTPClient :: forall (m :: * -> *).
MonadIO m =>
TokenProvider -> UCMVersion -> m AuthenticatedHttpClient
newAuthenticatedHTTPClient TokenProvider
tokenProvider UCMVersion
ucmVersion = IO AuthenticatedHttpClient -> m AuthenticatedHttpClient
forall a. IO a -> m a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO AuthenticatedHttpClient -> m AuthenticatedHttpClient)
-> IO AuthenticatedHttpClient -> m AuthenticatedHttpClient
forall a b. (a -> b) -> a -> b
$ do
  let managerSettings :: ManagerSettings
managerSettings =
        ManagerSettings
HTTP.tlsManagerSettings
          ManagerSettings
-> (ManagerSettings -> ManagerSettings) -> ManagerSettings
forall a b. a -> (a -> b) -> b
& (Request -> IO Request) -> ManagerSettings -> ManagerSettings
HTTP.addRequestMiddleware (TokenProvider -> Request -> IO Request
authMiddleware TokenProvider
tokenProvider)
          ManagerSettings
-> (ManagerSettings -> ManagerSettings) -> ManagerSettings
forall a b. a -> (a -> b) -> b
& UCMVersion -> ManagerSettings -> ManagerSettings
HTTP.setUserAgent (UCMVersion -> UCMVersion
HTTP.ucmUserAgent UCMVersion
ucmVersion)
  Manager -> AuthenticatedHttpClient
AuthenticatedHttpClient (Manager -> AuthenticatedHttpClient)
-> IO Manager -> IO AuthenticatedHttpClient
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> ManagerSettings -> IO Manager
forall (m :: * -> *). MonadIO m => ManagerSettings -> m Manager
HTTP.newTlsManagerWith ManagerSettings
managerSettings

-- | Adds Bearer tokens to requests according to their host.
-- If a CredentialFailure occurs (failure to refresh a token), auth is simply omitted,
-- and the request is likely to trigger a 401 response which the caller can detect and initiate a re-auth.
--
-- If a host isn't associated with any credentials auth is omitted.
authMiddleware :: TokenProvider -> (Request -> IO Request)
authMiddleware :: TokenProvider -> Request -> IO Request
authMiddleware TokenProvider
tokenProvider Request
req = do
  -- The http manager "may run this function multiple times" when preparing a request.
  -- We may wish to look into a better way to attach auth to our requests in middleware, but
  -- this is a simple fix that works for now.
  -- https://github.com/snoyberg/http-client/issues/350
  case HeaderName -> [(HeaderName, ByteString)] -> Maybe ByteString
forall a b. Eq a => a -> [(a, b)] -> Maybe b
Prelude.lookup (HeaderName
"Authorization") (Request -> [(HeaderName, ByteString)]
HTTP.requestHeaders Request
req) of
    Just ByteString
_ -> Request -> IO Request
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Request
req
    Maybe ByteString
Nothing -> do
      case URI -> Either UCMVersion CodeserverId
codeserverIdFromURI (URI -> Either UCMVersion CodeserverId)
-> URI -> Either UCMVersion CodeserverId
forall a b. (a -> b) -> a -> b
$ (Request -> URI
HTTP.getUri Request
req) of
        -- If we can't identify an appropriate codeserver we pass it through without any auth.
        Left UCMVersion
_ -> Request -> IO Request
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Request
req
        Right CodeserverId
codeserverHost -> do
          TokenProvider
tokenProvider CodeserverId
codeserverHost IO (Either CredentialFailure UCMVersion)
-> (Either CredentialFailure UCMVersion -> IO Request)
-> IO Request
forall a b. IO a -> (a -> IO b) -> IO b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \case
            Right UCMVersion
token -> do
              let newReq :: Request
newReq = ByteString -> Request -> Request
HTTP.applyBearerAuth (UCMVersion -> ByteString
Text.encodeUtf8 UCMVersion
token) Request
req
              Request -> IO Request
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Request
newReq
            Left CredentialFailure
_err -> do
              Request -> IO Request
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Request
req