diff --git a/api/openapi.json b/api/openapi.json index eb21723..6227099 100644 --- a/api/openapi.json +++ b/api/openapi.json @@ -1268,6 +1268,16 @@ } } } + }, + "410": { + "description": "Error response", + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/openApiErrorSchema" + } + } + } } } } diff --git a/internal/auth/auth.go b/internal/auth/auth.go index b31daa9..808ef5f 100644 --- a/internal/auth/auth.go +++ b/internal/auth/auth.go @@ -16,6 +16,7 @@ import ( "os/exec" "runtime" "strings" + "sync/atomic" "time" "github.com/Life-USTC/CLI/internal/config" @@ -254,6 +255,35 @@ func callbackRedirectURI(addr net.Addr) string { return fmt.Sprintf("http://%s/callback", addr.String()) } +type callbackResult struct { + code string + state string + err string +} + +func oauthCallbackHandler(results chan<- callbackResult) http.HandlerFunc { + var delivered atomic.Bool + return func(w http.ResponseWriter, r *http.Request) { + if !delivered.CompareAndSwap(false, true) { + http.Error(w, "Authentication callback was already received. You can close this tab.", http.StatusConflict) + return + } + q := r.URL.Query() + result := callbackResult{code: q.Get("code"), state: q.Get("state"), err: q.Get("error")} + select { + case results <- result: + default: + http.Error(w, "Authentication callback could not be delivered. Return to the terminal and retry.", http.StatusServiceUnavailable) + return + } + if result.err != "" { + _, _ = w.Write([]byte("

Authentication failed

You can close this tab.

")) + return + } + _, _ = w.Write([]byte("

Authentication successful!

You can close this tab and return to the terminal.

")) + } +} + // Login runs the full OAuth2 Authorization Code + PKCE flow. // Returns a credential to store. func Login(server string) (*config.Credential, error) { @@ -278,6 +308,7 @@ func Login(server string) (*config.Credential, error) { if err != nil { return nil, err } + defer func() { _ = listener.Close() }() redirectURI := callbackRedirectURI(listener.Addr()) // Register client @@ -314,28 +345,17 @@ func Login(server string) (*config.Credential, error) { oauth2.SetAuthURLParam("code_challenge_method", "S256"), ) - // Channel for callback result - type callbackResult struct { - code string - state string - err string - } ch := make(chan callbackResult, 1) mux := http.NewServeMux() - mux.HandleFunc("/callback", func(w http.ResponseWriter, r *http.Request) { - q := r.URL.Query() - if e := q.Get("error"); e != "" { - ch <- callbackResult{err: e} - _, _ = w.Write([]byte("

Authentication failed

You can close this tab.

")) - return - } - ch <- callbackResult{code: q.Get("code"), state: q.Get("state")} - _, _ = w.Write([]byte("

Authentication successful!

You can close this tab and return to the terminal.

")) - }) + mux.HandleFunc("/callback", oauthCallbackHandler(ch)) srv := &http.Server{Handler: mux} go func() { _ = srv.Serve(listener) }() - defer func() { _ = srv.Shutdown(context.Background()) }() + defer func() { + shutdownCtx, cancel := context.WithTimeout(context.Background(), 2*time.Second) + defer cancel() + _ = srv.Shutdown(shutdownCtx) + }() // Open browser fmt.Println() diff --git a/internal/auth/auth_test.go b/internal/auth/auth_test.go index ae0e23f..75af9f7 100644 --- a/internal/auth/auth_test.go +++ b/internal/auth/auth_test.go @@ -4,6 +4,8 @@ import ( "crypto/rand" "crypto/rsa" "net" + "net/http" + "net/http/httptest" "net/url" "strconv" "testing" @@ -13,6 +15,52 @@ import ( "github.com/go-jose/go-jose/v4/jwt" ) +func TestOAuthCallbackHandlerDeliversOnlyFirstRequest(t *testing.T) { + results := make(chan callbackResult, 1) + handler := oauthCallbackHandler(results) + + first := httptest.NewRecorder() + handler(first, httptest.NewRequest(http.MethodGet, "/callback?code=first&state=state-1", nil)) + if first.Code != http.StatusOK { + t.Fatalf("first status = %d", first.Code) + } + result := <-results + if result.code != "first" || result.state != "state-1" || result.err != "" { + t.Fatalf("result = %#v", result) + } + + repeated := httptest.NewRecorder() + handler(repeated, httptest.NewRequest(http.MethodGet, "/callback?code=second&state=state-2", nil)) + if repeated.Code != http.StatusConflict { + t.Fatalf("repeated status = %d, want %d", repeated.Code, http.StatusConflict) + } + select { + case extra := <-results: + t.Fatalf("unexpected repeated result: %#v", extra) + default: + } +} + +func TestOAuthCallbackHandlerDoesNotBlockWhenResultBufferIsFull(t *testing.T) { + results := make(chan callbackResult, 1) + results <- callbackResult{code: "occupied"} + handler := oauthCallbackHandler(results) + done := make(chan struct{}) + response := httptest.NewRecorder() + go func() { + handler(response, httptest.NewRequest(http.MethodGet, "/callback?error=denied", nil)) + close(done) + }() + select { + case <-done: + case <-time.After(time.Second): + t.Fatal("callback handler blocked on a full result channel") + } + if response.Code != http.StatusServiceUnavailable { + t.Fatalf("status = %d, want %d", response.Code, http.StatusServiceUnavailable) + } +} + func TestCallbackRedirectURIMatchesLoopbackListener(t *testing.T) { listener, err := net.Listen("tcp", "127.0.0.1:0") if err != nil { diff --git a/internal/openapi/client.gen.go b/internal/openapi/client.gen.go index 156de51..c20adfa 100644 --- a/internal/openapi/client.gen.go +++ b/internal/openapi/client.gen.go @@ -17833,6 +17833,7 @@ type WorkspaceCalendarFeedExportResponse struct { JSON401 *OpenApiErrorSchema JSON403 *OpenApiErrorSchema JSON404 *OpenApiErrorSchema + JSON410 *OpenApiErrorSchema } // Status returns HTTPResponse.Status @@ -21168,6 +21169,13 @@ func ParseWorkspaceCalendarFeedExportResponse(rsp *http.Response) (*WorkspaceCal } response.JSON404 = &dest + case strings.Contains(rsp.Header.Get("Content-Type"), "json") && rsp.StatusCode == 410: + var dest OpenApiErrorSchema + if err := json.Unmarshal(bodyBytes, &dest); err != nil { + return nil, err + } + response.JSON410 = &dest + } return response, nil