Cache Github status token

.github/workflows/build.yml

@@ -22,7 +22,7 @@ jobs:
         enable-stack: true
         stack-version: 'latest'
     - name: Cache
-      uses: actions/cache@v1
+      uses: actions/cache@v4
       env:
         cache-name: cache-stack
       with:

.github/workflows/release.yml

@@ -28,7 +28,7 @@ jobs:
         enable-stack: true
         stack-version: 'latest'
     - name: Cache
-      uses: actions/cache@v1
+      uses: actions/cache@v4
       env:
         cache-name: cache-stack
       with:

package.yaml

@@ -132,3 +132,6 @@ tests:
     - taskrunner
     - tasty
     - tasty-golden
+    - wai
+    - warp
+    - http-types

src/App.hs

@@ -1,5 +1,6 @@
 {-# LANGUAGE TemplateHaskell #-}
 {-# LANGUAGE DuplicateRecordFields #-}
+{-# LANGUAGE RecursiveDo #-}
 {-# OPTIONS_GHC -Wno-ambiguous-fields #-}
 
 module App where
@@ -39,6 +40,7 @@ import Utils
 import qualified RemoteCache
 import RemoteCache (getLatestBuildHash)
 import CommitStatus (updateCommitStatus, StatusRequest (..))
+import qualified CommitStatus
 import qualified System.Process as Process
 import Control.Monad.EarlyReturn (withEarlyReturn, earlyReturn)
 import Data.Time.Format.ISO8601 (iso8601Show)
@@ -124,27 +126,35 @@ main = do
     responsePipeReadFd <- handleToFd responsePipeRead
     hSetBuffering responsePipeWrite LineBuffering
 
-    parentEnv <- getEnvironment
-
-    cwd <- getCurrentDirectory
-
-    -- TODO: handle spawn error here
-    -- TODO: should we use withCreateProcess?
-    -- TODO: should we use delegate_ctlc or DIY? See https://hackage.haskell.org/package/process-1.6.20.0/docs/System-Process.html#g:4
-    -- -> We should DIY because we need to flush stream etc.
-    (Nothing, Just stdoutPipe, Just stderrPipe, processHandle) <- Process.createProcess
-      (proc args.cmd args.args) { std_in = UseHandle devnull, std_out = CreatePipe
-      , std_err = CreatePipe
-      , env=Just $ nubOrdOn fst $
-          [ ("BASH_FUNC_snapshot%%", "() {\n" <> $(embedStringFile "src/snapshot.sh") <> "\n}")
-          , ("_taskrunner_request_pipe", show requestPipeWriteFd)
-          , ("_taskrunner_response_pipe", show responsePipeReadFd)
-          ] <> parentEnv
-        }
+    -- Recursive: AppState is used before process is started (mostly for logging)
+    rec
+
+      appState <- AppState settings jobName buildId isToplevel <$> newIORef Nothing <*> newIORef Nothing <*> newIORef False <*> pure toplevelStderr <*> pure subprocessStderr <*> pure logFile
+        <*> newIORef Nothing
+
+      when (isToplevel && appState.settings.enableCommitStatus) do
+        -- Note: sets env for subprocesses, so has to be called before starting subprocess
+        void $ CommitStatus.getClient appState
 
-    (subprocessStderrRead, subprocessStderr) <- createPipe
+      parentEnv <- getEnvironment
+
+      cwd <- getCurrentDirectory
+
+      -- TODO: handle spawn error here
+      -- TODO: should we use withCreateProcess?
+      -- TODO: should we use delegate_ctlc or DIY? See https://hackage.haskell.org/package/process-1.6.20.0/docs/System-Process.html#g:4
+      -- -> We should DIY because we need to flush stream etc.
+      (Nothing, Just stdoutPipe, Just stderrPipe, processHandle) <- Process.createProcess
+        (proc args.cmd args.args) { std_in = UseHandle devnull, std_out = CreatePipe
+        , std_err = CreatePipe
+        , env=Just $ nubOrdOn fst $
+            [ ("BASH_FUNC_snapshot%%", "() {\n" <> $(embedStringFile "src/snapshot.sh") <> "\n}")
+            , ("_taskrunner_request_pipe", show requestPipeWriteFd)
+            , ("_taskrunner_response_pipe", show responsePipeReadFd)
+            ] <> parentEnv
+          }
 
-    appState <- AppState settings jobName buildId isToplevel <$> newIORef Nothing <*> newIORef Nothing <*> newIORef False <*> pure toplevelStderr <*> pure subprocessStderr <*> pure logFile
+      (subprocessStderrRead, subprocessStderr) <- createPipe
 
     logDebug appState $ "Running command: " <> show (args.cmd : args.args)
     logDebug appState $ "  buildId: " <> show buildId

src/CommitStatus.hs

@@ -11,12 +11,12 @@ import qualified Data.Text as T
 import qualified Data.Text.Encoding as TE
 import qualified Network.HTTP.Client as HTTP
 import Network.HTTP.Client.TLS (tlsManagerSettings)
-import System.Environment (getEnv)
+import System.Environment (getEnv, lookupEnv, setEnv)
 import Network.HTTP.Types.Status (Status(..))
 import Data.Aeson.Decoding (eitherDecode)
 import qualified Data.Text as Text
 import Utils (getCurrentCommit, logError, logDebug)
-import Types (AppState)
+import Types (AppState(..), GithubClient(..))
 
 -- Define the data types for the status update
 data StatusRequest = StatusRequest
@@ -33,70 +33,105 @@ newtype InstallationTokenResponse = InstallationTokenResponse
   } deriving (Show, Generic)
   deriving anyclass (FromJSON)
 
-updateCommitStatus :: MonadIO m => AppState -> StatusRequest -> m ()
-updateCommitStatus appState statusRequest = liftIO do
+getClient :: AppState -> IO GithubClient
+getClient appState = do
+  mClient <- readIORef appState.githubClient
+  case mClient of
+    Just client -> pure client
+    Nothing -> do
+      client <- initClient appState
+      writeIORef appState.githubClient $ Just client
+      pure client
+
+initClient :: AppState -> IO GithubClient
+initClient appState = do
   -- Load environment variables
+  apiUrl <- fromMaybe "https://api.github.com" <$> lookupEnv "GITHUB_API_URL"
   appId <- getEnv "GITHUB_APP_ID"
   installationId <- getEnv "GITHUB_INSTALLATION_ID"
   privateKeyStr <- getEnv "GITHUB_APP_PRIVATE_KEY"
   owner <- getEnv "GITHUB_REPOSITORY_OWNER"
   repo <- getEnv "GITHUB_REPOSITORY"
-
-  sha <- getCurrentCommit appState
-
-  let privateKeyBytes = encodeUtf8 $ Text.replace "|" "\n" $ toText privateKeyStr
-  let privateKey = fromMaybe (error "Invalid github key") $ readRsaSecret privateKeyBytes
-
-  -- Create the JWT token
-  now <- getPOSIXTime
-  let claims = mempty { iss = stringOrURI $ T.pack appId
-                   , iat = numericDate now
-                   , exp = numericDate (now + 5 * 60)
-                   }
-  let jwt = encodeSigned (EncodeRSAPrivateKey privateKey) (mempty { alg = Just RS256 }) claims
-
   -- Prepare the HTTP manager
   manager <- HTTP.newManager tlsManagerSettings
 
-  -- Get the installation access token
-  let installUrl = "https://api.github.com/app/installations/" ++ installationId ++ "/access_tokens"
-  initRequest <- HTTP.parseRequest installUrl
-  let request = initRequest
-                { HTTP.method = "POST"
-                , HTTP.requestHeaders =
-                    [ ("Authorization", "Bearer " <> TE.encodeUtf8 jwt)
-                    , ("Accept", "application/vnd.github.v3+json")
-                    , ("User-Agent", "restaumatic-bot")
-                    ]
-                }
-  response <- HTTP.httpLbs request manager
-  let mTokenResponse = eitherDecode @InstallationTokenResponse (HTTP.responseBody response)
-  case mTokenResponse of
-    Left err -> do
-      logError appState $ "CommitStatus: Failed to parse installation token response: " <> show err
-      logError appState $ "CommitStatus: Response: " <> decodeUtf8 response.responseBody
-      exitFailure
-    Right tokenResponse -> do
-      let accessToken = tokenResponse.token
+  let createToken = do
+        let privateKeyBytes = encodeUtf8 $ Text.replace "|" "\n" $ toText privateKeyStr
+        let privateKey = fromMaybe (error "Invalid github key") $ readRsaSecret privateKeyBytes
+
+        -- Create the JWT token
+        now <- getPOSIXTime
+        let claims = mempty { iss = stringOrURI $ T.pack appId
+                         , iat = numericDate now
+                         , exp = numericDate (now + 5 * 60)
+                         }
+        let jwt = encodeSigned (EncodeRSAPrivateKey privateKey) (mempty { alg = Just RS256 }) claims
 
-      -- Prepare the status update request
-      let statusUrl = "https://api.github.com/repos/" ++ owner ++ "/" ++ repo ++ "/statuses/" ++ toString sha
-      initStatusRequest <- HTTP.parseRequest statusUrl
-      let statusReq = initStatusRequest
+        -- Get the installation access token
+        let installUrl = apiUrl <> "/app/installations/" ++ installationId ++ "/access_tokens"
+        initRequest <- HTTP.parseRequest installUrl
+        let request = initRequest
                       { HTTP.method = "POST"
                       , HTTP.requestHeaders =
-                          [ ("Authorization", "Bearer " <> TE.encodeUtf8 accessToken)
+                          [ ("Authorization", "Bearer " <> TE.encodeUtf8 jwt)
                           , ("Accept", "application/vnd.github.v3+json")
-                          , ("Content-Type", "application/json")
                           , ("User-Agent", "restaumatic-bot")
                           ]
-                      , HTTP.requestBody = HTTP.RequestBodyLBS $ encode statusRequest
                       }
-      statusResponse <- HTTP.httpLbs statusReq manager
-      if statusResponse.responseStatus.statusCode == 201
-        then
-          logDebug appState "Commit status updated successfully"
-        else do
-          logError appState $ "CommitStatus: Failed to update commit status: " <> show statusResponse
-          logError appState $ "CommitStatus: Response: " <> decodeUtf8 response.responseBody
-          exitFailure
+        response <- HTTP.httpLbs request manager
+        let mTokenResponse = eitherDecode @InstallationTokenResponse (HTTP.responseBody response)
+        case mTokenResponse of
+          Left err -> do
+            logError appState $ "CommitStatus: Failed to parse installation token response: " <> show err
+            logError appState $ "CommitStatus: Response: " <> decodeUtf8 response.responseBody
+
+            -- FIXME: handle the error better
+            exitFailure
+          Right tokenResponse ->
+            pure tokenResponse.token
+
+  -- Try to read token from environment variable
+  -- Otherwise generate a new one, and set env for future uses (also in child processes)
+  accessToken <- lookupEnv "_taskrunner_github_access_token" >>= \case
+    Just token -> pure $ T.pack token
+    Nothing -> do
+      token <- createToken
+      setEnv "_taskrunner_github_access_token" $ T.unpack token
+      pure token
+
+  pure $ GithubClient { apiUrl = T.pack apiUrl
+                      , appId = T.pack appId
+                      , installationId = T.pack installationId
+                      , privateKey = T.pack privateKeyStr
+                      , owner = T.pack owner
+                      , repo = T.pack repo
+                      , manager = manager
+                      , accessToken = accessToken
+                      }
+
+updateCommitStatus :: MonadIO m => AppState -> StatusRequest -> m ()
+updateCommitStatus appState statusRequest = liftIO do
+  client <- getClient appState
+  sha <- getCurrentCommit appState
+
+  -- Prepare the status update request
+  let statusUrl = toString client.apiUrl <> "/repos/" ++ toString client.owner ++ "/" ++ toString client.repo ++ "/statuses/" ++ toString sha
+  initStatusRequest <- HTTP.parseRequest statusUrl
+  let statusReq = initStatusRequest
+                  { HTTP.method = "POST"
+                  , HTTP.requestHeaders =
+                      [ ("Authorization", "Bearer " <> TE.encodeUtf8 client.accessToken)
+                      , ("Accept", "application/vnd.github.v3+json")
+                      , ("Content-Type", "application/json")
+                      , ("User-Agent", "restaumatic-bot")
+                      ]
+                  , HTTP.requestBody = HTTP.RequestBodyLBS $ encode statusRequest
+                  }
+  statusResponse <- HTTP.httpLbs statusReq client.manager
+  if statusResponse.responseStatus.statusCode == 201
+    then
+      logDebug appState "Commit status updated successfully"
+    else do
+      logError appState $ "CommitStatus: Failed to update commit status: " <> show statusResponse
+      logError appState $ "CommitStatus: Response: " <> decodeUtf8 statusResponse.responseBody
+      exitFailure

src/Types.hs

@@ -3,6 +3,7 @@ module Types where
 import Universum
 import SnapshotCliArgs (SnapshotCliArgs)
 import Data.Aeson (FromJSON, ToJSON)
+import qualified Network.HTTP.Client as HTTP
 
 data Settings = Settings
   { stateDirectory :: FilePath
@@ -48,4 +49,19 @@ data AppState = AppState
   , toplevelStderr :: Handle
   , subprocessStderr :: Handle
   , logOutput :: Handle
+
+  -- | Lazily initialized Github client
+  , githubClient :: IORef (Maybe GithubClient)
+  }
+
+-- Unfortunately the type has to live there due to circular dependencies (AppState -> GithubClient -> AppState)
+data GithubClient = GithubClient
+  { apiUrl :: Text
+  , appId :: Text
+  , installationId :: Text
+  , privateKey :: Text
+  , owner :: Text
+  , repo :: Text
+  , manager :: HTTP.Manager
+  , accessToken :: Text
   }

taskrunner.cabal

@@ -208,6 +208,7 @@ test-suite taskrunner-test
   type: exitcode-stdio-1.0
   main-is: Spec.hs
   other-modules:
+      FakeGithubApi
       Paths_taskrunner
   autogen-modules:
       Paths_taskrunner
@@ -291,5 +292,7 @@ test-suite taskrunner-test
     , transformers
     , universum
     , unix
+    , wai
+    , warp
     , zstd
   default-language: Haskell2010

test/FakeGithubApi.hs

@@ -0,0 +1,72 @@
+{-# LANGUAGE OverloadedStrings #-}
+{-# LANGUAGE RecursiveDo #-}
+
+module FakeGithubApi (Server, start, stop, clearOutput, getOutput) where
+
+import Universum
+
+import Network.Wai
+import qualified Network.Wai.Handler.Warp as Warp
+import Network.HTTP.Types (status200, status201, status400, status404, methodPost)
+import Data.Aeson (encode, object, (.=))
+
+import Control.Concurrent (forkIO, ThreadId, killThread)
+
+-- Mock handler function
+app :: Server -> Application
+app server req respond = do
+    let path = pathInfo req
+    case path of
+        ["app", "installations", instId, "access_tokens"] ->
+            handleAccessTokenRequest server instId req respond
+        ["repos", owner, repo, "statuses", commitSha] ->
+            handleCommitStatusRequest server owner repo commitSha req respond
+        _ -> respond $ responseLBS status404 [] "Not Found"
+
+handleAccessTokenRequest :: Server -> Text -> Request -> (Response -> IO ResponseReceived) -> IO ResponseReceived
+handleAccessTokenRequest server instId req respond =
+    if requestMethod req == methodPost
+        then do
+          addOutput server $ "Requested access token for installation " <> instId
+          respond $ responseLBS status200 [("Content-Type", "application/json")]
+            (encode $ object ["token" .= ("mock-access-token" :: Text), "installation_id" .= instId])
+        else respond $ responseLBS status400 [] "Bad Request"
+
+handleCommitStatusRequest :: Server -> Text -> Text -> Text -> Request -> (Response -> IO ResponseReceived) -> IO ResponseReceived
+handleCommitStatusRequest server owner repo commitSha req respond =
+    if requestMethod req == methodPost
+        then do
+          body <- strictRequestBody req
+          -- Note: commit SHA omitted because it's nondeterministic
+          addOutput server $ "Updated commit status for " <> owner <> "/" <> repo <> " to " <> decodeUtf8 body
+          respond $ responseLBS status201 [("Content-Type", "application/json")]
+            (encode $ object ["state" .= ("success" :: Text), "sha" .= commitSha, "repository" .= repo, "owner" .= owner])
+        else respond $ responseLBS status400 [] "Bad Request"
+
+data Server = Server
+  { tid :: ThreadId
+  , output :: IORef [Text]
+  }
+
+start :: Int -> IO Server
+start port = do
+  started <- newEmptyMVar
+  output <- newIORef []
+  let settings = Warp.setPort port $ Warp.setBeforeMainLoop (putMVar started ()) Warp.defaultSettings
+  rec
+    let server = Server {tid, output}
+    tid <- forkIO $ Warp.runSettings settings $ app server
+  takeMVar started
+  pure server
+
+stop :: Server -> IO ()
+stop (Server {tid}) = killThread tid
+
+addOutput :: Server -> Text -> IO ()
+addOutput (Server {output}) msg = modifyIORef output (msg :)
+
+clearOutput :: Server -> IO ()
+clearOutput (Server {output}) = writeIORef output []
+
+getOutput :: Server -> IO [Text]
+getOutput (Server {output}) = reverse <$> readIORef output