-
Notifications
You must be signed in to change notification settings - Fork 342
/
client_cache.go
142 lines (125 loc) · 4.27 KB
/
client_cache.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
139
140
141
142
package azure
import (
"context"
"fmt"
"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/azcore/to"
"github.com/Azure/azure-sdk-for-go/sdk/azidentity"
"github.com/Azure/azure-sdk-for-go/sdk/storage/azblob/container"
"github.com/Azure/azure-sdk-for-go/sdk/storage/azblob/sas"
"github.com/Azure/azure-sdk-for-go/sdk/storage/azblob/service"
lru "github.com/hnlq715/golang-lru"
"github.com/puzpuzpuz/xsync"
"github.com/treeverse/lakefs/pkg/block/params"
)
type ClientCache struct {
serviceToClient *xsync.MapOf[string, *service.Client]
containerToClient *xsync.MapOf[string, *container.Client]
// udcCache - User Delegation Credential cache used to reduce POST requests while creating pre-signed URLs
udcCache *lru.ARCCache
params params.Azure
}
func NewCache(p params.Azure) (*ClientCache, error) {
l, err := lru.NewARC(udcCacheSize)
if err != nil {
return nil, err
}
return &ClientCache{
serviceToClient: xsync.NewMapOf[*service.Client](),
containerToClient: xsync.NewMapOf[*container.Client](),
udcCache: l,
params: p,
}, nil
}
func mapKey(storageAccount, containerName string) string {
return fmt.Sprintf("%s#%s", storageAccount, containerName)
}
func (c *ClientCache) NewContainerClient(storageAccount, containerName string) (*container.Client, error) {
key := mapKey(storageAccount, containerName)
var err error
cl, _ := c.containerToClient.LoadOrCompute(key, func() *container.Client {
var svc *service.Client
svc, err = c.NewServiceClient(storageAccount)
if err != nil {
return nil
}
return svc.NewContainerClient(containerName)
})
if err != nil {
return nil, err
}
return cl, nil
}
func (c *ClientCache) NewServiceClient(storageAccount string) (*service.Client, error) {
p := c.params
// Use StorageAccessKey to initialize storage account client only if it was provided for this given storage account
// Otherwise fall back to the default credentials
if p.StorageAccount != storageAccount {
p.StorageAccount = storageAccount
p.StorageAccessKey = ""
}
var err error
cl, _ := c.serviceToClient.LoadOrCompute(storageAccount, func() *service.Client {
var svc *service.Client
svc, err = BuildAzureServiceClient(p)
if err != nil {
return nil
}
return svc
})
if err != nil {
return nil, err
}
return cl, nil
}
func (c *ClientCache) NewUDC(ctx context.Context, storageAccount string, expiry *time.Time) (*service.UserDelegationCredential, error) {
// Otherwise assume using role based credentials and build signed URL using user delegation credentials
currentTime := time.Now().UTC().Add(-10 * time.Second)
// UDC expiry time of PreSignedExpiry + hour
udcExpiry := expiry.Add(time.Hour)
info := service.KeyInfo{
Start: to.Ptr(currentTime.UTC().Format(sas.TimeFormat)),
Expiry: to.Ptr(udcExpiry.Format(sas.TimeFormat)),
}
var udc *service.UserDelegationCredential
// Check udcCache
res, ok := c.udcCache.Get(storageAccount)
if !ok {
svc, err := c.NewServiceClient(storageAccount)
if err != nil {
return nil, err
}
udc, err = svc.GetUserDelegationCredential(ctx, info, nil)
if err != nil {
return nil, err
}
// UDC expires after PreSignedExpiry + hour but cache entry expires after an hour
c.udcCache.AddEx(storageAccount, udc, time.Hour)
} else {
udc = res.(*service.UserDelegationCredential)
}
return udc, nil
}
func BuildAzureServiceClient(params params.Azure) (*service.Client, error) {
var url string
if params.TestEndpointURL != "" { // For testing purposes - override default url template
url = params.TestEndpointURL
} else {
url = fmt.Sprintf(URLTemplate, params.StorageAccount)
}
options := service.ClientOptions{ClientOptions: azcore.ClientOptions{Retry: policy.RetryOptions{TryTimeout: params.TryTimeout}}}
if params.StorageAccessKey != "" {
cred, err := service.NewSharedKeyCredential(params.StorageAccount, params.StorageAccessKey)
if err != nil {
return nil, fmt.Errorf("invalid credentials: %w", err)
}
return service.NewClientWithSharedKeyCredential(url, cred, &options)
}
defaultCreds, err := azidentity.NewDefaultAzureCredential(nil)
if err != nil {
return nil, fmt.Errorf("missing credentials: %w", err)
}
return service.NewClient(url, defaultCreds, &options)
}