Skip to content

Commit 6bde1fe

Browse files
zylazyla
committed
Reuse cached token in subprocesses
1 parent 427a9ab commit 6bde1fe

5 files changed

+73
-56
lines changed

‎src/App.hs

Lines changed: 29 additions & 20 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,6 @@
11
{-# LANGUAGE TemplateHaskell #-}
22
{-# LANGUAGE DuplicateRecordFields #-}
3+
{-# LANGUAGE RecursiveDo #-}
34
{-# OPTIONS_GHC -Wno-ambiguous-fields #-}
45

56
module App where
@@ -39,6 +40,7 @@ import Utils
3940
import qualified RemoteCache
4041
import RemoteCache (getLatestBuildHash)
4142
import CommitStatus (updateCommitStatus, StatusRequest (..))
43+
import qualified CommitStatus
4244
import qualified System.Process as Process
4345
import Control.Monad.EarlyReturn (withEarlyReturn, earlyReturn)
4446
import Data.Time.Format.ISO8601 (iso8601Show)
@@ -124,28 +126,35 @@ main = do
124126
responsePipeReadFd <- handleToFd responsePipeRead
125127
hSetBuffering responsePipeWrite LineBuffering
126128

127-
parentEnv <- getEnvironment
128-
129-
cwd <- getCurrentDirectory
130-
131-
-- TODO: handle spawn error here
132-
-- TODO: should we use withCreateProcess?
133-
-- 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
134-
-- -> We should DIY because we need to flush stream etc.
135-
(Nothing, Just stdoutPipe, Just stderrPipe, processHandle) <- Process.createProcess
136-
(proc args.cmd args.args) { std_in = UseHandle devnull, std_out = CreatePipe
137-
, std_err = CreatePipe
138-
, env=Just $ nubOrdOn fst $
139-
[ ("BASH_FUNC_snapshot%%", "() {\n" <> $(embedStringFile "src/snapshot.sh") <> "\n}")
140-
, ("_taskrunner_request_pipe", show requestPipeWriteFd)
141-
, ("_taskrunner_response_pipe", show responsePipeReadFd)
142-
] <> parentEnv
143-
}
129+
-- Recursive: AppState is used before process is started (mostly for logging)
130+
rec
131+
132+
appState <- AppState settings jobName buildId isToplevel <$> newIORef Nothing <*> newIORef Nothing <*> newIORef False <*> pure toplevelStderr <*> pure subprocessStderr <*> pure logFile
133+
<*> newIORef Nothing
134+
135+
when isToplevel do
136+
-- Note: potentially sets env for subprocesses
137+
void $ CommitStatus.getClient appState
144138

145-
(subprocessStderrRead, subprocessStderr) <- createPipe
139+
parentEnv <- getEnvironment
140+
141+
cwd <- getCurrentDirectory
142+
143+
-- TODO: handle spawn error here
144+
-- TODO: should we use withCreateProcess?
145+
-- 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
146+
-- -> We should DIY because we need to flush stream etc.
147+
(Nothing, Just stdoutPipe, Just stderrPipe, processHandle) <- Process.createProcess
148+
(proc args.cmd args.args) { std_in = UseHandle devnull, std_out = CreatePipe
149+
, std_err = CreatePipe
150+
, env=Just $ nubOrdOn fst $
151+
[ ("BASH_FUNC_snapshot%%", "() {\n" <> $(embedStringFile "src/snapshot.sh") <> "\n}")
152+
, ("_taskrunner_request_pipe", show requestPipeWriteFd)
153+
, ("_taskrunner_response_pipe", show responsePipeReadFd)
154+
] <> parentEnv
155+
}
146156

147-
appState <- AppState settings jobName buildId isToplevel <$> newIORef Nothing <*> newIORef Nothing <*> newIORef False <*> pure toplevelStderr <*> pure subprocessStderr <*> pure logFile
148-
<*> newIORef Nothing
157+
(subprocessStderrRead, subprocessStderr) <- createPipe
149158

150159
logDebug appState $ "Running command: " <> show (args.cmd : args.args)
151160
logDebug appState $ " buildId: " <> show buildId

‎src/CommitStatus.hs

Lines changed: 40 additions & 31 deletions
Original file line numberDiff line numberDiff line change
@@ -11,7 +11,7 @@ import qualified Data.Text as T
1111
import qualified Data.Text.Encoding as TE
1212
import qualified Network.HTTP.Client as HTTP
1313
import Network.HTTP.Client.TLS (tlsManagerSettings)
14-
import System.Environment (getEnv, lookupEnv)
14+
import System.Environment (getEnv, lookupEnv, setEnv)
1515
import Network.HTTP.Types.Status (Status(..))
1616
import Data.Aeson.Decoding (eitherDecode)
1717
import qualified Data.Text as Text
@@ -55,40 +55,49 @@ initClient appState = do
5555
-- Prepare the HTTP manager
5656
manager <- HTTP.newManager tlsManagerSettings
5757

58-
let privateKeyBytes = encodeUtf8 $ Text.replace "|" "\n" $ toText privateKeyStr
59-
let privateKey = fromMaybe (error "Invalid github key") $ readRsaSecret privateKeyBytes
58+
let createToken = do
59+
let privateKeyBytes = encodeUtf8 $ Text.replace "|" "\n" $ toText privateKeyStr
60+
let privateKey = fromMaybe (error "Invalid github key") $ readRsaSecret privateKeyBytes
6061

61-
-- Create the JWT token
62-
now <- getPOSIXTime
63-
let claims = mempty { iss = stringOrURI $ T.pack appId
64-
, iat = numericDate now
65-
, exp = numericDate (now + 5 * 60)
66-
}
67-
let jwt = encodeSigned (EncodeRSAPrivateKey privateKey) (mempty { alg = Just RS256 }) claims
62+
-- Create the JWT token
63+
now <- getPOSIXTime
64+
let claims = mempty { iss = stringOrURI $ T.pack appId
65+
, iat = numericDate now
66+
, exp = numericDate (now + 5 * 60)
67+
}
68+
let jwt = encodeSigned (EncodeRSAPrivateKey privateKey) (mempty { alg = Just RS256 }) claims
6869

69-
-- Get the installation access token
70-
let installUrl = apiUrl <> "/app/installations/" ++ installationId ++ "/access_tokens"
71-
initRequest <- HTTP.parseRequest installUrl
72-
let request = initRequest
73-
{ HTTP.method = "POST"
74-
, HTTP.requestHeaders =
75-
[ ("Authorization", "Bearer " <> TE.encodeUtf8 jwt)
76-
, ("Accept", "application/vnd.github.v3+json")
77-
, ("User-Agent", "restaumatic-bot")
78-
]
79-
}
80-
response <- HTTP.httpLbs request manager
81-
let mTokenResponse = eitherDecode @InstallationTokenResponse (HTTP.responseBody response)
82-
accessToken <- case mTokenResponse of
83-
Left err -> do
84-
logError appState $ "CommitStatus: Failed to parse installation token response: " <> show err
85-
logError appState $ "CommitStatus: Response: " <> decodeUtf8 response.responseBody
70+
-- Get the installation access token
71+
let installUrl = apiUrl <> "/app/installations/" ++ installationId ++ "/access_tokens"
72+
initRequest <- HTTP.parseRequest installUrl
73+
let request = initRequest
74+
{ HTTP.method = "POST"
75+
, HTTP.requestHeaders =
76+
[ ("Authorization", "Bearer " <> TE.encodeUtf8 jwt)
77+
, ("Accept", "application/vnd.github.v3+json")
78+
, ("User-Agent", "restaumatic-bot")
79+
]
80+
}
81+
response <- HTTP.httpLbs request manager
82+
let mTokenResponse = eitherDecode @InstallationTokenResponse (HTTP.responseBody response)
83+
case mTokenResponse of
84+
Left err -> do
85+
logError appState $ "CommitStatus: Failed to parse installation token response: " <> show err
86+
logError appState $ "CommitStatus: Response: " <> decodeUtf8 response.responseBody
8687

87-
-- FIXME: handle the error better
88-
exitFailure
89-
Right tokenResponse ->
90-
pure tokenResponse.token
88+
-- FIXME: handle the error better
89+
exitFailure
90+
Right tokenResponse ->
91+
pure tokenResponse.token
9192

93+
-- Try to read token from environment variable
94+
-- Otherwise generate a new one, and set env for future uses (also in child processes)
95+
accessToken <- lookupEnv "_taskrunner_github_access_token" >>= \case
96+
Just token -> pure $ T.pack token
97+
Nothing -> do
98+
token <- createToken
99+
setEnv "_taskrunner_github_access_token" $ T.unpack token
100+
pure token
92101

93102
pure $ GithubClient { apiUrl = T.pack apiUrl
94103
, appId = T.pack appId

‎test/FakeGithubApi.hs

Lines changed: 4 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -49,12 +49,13 @@ data Server = Server
4949
}
5050

5151
start :: Int -> IO Server
52-
start port = mdo
52+
start port = do
5353
started <- newEmptyMVar
5454
output <- newIORef []
5555
let settings = Warp.setPort port $ Warp.setBeforeMainLoop (putMVar started ()) Warp.defaultSettings
56-
tid <- forkIO $ Warp.runSettings settings $ app server
57-
let server = Server {tid, output}
56+
rec
57+
let server = Server {tid, output}
58+
tid <- forkIO $ Warp.runSettings settings $ app server
5859
takeMVar started
5960
pure server
6061

‎test/t/github-commit-status-failure-nested.out

Lines changed: 0 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -4,6 +4,5 @@
44
Requested access token for installation 123
55
Updated commit status for fakeowner/fakerepo to {"context":"othertask","description":"not cached","state":"pending","target_url":null}
66
Updated commit status for fakeowner/fakerepo to {"context":"othertask","description":null,"state":"failure","target_url":null}
7-
Requested access token for installation 123
87
Updated commit status for fakeowner/fakerepo to {"context":"mytask","description":null,"state":"failure","target_url":null}
98
-- exit code: 1

‎test/t/github-commit-status-nested.out

Lines changed: 0 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -3,7 +3,6 @@
33
-- github:
44
Requested access token for installation 123
55
Updated commit status for fakeowner/fakerepo to {"context":"mytask","description":"not cached","state":"pending","target_url":null}
6-
Requested access token for installation 123
76
Updated commit status for fakeowner/fakerepo to {"context":"othertask","description":"not cached","state":"pending","target_url":null}
87
Updated commit status for fakeowner/fakerepo to {"context":"othertask","description":null,"state":"success","target_url":null}
98
Updated commit status for fakeowner/fakerepo to {"context":"mytask","description":null,"state":"success","target_url":null}

0 commit comments

Comments
 (0)