forked from Azure/azure-kusto-go
-
Notifications
You must be signed in to change notification settings - Fork 0
/
tokenprovider.go
123 lines (104 loc) · 3.69 KB
/
tokenprovider.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
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
package kusto
import (
"context"
"fmt"
"github.com/andyliuliming/azure-kusto-go/kusto/utils"
"net/http"
"strings"
"sync/atomic"
"github.com/Azure/azure-sdk-for-go/sdk/azcore"
"github.com/Azure/azure-sdk-for-go/sdk/azcore/policy"
)
type TokenProvider struct {
tokenCred azcore.TokenCredential //Holds the received token credential as per the authorization
tokenScheme string //Contains token scheme for tokenprovider
customToken string //Holds the custom auth token to be used for authorization
initOnce utils.OnceWithInit[*tokenWrapperResult] //To ensure tokenprovider will be initialized only once while aquiring token
scopes []string //Contains scopes of the auth token
http atomic.Value //Contains the http client to be used for token provider
}
// tokenProvider need to be received as reference, to reflect updations to the structs
func (tkp *TokenProvider) AcquireToken(ctx context.Context) (string, string, error) {
if !isEmpty(tkp.customToken) {
return tkp.customToken, tkp.tokenScheme, nil
}
if tkp.initOnce != nil {
_, err := tkp.initOnce.DoWithInit()
if err != nil {
return "", "", err
}
}
if tkp.tokenCred != nil {
token, err := tkp.tokenCred.GetToken(ctx, policy.TokenRequestOptions{Scopes: tkp.scopes})
if err != nil {
return "", "", err
}
return token.Token, tkp.tokenScheme, nil
}
return "", "", fmt.Errorf("Error: No token info present in token provider")
}
func (tkp *TokenProvider) AuthorizationRequired() bool {
return !(tkp.initOnce == nil && tkp.tokenCred == nil && isEmpty(tkp.customToken))
}
type tokenWrapperResult struct {
credential azcore.TokenCredential
scopes []string
}
func (tkp *TokenProvider) setInit(kcsb *ConnectionStringBuilder, f func(*CloudInfo, *azcore.ClientOptions, string) (azcore.TokenCredential, error)) {
tkp.initOnce = utils.NewOnceWithInit(func() (*tokenWrapperResult, error) {
wrapper, err := tokenWrapper(kcsb, func() *http.Client { return tkp.http.Load().(*http.Client) }, f)
if err != nil {
return nil, err
}
tkp.tokenCred = wrapper.credential
tkp.scopes = wrapper.scopes
return wrapper, err
})
}
func (tkp *TokenProvider) SetHttp(http *http.Client) {
tkp.http.Store(http)
}
func tokenWrapper(kcsb *ConnectionStringBuilder, http func() *http.Client, f func(*CloudInfo, *azcore.ClientOptions, string) (azcore.TokenCredential, error)) (*tokenWrapperResult,
error) {
ci, cliOpts, appClientId, err := getCommonCloudInfo(kcsb, http)
if err != nil {
return nil, err
}
credential, err := f(ci, cliOpts, appClientId)
if err != nil {
return nil, err
}
resourceURI := ci.KustoServiceResourceID
if ci.LoginMfaRequired {
resourceURI = strings.Replace(resourceURI, ".kusto.", ".kustomfa.", 1)
}
scopes := []string{fmt.Sprintf("%s/.default", resourceURI)}
return &tokenWrapperResult{
credential: credential,
scopes: scopes,
}, nil
}
func getCommonCloudInfo(kcsb *ConnectionStringBuilder, http func() *http.Client) (*CloudInfo, *azcore.ClientOptions, string, error) {
client := http()
if http == nil {
return nil, nil, "", fmt.Errorf("error: No http client provided")
}
cloud, err := GetMetadata(kcsb.DataSource, client)
if err != nil {
return nil, nil, "", err
}
cliOpts := kcsb.ClientOptions
appClientId := kcsb.ApplicationClientId
if cliOpts == nil {
cliOpts = &azcore.ClientOptions{
Transport: client,
}
}
if isEmpty(cliOpts.Cloud.ActiveDirectoryAuthorityHost) {
cliOpts.Cloud.ActiveDirectoryAuthorityHost = cloud.LoginEndpoint
}
if isEmpty(appClientId) {
appClientId = cloud.KustoClientAppID
}
return &cloud, cliOpts, appClientId, nil
}