snapshot of https://github.com/ory/keto.git (master, 2026-09-16; blobs > 131072B omitted)

oryx/pagination/keysetpagination_v2/page_token.go

@@ -0,0 +1,245 @@
+// Copyright © 2023 Ory Corp
+// SPDX-License-Identifier: Apache-2.0
+
+package keysetpagination
+
+import (
+	"crypto/rand"
+	"database/sql"
+	"encoding/base64"
+	"encoding/json"
+	"time"
+
+	"github.com/gofrs/uuid"
+	"github.com/pkg/errors"
+
+	"github.com/ory/herodot"
+
+	"github.com/ory/x/aead"
+)
+
+// pageTokenContext is passed as additional authenticated data to the AEAD.
+// It binds the ciphertext to its purpose: even if the encryption key is
+// reused in another context (for example when it is derived from a shared
+// system secret), ciphertexts from that context are rejected as page tokens
+// and vice versa.
+const pageTokenContext = "ory/keysetpagination_v2/page_token"
+
+// fallbackEncryptionKey seals page tokens when no pagination secrets are
+// configured. It is well known, so those tokens are only obfuscated, not
+// authenticated. Configure pagination secrets to authenticate tokens.
+var fallbackEncryptionKey [32]byte
+
+type (
+	PageToken struct {
+		testNow func() time.Time
+		cols    []Column
+	}
+	jsonPageToken = struct {
+		ExpiresAt time.Time    `json:"e"`
+		Cols      []jsonColumn `json:"c"`
+	}
+
+	jsonColumn = struct {
+		Name          string `json:"n"`
+		Order         Order  `json:"o,omitempty"`
+		Nullable      bool   `json:"nl,omitempty"`
+		HasConstraint bool   `json:"hc,omitempty"`
+
+		ValueAny   any        `json:"v,omitempty"`
+		ValueTime  *time.Time `json:"vt,omitempty"`
+		ValueUUID  *uuid.UUID `json:"vu,omitempty"`
+		ValueInt64 *int64     `json:"vi,omitempty"`
+		ValueNull  bool       `json:"vn,omitempty"`
+	}
+	Column struct {
+		Name     string
+		Order    Order
+		Value    any
+		Nullable bool
+		// HasConstraint marks if the column is already constrained by WHERE condition.
+		HasConstraint bool
+	}
+)
+
+func (t PageToken) Columns() []Column { return t.cols }
+
+// Encrypt encrypts the page token using the first key in the provided keyset.
+// It uses a fallback key if no keys are provided.
+func (t PageToken) Encrypt(keys [][32]byte) string {
+	key := fallbackEncryptionKey
+	if len(keys) > 0 {
+		key = keys[0]
+	}
+	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) {
+	now := time.Now
+	if t.testNow != nil {
+		now = t.testNow
+	}
+	toEncode := jsonPageToken{
+		ExpiresAt: now().Add(time.Hour).UTC(),
+		Cols:      make([]jsonColumn, len(t.cols)),
+	}
+	for i, col := range t.cols {
+		toEncode.Cols[i] = jsonColumn{
+			Name:          col.Name,
+			Order:         col.Order,
+			Nullable:      col.Nullable,
+			HasConstraint: col.HasConstraint,
+		}
+		switch v := col.Value.(type) {
+		case time.Time:
+			toEncode.Cols[i].ValueTime = new(v)
+		case uuid.UUID:
+			toEncode.Cols[i].ValueUUID = new(v)
+		case uuid.NullUUID:
+			if v.Valid {
+				toEncode.Cols[i].ValueUUID = new(v.UUID)
+			} else {
+				toEncode.Cols[i].ValueNull = true
+			}
+		case sql.NullString:
+			if v.Valid {
+				toEncode.Cols[i].ValueAny = new(v.String)
+			} else {
+				toEncode.Cols[i].ValueNull = true
+			}
+		case int64:
+			toEncode.Cols[i].ValueInt64 = new(v)
+		case sql.NullInt64:
+			if v.Valid {
+				toEncode.Cols[i].ValueInt64 = new(v.Int64)
+			} else {
+				toEncode.Cols[i].ValueNull = true
+			}
+		default:
+			toEncode.Cols[i].ValueAny = v
+		}
+	}
+	return json.Marshal(toEncode)
+}
+
+func ErrPageTokenExpired() *herodot.DefaultError {
+	return herodot.ErrBadRequest().WithError("page token expired, do not persist page tokens")
+}
+
+func (t *PageToken) UnmarshalJSON(data []byte) error {
+	rawToken := jsonPageToken{}
+	if err := json.Unmarshal(data, &rawToken); err != nil {
+		return err
+	}
+	t.cols = make([]Column, len(rawToken.Cols))
+	for i, col := range rawToken.Cols {
+		t.cols[i] = Column{
+			Name:          col.Name,
+			Order:         col.Order,
+			Nullable:      col.Nullable,
+			HasConstraint: col.HasConstraint,
+		}
+		switch {
+		case col.ValueNull:
+			t.cols[i].Value = nil
+		case col.ValueAny != nil:
+			t.cols[i].Value = col.ValueAny
+		case col.ValueUUID != nil:
+			t.cols[i].Value = *col.ValueUUID
+		case col.ValueTime != nil:
+			t.cols[i].Value = *col.ValueTime
+		case col.ValueInt64 != nil:
+			t.cols[i].Value = *col.ValueInt64
+		}
+	}
+
+	now := time.Now
+	if t.testNow != nil {
+		now = t.testNow
+	}
+	if rawToken.ExpiresAt.Before(now().UTC()) {
+		return errors.WithStack(ErrPageTokenExpired())
+	}
+	return nil
+}
+
+func NewPageToken(cols ...Column) PageToken { return PageToken{cols: cols} }
+
+func (t *PageToken) encrypt(key [32]byte) (string, error) {
+	raw, err := json.Marshal(t)
+	if err != nil {
+		return "", errors.Wrap(err, "cannot marshal page token")
+	}
+
+	a, err := aead.New(key)
+	if err != nil {
+		return "", errors.Wrap(err, "cannot create AEAD")
+	}
+
+	// The nonce is prepended to the ciphertext. AEADs that manage the nonce
+	// internally report a nonce size of zero, so this also covers them.
+	nonce := make([]byte, a.NonceSize())
+	if _, err := rand.Read(nonce); err != nil {
+		return "", errors.Wrap(err, "cannot generate nonce")
+	}
+
+	return base64.URLEncoding.EncodeToString(a.Seal(nonce, nonce, raw, []byte(pageTokenContext))), 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 {
+		return errors.WithStack(ErrInvalidPaginationToken())
+	}
+
+	dec, err := openAEAD(key, raw)
+	if err != nil {
+		// Tokens issued before the switch to a context-bound AEAD are sealed
+		// with NaCl secretbox. Remove this fallback once all tokens issued by
+		// the previous version have expired.
+		var ok bool
+		dec, ok = aead.OpenLegacySecretbox(key, raw)
+		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
+}
+
+func openAEAD(key [32]byte, raw []byte) ([]byte, error) {
+	a, err := aead.New(key)
+	if err != nil {
+		return nil, errors.Wrap(err, "cannot create AEAD")
+	}
+	if len(raw) < a.NonceSize() {
+		return nil, errors.New("ciphertext too short")
+	}
+	nonce, ciphertext := raw[:a.NonceSize()], raw[a.NonceSize():]
+	bs, err := a.Open(nil, nonce, ciphertext, []byte(pageTokenContext))
+	if err != nil {
+		return nil, errors.Wrap(err, "cannot open AEAD")
+	}
+
+	return bs, nil
+}