Skip to content

Commit 40e4678

Browse files
zepatrikzepatrik
authored
chore(pagination): improve encryption key handling (#870)
1 parent 3729901 commit 40e4678

5 files changed

Lines changed: 66 additions & 28 deletions

File tree

‎pagination/keysetpagination_v2/page_token.go‎

Lines changed: 10 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -29,8 +29,16 @@ type (
2929
}
3030
)
3131

32-
func (t PageToken) Columns() []Column { return t.cols }
33-
func (t PageToken) Encrypt(key *[32]byte) string { return hyrumtoken.Marshal(key, t) }
32+
func (t PageToken) Columns() []Column { return t.cols }
33+
34+
// Encrypt encrypts the page token using the first key in the provided keyset.
35+
// It panics if no keys are provided.
36+
func (t PageToken) Encrypt(keys [][32]byte) string {
37+
if len(keys) == 0 {
38+
panic("keyset pagination: cannot encrypt page token with no keys")
39+
}
40+
return hyrumtoken.Marshal(&keys[0], t)
41+
}
3442

3543
func (t PageToken) MarshalJSON() ([]byte, error) {
3644
now := time.Now

‎pagination/keysetpagination_v2/page_token_test.go‎

Lines changed: 23 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -39,3 +39,26 @@ func TestPageToken(t *testing.T) {
3939
assert.ErrorIs(t, decodedToken.UnmarshalJSON(raw), ErrPageTokenExpired)
4040
})
4141
}
42+
43+
func TestPageToken_Encrypt(t *testing.T) {
44+
t.Parallel()
45+
46+
keys := [][32]byte{{1, 2, 3}, {4, 5, 6}}
47+
token := NewPageToken(Column{Name: "id", Value: "token"})
48+
49+
t.Run("encrypts with the first key", func(t *testing.T) {
50+
encrypted := token.Encrypt(keys)
51+
52+
decrypted, err := ParsePageToken(keys[:1], encrypted)
53+
require.NoError(t, err)
54+
assert.Equal(t, token, decrypted)
55+
56+
_, err = ParsePageToken(keys[1:], encrypted)
57+
assert.ErrorContains(t, err, "decrypt token")
58+
})
59+
60+
t.Run("panics with no keys", func(t *testing.T) {
61+
assert.PanicsWithValue(t, "keyset pagination: cannot encrypt page token with no keys", func() { token.Encrypt(nil) })
62+
assert.PanicsWithValue(t, "keyset pagination: cannot encrypt page token with no keys", func() { token.Encrypt([][32]byte{}) })
63+
})
64+
}

‎pagination/keysetpagination_v2/parse_header_test.go‎

Lines changed: 3 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -18,14 +18,13 @@ func TestParseHeader(t *testing.T) {
1818

1919
u, err := url.Parse("https://www.ory.sh/")
2020
require.NoError(t, err)
21-
key := [32]byte{1, 2, 3}
22-
keys := [][32]byte{key}
21+
keys := [][32]byte{{1, 2, 3}}
2322
defaultToken, nextToken := NewPageToken(Column{Name: "id", Value: "default"}), NewPageToken(Column{Name: "id", Value: "next"})
2423

2524
t.Run("has next page", func(t *testing.T) {
2625
p := NewPaginator(WithSize(2), WithDefaultToken(defaultToken), WithToken(nextToken))
2726
r := httptest.NewRecorder()
28-
SetLinkHeader(r, &key, u, p)
27+
SetLinkHeader(r, keys, u, p)
2928

3029
first, next, isLast := ParseHeader(&http.Response{Header: r.Header()})
3130
require.NotEqual(t, first, next, r.Header())
@@ -43,7 +42,7 @@ func TestParseHeader(t *testing.T) {
4342
t.Run("is last page", func(t *testing.T) {
4443
p := NewPaginator(WithSize(2), WithDefaultToken(defaultToken), WithToken(nextToken), withIsLast(true))
4544
r := httptest.NewRecorder()
46-
SetLinkHeader(r, &key, u, p)
45+
SetLinkHeader(r, keys, u, p)
4746

4847
first, next, isLast := ParseHeader(&http.Response{Header: r.Header()})
4948
assert.Empty(t, next, r.Header())

‎pagination/keysetpagination_v2/request_params.go‎

Lines changed: 10 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -68,11 +68,11 @@ type ResponseHeaders struct {
6868

6969
// SetLinkHeader adds the Link header for the page encoded by the paginator.
7070
// It contains links to the first and next page, if one exists.
71-
func SetLinkHeader(w http.ResponseWriter, key *[32]byte, u *url.URL, p *Paginator) {
71+
func SetLinkHeader(w http.ResponseWriter, keys [][32]byte, u *url.URL, p *Paginator) {
7272
size := p.Size()
73-
link := []string{linkPart(u, "first", p.DefaultToken().Encrypt(key), size)}
73+
link := []string{linkPart(u, "first", p.DefaultToken().Encrypt(keys), size)}
7474
if !p.isLast {
75-
link = append(link, linkPart(u, "next", p.PageToken().Encrypt(key), size))
75+
link = append(link, linkPart(u, "next", p.PageToken().Encrypt(keys), size))
7676
}
7777
w.Header().Set("Link", strings.Join(link, ","))
7878
}
@@ -109,9 +109,14 @@ func ParseQueryParams(keys [][32]byte, q url.Values) ([]Option, error) {
109109
return opts, nil
110110
}
111111

112+
// ParsePageToken parses a page token from the given raw string using the provided keys.
113+
// It panics if no keys are provided.
112114
func ParsePageToken(keys [][32]byte, raw string) (t PageToken, err error) {
113-
for _, key := range keys {
114-
err = errors.WithStack(hyrumtoken.Unmarshal(&key, raw, &t))
115+
if len(keys) == 0 {
116+
panic("keysetpagination: cannot parse page token with no keys")
117+
}
118+
for i := range keys {
119+
err = errors.WithStack(hyrumtoken.Unmarshal(&keys[i], raw, &t))
115120
if err == nil {
116121
return
117122
}

‎pagination/keysetpagination_v2/request_params_test.go‎

Lines changed: 20 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -17,7 +17,7 @@ import (
1717
func TestSetLinkHeader(t *testing.T) {
1818
t.Parallel()
1919

20-
key := [32]byte{1, 2, 3}
20+
keys := [][32]byte{{1, 2, 3}}
2121
defaultToken, nextToken := NewPageToken(Column{Name: "id", Value: "default"}), NewPageToken(Column{Name: "id", Value: "next"})
2222
opts := []Option{WithSize(2), WithDefaultToken(defaultToken), WithToken(nextToken)}
2323

@@ -30,7 +30,7 @@ func TestSetLinkHeader(t *testing.T) {
3030
assert.Equal(t, "https", u.Scheme)
3131
assert.Equal(t, "ory.sh", u.Host)
3232
raw := u.Query().Get("page_token")
33-
token, err := ParsePageToken([][32]byte{key}, raw)
33+
token, err := ParsePageToken(keys, raw)
3434
require.NoError(t, err)
3535
return token
3636
}
@@ -39,7 +39,7 @@ func TestSetLinkHeader(t *testing.T) {
3939
r := httptest.NewRecorder()
4040
p := NewPaginator(opts...)
4141

42-
SetLinkHeader(r, &key, u, p)
42+
SetLinkHeader(r, keys, u, p)
4343

4444
assert.Len(t, r.Result().Header.Values("link"), 1, "make sure we send one header with multiple comma-separated values rather than multiple headers")
4545
links := link.ParseResponse(r.Result())
@@ -55,7 +55,7 @@ func TestSetLinkHeader(t *testing.T) {
5555
r := httptest.NewRecorder()
5656
p := NewPaginator(append(opts, withIsLast(true))...)
5757

58-
SetLinkHeader(r, &key, u, p)
58+
SetLinkHeader(r, keys, u, p)
5959

6060
assert.Len(t, r.Result().Header.Values("link"), 1, "make sure we send one header with multiple comma-separated values rather than multiple headers")
6161
links := link.ParseResponse(r.Result())
@@ -73,20 +73,23 @@ func TestParsePageToken(t *testing.T) {
7373
keys := [][32]byte{{1, 2, 3}, {4, 5, 6}}
7474

7575
expectedToken := NewPageToken(Column{Name: "id", Value: "token"}, Column{Name: "name", Order: OrderDescending, Value: "test"})
76+
encryptedToken := expectedToken.Encrypt(keys)
7677

7778
t.Run("with valid key", func(t *testing.T) {
78-
for i, key := range keys {
79-
encoded := expectedToken.Encrypt(&key)
80-
token, err := ParsePageToken(keys, encoded)
81-
require.NoErrorf(t, err, "%d", i)
82-
assert.Equal(t, expectedToken, token)
83-
}
79+
token, err := ParsePageToken(keys, encryptedToken)
80+
require.NoError(t, err)
81+
assert.Equal(t, expectedToken, token)
82+
})
83+
84+
t.Run("with rotated key", func(t *testing.T) {
85+
encryptedToken := expectedToken.Encrypt(keys[1:])
86+
token, err := ParsePageToken(keys, encryptedToken)
87+
require.NoError(t, err)
88+
assert.Equal(t, expectedToken, token)
8489
})
8590

8691
t.Run("with invalid key", func(t *testing.T) {
87-
invalidKey := [32]byte{7, 8, 9}
88-
encoded := expectedToken.Encrypt(&invalidKey)
89-
token, err := ParsePageToken(keys, encoded)
92+
token, err := ParsePageToken([][32]byte{{7, 8, 9}}, encryptedToken)
9093
require.ErrorContains(t, err, "decrypt token")
9194
assert.Zero(t, token)
9295
})
@@ -98,7 +101,7 @@ func TestParse(t *testing.T) {
98101
keys := [][32]byte{{1, 2, 3}}
99102
token := NewPageToken(Column{Name: "id", Value: "token"}, Column{Name: "name", Order: OrderDescending, Value: "test"})
100103
defaultToken := NewPageToken(Column{Name: "id", Value: "default"}, Column{Name: "name", Order: OrderDescending, Value: "default name"})
101-
encodedToken := token.Encrypt(&keys[0])
104+
encryptedToken := token.Encrypt(keys)
102105

103106
for _, tc := range []struct {
104107
name string
@@ -114,7 +117,7 @@ func TestParse(t *testing.T) {
114117
},
115118
{
116119
name: "with page token",
117-
q: url.Values{"page_token": {encodedToken}},
120+
q: url.Values{"page_token": {encryptedToken}},
118121
expectedSize: DefaultSize,
119122
expectedToken: token,
120123
},
@@ -126,7 +129,7 @@ func TestParse(t *testing.T) {
126129
},
127130
{
128131
name: "with page size and page token",
129-
q: url.Values{"page_size": {"123"}, "page_token": {encodedToken}},
132+
q: url.Values{"page_size": {"123"}, "page_token": {encryptedToken}},
130133
expectedSize: 123,
131134
expectedToken: token,
132135
},
@@ -158,7 +161,7 @@ func TestParse(t *testing.T) {
158161
assert.Equal(t, defaultToken, paginator.PageToken())
159162
assert.Equal(t, DefaultSize, paginator.Size())
160163

161-
opts, err = ParseQueryParams(keys, url.Values{"page_token": {"", encodedToken, ""}, "page_size": {"", "123", ""}})
164+
opts, err = ParseQueryParams(keys, url.Values{"page_token": {"", encryptedToken, ""}, "page_size": {"", "123", ""}})
162165
require.NoError(t, err)
163166
paginator = NewPaginator(append(opts, WithDefaultToken(defaultToken))...)
164167
assert.Equal(t, token, paginator.PageToken())

0 commit comments

Comments
 (0)