-
Notifications
You must be signed in to change notification settings - Fork 6
/
akams.go
274 lines (254 loc) · 7.83 KB
/
akams.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
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
// Copyright (c) Microsoft Corporation.
// Licensed under the MIT License.
// akams package provides a client for the aka.ms API.
// See https://aka.ms/akaapi for more information.
package akams
import (
"bytes"
"context"
"encoding/json"
"fmt"
"io"
"net/http"
"net/url"
)
const Scope = "https://microsoft.onmicrosoft.com/redirectionapi"
// Host is the host identifier.
type Host string
const (
HostAkaMs Host = "1"
HostGoMicrosoftCOM Host = "2"
HostSpoMs Host = "3"
HostOfficeCOM Host = "4"
HostOffice365COM Host = "5"
HostO365COM Host = "6"
HostMicrosoft365COM Host = "7"
)
const (
defaultBulkSize = 300
defaultMaxSizeBytes = 50_000
)
// ResponseError is an error returned by the aka.ms API.
type ResponseError struct {
StatusCode int
Body string
}
func (e *ResponseError) Error() string {
return fmt.Sprintf("request failed: %d\n%s", e.StatusCode, e.Body)
}
// Client is a client for the aka.ms API.
type Client struct {
baseURL *url.URL
httpClient *http.Client
bulkSize int
maxSizeBytes int
}
// NewClient creates a new [Client].
func NewClient(tenant string, httpClient *http.Client) (*Client, error) {
const apiProdBaseUrl = "https://redirectionapi.trafficmanager.net/api"
return NewClientCustom(apiProdBaseUrl, HostAkaMs, tenant, httpClient)
}
// NewClientCustom creates a new [Client] with a custom API base URL and host.
func NewClientCustom(apiBaseURL string, host Host, tenant string, httpClient *http.Client) (*Client, error) {
if httpClient == nil {
httpClient = &http.Client{}
}
baseURL, err := url.Parse(apiBaseURL)
if err != nil {
return nil, err
}
baseURL = baseURL.JoinPath("aka", string(host), tenant, "/")
return &Client{baseURL: baseURL, httpClient: httpClient, bulkSize: defaultBulkSize, maxSizeBytes: defaultMaxSizeBytes}, nil
}
// SetBulkLimit sets the maximum number of items and the maximum size in bytes
// for bulk operations. The default is 300 items and 50_000 bytes.
// Setting a value of 0 or negative for bulkSize or maxSizeBytes will reset the limit to the default.
func (c *Client) SetBulkLimit(bulkSize int, maxSizeBytes int) {
if bulkSize <= 0 {
bulkSize = defaultBulkSize
} else {
c.bulkSize = bulkSize
}
if maxSizeBytes <= 0 {
maxSizeBytes = defaultMaxSizeBytes
} else {
c.maxSizeBytes = maxSizeBytes
}
}
// CreateBulk creates multiple links in bulk.
func (c *Client) CreateBulk(ctx context.Context, links []CreateLinkRequest) error {
return chunkSlice(links, c.bulkSize, c.maxSizeBytes, func(r io.Reader) error {
req, err := c.newRequest(ctx, http.MethodPost, "bulk", r)
if err != nil {
return err
}
resp, err := c.do(req)
if err != nil {
return err
}
defer resp.Body.Close()
if resp.StatusCode != http.StatusOK {
return c.reqError(resp)
}
return nil
})
}
// UpdateBulk updates multiple links in bulk.
func (c *Client) UpdateBulk(ctx context.Context, links []UpdateLinkRequest) error {
return chunkSlice(links, c.bulkSize, c.maxSizeBytes, func(r io.Reader) error {
req, err := c.newRequest(ctx, http.MethodPut, "bulk", r)
if err != nil {
return err
}
resp, err := c.do(req)
if err != nil {
return err
}
defer resp.Body.Close()
switch resp.StatusCode {
case http.StatusAccepted, http.StatusNoContent, http.StatusNotFound:
// success
default:
return c.reqError(resp)
}
return nil
})
}
// CreateOrUpdateBulk creates or updates multiple links in bulk.
// If a link already exists, it will be updated.
// If a link does not exist, it will be created.
// If any link fails to be created or updated, the function will return an error.
func (c *Client) CreateOrUpdateBulk(ctx context.Context, links []CreateLinkRequest) error {
// First try to create all links.
err := c.CreateBulk(ctx, links)
if err == nil {
// All links were created successfully.
return nil
}
// Bad request error is returned when some links already exist.
if e, ok := err.(*ResponseError); !ok || e.StatusCode != http.StatusBadRequest {
return err
}
// We need to identify which links already exist and which don't.
toCreate := make([]CreateLinkRequest, 0, len(links))
toUpdate := make([]UpdateLinkRequest, 0, len(links))
for _, l := range links {
exists, err := c.exists(ctx, l.ShortURL)
if err != nil {
return err
}
if !exists {
toCreate = append(toCreate, l)
} else {
toUpdate = append(toUpdate, l.ToUpdateLinkRequest())
}
}
// Create the links that don't exist and update the ones that do.
if len(toCreate) != 0 {
if err := c.CreateBulk(ctx, toCreate); err != nil {
return err
}
}
if len(toUpdate) != 0 {
if err := c.UpdateBulk(ctx, toUpdate); err != nil {
return err
}
}
return nil
}
func (c *Client) exists(ctx context.Context, shortURL string) (bool, error) {
req, err := c.newRequest(ctx, http.MethodGet, shortURL, nil)
if err != nil {
return false, err
}
resp, err := c.do(req)
if err != nil {
return false, err
}
defer resp.Body.Close()
switch resp.StatusCode {
case http.StatusOK:
return true, nil
case http.StatusNotFound:
return false, nil
default:
return false, c.reqError(resp)
}
}
func (c *Client) reqError(resp *http.Response) error {
// Try to preserve the response body. It may have important context to fix the issue.
body, _ := io.ReadAll(resp.Body)
return &ResponseError{StatusCode: resp.StatusCode, Body: string(body)}
}
func (c *Client) newRequest(ctx context.Context, method string, urlStr string, body io.Reader) (*http.Request, error) {
u := c.baseURL.JoinPath(urlStr)
req, err := http.NewRequestWithContext(ctx, method, u.String(), body)
if err != nil {
return nil, fmt.Errorf("failed to create request: %v", err)
}
req.Header.Set("Accept", "application/json")
if body != nil {
req.Header.Set("Content-Type", "application/json")
}
return req, nil
}
func (c *Client) do(req *http.Request) (*http.Response, error) {
resp, err := c.httpClient.Do(req)
if err != nil {
return nil, fmt.Errorf("request failed: %v", err)
}
return resp, nil
}
// chunkSlice chunks s, encodes each chunk into JSON,
// and calls fn with the encoded chunk.
// The maximum encoded size of each chunk is limited to maxSizeBytes,
// and each chunk has at most bulkSize items.
// The order of the items is preserved.
// If the size of an item is larger than maxSizeBytes, an error will be returned.
// If fn returns an error, the function will stop and return the error.
func chunkSlice[T any](s []T, bulkSize int, maxSizeBytes int, fn func(io.Reader) error) error {
var buf bytes.Buffer
buf.WriteByte('[')
if len(s) == 0 {
buf.WriteByte(']')
return fn(&buf)
}
// We try to convert the slice to JSON in chunks to avoid hitting the maximum request size.
// The end result have the same encoding as if we had used json.Marshal.
enc := json.NewEncoder(&buf)
enc.SetEscapeHTML(false)
var callFn bool
for i := 0; i < len(s); {
lastSize := buf.Len() // keep in case we need to rewind.
if err := enc.Encode(s[i]); err != nil {
return err
}
buf.Truncate(buf.Len() - 1) // Remove the trailing newline added by enc.Encode.
buf.WriteByte(',') // Add a comma to separate items.
if buf.Len() > maxSizeBytes {
// The last item was too big.
if size := buf.Len() - lastSize; size > maxSizeBytes {
// The last item is too big to fit in a chunk.
return fmt.Errorf("item %d is too large: %d bytes > %d byte maximum", i, size, maxSizeBytes)
}
// Rewind and call fn.
buf.Truncate(lastSize)
callFn = true
} else {
// The last item fits, continue.
i++
// Call fn if we reached the end or the bulk size.
callFn = i == len(s) || i%bulkSize == 0
}
if callFn {
buf.Truncate(buf.Len() - 1) // Remove the trailing comma.
buf.WriteByte(']') // Close the JSON array.
if err := fn(&buf); err != nil {
return err
}
buf.Reset() // Reset the buffer for the next chunk.
buf.WriteByte('[') // Open a new JSON array.
}
}
return nil
}