Extract state from authURL

code-asher · code-asher · commit f862651550ef · 2024-02-08T18:28:23.000-09:00
diff --git a/coderd/coderdtest/oidctest/helper.go b/coderd/coderdtest/oidctest/helper.go
@@ -122,10 +122,13 @@ func (h *LoginHelper) ForceRefresh(t *testing.T, db database.Store, user *coders
 // unit tests, it's easier to skip this step sometimes. It does make an actual
 // request to the IDP, so it should be equivalent to doing this "manually" with
 // actual requests.
-//
-// TODO: Is state param optional? Can we grab it from the authURL?
-func OAuth2GetCode(authURL string, state string, doRequest func(req *http.Request) (*http.Response, error)) (string, error) {
-	r, err := http.NewRequestWithContext(context.Background(), http.MethodGet, authURL, nil)
+func OAuth2GetCode(rawAuthURL string, doRequest func(req *http.Request) (*http.Response, error)) (string, error) {
+	authURL, err := url.Parse(rawAuthURL)
+	if err != nil {
+		return "", xerrors.Errorf("failed to parse auth URL: %w", err)
+	}
+
+	r, err := http.NewRequestWithContext(context.Background(), http.MethodGet, rawAuthURL, nil)
 	if err != nil {
 		return "", xerrors.Errorf("failed to create auth request: %w", err)
 	}
@@ -156,6 +159,7 @@ func OAuth2GetCode(authURL string, state string, doRequest func(req *http.Reques
 		return "", xerrors.Errorf("expected code in redirect location")
 	}
 
+	state := authURL.Query().Get("state")
 	newState := toURL.Query().Get("state")
 	if newState != state {
 		return "", xerrors.Errorf("expected state %q, got %q", state, newState)
diff --git a/coderd/coderdtest/oidctest/idp.go b/coderd/coderdtest/oidctest/idp.go
@@ -539,7 +539,7 @@ func (f *FakeIDP) CreateAuthCode(t testing.TB, state string) string {
 	// it expects some claims to be present.
 	f.stateToIDTokenClaims.Store(state, jwt.MapClaims{})
 
-	code, err := OAuth2GetCode(f.cfg.AuthCodeURL(state), state, func(req *http.Request) (*http.Response, error) {
+	code, err := OAuth2GetCode(f.cfg.AuthCodeURL(state), func(req *http.Request) (*http.Response, error) {
 		rw := httptest.NewRecorder()
 		f.handler.ServeHTTP(rw, req)
 		resp := rw.Result()
diff --git a/enterprise/coderd/oauth2_test.go b/enterprise/coderd/oauth2_test.go
@@ -945,7 +945,6 @@ func authorizationFlow(ctx context.Context, client *codersdk.Client, cfg *oauth2
 	state := uuid.NewString()
 	return oidctest.OAuth2GetCode(
 		cfg.AuthCodeURL(state),
-		state,
 		func(req *http.Request) (*http.Response, error) {
 			// TODO: Would be better if client had a .Do() method.
 			// TODO: Is this the best way to handle redirects?
