Retool Credential manager to use a singleton

Author: ChrisPenner

hie.yaml

@@ -42,6 +42,10 @@ cradle:
     - path: "lib/unison-hashing/src"
       component: "unison-hashing:lib"
 
+    - path: "lib/unison-credentials/src"
+      component: "unison-credentials:lib"
+
+
     - path: "lib/unison-prelude/src"
       component: "unison-prelude:lib"
 
@@ -98,7 +102,7 @@ cradle:
 
     - path: "parser-typechecker/tests"
       component: "unison-parser-typechecker:test:parser-typechecker-tests"
-      
+
     - path: "unison-runtime/src"
       component: "unison-runtime:lib"
 

lib/unison-credentials/src/Unison/Auth/CredentialFile.hs

@@ -6,7 +6,6 @@ import Data.Aeson qualified as Aeson
 import System.FilePath (takeDirectory, (</>))
 import System.IO.LockFile
 import Unison.Auth.Types
-import Unison.Debug qualified as Debug
 import Unison.Prelude
 import UnliftIO.Directory
 
@@ -26,26 +25,26 @@ getCredentialJSONFilePath = do
 
 -- | Atomically update the credential storage file.
 -- Creates an empty file automatically if one doesn't exist.
-atomicallyModifyCredentialsFile :: (MonadIO m) => (Credentials -> Credentials) -> m Credentials
-atomicallyModifyCredentialsFile f = liftIO $ do
-  credentialJSONPath <- getCredentialJSONFilePath
-  doesFileExist credentialJSONPath >>= \case
+atomicallyModifyCredentialsFile :: (MonadUnliftIO m) => (Credentials -> m (Credentials, r)) -> m r
+atomicallyModifyCredentialsFile f = do
+  credentialJSONPath <- liftIO $ getCredentialJSONFilePath
+  liftIO (doesFileExist credentialJSONPath) >>= \case
     True -> pure ()
-    False -> do
+    False -> liftIO $ do
       createDirectoryIfMissing True $ takeDirectory credentialJSONPath
       Aeson.encodeFile credentialJSONPath emptyCredentials
 
-  withLockFile lockfileConfig (withLockExt credentialJSONPath) $ do
+  toIO <- askRunInIO
+  liftIO $ withLockFile lockfileConfig (withLockExt credentialJSONPath) $ toIO $ do
     credentials <-
-      Aeson.eitherDecodeFileStrict credentialJSONPath >>= \case
+      liftIO (Aeson.eitherDecodeFileStrict credentialJSONPath) >>= \case
         -- If something goes wrong, just wipe the credentials file so we're in a clean slate.
         -- In the worst case the user will simply need to log in again.
-        Left err -> do
-          Debug.debugM Debug.Auth "Error decoding credentials file" err
-          Aeson.encodeFile credentialJSONPath emptyCredentials
+        Left _err -> do
+          liftIO $ Aeson.encodeFile credentialJSONPath emptyCredentials
           pure emptyCredentials
         Right creds -> pure creds
-    let newCredentials = f credentials
+    (newCredentials, r) <- f credentials
     when (newCredentials /= credentials) $ do
-      Aeson.encodeFile credentialJSONPath $ newCredentials
-    pure newCredentials
+      liftIO $ Aeson.encodeFile credentialJSONPath newCredentials
+    pure r

lib/unison-credentials/src/Unison/Auth/CredentialManager.hs

@@ -4,7 +4,7 @@ module Unison.Auth.CredentialManager
   ( saveCredentials,
     CredentialManager,
     newCredentialManager,
-    getCredentials,
+    getCodeserverCredentials,
     getOrCreatePersonalKey,
     isExpired,
   )
@@ -13,9 +13,11 @@ where
 import Control.Monad.Trans.Except
 import Data.Map qualified as Map
 import Data.Time.Clock (addUTCTime, diffUTCTime, getCurrentTime)
-import Unison.Auth.CredentialFile
+import System.IO.Unsafe (unsafePerformIO)
+import Unison.Auth.CredentialFile qualified as CF
 import Unison.Auth.PersonalKey (PersonalPrivateKey, generatePersonalKey)
-import Unison.Auth.Types
+import Unison.Auth.Types hiding (getCodeserverCredentials)
+import Unison.Auth.Types qualified as Auth
 import Unison.Prelude
 import Unison.Share.Types (CodeserverId)
 import UnliftIO qualified
@@ -25,46 +27,55 @@ import UnliftIO qualified
 -- Note: Currently the in-memory cache is _not_ updated if a different UCM updates
 -- the credentials file, however this shouldn't pose any problems, since auth will still
 -- be refreshed if we encounter any auth failures on requests.
-newtype CredentialManager = CredentialManager (UnliftIO.MVar Credentials)
+newtype CredentialManager = CredentialManager (UnliftIO.MVar (Maybe Credentials {- Credentials may or may not be initialized -}))
+
+-- | A global CredentialManager instance/singleton.
+globalCredentialsManager :: CredentialManager
+globalCredentialsManager = unsafePerformIO do
+  CredentialManager <$> UnliftIO.newMVar Nothing
+{-# NOINLINE globalCredentialsManager #-}
 
 -- | Fetches the user's personal key from the active profile, if it exists.
 -- Otherwise it creates a new personal key, saves it to the active profile, and returns it.
 getOrCreatePersonalKey :: (MonadUnliftIO m) => CredentialManager -> m PersonalPrivateKey
-getOrCreatePersonalKey credMan@(CredentialManager credsVar) = do
-  Credentials {activeProfile, personalKeys} <- liftIO (UnliftIO.readMVar credsVar)
-  case Map.lookup activeProfile personalKeys of
-    Just pk -> pure pk
-    Nothing -> do
-      pk <- generatePersonalKey
-      _ <- modifyCredentials credMan $ \creds ->
-        creds {personalKeys = Map.insert activeProfile pk creds.personalKeys}
-      pure pk
+getOrCreatePersonalKey credMan = do
+  modifyCredentials credMan \creds@(Credentials {activeProfile, personalKeys}) -> do
+    case Map.lookup activeProfile personalKeys of
+      Just pk -> pure (creds, pk)
+      Nothing -> do
+        pk <- generatePersonalKey
+        pure (creds {personalKeys = Map.insert activeProfile pk personalKeys}, pk)
 
 -- | Saves credentials to the active profile.
 saveCredentials :: (UnliftIO.MonadUnliftIO m) => CredentialManager -> CodeserverId -> CodeserverCredentials -> m ()
 saveCredentials credManager aud creds = do
-  void . modifyCredentials credManager $ setCodeserverCredentials aud creds
+  void . modifyCredentials credManager $ \cf -> pure (setCodeserverCredentials aud creds cf, ())
 
 -- | Atomically update the credential storage file, and update the in-memory cache.
-modifyCredentials :: (UnliftIO.MonadUnliftIO m) => CredentialManager -> (Credentials -> Credentials) -> m Credentials
+modifyCredentials :: (UnliftIO.MonadUnliftIO m) => CredentialManager -> (Credentials -> m (Credentials, r)) -> m r
 modifyCredentials (CredentialManager credsVar) f = do
   UnliftIO.modifyMVar credsVar $ \_ -> do
-    newCreds <- atomicallyModifyCredentialsFile f
-    pure (newCreds, newCreds)
+    (creds, r) <- CF.atomicallyModifyCredentialsFile (f >=> \(creds', r') -> pure (creds', (creds', r')))
+    pure (Just creds, r)
+
+readCredentials :: (UnliftIO.MonadUnliftIO m) => CredentialManager -> m Credentials
+readCredentials (CredentialManager credsVar) = do
+  UnliftIO.modifyMVar credsVar $ \mayCreds -> case mayCreds of
+    Just creds -> pure (mayCreds, creds)
+    Nothing -> do
+      creds <- CF.atomicallyModifyCredentialsFile \c -> pure (c, c)
+      pure (Just creds, creds)
 
-getCredentials :: (MonadIO m) => CredentialManager -> CodeserverId -> m (Either CredentialFailure CodeserverCredentials)
-getCredentials (CredentialManager credsVar) aud = runExceptT do
-  creds <- lift (UnliftIO.readMVar credsVar)
-  codeserverCreds <- except (getCodeserverCredentials aud creds)
+getCodeserverCredentials :: (MonadIO m) => CredentialManager -> CodeserverId -> m (Either CredentialFailure CodeserverCredentials)
+getCodeserverCredentials credMan aud = runExceptT do
+  creds <- liftIO $ readCredentials credMan
+  codeserverCreds <- except (Auth.getCodeserverCredentials aud creds)
   lift (isExpired codeserverCreds) >>= \case
     True -> throwE (ReauthRequired aud)
     False -> pure codeserverCreds
 
-newCredentialManager :: (MonadIO m) => m CredentialManager
-newCredentialManager = do
-  credentials <- atomicallyModifyCredentialsFile id
-  credentialsVar <- UnliftIO.newMVar credentials
-  pure (CredentialManager credentialsVar)
+newCredentialManager :: CredentialManager
+newCredentialManager = globalCredentialsManager
 
 -- | Checks whether CodeserverCredentials are expired.
 isExpired :: (MonadIO m) => CodeserverCredentials -> m Bool

unison-cli/src/Unison/Auth/Tokens.hs

@@ -37,7 +37,7 @@ newTokenProvider manager host = UnliftIO.try @_ @CredentialFailure $ do
       -- If the access token is provided via environment variable, we don't need to refresh it.
       pure accessToken
     Nothing -> do
-      creds@CodeserverCredentials {tokens, discoveryURI} <- throwEitherM $ getCredentials manager host
+      creds@CodeserverCredentials {tokens, discoveryURI} <- throwEitherM $ getCodeserverCredentials manager host
       let Tokens {accessToken = currentAccessToken} = tokens
       expired <- isExpired creds
       if expired

unison-cli/src/Unison/Codebase/Editor/HandleInput/AuthLogin.hs

@@ -23,7 +23,7 @@ import Network.Wai
 import Network.Wai qualified as Wai
 import Network.Wai.Handler.Warp qualified as Warp
 import U.Codebase.Sqlite.Queries qualified as Q
-import Unison.Auth.CredentialManager (getCredentials, saveCredentials)
+import Unison.Auth.CredentialManager (getCodeserverCredentials, saveCredentials)
 import Unison.Auth.Discovery (discoveryURIForCodeserver, fetchDiscoveryDoc)
 import Unison.Auth.Types
   ( Code,
@@ -55,7 +55,7 @@ ucmOAuthClientID = "ucm"
 ensureAuthenticatedWithCodeserver :: CodeserverURI -> Cli UserInfo
 ensureAuthenticatedWithCodeserver codeserverURI = do
   Cli.Env {credentialManager} <- ask
-  getCredentials credentialManager (codeserverIdFromCodeserverURI codeserverURI) >>= \case
+  getCodeserverCredentials credentialManager (codeserverIdFromCodeserverURI codeserverURI) >>= \case
     Right (CodeserverCredentials {userInfo}) -> pure userInfo
     Left _ -> authLogin codeserverURI