Skip to content

Commit dbcbe48

Browse files
authored
Merge pull request #641 from router-for-me/url-OAuth-add-ter
OAuth and management
2 parents 6390886 + 781bc15 commit dbcbe48

File tree

15 files changed

+738
-221
lines changed

15 files changed

+738
-221
lines changed

internal/api/handlers/management/auth_files.go

Lines changed: 59 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -197,6 +197,19 @@ func stopCallbackForwarder(port int) {
197197
stopForwarderInstance(port, forwarder)
198198
}
199199

200+
func stopCallbackForwarderInstance(port int, forwarder *callbackForwarder) {
201+
if forwarder == nil {
202+
return
203+
}
204+
callbackForwardersMu.Lock()
205+
if current := callbackForwarders[port]; current == forwarder {
206+
delete(callbackForwarders, port)
207+
}
208+
callbackForwardersMu.Unlock()
209+
210+
stopForwarderInstance(port, forwarder)
211+
}
212+
200213
func stopForwarderInstance(port int, forwarder *callbackForwarder) {
201214
if forwarder == nil || forwarder.server == nil {
202215
return
@@ -785,14 +798,16 @@ func (h *Handler) RequestAnthropicToken(c *gin.Context) {
785798
RegisterOAuthSession(state, "anthropic")
786799

787800
isWebUI := isWebUIRequest(c)
801+
var forwarder *callbackForwarder
788802
if isWebUI {
789803
targetURL, errTarget := h.managementCallbackURL("/anthropic/callback")
790804
if errTarget != nil {
791805
log.WithError(errTarget).Error("failed to compute anthropic callback target")
792806
c.JSON(http.StatusInternalServerError, gin.H{"error": "callback server unavailable"})
793807
return
794808
}
795-
if _, errStart := startCallbackForwarder(anthropicCallbackPort, "anthropic", targetURL); errStart != nil {
809+
var errStart error
810+
if forwarder, errStart = startCallbackForwarder(anthropicCallbackPort, "anthropic", targetURL); errStart != nil {
796811
log.WithError(errStart).Error("failed to start anthropic callback forwarder")
797812
c.JSON(http.StatusInternalServerError, gin.H{"error": "failed to start callback server"})
798813
return
@@ -801,14 +816,17 @@ func (h *Handler) RequestAnthropicToken(c *gin.Context) {
801816

802817
go func() {
803818
if isWebUI {
804-
defer stopCallbackForwarder(anthropicCallbackPort)
819+
defer stopCallbackForwarderInstance(anthropicCallbackPort, forwarder)
805820
}
806821

807822
// Helper: wait for callback file
808823
waitFile := filepath.Join(h.cfg.AuthDir, fmt.Sprintf(".oauth-anthropic-%s.oauth", state))
809824
waitForFile := func(path string, timeout time.Duration) (map[string]string, error) {
810825
deadline := time.Now().Add(timeout)
811826
for {
827+
if !IsOAuthSessionPending(state, "anthropic") {
828+
return nil, errOAuthSessionNotPending
829+
}
812830
if time.Now().After(deadline) {
813831
SetOAuthSessionError(state, "Timeout waiting for OAuth callback")
814832
return nil, fmt.Errorf("timeout waiting for OAuth callback")
@@ -828,6 +846,9 @@ func (h *Handler) RequestAnthropicToken(c *gin.Context) {
828846
// Wait up to 5 minutes
829847
resultMap, errWait := waitForFile(waitFile, 5*time.Minute)
830848
if errWait != nil {
849+
if errors.Is(errWait, errOAuthSessionNotPending) {
850+
return
851+
}
831852
authErr := claude.NewAuthenticationError(claude.ErrCallbackTimeout, errWait)
832853
log.Error(claude.GetUserFriendlyMessage(authErr))
833854
return
@@ -933,6 +954,7 @@ func (h *Handler) RequestAnthropicToken(c *gin.Context) {
933954
}
934955
fmt.Println("You can now use Claude services through this CLI")
935956
CompleteOAuthSession(state)
957+
CompleteOAuthSessionsByProvider("anthropic")
936958
}()
937959

938960
c.JSON(200, gin.H{"status": "ok", "url": authURL, "state": state})
@@ -968,14 +990,16 @@ func (h *Handler) RequestGeminiCLIToken(c *gin.Context) {
968990
RegisterOAuthSession(state, "gemini")
969991

970992
isWebUI := isWebUIRequest(c)
993+
var forwarder *callbackForwarder
971994
if isWebUI {
972995
targetURL, errTarget := h.managementCallbackURL("/google/callback")
973996
if errTarget != nil {
974997
log.WithError(errTarget).Error("failed to compute gemini callback target")
975998
c.JSON(http.StatusInternalServerError, gin.H{"error": "callback server unavailable"})
976999
return
9771000
}
978-
if _, errStart := startCallbackForwarder(geminiCallbackPort, "gemini", targetURL); errStart != nil {
1001+
var errStart error
1002+
if forwarder, errStart = startCallbackForwarder(geminiCallbackPort, "gemini", targetURL); errStart != nil {
9791003
log.WithError(errStart).Error("failed to start gemini callback forwarder")
9801004
c.JSON(http.StatusInternalServerError, gin.H{"error": "failed to start callback server"})
9811005
return
@@ -984,7 +1008,7 @@ func (h *Handler) RequestGeminiCLIToken(c *gin.Context) {
9841008

9851009
go func() {
9861010
if isWebUI {
987-
defer stopCallbackForwarder(geminiCallbackPort)
1011+
defer stopCallbackForwarderInstance(geminiCallbackPort, forwarder)
9881012
}
9891013

9901014
// Wait for callback file written by server route
@@ -993,6 +1017,9 @@ func (h *Handler) RequestGeminiCLIToken(c *gin.Context) {
9931017
deadline := time.Now().Add(5 * time.Minute)
9941018
var authCode string
9951019
for {
1020+
if !IsOAuthSessionPending(state, "gemini") {
1021+
return
1022+
}
9961023
if time.Now().After(deadline) {
9971024
log.Error("oauth flow timed out")
9981025
SetOAuthSessionError(state, "OAuth flow timed out")
@@ -1093,7 +1120,9 @@ func (h *Handler) RequestGeminiCLIToken(c *gin.Context) {
10931120

10941121
// Initialize authenticated HTTP client via GeminiAuth to honor proxy settings
10951122
gemAuth := geminiAuth.NewGeminiAuth()
1096-
gemClient, errGetClient := gemAuth.GetAuthenticatedClient(ctx, &ts, h.cfg, true)
1123+
gemClient, errGetClient := gemAuth.GetAuthenticatedClient(ctx, &ts, h.cfg, &geminiAuth.WebLoginOptions{
1124+
NoBrowser: true,
1125+
})
10971126
if errGetClient != nil {
10981127
log.Errorf("failed to get authenticated client: %v", errGetClient)
10991128
SetOAuthSessionError(state, "Failed to get authenticated client")
@@ -1166,6 +1195,7 @@ func (h *Handler) RequestGeminiCLIToken(c *gin.Context) {
11661195
}
11671196

11681197
CompleteOAuthSession(state)
1198+
CompleteOAuthSessionsByProvider("gemini")
11691199
fmt.Printf("You can now use Gemini CLI services through this CLI; token saved to %s\n", savedPath)
11701200
}()
11711201

@@ -1207,14 +1237,16 @@ func (h *Handler) RequestCodexToken(c *gin.Context) {
12071237
RegisterOAuthSession(state, "codex")
12081238

12091239
isWebUI := isWebUIRequest(c)
1240+
var forwarder *callbackForwarder
12101241
if isWebUI {
12111242
targetURL, errTarget := h.managementCallbackURL("/codex/callback")
12121243
if errTarget != nil {
12131244
log.WithError(errTarget).Error("failed to compute codex callback target")
12141245
c.JSON(http.StatusInternalServerError, gin.H{"error": "callback server unavailable"})
12151246
return
12161247
}
1217-
if _, errStart := startCallbackForwarder(codexCallbackPort, "codex", targetURL); errStart != nil {
1248+
var errStart error
1249+
if forwarder, errStart = startCallbackForwarder(codexCallbackPort, "codex", targetURL); errStart != nil {
12181250
log.WithError(errStart).Error("failed to start codex callback forwarder")
12191251
c.JSON(http.StatusInternalServerError, gin.H{"error": "failed to start callback server"})
12201252
return
@@ -1223,14 +1255,17 @@ func (h *Handler) RequestCodexToken(c *gin.Context) {
12231255

12241256
go func() {
12251257
if isWebUI {
1226-
defer stopCallbackForwarder(codexCallbackPort)
1258+
defer stopCallbackForwarderInstance(codexCallbackPort, forwarder)
12271259
}
12281260

12291261
// Wait for callback file
12301262
waitFile := filepath.Join(h.cfg.AuthDir, fmt.Sprintf(".oauth-codex-%s.oauth", state))
12311263
deadline := time.Now().Add(5 * time.Minute)
12321264
var code string
12331265
for {
1266+
if !IsOAuthSessionPending(state, "codex") {
1267+
return
1268+
}
12341269
if time.Now().After(deadline) {
12351270
authErr := codex.NewAuthenticationError(codex.ErrCallbackTimeout, fmt.Errorf("timeout waiting for OAuth callback"))
12361271
log.Error(codex.GetUserFriendlyMessage(authErr))
@@ -1346,6 +1381,7 @@ func (h *Handler) RequestCodexToken(c *gin.Context) {
13461381
}
13471382
fmt.Println("You can now use Codex services through this CLI")
13481383
CompleteOAuthSession(state)
1384+
CompleteOAuthSessionsByProvider("codex")
13491385
}()
13501386

13511387
c.JSON(200, gin.H{"status": "ok", "url": authURL, "state": state})
@@ -1391,14 +1427,16 @@ func (h *Handler) RequestAntigravityToken(c *gin.Context) {
13911427
RegisterOAuthSession(state, "antigravity")
13921428

13931429
isWebUI := isWebUIRequest(c)
1430+
var forwarder *callbackForwarder
13941431
if isWebUI {
13951432
targetURL, errTarget := h.managementCallbackURL("/antigravity/callback")
13961433
if errTarget != nil {
13971434
log.WithError(errTarget).Error("failed to compute antigravity callback target")
13981435
c.JSON(http.StatusInternalServerError, gin.H{"error": "callback server unavailable"})
13991436
return
14001437
}
1401-
if _, errStart := startCallbackForwarder(antigravityCallbackPort, "antigravity", targetURL); errStart != nil {
1438+
var errStart error
1439+
if forwarder, errStart = startCallbackForwarder(antigravityCallbackPort, "antigravity", targetURL); errStart != nil {
14021440
log.WithError(errStart).Error("failed to start antigravity callback forwarder")
14031441
c.JSON(http.StatusInternalServerError, gin.H{"error": "failed to start callback server"})
14041442
return
@@ -1407,13 +1445,16 @@ func (h *Handler) RequestAntigravityToken(c *gin.Context) {
14071445

14081446
go func() {
14091447
if isWebUI {
1410-
defer stopCallbackForwarder(antigravityCallbackPort)
1448+
defer stopCallbackForwarderInstance(antigravityCallbackPort, forwarder)
14111449
}
14121450

14131451
waitFile := filepath.Join(h.cfg.AuthDir, fmt.Sprintf(".oauth-antigravity-%s.oauth", state))
14141452
deadline := time.Now().Add(5 * time.Minute)
14151453
var authCode string
14161454
for {
1455+
if !IsOAuthSessionPending(state, "antigravity") {
1456+
return
1457+
}
14171458
if time.Now().After(deadline) {
14181459
log.Error("oauth flow timed out")
14191460
SetOAuthSessionError(state, "OAuth flow timed out")
@@ -1576,6 +1617,7 @@ func (h *Handler) RequestAntigravityToken(c *gin.Context) {
15761617
}
15771618

15781619
CompleteOAuthSession(state)
1620+
CompleteOAuthSessionsByProvider("antigravity")
15791621
fmt.Printf("Authentication successful! Token saved to %s\n", savedPath)
15801622
if projectID != "" {
15811623
fmt.Printf("Using GCP project: %s\n", projectID)
@@ -1653,14 +1695,16 @@ func (h *Handler) RequestIFlowToken(c *gin.Context) {
16531695
RegisterOAuthSession(state, "iflow")
16541696

16551697
isWebUI := isWebUIRequest(c)
1698+
var forwarder *callbackForwarder
16561699
if isWebUI {
16571700
targetURL, errTarget := h.managementCallbackURL("/iflow/callback")
16581701
if errTarget != nil {
16591702
log.WithError(errTarget).Error("failed to compute iflow callback target")
16601703
c.JSON(http.StatusInternalServerError, gin.H{"status": "error", "error": "callback server unavailable"})
16611704
return
16621705
}
1663-
if _, errStart := startCallbackForwarder(iflowauth.CallbackPort, "iflow", targetURL); errStart != nil {
1706+
var errStart error
1707+
if forwarder, errStart = startCallbackForwarder(iflowauth.CallbackPort, "iflow", targetURL); errStart != nil {
16641708
log.WithError(errStart).Error("failed to start iflow callback forwarder")
16651709
c.JSON(http.StatusInternalServerError, gin.H{"status": "error", "error": "failed to start callback server"})
16661710
return
@@ -1669,14 +1713,17 @@ func (h *Handler) RequestIFlowToken(c *gin.Context) {
16691713

16701714
go func() {
16711715
if isWebUI {
1672-
defer stopCallbackForwarder(iflowauth.CallbackPort)
1716+
defer stopCallbackForwarderInstance(iflowauth.CallbackPort, forwarder)
16731717
}
16741718
fmt.Println("Waiting for authentication...")
16751719

16761720
waitFile := filepath.Join(h.cfg.AuthDir, fmt.Sprintf(".oauth-iflow-%s.oauth", state))
16771721
deadline := time.Now().Add(5 * time.Minute)
16781722
var resultMap map[string]string
16791723
for {
1724+
if !IsOAuthSessionPending(state, "iflow") {
1725+
return
1726+
}
16801727
if time.Now().After(deadline) {
16811728
SetOAuthSessionError(state, "Authentication failed")
16821729
fmt.Println("Authentication failed: timeout waiting for callback")
@@ -1743,6 +1790,7 @@ func (h *Handler) RequestIFlowToken(c *gin.Context) {
17431790
}
17441791
fmt.Println("You can now use iFlow services through this CLI")
17451792
CompleteOAuthSession(state)
1793+
CompleteOAuthSessionsByProvider("iflow")
17461794
}()
17471795

17481796
c.JSON(http.StatusOK, gin.H{"status": "ok", "url": authURL, "state": state})

0 commit comments

Comments
 (0)