Skip to content

Commit cd9ce0e

Browse files
ChrisPennerChrisPenner
committed
Pass Token Provider
1 parent ca4bb33 commit cd9ce0e

File tree

10 files changed

+73
-33
lines changed

10 files changed

+73
-33
lines changed

‎parser-typechecker/src/Unison/Codebase/SqliteCodebase/Migrations.hs‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -89,7 +89,8 @@ migrations regionVar getDeclType termBuffer declBuffer rootCodebasePath =
8989
sqlMigration 19 Q.addMergeBranchTables,
9090
sqlMigration 20 Q.addUpdateBranchTable,
9191
sqlMigration 21 Q.addDerivedDependentsByDependencyIndex,
92-
sqlMigration 22 Q.addUpgradeBranchTable
92+
sqlMigration 22 Q.addUpgradeBranchTable,
93+
sqlMigration 23 Q.addSyncV3TempTables
9394
]
9495
where
9596
runT :: Sqlite.Transaction () -> Sqlite.Connection -> IO ()

‎parser-typechecker/src/Unison/Codebase/SqliteCodebase/Operations.hs‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -85,6 +85,7 @@ createSchema = do
8585
Q.addUpdateBranchTable
8686
Q.addDerivedDependentsByDependencyIndex
8787
Q.addUpgradeBranchTable
88+
Q.addSyncV3TempTables
8889
(_, emptyCausalHashId) <- emptyCausalHash
8990
(_, ProjectBranchRow {projectId, branchId}) <-
9091
insertProjectAndBranch scratchProjectName scratchBranchName emptyCausalHashId

‎unison-cli/src/Unison/Cli/Monad.hs‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -71,6 +71,7 @@ import U.Codebase.Sqlite.DbId (ProjectBranchId, ProjectId)
7171
import U.Codebase.Sqlite.Queries qualified as Q
7272
import Unison.Auth.CredentialManager (CredentialManager)
7373
import Unison.Auth.HTTPClient (AuthenticatedHttpClient)
74+
import Unison.Auth.Tokens (TokenProvider)
7475
import Unison.Codebase (Codebase)
7576
import Unison.Codebase qualified as Codebase
7677
import Unison.Codebase.Editor.Input (Input)
@@ -158,6 +159,9 @@ type SourceName = Text
158159
-- Get the environment with 'ask'.
159160
data Env = Env
160161
{ authHTTPClient :: AuthenticatedHttpClient,
162+
-- | How to get auth tokens for a given codeserver.
163+
-- Using AuthenticatedHttpClient takes care of this, but websocket connection need to provide auth headers manually.
164+
tokenProvider :: TokenProvider,
161165
codebase :: Codebase IO Symbol Ann,
162166
credentialManager :: CredentialManager,
163167
-- | Generate a unique name.

‎unison-cli/src/Unison/Codebase/Transcript/Runner.hs‎

Lines changed: 7 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -96,7 +96,9 @@ withRunner ::
9696
m r
9797
withRunner isTest verbosity ucmVersion action = do
9898
credMan <- AuthN.newCredentialManager
99-
authenticatedHTTPClient <- initTranscriptAuthenticatedHTTPClient credMan
99+
let tokenProvider :: AuthN.TokenProvider
100+
tokenProvider = AuthN.newTokenProvider credMan
101+
authenticatedHTTPClient <- AuthN.newAuthenticatedHTTPClient tokenProvider ucmVersion
100102

101103
-- If we're in a transcript test, configure the environment to use a non-existent fzf binary
102104
-- so that errors are consistent.
@@ -130,6 +132,7 @@ withRunner isTest verbosity ucmVersion action = do
130132
ucmVersion
131133
baseUrlText
132134
authenticatedHTTPClient
135+
tokenProvider
133136
credMan
134137
stanzas
135138
where
@@ -138,11 +141,6 @@ withRunner isTest verbosity ucmVersion action = do
138141
RTI.withRuntime False RTI.Persistent ucmVersion \runtime ->
139142
RTI.withRuntime True RTI.Persistent ucmVersion \sbRuntime ->
140143
action runtime sbRuntime
141-
initTranscriptAuthenticatedHTTPClient :: AuthN.CredentialManager -> m AuthN.AuthenticatedHttpClient
142-
initTranscriptAuthenticatedHTTPClient credMan = liftIO $ do
143-
let tokenProvider :: AuthN.TokenProvider
144-
tokenProvider = AuthN.newTokenProvider credMan
145-
AuthN.newAuthenticatedHTTPClient tokenProvider ucmVersion
146144

147145
isGeneratedBlock :: ProcessedBlock -> Bool
148146
isGeneratedBlock = generated . getCommonInfoTags
@@ -157,10 +155,11 @@ run ::
157155
UCMVersion ->
158156
Text ->
159157
AuthN.AuthenticatedHttpClient ->
158+
AuthN.TokenProvider ->
160159
AuthN.CredentialManager ->
161160
Transcript ->
162161
IO (Either Error Transcript)
163-
run isTest verbosity codebase runtime sbRuntime ucmVersion baseURL authenticatedHTTPClient credMan transcript = UnliftIO.try do
162+
run isTest verbosity codebase runtime sbRuntime ucmVersion baseURL authenticatedHTTPClient tokenProvider credMan transcript = UnliftIO.try do
164163
let behaviors = extractBehaviors $ settings transcript
165164
let stanzas' = stanzas transcript
166165
httpManager <- HTTP.newManager HTTP.defaultManagerSettings
@@ -518,6 +517,7 @@ run isTest verbosity codebase runtime sbRuntime ucmVersion baseURL authenticated
518517
let env =
519518
Cli.Env
520519
{ authHTTPClient = authenticatedHTTPClient,
520+
tokenProvider,
521521
codebase,
522522
credentialManager = credMan,
523523
generateUniqueName = do

‎unison-cli/src/Unison/CommandLine/Main.hs‎

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -26,6 +26,7 @@ import U.Codebase.Sqlite.Queries qualified as Queries
2626
import Unison.Auth.CredentialManager qualified as AuthN
2727
import Unison.Auth.HTTPClient (AuthenticatedHttpClient)
2828
import Unison.Auth.HTTPClient qualified as AuthN
29+
import Unison.Auth.Tokens (TokenProvider)
2930
import Unison.Cli.Monad qualified as Cli
3031
import Unison.Cli.Pretty qualified as P
3132
import Unison.Cli.ProjectUtils qualified as ProjectUtils
@@ -146,11 +147,12 @@ main ::
146147
Maybe Server.BaseUrl ->
147148
UCMVersion ->
148149
AuthN.AuthenticatedHttpClient ->
150+
TokenProvider ->
149151
AuthN.CredentialManager ->
150152
(PP.ProjectPathIds -> IO ()) ->
151153
ShouldWatchFiles ->
152154
IO ()
153-
main dir welcome ppIds initialInputs runtime sbRuntime codebase serverBaseUrl ucmVersion authHTTPClient credentialManager lspCheckForChanges shouldWatchFiles = do
155+
main dir welcome ppIds initialInputs runtime sbRuntime codebase serverBaseUrl ucmVersion authHTTPClient tokenProvider credentialManager lspCheckForChanges shouldWatchFiles = do
154156
-- we don't like FSNotify's debouncing (it seems to drop later events)
155157
-- so we will be doing our own instead
156158
let config = FSNotify.defaultConfig
@@ -288,6 +290,7 @@ main dir welcome ppIds initialInputs runtime sbRuntime codebase serverBaseUrl uc
288290
{ authHTTPClient,
289291
codebase,
290292
credentialManager,
293+
tokenProvider,
291294
loadSource = loadSourceFile,
292295
lspCheckForChanges,
293296
writeSource,

‎unison-cli/src/Unison/MCP/Cli.hs‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -102,6 +102,7 @@ cliToMCP projCtx cli = do
102102
let cliEnv =
103103
Cli.Env
104104
{ authHTTPClient = authenticatedHTTPClient,
105+
tokenProvider,
105106
codebase,
106107
credentialManager = credMan,
107108
generateUniqueName = do

‎unison-cli/src/Unison/Main.hs‎

Lines changed: 13 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -54,6 +54,7 @@ import Text.Megaparsec qualified as MP
5454
import U.Codebase.Sqlite.Queries qualified as Queries
5555
import Unison.Auth.CredentialManager qualified as AuthN
5656
import Unison.Auth.HTTPClient qualified as AuthN
57+
import Unison.Auth.Tokens (TokenProvider)
5758
import Unison.Auth.Tokens qualified as AuthN
5859
import Unison.Cli.ProjectUtils qualified as ProjectUtils
5960
import Unison.Codebase (Codebase, CodebasePath)
@@ -185,7 +186,8 @@ main version = do
185186
let serverUrl = Nothing
186187
let ucmVersion = Version.gitDescribeWithDate version
187188
credMan <- liftIO $ AuthN.newCredentialManager
188-
authenticatedHTTPClient <- initTranscriptAuthenticatedHTTPClient ucmVersion credMan
189+
let tokenProvider = AuthN.newTokenProvider credMan
190+
authenticatedHTTPClient <- AuthN.newAuthenticatedHTTPClient tokenProvider ucmVersion
189191
startProjectPath <- Codebase.runTransaction theCodebase Codebase.expectCurrentProjectPath
190192
launch
191193
version
@@ -195,6 +197,7 @@ main version = do
195197
theCodebase
196198
[Left fileEvent, Right $ Input.ExecuteI NoProf mainName args, Right Input.QuitI]
197199
authenticatedHTTPClient
200+
tokenProvider
198201
credMan
199202
serverUrl
200203
(PP.toIds startProjectPath)
@@ -213,7 +216,8 @@ main version = do
213216
let serverUrl = Nothing
214217
let ucmVersion = Version.gitDescribeWithDate version
215218
credMan <- liftIO $ AuthN.newCredentialManager
216-
authenticatedHTTPClient <- initTranscriptAuthenticatedHTTPClient ucmVersion credMan
219+
let tokenProvider = AuthN.newTokenProvider credMan
220+
authenticatedHTTPClient <- AuthN.newAuthenticatedHTTPClient tokenProvider ucmVersion
217221
startProjectPath <- Codebase.runTransaction theCodebase Codebase.expectCurrentProjectPath
218222
launch
219223
version
@@ -223,6 +227,7 @@ main version = do
223227
theCodebase
224228
[Left fileEvent, Right $ Input.ExecuteI NoProf mainName args, Right Input.QuitI]
225229
authenticatedHTTPClient
230+
tokenProvider
226231
credMan
227232
serverUrl
228233
(PP.toIds startProjectPath)
@@ -330,7 +335,8 @@ main version = do
330335
let isTest = False
331336
let ucmVersion = Version.gitDescribeWithDate version
332337
credMan <- liftIO $ AuthN.newCredentialManager
333-
authenticatedHTTPClient <- initTranscriptAuthenticatedHTTPClient ucmVersion credMan
338+
let tokenProvider = AuthN.newTokenProvider credMan
339+
authenticatedHTTPClient <- AuthN.newAuthenticatedHTTPClient tokenProvider ucmVersion
334340
mcpServerConfig <- MCP.initServer theCodebase runtime sbRuntime (Just currentDir) ucmVersion authenticatedHTTPClient
335341
Server.startServer
336342
isTest
@@ -374,6 +380,7 @@ main version = do
374380
theCodebase
375381
[]
376382
authenticatedHTTPClient
383+
tokenProvider
377384
credMan
378385
mayBaseUrl
379386
(PP.toIds startingProjectPath)
@@ -596,14 +603,15 @@ launch ::
596603
Codebase.Codebase IO Symbol Ann ->
597604
[Either Input.Event Input.Input] ->
598605
AuthN.AuthenticatedHttpClient ->
606+
TokenProvider ->
599607
AuthN.CredentialManager ->
600608
Maybe Server.BaseUrl ->
601609
PP.ProjectPathIds ->
602610
InitResult ->
603611
(PP.ProjectPathIds -> IO ()) ->
604612
CommandLine.ShouldWatchFiles ->
605613
IO ()
606-
launch version dir runtime sbRuntime codebase inputs authenticatedHTTPClient credMan serverBaseUrl startingPath initResult lspCheckForChanges shouldWatchFiles = do
614+
launch version dir runtime sbRuntime codebase inputs authenticatedHTTPClient tokenProvider credMan serverBaseUrl startingPath initResult lspCheckForChanges shouldWatchFiles = do
607615
showWelcomeHint <- Codebase.runTransaction codebase Queries.doProjectsExist
608616
let isNewCodebase = case initResult of
609617
CreatedCodebase -> NewlyCreatedCodebase
@@ -621,6 +629,7 @@ launch version dir runtime sbRuntime codebase inputs authenticatedHTTPClient cre
621629
serverBaseUrl
622630
ucmVersion
623631
authenticatedHTTPClient
632+
tokenProvider
624633
credMan
625634
lspCheckForChanges
626635
shouldWatchFiles

‎unison-cli/src/Unison/Share/SyncV3.hs‎

Lines changed: 10 additions & 20 deletions
Original file line numberDiff line numberDiff line change
@@ -3,18 +3,17 @@ module Unison.Share.SyncV3
33
)
44
where
55

6+
import Network.Socket (withSocketsDo)
67
import Control.Arrow ((&&&))
78
import Control.Monad.Reader
89
import Data.Set qualified as Set
9-
import Data.Set.Lens qualified as Lens
10+
import Data.Text.Encoding as Text
1011
import GHC.Natural
1112
import Ki qualified
1213
import Network.WebSockets qualified as WS
1314
import U.Codebase.HashTags
1415
import U.Codebase.Sqlite.DbId
15-
import U.Codebase.Sqlite.Entity qualified as Entity
1616
import U.Codebase.Sqlite.Queries qualified as Q
17-
import U.Codebase.Sqlite.TempEntity (TempEntity)
1817
import U.Codebase.Sqlite.V2.HashHandle (v2HashHandle)
1918
import Unison.Cli.Monad
2019
import Unison.Cli.Monad qualified as Cli
@@ -27,9 +26,11 @@ import Unison.Server.Orphans ()
2726
import Unison.Share.API.Hash qualified as Share
2827
import Unison.Share.Codeserver qualified as Codeserver
2928
import Unison.Share.Sync.Types qualified as Sync
29+
import Unison.Share.Types
3030
import Unison.Sync.Common qualified as Sync
3131
import Unison.SyncV3.Types
3232
import Unison.SyncV3.Types as SyncV3
33+
import Unison.SyncV3.Utils (tempEntityDependencies)
3334
import Unison.Util.Servant.CBOR qualified as CBOR
3435
import Unison.Util.Websockets (Queues (..), withQueues)
3536
import UnliftIO.STM
@@ -58,14 +59,16 @@ syncFromCodeserver ::
5859
Share.HashJWT ->
5960
Cli (Either (Sync.SyncError SyncV3.SyncError) (CausalHash, CausalHashId))
6061
syncFromCodeserver _shouldValidate codeserver branchRef hashJwt = do
61-
Cli.Env {codebase} <- ask
62+
Cli.Env {codebase, tokenProvider} <- ask
6263
let host = Codeserver.codeserverRegName codeserver
6364
let syncV3Path = "/ucm/v3/sync/download"
6465
let rootCausalHash = Share.hashJWTHash hashJwt
6566
-- Enable compression
6667
let connectionOptions = WS.defaultConnectionOptions {WS.connectionCompressionOptions = WS.PermessageDeflateCompression WS.defaultPermessageDeflate}
67-
-- TODO: Add authentication headers manually.
68-
let headers = []
68+
headers <-
69+
(liftIO (tokenProvider (codeserverIdFromCodeserverURI codeserver))) <&> \case
70+
Left {} -> []
71+
Right token -> [("Authorization", "Bearer " <> Text.encodeUtf8 token)]
6972
let runner = case Codeserver.codeserverScheme codeserver of
7073
Codeserver.Https ->
7174
let tlsPort = 443
@@ -76,7 +79,7 @@ syncFromCodeserver _shouldValidate codeserver branchRef hashJwt = do
7679
port = maybe tlsPort id $ (Codeserver.codeserverPort) codeserver
7780
in WS.runClientWith host port
7881
Debug.debugLogM Debug.Temp "Obtaining Connection"
79-
liftIO $ (runner syncV3Path connectionOptions headers) \conn -> do
82+
liftIO $ withSocketsDo $ (runner syncV3Path connectionOptions headers) \conn -> do
8083
Debug.debugLogM Debug.Temp "Obtained Connection"
8184
withQueues inputBuffer outputBuffer conn $ \queues@Queues {send} -> do
8285
Debug.debugLogM Debug.Temp "Obtained Queues"
@@ -216,16 +219,3 @@ flushTemp codebase rootCausalHash = do
216219
loop
217220
loop
218221
Q.expectCausalHashIdByCausalHash (Sync.hash32ToCausalHash rootCausalHash)
219-
220-
tempEntityDependencies :: TempEntity -> Set (EntityKind, Hash32)
221-
tempEntityDependencies entity = do
222-
let componentDeps = Lens.setOf Entity.defns_ entity
223-
patchDeps = Lens.setOf Entity.patches_ entity
224-
branchHashes = Lens.setOf Entity.branchHashes_ entity <> Lens.setOf Entity.branches_ entity
225-
causalHashes = Lens.setOf Entity.causalHashes_ entity
226-
in Set.unions
227-
[ Set.map (DefnComponentEntity,) componentDeps,
228-
Set.map (PatchEntity,) patchDeps,
229-
Set.map (NamespaceEntity,) branchHashes,
230-
Set.map (CausalEntity,) causalHashes
231-
]

‎unison-share-api/src/Unison/SyncV3/Utils.hs‎

Lines changed: 30 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,30 @@
1+
module Unison.SyncV3.Utils (tempEntityDependencies, entityDependencies) where
2+
3+
import Data.Set (Set)
4+
import Data.Set qualified as Set
5+
import Data.Set.Lens qualified as Lens
6+
import U.Codebase.Sqlite.Entity qualified as Entity
7+
import U.Codebase.Sqlite.TempEntity
8+
import Unison.Hash32 (Hash32)
9+
import Unison.SyncV3.Types
10+
import Unison.Util.Servant.CBOR qualified as CBOR
11+
12+
tempEntityDependencies :: TempEntity -> Set (EntityKind, Hash32)
13+
tempEntityDependencies entity = do
14+
let componentDeps = Lens.setOf Entity.defns_ entity
15+
patchDeps = Lens.setOf Entity.patches_ entity
16+
branchHashes = Lens.setOf Entity.branchHashes_ entity <> Lens.setOf Entity.branches_ entity
17+
causalHashes = Lens.setOf Entity.causalHashes_ entity
18+
in Set.unions
19+
[ Set.map (DefnComponentEntity,) componentDeps,
20+
Set.map (PatchEntity,) patchDeps,
21+
Set.map (NamespaceEntity,) branchHashes,
22+
Set.map (CausalEntity,) causalHashes
23+
]
24+
25+
entityDependencies :: Entity hash text -> Set (EntityKind, Hash32)
26+
entityDependencies Entity {entityData} = do
27+
case (CBOR.deserialiseOrFailCBORBytes $ entityData) of
28+
-- TODO: proper error handling
29+
Left err -> error $ show err
30+
Right tempEntity -> tempEntityDependencies tempEntity

‎unison-share-api/unison-share-api.cabal‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -53,6 +53,7 @@ library
5353
Unison.SyncV2.API
5454
Unison.SyncV2.Types
5555
Unison.SyncV3.Types
56+
Unison.SyncV3.Utils
5657
Unison.Util.Find
5758
Unison.Util.Servant.CBOR
5859
Unison.Util.Websockets

0 commit comments

Comments
 (0)