this repo has no description
iterate on Body / GetBody wrangling
+130
-30
+4
-11
atproto/client/api_request.go
+4
-11
atproto/client/api_request.go
···
1
1
package client
2
2
3
3
import (
4
-
"bytes"
5
4
"context"
6
5
"fmt"
7
6
"io"
···
27
26
Endpoint syntax.NSID
28
27
29
28
// Optional request body (may be nil). If this is provided, then 'Content-Type' header should be specified
30
-
Body io.ReadCloser
29
+
Body io.Reader
31
30
32
31
// Optional function to return new reader for request body; used for retries. strongly recommended if Body is defined. Body still needs to be defined, even if this function is provided.
33
32
GetBody func() (io.ReadCloser, error)
···
52
51
53
52
// logic to turn "whatever io.Reader we are handed" in to something relatively re-tryable (using GetBody)
54
53
if body != nil {
54
+
// NOTE: http.NewRequestWithContext already handles GetBody() as well as ContentLength for specific types like bytes.Buffer and strings.Reader. We just want to add io.Seeker here, for things like files-on-disk.
55
55
switch v := body.(type) {
56
-
case *bytes.Buffer:
57
-
req.Body = io.NopCloser(v)
58
-
req.GetBody = func() (io.ReadCloser, error) {
59
-
return io.NopCloser(v), nil
60
-
}
61
56
case io.Seeker:
62
57
req.Body = io.NopCloser(body)
63
58
req.GetBody = func() (io.ReadCloser, error) {
64
59
v.Seek(0, 0)
65
60
return io.NopCloser(body), nil
66
61
}
67
-
case io.ReadCloser:
68
-
req.Body = v
69
-
case io.Reader:
70
-
req.Body = io.NopCloser(body)
62
+
default:
63
+
req.Body = body
71
64
}
72
65
}
73
66
return &req
+49
-17
atproto/client/password_auth.go
+49
-17
atproto/client/password_auth.go
···
14
14
)
15
15
16
16
type PasswordAuth struct {
17
-
Session SessionData
17
+
Session PasswordSessionData
18
18
// TODO: RefreshCallback
19
19
20
-
lk sync.Mutex
20
+
// lock which protects concurrent access to AccessToken and RefreshToken in session data
21
+
lk sync.RWMutex
22
+
}
23
+
24
+
type PasswordSessionData struct {
25
+
AccessToken string `json:"access_token"`
26
+
RefreshToken string `json:"refresh_token"`
27
+
AccountDID syntax.DID `json:"account_did"`
28
+
Host string `json:"host"`
21
29
}
22
30
23
-
type SessionData struct {
24
-
AccessToken string
25
-
RefreshToken string
26
-
AccountDID syntax.DID
27
-
Host string
31
+
func (sd *PasswordSessionData) Clone() PasswordSessionData {
32
+
return PasswordSessionData{
33
+
AccessToken: sd.AccessToken,
34
+
RefreshToken: sd.RefreshToken,
35
+
AccountDID: sd.AccountDID,
36
+
Host: sd.Host,
37
+
}
28
38
}
29
39
30
40
func (a *PasswordAuth) DoWithAuth(c *http.Client, req *http.Request) (*http.Response, error) {
31
-
req.Header.Set("Authorization", "Bearer "+a.Session.AccessToken)
41
+
accessToken, refreshToken := a.GetTokens()
42
+
req.Header.Set("Authorization", "Bearer "+accessToken)
32
43
resp, err := c.Do(req)
33
44
if err != nil {
34
45
return nil, err
···
50
61
}
51
62
52
63
// ok, we had an expired token, try a refresh
53
-
if err := a.Refresh(req.Context(), c); err != nil {
64
+
if err := a.Refresh(req.Context(), c, refreshToken); err != nil {
54
65
return nil, err
55
66
}
56
67
···
62
73
}
63
74
}
64
75
65
-
retry.Header.Set("Authorization", "Bearer "+a.Session.AccessToken)
76
+
accessToken, _ = a.GetTokens()
77
+
78
+
retry.Header.Set("Authorization", "Bearer "+accessToken)
66
79
retryResp, err := c.Do(retry)
67
80
if err != nil {
68
81
return nil, err
···
71
84
return retryResp, err
72
85
}
73
86
87
+
// Returns current access and refresh tokens (take a read-lock on session data)
88
+
func (a *PasswordAuth) GetTokens() (string, string) {
89
+
a.lk.RLock()
90
+
defer a.lk.RUnlock()
91
+
return a.Session.AccessToken, a.Session.RefreshToken
92
+
}
93
+
94
+
// Refreshes auth tokens (takes a write-lock on session data).
95
+
//
96
+
// `priorRefreshToken` argument is used to check if a concurrent refresh already took place.
97
+
//
74
98
// TODO: need a "Logout" method as well? which takes the refresh token (not access token)
75
-
func (a *PasswordAuth) Refresh(ctx context.Context, c *http.Client) error {
76
-
77
-
prior := a.Session.RefreshToken
99
+
func (a *PasswordAuth) Refresh(ctx context.Context, c *http.Client, priorRefreshToken string) error {
78
100
79
101
a.lk.Lock()
80
102
defer a.lk.Unlock()
81
103
82
-
// XXX: basic concurrency check: if refresh token already changed, can bail here. should probably handle this better (accept refresh token as input?)
83
-
if prior != a.Session.RefreshToken {
104
+
// basic concurrency check: if refresh token already changed, can bail here (releasing lock)
105
+
if priorRefreshToken != "" && priorRefreshToken != a.Session.RefreshToken {
84
106
return nil
85
107
}
86
108
···
89
111
if err != nil {
90
112
return err
91
113
}
92
-
// TODO: this doesn't inherit User-Agent header
114
+
// NOTE: could try to pull User-Agent from a request and pass that through to here
93
115
req.Header.Set("User-Agent", "indigo-sdk")
94
116
95
117
// NOTE: using refresh token here, not access token
···
158
180
}
159
181
160
182
ra := PasswordAuth{
161
-
Session: SessionData{
183
+
Session: PasswordSessionData{
162
184
AccessToken: out.AccessJwt,
163
185
RefreshToken: out.RefreshJwt,
164
186
AccountDID: ident.DID,
···
169
191
c.AccountDID = &ident.DID
170
192
return c, nil
171
193
}
194
+
195
+
func ResumePasswordSession(data PasswordSessionData) *APIClient {
196
+
c := NewAPIClient(data.Host)
197
+
ra := PasswordAuth{
198
+
Session: data,
199
+
}
200
+
c.Auth = &ra
201
+
c.AccountDID = &data.AccountDID
202
+
return c
203
+
}
+73
-2
atproto/client/password_auth_test.go
+73
-2
atproto/client/password_auth_test.go
···
1
1
package client
2
2
3
3
import (
4
+
"bytes"
4
5
"context"
5
6
"encoding/json"
6
7
"fmt"
8
+
"io"
7
9
"net/http"
8
10
"net/http/httptest"
11
+
"os"
9
12
"strings"
10
13
"testing"
11
14
···
59
62
"refreshJwt": "refresh1",
60
63
})
61
64
return
62
-
case "/xrpc/com.example.get":
65
+
case "/xrpc/com.example.get", "/xrpc/com.example.post":
63
66
hdr := r.Header.Get("Authorization")
64
-
if hdr == "Bearer access1" {
67
+
if hdr == "Bearer access1" || hdr == "Bearer access2" {
65
68
w.Header().Set("Content-Type", "application/json")
66
69
fmt.Fprintln(w, "{\"status\":\"success\"}")
67
70
return
···
117
120
})
118
121
119
122
{
123
+
// simple GET requests, with token expire/retry
120
124
c, err := LoginWithPassword(ctx, &dir, syntax.Handle("user1.example.com").AtIdentifier(), "password1", "")
121
125
require.NoError(err)
122
126
err = c.Get(ctx, syntax.NSID("com.example.get"), nil, nil)
123
127
assert.NoError(err)
124
128
err = c.Get(ctx, syntax.NSID("com.example.expire"), nil, nil)
125
129
assert.NoError(err)
130
+
}
131
+
132
+
{
133
+
// simple POST request, with token expire/retry
134
+
c, err := LoginWithPassword(ctx, &dir, syntax.Handle("user1.example.com").AtIdentifier(), "password1", "")
135
+
require.NoError(err)
136
+
body := map[string]any{
137
+
"a": 123,
138
+
"b": "hello",
139
+
}
140
+
var out json.RawMessage
141
+
err = c.Post(ctx, syntax.NSID("com.example.post"), body, &out)
142
+
assert.NoError(err)
143
+
err = c.Post(ctx, syntax.NSID("com.example.expire"), body, &out)
144
+
assert.NoError(err)
145
+
}
146
+
147
+
{
148
+
// POST with bytes.Buffer body
149
+
c, err := LoginWithPassword(ctx, &dir, syntax.Handle("user1.example.com").AtIdentifier(), "password1", "")
150
+
require.NoError(err)
151
+
body := bytes.NewBufferString("some text")
152
+
req := NewAPIRequest(MethodProcedure, syntax.NSID("com.example.expire"), body)
153
+
req.Headers.Set("Content-Type", "text/plain")
154
+
resp, err := c.Do(ctx, req)
155
+
require.NoError(err)
156
+
assert.Equal(200, resp.StatusCode)
157
+
}
158
+
159
+
{
160
+
// POST with file on disk (can seek and retry)
161
+
c, err := LoginWithPassword(ctx, &dir, syntax.Handle("user1.example.com").AtIdentifier(), "password1", "")
162
+
require.NoError(err)
163
+
f, err := os.Open("testdata/body.json")
164
+
require.NoError(err)
165
+
req := NewAPIRequest(MethodProcedure, syntax.NSID("com.example.expire"), f)
166
+
req.Headers.Set("Content-Type", "application/json")
167
+
resp, err := c.Do(ctx, req)
168
+
require.NoError(err)
169
+
assert.Equal(200, resp.StatusCode)
170
+
}
171
+
172
+
{
173
+
// POST with pipe reader (can *not* retry)
174
+
c, err := LoginWithPassword(ctx, &dir, syntax.Handle("user1.example.com").AtIdentifier(), "password1", "")
175
+
require.NoError(err)
176
+
r1, w1 := io.Pipe()
177
+
go func() {
178
+
fmt.Fprintf(w1, "some data")
179
+
w1.Close()
180
+
}()
181
+
req1 := NewAPIRequest(MethodProcedure, syntax.NSID("com.example.post"), r1)
182
+
req1.Headers.Set("Content-Type", "text/plain")
183
+
resp, err := c.Do(ctx, req1)
184
+
require.NoError(err)
185
+
assert.Equal(200, resp.StatusCode)
186
+
187
+
// expect this to fail (can't re-read from Pipe)
188
+
r2, w2 := io.Pipe()
189
+
go func() {
190
+
fmt.Fprintf(w2, "some data")
191
+
w2.Close()
192
+
}()
193
+
req2 := NewAPIRequest(MethodProcedure, syntax.NSID("com.example.expire"), r2)
194
+
req2.Headers.Set("Content-Type", "text/plain")
195
+
_, err = c.Do(ctx, req2)
196
+
assert.Error(err)
126
197
}
127
198
}