-
Notifications
You must be signed in to change notification settings - Fork 365
/
m2m.go
74 lines (68 loc) · 2 KB
/
m2m.go
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
package common
import (
"context"
"encoding/json"
"errors"
"fmt"
"io"
"log"
"net/http"
"golang.org/x/oauth2"
"golang.org/x/oauth2/clientcredentials"
)
type oauthAuthorizationServer struct {
AuthorizationEndpoint string `json:"authorization_endpoint"`
TokenEndpoint string `json:"token_endpoint"`
}
var errNotAvailable = errors.New("not available")
func (c *DatabricksClient) getOAuthEndpoints() (*oauthAuthorizationServer, error) {
err := c.fixHost()
if err != nil {
return nil, fmt.Errorf("host: %w", err)
}
oidc := fmt.Sprintf("%s/oidc/.well-known/oauth-authorization-server", c.Host)
oidcResponse, err := http.Get(oidc)
if err != nil {
return nil, errNotAvailable
}
if oidcResponse.Body == nil {
return nil, fmt.Errorf("fetch .well-known: empty body")
}
defer oidcResponse.Body.Close()
raw, err := io.ReadAll(oidcResponse.Body)
if err != nil {
return nil, fmt.Errorf("read .well-known: %w", err)
}
var oauthEndpoints oauthAuthorizationServer
err = json.Unmarshal(raw, &oauthEndpoints)
if err != nil {
return nil, fmt.Errorf("parse .well-known: %w", err)
}
return &oauthEndpoints, nil
}
func (c *DatabricksClient) configureWithOAuthM2M(
ctx context.Context) (func(r *http.Request) error, error) {
if !c.IsAws() || c.ClientID == "" || c.ClientSecret == "" || c.Host == "" {
return nil, nil
}
// workaround for accounts endpoint not having yet a well-known OIDC alias
if c.TokenEndpoint == "" {
endpoints, err := c.getOAuthEndpoints()
if err == errNotAvailable {
return nil, nil
}
if err != nil {
return nil, fmt.Errorf("databricks oauth: %w", err)
}
c.TokenEndpoint = endpoints.TokenEndpoint
}
log.Printf("[INFO] Generating Databricks OAuth token for Service Principal (%s)", c.ClientID)
ts := (&clientcredentials.Config{
ClientID: c.ClientID,
ClientSecret: c.ClientSecret,
AuthStyle: oauth2.AuthStyleInHeader,
TokenURL: c.TokenEndpoint,
Scopes: []string{"all-apis"},
}).TokenSource(ctx)
return newOidcAuthorizerWithJustBearer(ts), nil
}