Skip to content

Commit 07990a0

Browse files
EmyrkEmyrk
committed
Work on instrumenting related oauth2 calls
1 parent b4410df commit 07990a0

File tree

4 files changed

+60
-24
lines changed

4 files changed

+60
-24
lines changed

‎coderd/coderdtest/oidctest/idp.go

Lines changed: 4 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -421,18 +421,19 @@ func (f *FakeIDP) CreateAuthCode(t testing.TB, state string, opts ...func(r *htt
421421
rw := httptest.NewRecorder()
422422
f.handler.ServeHTTP(rw, r)
423423
resp := rw.Result()
424+
defer resp.Body.Close()
424425

425426
require.Equal(t, http.StatusTemporaryRedirect, resp.StatusCode, "expected redirect")
426427
to := resp.Header.Get("Location")
427428
require.NotEmpty(t, to, "expected redirect location")
428429

429-
toUrl, err := url.Parse(to)
430+
toURL, err := url.Parse(to)
430431
require.NoError(t, err, "failed to parse redirect location")
431432

432-
code := toUrl.Query().Get("code")
433+
code := toURL.Query().Get("code")
433434
require.NotEmpty(t, code, "expected code in redirect location")
434435

435-
newState := toUrl.Query().Get("state")
436+
newState := toURL.Query().Get("state")
436437
require.Equal(t, state, newState, "expected state to match")
437438

438439
return code

‎coderd/externalauth/externalauth.go

Lines changed: 18 additions & 19 deletions
Original file line numberDiff line numberDiff line change
@@ -29,7 +29,7 @@ import (
2929

3030
// Config is used for authentication for Git operations.
3131
type Config struct {
32-
promoauth.OAuth2Config
32+
promoauth.InstrumentedeOAuth2Config
3333
// ID is a unique identifier for the authenticator.
3434
ID string
3535
// Type is the type of provider.
@@ -187,12 +187,8 @@ func (c *Config) ValidateToken(ctx context.Context, token string) (bool, *coders
187187
return false, nil, err
188188
}
189189

190-
cli := http.DefaultClient
191-
if v, ok := ctx.Value(oauth2.HTTPClient).(*http.Client); ok {
192-
cli = v
193-
}
194190
req.Header.Set("Authorization", fmt.Sprintf("Bearer %s", token))
195-
res, err := cli.Do(req)
191+
res, err := c.InstrumentedeOAuth2Config.Do(ctx, "ValidateToken", req)
196192
if err != nil {
197193
return false, nil, err
198194
}
@@ -242,7 +238,7 @@ func (c *Config) AppInstallations(ctx context.Context, token string) ([]codersdk
242238
return nil, false, err
243239
}
244240
req.Header.Set("Authorization", fmt.Sprintf("Bearer %s", token))
245-
res, err := http.DefaultClient.Do(req)
241+
res, err := c.InstrumentedeOAuth2Config.Do(ctx, "AppInstallations", req)
246242
if err != nil {
247243
return nil, false, err
248244
}
@@ -282,6 +278,8 @@ func (c *Config) AppInstallations(ctx context.Context, token string) ([]codersdk
282278
}
283279

284280
type DeviceAuth struct {
281+
// Cfg is provided for the http client method.
282+
Cfg promoauth.InstrumentedeOAuth2Config
285283
ClientID string
286284
TokenURL string
287285
Scopes []string
@@ -302,8 +300,8 @@ func (c *DeviceAuth) AuthorizeDevice(ctx context.Context) (*codersdk.ExternalAut
302300
if err != nil {
303301
return nil, err
304302
}
303+
resp, err := c.Cfg.Do(ctx, "AuthorizeDevice", req)
305304
req.Header.Set("Accept", "application/json")
306-
resp, err := http.DefaultClient.Do(req)
307305
if err != nil {
308306
return nil, err
309307
}
@@ -458,24 +456,25 @@ func ConvertConfig(instrument *promoauth.Factory, entries []codersdk.ExternalAut
458456
}
459457

460458
cfg := &Config{
461-
OAuth2Config: instrument.New(entry.ID, oauthConfig),
462-
ID: entry.ID,
463-
Regex: regex,
464-
Type: entry.Type,
465-
NoRefresh: entry.NoRefresh,
466-
ValidateURL: entry.ValidateURL,
467-
AppInstallationsURL: entry.AppInstallationsURL,
468-
AppInstallURL: entry.AppInstallURL,
469-
DisplayName: entry.DisplayName,
470-
DisplayIcon: entry.DisplayIcon,
471-
ExtraTokenKeys: entry.ExtraTokenKeys,
459+
InstrumentedeOAuth2Config: instrument.New(entry.ID, oauthConfig),
460+
ID: entry.ID,
461+
Regex: regex,
462+
Type: entry.Type,
463+
NoRefresh: entry.NoRefresh,
464+
ValidateURL: entry.ValidateURL,
465+
AppInstallationsURL: entry.AppInstallationsURL,
466+
AppInstallURL: entry.AppInstallURL,
467+
DisplayName: entry.DisplayName,
468+
DisplayIcon: entry.DisplayIcon,
469+
ExtraTokenKeys: entry.ExtraTokenKeys,
472470
}
473471

474472
if entry.DeviceFlow {
475473
if entry.DeviceCodeURL == "" {
476474
return nil, xerrors.Errorf("external auth provider %q: device auth url must be provided", entry.ID)
477475
}
478476
cfg.DeviceAuth = &DeviceAuth{
477+
Cfg: cfg,
479478
ClientID: entry.ClientID,
480479
TokenURL: oc.Endpoint.TokenURL,
481480
Scopes: entry.Scopes,

‎coderd/promoauth/oauth2.go

Lines changed: 37 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -18,6 +18,24 @@ type OAuth2Config interface {
1818
TokenSource(context.Context, *oauth2.Token) oauth2.TokenSource
1919
}
2020

21+
type InstrumentedeOAuth2Config interface {
22+
OAuth2Config
23+
24+
// Do is provided as a convience method to make a request with the oauth2 client.
25+
// It mirrors `http.Client.Do`.
26+
// We need this because Coder adds some extra functionality to
27+
// oauth clients such as the `ValidateToken()` method.
28+
Do(ctx context.Context, source string, req *http.Request) (*http.Response, error)
29+
}
30+
31+
type HTTPDo interface {
32+
// Do is provided as a convience method to make a request with the oauth2 client.
33+
// It mirrors `http.Client.Do`.
34+
// We need this because Coder adds some extra functionality to
35+
// oauth clients such as the `ValidateToken()` method.
36+
Do(ctx context.Context, source string, req *http.Request) (*http.Response, error)
37+
}
38+
2139
var _ OAuth2Config = (*Config)(nil)
2240

2341
type Factory struct {
@@ -65,6 +83,11 @@ type Config struct {
6583
metrics *metrics
6684
}
6785

86+
func (c *Config) Do(ctx context.Context, source string, req *http.Request) (*http.Response, error) {
87+
cli := c.oauthHTTPClient(ctx, source)
88+
return cli.Do(req)
89+
}
90+
6891
func (c *Config) AuthCodeURL(state string, opts ...oauth2.AuthCodeOption) string {
6992
// No external requests are made when constructing the auth code url.
7093
return c.underlying.AuthCodeURL(state, opts...)
@@ -86,6 +109,10 @@ func (c *Config) TokenSource(ctx context.Context, token *oauth2.Token) oauth2.To
86109
// source that will make a network request when the 'Token' method is called on
87110
// it if the token is expired.
88111
func (c *Config) wrapClient(ctx context.Context, source string) context.Context {
112+
return context.WithValue(ctx, oauth2.HTTPClient, c.oauthHTTPClient(ctx, source))
113+
}
114+
115+
func (c *Config) oauthHTTPClient(ctx context.Context, source string) *http.Client {
89116
cli := &http.Client{}
90117

91118
// Check if the context has an http client already.
@@ -95,7 +122,7 @@ func (c *Config) wrapClient(ctx context.Context, source string) context.Context
95122

96123
// The new tripper will instrument every request made by the oauth2 client.
97124
cli.Transport = newInstrumentedTripper(c, source, cli.Transport)
98-
return context.WithValue(ctx, oauth2.HTTPClient, cli)
125+
return cli
99126
}
100127

101128
type instrumentedTripper struct {
@@ -133,5 +160,14 @@ func (i *instrumentedTripper) RoundTrip(r *http.Request) (*http.Response, error)
133160
"source": i.source,
134161
"status_code": fmt.Sprintf("%d", statusCode),
135162
}).Inc()
163+
if err == nil {
164+
fmt.Println(map[string]string{
165+
"limit": resp.Header.Get("x-ratelimit-limit"),
166+
"remain": resp.Header.Get("x-ratelimit-remaining"),
167+
"used": resp.Header.Get("x-ratelimit-used"),
168+
"reset": resp.Header.Get("x-ratelimit-reset"),
169+
"resource": resp.Header.Get("x-ratelimit-resource"),
170+
})
171+
}
136172
return resp, err
137173
}

‎coderd/workspaceagents.go

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -2100,7 +2100,7 @@ func (api *API) workspaceAgentsExternalAuth(rw http.ResponseWriter, r *http.Requ
21002100
})
21012101
return
21022102
}
2103-
httpapi.Write(ctx, rw, http.StatusOK, resp)
2103+
httpapi.Write(ctx, rw, http.StatusInternalServerError, resp)
21042104
return
21052105
}
21062106
}

0 commit comments

Comments
 (0)