Work on instrumenting related oauth2 calls

coder · Emyrk · Jan 10, 2024 · Jan 9, 2024 · Jan 9, 2024 · Jan 9, 2024
commit 3377a9b20e28c8fb6bb6dd274c5312d763b7a801
diff --git a/coderd/coderdtest/oidctest/idp.go b/coderd/coderdtest/oidctest/idp.go
@@ -421,18 +421,19 @@ func (f *FakeIDP) CreateAuthCode(t testing.TB, state string, opts ...func(r *htt
 	rw := httptest.NewRecorder()
 	f.handler.ServeHTTP(rw, r)
 	resp := rw.Result()
+	defer resp.Body.Close()
 
 	require.Equal(t, http.StatusTemporaryRedirect, resp.StatusCode, "expected redirect")
 	to := resp.Header.Get("Location")
 	require.NotEmpty(t, to, "expected redirect location")
 
-	toUrl, err := url.Parse(to)
+	toURL, err := url.Parse(to)
 	require.NoError(t, err, "failed to parse redirect location")
 
-	code := toUrl.Query().Get("code")
+	code := toURL.Query().Get("code")
 	require.NotEmpty(t, code, "expected code in redirect location")
 
-	newState := toUrl.Query().Get("state")
+	newState := toURL.Query().Get("state")
 	require.Equal(t, state, newState, "expected state to match")
 
 	return code

diff --git a/coderd/externalauth/externalauth.go b/coderd/externalauth/externalauth.go
@@ -29,7 +29,7 @@ import (
 
 // Config is used for authentication for Git operations.
 type Config struct {
-	promoauth.OAuth2Config
+	promoauth.InstrumentedeOAuth2Config
 	// ID is a unique identifier for the authenticator.
 	ID string
 	// Type is the type of provider.
@@ -187,12 +187,8 @@ func (c *Config) ValidateToken(ctx context.Context, token string) (bool, *coders
 		return false, nil, err
 	}
 
-	cli := http.DefaultClient
-	if v, ok := ctx.Value(oauth2.HTTPClient).(*http.Client); ok {
-		cli = v
-	}
 	req.Header.Set("Authorization", fmt.Sprintf("Bearer %s", token))
-	res, err := cli.Do(req)
+	res, err := c.InstrumentedeOAuth2Config.Do(ctx, "ValidateToken", req)
 	if err != nil {
 		return false, nil, err
 	}
@@ -242,7 +238,7 @@ func (c *Config) AppInstallations(ctx context.Context, token string) ([]codersdk
 		return nil, false, err
 	}
 	req.Header.Set("Authorization", fmt.Sprintf("Bearer %s", token))
-	res, err := http.DefaultClient.Do(req)
+	res, err := c.InstrumentedeOAuth2Config.Do(ctx, "AppInstallations", req)
 	if err != nil {
 		return nil, false, err
 	}
@@ -282,6 +278,8 @@ func (c *Config) AppInstallations(ctx context.Context, token string) ([]codersdk
 }
 
 type DeviceAuth struct {
+	// Cfg is provided for the http client method.
+	Cfg      promoauth.InstrumentedeOAuth2Config
 	ClientID string
 	TokenURL string
 	Scopes   []string
@@ -302,8 +300,8 @@ func (c *DeviceAuth) AuthorizeDevice(ctx context.Context) (*codersdk.ExternalAut
 	if err != nil {
 		return nil, err
 	}
+	resp, err := c.Cfg.Do(ctx, "AuthorizeDevice", req)
 	req.Header.Set("Accept", "application/json")
-	resp, err := http.DefaultClient.Do(req)
 	if err != nil {
 		return nil, err
 	}
@@ -458,24 +456,25 @@ func ConvertConfig(instrument *promoauth.Factory, entries []codersdk.ExternalAut
 		}
 
 		cfg := &Config{
-			OAuth2Config:        instrument.New(entry.ID, oauthConfig),
-			ID:                  entry.ID,
-			Regex:               regex,
-			Type:                entry.Type,
-			NoRefresh:           entry.NoRefresh,
-			ValidateURL:         entry.ValidateURL,
-			AppInstallationsURL: entry.AppInstallationsURL,
-			AppInstallURL:       entry.AppInstallURL,
-			DisplayName:         entry.DisplayName,
-			DisplayIcon:         entry.DisplayIcon,
-			ExtraTokenKeys:      entry.ExtraTokenKeys,
+			InstrumentedeOAuth2Config: instrument.New(entry.ID, oauthConfig),
+			ID:                        entry.ID,
+			Regex:                     regex,
+			Type:                      entry.Type,
+			NoRefresh:                 entry.NoRefresh,
+			ValidateURL:               entry.ValidateURL,
+			AppInstallationsURL:       entry.AppInstallationsURL,
+			AppInstallURL:             entry.AppInstallURL,
+			DisplayName:               entry.DisplayName,
+			DisplayIcon:               entry.DisplayIcon,
+			ExtraTokenKeys:            entry.ExtraTokenKeys,
 		}
 
 		if entry.DeviceFlow {
 			if entry.DeviceCodeURL == "" {
 				return nil, xerrors.Errorf("external auth provider %q: device auth url must be provided", entry.ID)
 			}
 			cfg.DeviceAuth = &DeviceAuth{
+				Cfg:      cfg,
 				ClientID: entry.ClientID,
 				TokenURL: oc.Endpoint.TokenURL,
 				Scopes:   entry.Scopes,

diff --git a/coderd/promoauth/oauth2.go b/coderd/promoauth/oauth2.go
@@ -18,6 +18,24 @@ type OAuth2Config interface {
 	TokenSource(context.Context, *oauth2.Token) oauth2.TokenSource
 }
 
+type InstrumentedeOAuth2Config interface {
+	OAuth2Config
+
+	// Do is provided as a convience method to make a request with the oauth2 client.
+	// It mirrors `http.Client.Do`.
+	// We need this because Coder adds some extra functionality to
+	// oauth clients such as the `ValidateToken()` method.
+	Do(ctx context.Context, source string, req *http.Request) (*http.Response, error)
+}
+
+type HTTPDo interface {
+	// Do is provided as a convience method to make a request with the oauth2 client.
+	// It mirrors `http.Client.Do`.
+	// We need this because Coder adds some extra functionality to
+	// oauth clients such as the `ValidateToken()` method.
+	Do(ctx context.Context, source string, req *http.Request) (*http.Response, error)
+}
+
 var _ OAuth2Config = (*Config)(nil)
 
 type Factory struct {
@@ -65,6 +83,11 @@ type Config struct {
 	metrics    *metrics
 }
 
+func (c *Config) Do(ctx context.Context, source string, req *http.Request) (*http.Response, error) {
+	cli := c.oauthHTTPClient(ctx, source)
+	return cli.Do(req)
+}
+
 func (c *Config) AuthCodeURL(state string, opts ...oauth2.AuthCodeOption) string {
 	// No external requests are made when constructing the auth code url.
 	return c.underlying.AuthCodeURL(state, opts...)
@@ -86,6 +109,10 @@ func (c *Config) TokenSource(ctx context.Context, token *oauth2.Token) oauth2.To
 // source that will make a network request when the 'Token' method is called on
 // it if the token is expired.
 func (c *Config) wrapClient(ctx context.Context, source string) context.Context {
+	return context.WithValue(ctx, oauth2.HTTPClient, c.oauthHTTPClient(ctx, source))
+}
+
+func (c *Config) oauthHTTPClient(ctx context.Context, source string) *http.Client {
 	cli := &http.Client{}
 
 	// Check if the context has an http client already.
@@ -95,7 +122,7 @@ func (c *Config) wrapClient(ctx context.Context, source string) context.Context
 
 	// The new tripper will instrument every request made by the oauth2 client.
 	cli.Transport = newInstrumentedTripper(c, source, cli.Transport)
-	return context.WithValue(ctx, oauth2.HTTPClient, cli)
+	return cli
 }
 
 type instrumentedTripper struct {
@@ -133,5 +160,14 @@ func (i *instrumentedTripper) RoundTrip(r *http.Request) (*http.Response, error)
 		"source":      i.source,
 		"status_code": fmt.Sprintf("%d", statusCode),
 	}).Inc()
+	if err == nil {
+		fmt.Println(map[string]string{
+			"limit":    resp.Header.Get("x-ratelimit-limit"),
+			"remain":   resp.Header.Get("x-ratelimit-remaining"),
+			"used":     resp.Header.Get("x-ratelimit-used"),
+			"reset":    resp.Header.Get("x-ratelimit-reset"),
+			"resource": resp.Header.Get("x-ratelimit-resource"),
+		})
+	}
 	return resp, err
 }
diff --git a/coderd/workspaceagents.go b/coderd/workspaceagents.go
@@ -2100,7 +2100,7 @@ func (api *API) workspaceAgentsExternalAuth(rw http.ResponseWriter, r *http.Requ
 				})
 				return
 			}
-			httpapi.Write(ctx, rw, http.StatusOK, resp)
+			httpapi.Write(ctx, rw, http.StatusInternalServerError, resp)
 			return
 		}
 	}
-Original file line number
+Diff line change
@@ Expand Up @@
     				})
     				return
     			}
-    			httpapi.Write(ctx, rw, http.StatusOK, resp)
+    			httpapi.Write(ctx, rw, http.StatusInternalServerError, resp)
     			return
     		}
     	}
@@ Expand Down @@