-
Notifications
You must be signed in to change notification settings - Fork 822
/
token_provider.go
138 lines (110 loc) · 4.26 KB
/
token_provider.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
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
// Copyright (c) Microsoft Corporation. All rights reserved.
// Licensed under the MIT License.
package sbauth
import (
"context"
"strconv"
"time"
"github.com/Azure/azure-sdk-for-go/sdk/azcore"
"github.com/Azure/azure-sdk-for-go/sdk/azcore/policy"
"github.com/Azure/azure-sdk-for-go/sdk/messaging/azeventhubs/internal/auth"
"github.com/Azure/azure-sdk-for-go/sdk/messaging/azeventhubs/internal/exported"
"github.com/Azure/azure-sdk-for-go/sdk/messaging/azeventhubs/internal/sas"
)
// TokenProvider handles access tokens and expiration calculation for SAS
// keys (via connection strings) or TokenCredentials from Azure Identity.
type TokenProvider struct {
tokenCred azcore.TokenCredential
sasTokenProvider *sas.TokenProvider
// InsecureDisableTLS disables TLS. This is only used if the user is connecting to localhost
// and is using an emulator connection string. See [ConnectionStringProperties.Emulator] for
// details.
InsecureDisableTLS bool
}
// NewTokenProvider creates a tokenProvider from azcore.TokenCredential.
func NewTokenProvider(tokenCredential azcore.TokenCredential) *TokenProvider {
return &TokenProvider{tokenCred: tokenCredential}
}
// NewTokenProviderWithConnectionString creates a tokenProvider from a connection string.
func NewTokenProviderWithConnectionString(props exported.ConnectionStringProperties) (*TokenProvider, error) {
// NOTE: this is the value we've been using since forever. AFAIK, it's arbitrary.
const defaultTokenExpiry = 2 * time.Hour
var authOption sas.TokenProviderOption
if props.SharedAccessSignature == nil {
authOption = sas.TokenProviderWithKey(*props.SharedAccessKeyName, *props.SharedAccessKey, defaultTokenExpiry)
} else {
authOption = sas.TokenProviderWithSAS(*props.SharedAccessSignature)
}
provider, err := sas.NewTokenProvider(authOption)
if err != nil {
return nil, err
}
return &TokenProvider{sasTokenProvider: provider, InsecureDisableTLS: props.Emulator}, nil
}
// singleUseTokenProvider allows you to wrap an *auth.Token so it can be used
// with functions that require a TokenProvider, but only actually should get
// a single token (like cbs.NegotiateClaim)
type singleUseTokenProvider auth.Token
// GetToken will return this token.
// This function makes us compatible with auth.TokenProvider.
func (tp *singleUseTokenProvider) GetToken(uri string) (*auth.Token, error) {
return (*auth.Token)(tp), nil
}
// GetToken will retrieve a new token.
// This function makes us compatible with auth.TokenProvider.
func (tp *TokenProvider) GetToken(uri string) (*auth.Token, error) {
token, _, err := tp.getTokenImpl(uri)
return token, err
}
// GetToken returns a token (that is compatible as an auth.TokenProvider) and
// the calculated time when you should renew your token.
func (tp *TokenProvider) GetTokenAsTokenProvider(uri string) (*singleUseTokenProvider, time.Time, error) {
token, renewAt, err := tp.getTokenImpl(uri)
if err != nil {
return nil, time.Time{}, err
}
return (*singleUseTokenProvider)(token), renewAt, nil
}
func (tp *TokenProvider) getTokenImpl(uri string) (*auth.Token, time.Time, error) {
if tp.sasTokenProvider != nil {
return tp.getSASToken(uri)
} else {
return tp.getAZCoreToken()
}
}
func (tpa *TokenProvider) getAZCoreToken() (*auth.Token, time.Time, error) {
// not sure if URI plays in here.
accessToken, err := tpa.tokenCred.GetToken(context.TODO(), policy.TokenRequestOptions{
Scopes: []string{
"https://eventhubs.azure.net//.default",
},
})
if err != nil {
return nil, time.Time{}, err
}
authToken := &auth.Token{
TokenType: auth.CBSTokenTypeJWT,
Token: accessToken.Token,
Expiry: strconv.FormatInt(accessToken.ExpiresOn.Unix(), 10),
}
return authToken,
accessToken.ExpiresOn,
nil
}
func (tpa *TokenProvider) getSASToken(uri string) (*auth.Token, time.Time, error) {
authToken, err := tpa.sasTokenProvider.GetToken(uri)
if err != nil {
return nil, time.Time{}, err
}
// we can ignore the error here since we did the string-izing of the time
// in the first place.
var expiryTime time.Time
if authToken.Expiry != "0" {
// TODO: I'd like to just use the actual Expiry time we generated
// Filed here https://github.com/Azure/azure-sdk-for-go/issues/20468
expiryTime = time.Now().Add(time.Minute * 15)
}
return authToken,
expiryTime,
nil
}