Skip to content

Commit

Permalink
Move external backend key logic to receiver methods of backend.Key (#…
Browse files Browse the repository at this point in the history
…45985) (#46196)

There are a number of places that the backend.Key is assumed to
be a []byte and is passed to various functions from the `bytes`
package. This creates equivalent functionality on the backend.Key
type so that all assumptions about the backend key type can be
removed.
  • Loading branch information
rosstimothy authored Sep 5, 2024
1 parent 3a64eae commit c8c36df
Show file tree
Hide file tree
Showing 22 changed files with 649 additions and 119 deletions.
27 changes: 2 additions & 25 deletions lib/backend/backend.go
Original file line number Diff line number Diff line change
Expand Up @@ -18,12 +18,10 @@ limitations under the License.
package backend

import (
"bytes"
"context"
"fmt"
"io"
"sort"
"strings"
"time"

"github.com/gravitational/trace"
Expand Down Expand Up @@ -198,7 +196,7 @@ type Watch struct {
// String returns a user-friendly description
// of the watcher
func (w *Watch) String() string {
return fmt.Sprintf("Watcher(name=%v, prefixes=%v)", w.Name, string(bytes.Join(w.Prefixes, []byte(", "))))
return fmt.Sprintf("Watcher(name=%v, prefixes=%v)", w.Name, w.Prefixes)
}

// Watcher returns watcher
Expand Down Expand Up @@ -365,7 +363,7 @@ func (it Items) Swap(i, j int) {

// Less is part of sort.Interface.
func (it Items) Less(i, j int) bool {
return bytes.Compare(it[i].Key, it[j].Key) < 0
return it[i].Key.Compare(it[j].Key) < 0
}

// TTL returns TTL in duration units, rounds up to one second
Expand Down Expand Up @@ -415,24 +413,3 @@ func (p earliest) Less(i, j int) bool {
func (p earliest) Swap(i, j int) {
p[i], p[j] = p[j], p[i]
}

// Separator is used as a separator between key parts
const Separator = '/'

// NewKey joins parts into path separated by Separator,
// makes sure path always starts with Separator ("/")
func NewKey(parts ...string) Key {
return internalKey("", parts...)
}

// ExactKey is like Key, except a Separator is appended to the result
// path of Key. This is to ensure range matching of a path will only
// math child paths and not other paths that have the resulting path
// as a prefix.
func ExactKey(parts ...string) Key {
return append(NewKey(parts...), Separator)
}

func internalKey(internalPrefix string, parts ...string) Key {
return Key(strings.Join(append([]string{internalPrefix}, parts...), string(Separator)))
}
2 changes: 1 addition & 1 deletion lib/backend/buffer.go
Original file line number Diff line number Diff line change
Expand Up @@ -319,7 +319,7 @@ type BufferWatcher struct {
// String returns user-friendly representation
// of the buffer watcher
func (w *BufferWatcher) String() string {
return fmt.Sprintf("Watcher(name=%v, prefixes=%v, capacity=%v, size=%v)", w.Name, string(bytes.Join(w.Prefixes, []byte(", "))), w.capacity, len(w.eventsC))
return fmt.Sprintf("Watcher(name=%v, prefixes=%v, capacity=%v, size=%v)", w.Name, w.Prefixes, w.capacity, len(w.eventsC))
}

// Events returns events channel. This method performs internal work and should be re-called after each event
Expand Down
3 changes: 1 addition & 2 deletions lib/backend/dynamo/dynamodbbk.go
Original file line number Diff line number Diff line change
Expand Up @@ -17,7 +17,6 @@ limitations under the License.
package dynamo

import (
"bytes"
"context"
"net/http"
"sort"
Expand Down Expand Up @@ -517,7 +516,7 @@ func (b *Backend) CompareAndSwap(ctx context.Context, expected backend.Item, rep
if len(replaceWith.Key) == 0 {
return nil, trace.BadParameter("missing parameter Key")
}
if !bytes.Equal(expected.Key, replaceWith.Key) {
if expected.Key.Compare(replaceWith.Key) != 0 {
return nil, trace.BadParameter("expected and replaceWith keys should match")
}
r := record{
Expand Down
6 changes: 2 additions & 4 deletions lib/backend/etcdbk/etcd.go
Original file line number Diff line number Diff line change
Expand Up @@ -18,9 +18,7 @@ limitations under the License.
package etcdbk

import (
"bytes"
"context"
"crypto/subtle"
"crypto/tls"
"crypto/x509"
"encoding/base64"
Expand Down Expand Up @@ -736,7 +734,7 @@ func (b *EtcdBackend) CompareAndSwap(ctx context.Context, expected backend.Item,
if len(replaceWith.Key) == 0 {
return nil, trace.BadParameter("missing parameter Key")
}
if subtle.ConstantTimeCompare(expected.Key, replaceWith.Key) != 1 {
if expected.Key.Compare(replaceWith.Key) != 0 {
return nil, trace.BadParameter("expected and replaceWith keys should match")
}
var opts []clientv3.OpOption
Expand Down Expand Up @@ -1019,7 +1017,7 @@ func fromType(eventType mvccpb.Event_EventType) types.OpType {
}

func (b *EtcdBackend) trimPrefix(in backend.Key) backend.Key {
return bytes.TrimPrefix(in, backend.Key(b.cfg.Key))
return in.TrimPrefix(backend.Key(b.cfg.Key))
}

func (b *EtcdBackend) prependPrefix(in backend.Key) string {
Expand Down
2 changes: 1 addition & 1 deletion lib/backend/firestore/firestorebk.go
Original file line number Diff line number Diff line change
Expand Up @@ -516,7 +516,7 @@ func (b *Backend) CompareAndSwap(ctx context.Context, expected backend.Item, rep
if len(replaceWith.Key) == 0 {
return nil, trace.BadParameter("missing parameter Key")
}
if !bytes.Equal(expected.Key, replaceWith.Key) {
if expected.Key.Compare(replaceWith.Key) != 0 {
return nil, trace.BadParameter("expected and replaceWith keys should match")
}

Expand Down
90 changes: 89 additions & 1 deletion lib/backend/key.go
Original file line number Diff line number Diff line change
Expand Up @@ -16,5 +16,93 @@

package backend

import (
"bytes"
"fmt"
"strings"
)

// Key is the unique identifier for an [Item].
type Key = []byte
type Key []byte

// Separator is used as a separator between key parts
const Separator = '/'

// NewKey joins parts into path separated by Separator,
// makes sure path always starts with Separator ("/")
func NewKey(parts ...string) Key {
return internalKey("", parts...)
}

// ExactKey is like Key, except a Separator is appended to the result
// path of Key. This is to ensure range matching of a path will only
// math child paths and not other paths that have the resulting path
// as a prefix.
func ExactKey(parts ...string) Key {
return append(NewKey(parts...), Separator)
}

func internalKey(internalPrefix string, parts ...string) Key {
return Key(strings.Join(append([]string{internalPrefix}, parts...), string(Separator)))
}

// String returns the textual representation of the key with
// each component concatenated together via the [Separator].
func (k Key) String() string {
return string(k)
}

// HasPrefix reports whether the key begins with prefix.
func (k Key) HasPrefix(prefix Key) bool {
return bytes.HasPrefix(k, prefix)
}

// TrimPrefix returns the key without the provided leading prefix string.
// If the key doesn't start with prefix, it is returned unchanged.
func (k Key) TrimPrefix(prefix Key) Key {
return bytes.TrimPrefix(k, prefix)
}

func (k Key) PrependPrefix(p Key) Key {
return append(p, k...)
}

// HasSuffix reports whether the key ends with suffix.
func (k Key) HasSuffix(suffix Key) bool {
return bytes.HasSuffix(k, suffix)
}

// TrimSuffix returns the key without the provided trailing suffix string.
// If the key doesn't end with suffix, it is returned unchanged.
func (k Key) TrimSuffix(suffix Key) Key {
return bytes.TrimSuffix(k, suffix)
}

func (k Key) Components() [][]byte {
if len(k) == 0 {
return nil
}

sep := []byte{Separator}
return bytes.Split(bytes.TrimPrefix(k, sep), sep)
}

func (k Key) Compare(o Key) int {
return bytes.Compare(k, o)
}

// Scan implement sql.Scanner, allowing a [Key] to
// be directly retrieved from sql backends without
// an intermediary object.
func (k *Key) Scan(scan any) error {
switch key := scan.(type) {
case []byte:
*k = bytes.Clone(key)
case string:
*k = []byte(strings.Clone(key))
default:
return fmt.Errorf("invalid Key type %T", scan)
}

return nil
}
Loading

0 comments on commit c8c36df

Please sign in to comment.