fix: meaningful error and panic in pagination tokens

GitOrigin-RevId: 5e7fa89b3380c4052c8d7f9f850abf375564d2ed

Author: zepatrik

pagination/keysetpagination_v2/page_token.go

@@ -4,15 +4,17 @@
 package keysetpagination
 
 import (
+	"crypto/rand"
 	"database/sql"
+	"encoding/base64"
 	"encoding/json"
+	"io"
 	"time"
 
 	"github.com/gofrs/uuid"
-	"github.com/pkg/errors"
-	"github.com/ssoready/hyrumtoken"
-
 	"github.com/ory/herodot"
+	"github.com/pkg/errors"
+	"golang.org/x/crypto/nacl/secretbox"
 )
 
 var fallbackEncryptionKey = &[32]byte{}
@@ -58,7 +60,16 @@ func (t PageToken) Encrypt(keys [][32]byte) string {
 	if len(keys) > 0 {
 		key = &keys[0]
 	}
-	return hyrumtoken.Marshal(key, t)
+	enc, err := t.encrypt(key)
+	if err != nil {
+		// This should basically never happen, only if reading from the random source or marshaling the token fails.
+		// In both cases, we have a bigger problem than just not being able to generate the page token.
+		// Therefore, if we do get an error, there is no point in returning it to the client,
+		// as we already have a working result set. With this string, the next page will return an error,
+		// but that is better than breaking the current page.
+		return "internal error: failed to generate page token"
+	}
+	return enc
 }
 
 func (t PageToken) MarshalJSON() ([]byte, error) {
@@ -109,7 +120,7 @@ func (t PageToken) MarshalJSON() ([]byte, error) {
 	return json.Marshal(toEncode)
 }
 
-var ErrPageTokenExpired = herodot.ErrBadRequest.WithReason("page token expired, do not persist page tokens")
+var ErrPageTokenExpired = herodot.ErrBadRequest.WithError("page token expired, do not persist page tokens")
 
 func (t *PageToken) UnmarshalJSON(data []byte) error {
 	rawToken := jsonPageToken{}
@@ -149,3 +160,46 @@ func (t *PageToken) UnmarshalJSON(data []byte) error {
 }
 
 func NewPageToken(cols ...Column) PageToken { return PageToken{cols: cols} }
+
+func (t *PageToken) encrypt(key *[32]byte) (string, error) {
+	var nonce [24]byte
+	if _, err := io.ReadFull(rand.Reader, nonce[:]); err != nil {
+		return "", errors.Wrap(err, "cannot seed nonce")
+	}
+
+	raw, err := json.Marshal(t)
+	if err != nil {
+		return "", errors.Wrap(err, "cannot marshal page token")
+	}
+
+	enc := secretbox.Seal(nonce[:], raw, &nonce, key)
+	return base64.URLEncoding.EncodeToString(enc), nil
+}
+
+func (t *PageToken) decrypt(key *[32]byte, s string) error {
+	if s == "" {
+		return errors.WithStack(ErrInvalidPaginationToken)
+	}
+
+	raw, err := base64.URLEncoding.DecodeString(s)
+	if err != nil || len(raw) < 24 {
+		return errors.WithStack(ErrInvalidPaginationToken)
+	}
+
+	var nonce [24]byte
+	copy(nonce[:], raw[:24])
+
+	dec, ok := secretbox.Open(nil, raw[24:], &nonce, key)
+	if !ok {
+		return errors.WithStack(ErrInvalidPaginationToken)
+	}
+
+	if err := json.Unmarshal(dec, t); err != nil {
+		if errors.As(err, new(*herodot.DefaultError)) {
+			return err
+		}
+		return errors.WithStack(herodot.ErrInternalServerError.WithReason("unable to unmarshal page token").WithDebug(err.Error()))
+	}
+
+	return nil
+}

pagination/keysetpagination_v2/paginator.go

@@ -12,7 +12,7 @@ import (
 	"github.com/ory/herodot"
 )
 
-var ErrInvalidPaginationToken = herodot.ErrBadRequest.WithReason("invalid pagination token")
+var ErrInvalidPaginationToken = herodot.ErrBadRequest.WithError("invalid pagination token")
 
 type (
 	Paginator struct {

pagination/keysetpagination_v2/request_params.go

@@ -12,7 +12,6 @@ import (
 	"strings"
 
 	"github.com/pkg/errors"
-	"github.com/ssoready/hyrumtoken"
 )
 
 // Pagination Request Parameters
@@ -110,15 +109,18 @@ func ParseQueryParams(keys [][32]byte, q url.Values) ([]Option, error) {
 }
 
 // ParsePageToken parses a page token from the given raw string using the provided keys.
-// It panics if no keys are provided.
+// If decryption fails with all provided keys (including when no keys are given), it
+// falls back to using the fallbackEncryptionKey and returns an error if that also fails.
 func ParsePageToken(keys [][32]byte, raw string) (t PageToken, err error) {
 	for i := range keys {
-		err = errors.WithStack(hyrumtoken.Unmarshal(&keys[i], raw, &t))
-		if err == nil {
-			return
+		err = errors.WithStack(t.decrypt(&keys[i], raw))
+		if errors.Is(err, ErrInvalidPaginationToken) {
+			continue
 		}
+		// either we successfully decrypted the token, or we got an error that is not ErrInvalidPaginationToken, in both cases we should return immediately
+		return t, err
 	}
 	// as a last resort, try the fallback key
-	err = hyrumtoken.Unmarshal(fallbackEncryptionKey, raw, &t)
+	err = t.decrypt(fallbackEncryptionKey, raw)
 	return t, errors.WithStack(err)
 }