rls: Implementation of RLS key builders. (#3344)
This commit is contained in:
Easwar Swaminathan
committed by
GitHub
GitHub
parent
979f0a2f08
commit
d76e57cb09
174
balancer/rls/internal/keys/builder.go
Normal file
174
balancer/rls/internal/keys/builder.go
Normal file
@ -0,0 +1,174 @@
|
||||
// +build !appengine
|
||||
|
||||
/*
|
||||
*
|
||||
* Copyright 2020 gRPC authors.
|
||||
*
|
||||
* Licensed under the Apache License, Version 2.0 (the "License");
|
||||
* you may not use this file except in compliance with the License.
|
||||
* You may obtain a copy of the License at
|
||||
*
|
||||
* http://www.apache.org/licenses/LICENSE-2.0
|
||||
*
|
||||
* Unless required by applicable law or agreed to in writing, software
|
||||
* distributed under the License is distributed on an "AS IS" BASIS,
|
||||
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
* See the License for the specific language governing permissions and
|
||||
* limitations under the License.
|
||||
*
|
||||
*/
|
||||
|
||||
// Package keys provides functionality required to build RLS request keys.
|
||||
package keys
|
||||
|
||||
import (
|
||||
"errors"
|
||||
"fmt"
|
||||
"sort"
|
||||
"strings"
|
||||
|
||||
rlspb "google.golang.org/grpc/balancer/rls/internal/proto/grpc_lookup_v1"
|
||||
"google.golang.org/grpc/metadata"
|
||||
)
|
||||
|
||||
// BuilderMap provides a mapping from a request path to the key builder to be
|
||||
// used for that path.
|
||||
// The BuilderMap is constructed by parsing the RouteLookupConfig received by
|
||||
// the RLS balancer as part of its ServiceConfig, and is used by the picker in
|
||||
// the data path to build the RLS keys to be used for a given request.
|
||||
type BuilderMap map[string]builder
|
||||
|
||||
// MakeBuilderMap parses the provided RouteLookupConfig proto and returns a map
|
||||
// from paths to key builders.
|
||||
//
|
||||
// The following conditions are validated, and an error is returned if any of
|
||||
// them is not met:
|
||||
// grpc_keybuilders field
|
||||
// * must have at least one entry
|
||||
// * must not have two entries with the same Name
|
||||
// * must not have any entry with a Name with the service field unset or empty
|
||||
// * must not have any entries without a Name
|
||||
// * must not have a headers entry that has required_match set
|
||||
// * must not have two headers entries with the same key within one entry
|
||||
func MakeBuilderMap(cfg *rlspb.RouteLookupConfig) (BuilderMap, error) {
|
||||
kbs := cfg.GetGrpcKeybuilders()
|
||||
if len(kbs) == 0 {
|
||||
return nil, errors.New("rls: RouteLookupConfig does not contain any GrpcKeyBuilder")
|
||||
}
|
||||
|
||||
bm := make(map[string]builder)
|
||||
for _, kb := range kbs {
|
||||
var matchers []matcher
|
||||
seenKeys := make(map[string]bool)
|
||||
for _, h := range kb.GetHeaders() {
|
||||
if h.GetRequiredMatch() {
|
||||
return nil, fmt.Errorf("rls: GrpcKeyBuilder in RouteLookupConfig has required_match field set {%+v}", kbs)
|
||||
}
|
||||
key := h.GetKey()
|
||||
if seenKeys[key] {
|
||||
return nil, fmt.Errorf("rls: GrpcKeyBuilder in RouteLookupConfig contains repeated Key field in headers {%+v}", kbs)
|
||||
}
|
||||
seenKeys[key] = true
|
||||
matchers = append(matchers, matcher{key: h.GetKey(), names: h.GetNames()})
|
||||
}
|
||||
b := builder{matchers: matchers}
|
||||
|
||||
names := kb.GetNames()
|
||||
if len(names) == 0 {
|
||||
return nil, fmt.Errorf("rls: GrpcKeyBuilder in RouteLookupConfig does not contain any Name {%+v}", kbs)
|
||||
}
|
||||
for _, name := range names {
|
||||
if name.GetService() == "" {
|
||||
return nil, fmt.Errorf("rls: GrpcKeyBuilder in RouteLookupConfig contains a Name field with no Service {%+v}", kbs)
|
||||
}
|
||||
if strings.Contains(name.GetMethod(), `/`) {
|
||||
return nil, fmt.Errorf("rls: GrpcKeyBuilder in RouteLookupConfig contains a method with a slash {%+v}", kbs)
|
||||
}
|
||||
path := "/" + name.GetService() + "/" + name.GetMethod()
|
||||
if _, ok := bm[path]; ok {
|
||||
return nil, fmt.Errorf("rls: GrpcKeyBuilder in RouteLookupConfig contains repeated Name field {%+v}", kbs)
|
||||
}
|
||||
bm[path] = b
|
||||
}
|
||||
}
|
||||
return bm, nil
|
||||
}
|
||||
|
||||
// KeyMap represents the RLS keys to be used for a request.
|
||||
type KeyMap struct {
|
||||
// Map is the representation of an RLS key as a Go map. This is used when
|
||||
// an actual RLS request is to be sent out on the wire, since the
|
||||
// RouteLookupRequest proto expects a Go map.
|
||||
Map map[string]string
|
||||
// Str is the representation of an RLS key as a string, sorted by keys.
|
||||
// Since the RLS keys are part of the cache key in the request cache
|
||||
// maintained by the RLS balancer, and Go maps cannot be used as keys for
|
||||
// Go maps (the cache is implemented as a map), we need a stringified
|
||||
// version of it.
|
||||
Str string
|
||||
}
|
||||
|
||||
// RLSKey builds the RLS keys to be used for the given request, identified by
|
||||
// the request path and the request headers stored in metadata.
|
||||
func (bm BuilderMap) RLSKey(md metadata.MD, path string) KeyMap {
|
||||
b, ok := bm[path]
|
||||
if !ok {
|
||||
i := strings.LastIndex(path, "/")
|
||||
b, ok = bm[path[:i+1]]
|
||||
if !ok {
|
||||
return KeyMap{}
|
||||
}
|
||||
}
|
||||
return b.keys(md)
|
||||
}
|
||||
|
||||
// builder provides the actual functionality of building RLS keys. These are
|
||||
// stored in the BuilderMap.
|
||||
// While processing a pick, the picker looks in the BuilderMap for the
|
||||
// appropriate builder to be used for the given RPC. For each of the matchers
|
||||
// in the found builder, we iterate over the list of request headers (available
|
||||
// as metadata in the context). Once a header matches one of the names in the
|
||||
// matcher, we set the value of the header in the keyMap (with the key being
|
||||
// the one found in the matcher) and move on to the next matcher. If no
|
||||
// KeyBuilder was found in the map, or no header match was found, an empty
|
||||
// keyMap is returned.
|
||||
type builder struct {
|
||||
matchers []matcher
|
||||
}
|
||||
|
||||
// matcher helps extract a key from request headers based on a given name.
|
||||
type matcher struct {
|
||||
// The key used in the keyMap sent as part of the RLS request.
|
||||
key string
|
||||
// List of header names which can supply the value for this key.
|
||||
names []string
|
||||
}
|
||||
|
||||
func (b builder) keys(md metadata.MD) KeyMap {
|
||||
kvMap := make(map[string]string)
|
||||
for _, m := range b.matchers {
|
||||
for _, name := range m.names {
|
||||
if vals := md.Get(name); vals != nil {
|
||||
kvMap[m.key] = strings.Join(vals, ",")
|
||||
break
|
||||
}
|
||||
}
|
||||
}
|
||||
return KeyMap{Map: kvMap, Str: mapToString(kvMap)}
|
||||
}
|
||||
|
||||
func mapToString(kv map[string]string) string {
|
||||
var keys []string
|
||||
for k := range kv {
|
||||
keys = append(keys, k)
|
||||
}
|
||||
sort.Strings(keys)
|
||||
var sb strings.Builder
|
||||
for i, k := range keys {
|
||||
if i != 0 {
|
||||
fmt.Fprint(&sb, ",")
|
||||
}
|
||||
fmt.Fprintf(&sb, "%s=%s", k, kv[k])
|
||||
}
|
||||
return sb.String()
|
||||
}
|
||||
332
balancer/rls/internal/keys/builder_test.go
Normal file
332
balancer/rls/internal/keys/builder_test.go
Normal file
@ -0,0 +1,332 @@
|
||||
// +build !appengine
|
||||
|
||||
/*
|
||||
*
|
||||
* Copyright 2020 gRPC authors.
|
||||
*
|
||||
* Licensed under the Apache License, Version 2.0 (the "License");
|
||||
* you may not use this file except in compliance with the License.
|
||||
* You may obtain a copy of the License at
|
||||
*
|
||||
* http://www.apache.org/licenses/LICENSE-2.0
|
||||
*
|
||||
* Unless required by applicable law or agreed to in writing, software
|
||||
* distributed under the License is distributed on an "AS IS" BASIS,
|
||||
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
* See the License for the specific language governing permissions and
|
||||
* limitations under the License.
|
||||
*
|
||||
*/
|
||||
|
||||
package keys
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"strings"
|
||||
"testing"
|
||||
|
||||
"github.com/google/go-cmp/cmp"
|
||||
rlspb "google.golang.org/grpc/balancer/rls/internal/proto/grpc_lookup_v1"
|
||||
"google.golang.org/grpc/metadata"
|
||||
)
|
||||
|
||||
var (
|
||||
goodKeyBuilder1 = &rlspb.GrpcKeyBuilder{
|
||||
Names: []*rlspb.GrpcKeyBuilder_Name{
|
||||
{Service: "gFoo"},
|
||||
},
|
||||
Headers: []*rlspb.NameMatcher{
|
||||
{Key: "k1", Names: []string{"n1"}},
|
||||
{Key: "k2", Names: []string{"n1"}},
|
||||
},
|
||||
}
|
||||
goodKeyBuilder2 = &rlspb.GrpcKeyBuilder{
|
||||
Names: []*rlspb.GrpcKeyBuilder_Name{
|
||||
{Service: "gBar", Method: "method1"},
|
||||
{Service: "gFoobar"},
|
||||
},
|
||||
Headers: []*rlspb.NameMatcher{
|
||||
{Key: "k1", Names: []string{"n1", "n2"}},
|
||||
},
|
||||
}
|
||||
)
|
||||
|
||||
func TestMakeBuilderMap(t *testing.T) {
|
||||
wantBuilderMap1 := map[string]builder{
|
||||
"/gFoo/": {matchers: []matcher{{key: "k1", names: []string{"n1"}}, {key: "k2", names: []string{"n1"}}}},
|
||||
}
|
||||
wantBuilderMap2 := map[string]builder{
|
||||
"/gFoo/": {matchers: []matcher{{key: "k1", names: []string{"n1"}}, {key: "k2", names: []string{"n1"}}}},
|
||||
"/gBar/method1": {matchers: []matcher{{key: "k1", names: []string{"n1", "n2"}}}},
|
||||
"/gFoobar/": {matchers: []matcher{{key: "k1", names: []string{"n1", "n2"}}}},
|
||||
}
|
||||
|
||||
tests := []struct {
|
||||
desc string
|
||||
cfg *rlspb.RouteLookupConfig
|
||||
wantBuilderMap BuilderMap
|
||||
}{
|
||||
{
|
||||
desc: "One good GrpcKeyBuilder",
|
||||
cfg: &rlspb.RouteLookupConfig{
|
||||
GrpcKeybuilders: []*rlspb.GrpcKeyBuilder{goodKeyBuilder1},
|
||||
},
|
||||
wantBuilderMap: wantBuilderMap1,
|
||||
},
|
||||
{
|
||||
desc: "Two good GrpcKeyBuilders",
|
||||
cfg: &rlspb.RouteLookupConfig{
|
||||
GrpcKeybuilders: []*rlspb.GrpcKeyBuilder{goodKeyBuilder1, goodKeyBuilder2},
|
||||
},
|
||||
wantBuilderMap: wantBuilderMap2,
|
||||
},
|
||||
}
|
||||
|
||||
for _, test := range tests {
|
||||
t.Run(test.desc, func(t *testing.T) {
|
||||
builderMap, err := MakeBuilderMap(test.cfg)
|
||||
if err != nil || !cmp.Equal(builderMap, test.wantBuilderMap, cmp.AllowUnexported(builder{}, matcher{})) {
|
||||
t.Errorf("MakeBuilderMap(%+v) returned {%v, %v}, want: {%v, nil}", test.cfg, builderMap, err, test.wantBuilderMap)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestMakeBuilderMapErrors(t *testing.T) {
|
||||
emptyServiceKeyBuilder := &rlspb.GrpcKeyBuilder{
|
||||
Names: []*rlspb.GrpcKeyBuilder_Name{
|
||||
{Service: "bFoo", Method: "method1"},
|
||||
{Service: "bBar"},
|
||||
{Method: "method1"},
|
||||
},
|
||||
Headers: []*rlspb.NameMatcher{{Key: "k1", Names: []string{"n1", "n2"}}},
|
||||
}
|
||||
requiredMatchKeyBuilder := &rlspb.GrpcKeyBuilder{
|
||||
Names: []*rlspb.GrpcKeyBuilder_Name{{Service: "bFoo", Method: "method1"}},
|
||||
Headers: []*rlspb.NameMatcher{{Key: "k1", Names: []string{"n1", "n2"}, RequiredMatch: true}},
|
||||
}
|
||||
repeatedHeadersKeyBuilder := &rlspb.GrpcKeyBuilder{
|
||||
Names: []*rlspb.GrpcKeyBuilder_Name{
|
||||
{Service: "gBar", Method: "method1"},
|
||||
{Service: "gFoobar"},
|
||||
},
|
||||
Headers: []*rlspb.NameMatcher{
|
||||
{Key: "k1", Names: []string{"n1", "n2"}},
|
||||
{Key: "k1", Names: []string{"n1", "n2"}},
|
||||
},
|
||||
}
|
||||
methodNameWithSlashKeyBuilder := &rlspb.GrpcKeyBuilder{
|
||||
Names: []*rlspb.GrpcKeyBuilder_Name{{Service: "gBar", Method: "method1/foo"}},
|
||||
Headers: []*rlspb.NameMatcher{{Key: "k1", Names: []string{"n1", "n2"}}},
|
||||
}
|
||||
|
||||
tests := []struct {
|
||||
desc string
|
||||
cfg *rlspb.RouteLookupConfig
|
||||
wantErrPrefix string
|
||||
}{
|
||||
{
|
||||
desc: "No GrpcKeyBuilder",
|
||||
cfg: &rlspb.RouteLookupConfig{},
|
||||
wantErrPrefix: "rls: RouteLookupConfig does not contain any GrpcKeyBuilder",
|
||||
},
|
||||
{
|
||||
desc: "Two GrpcKeyBuilders with same Name",
|
||||
cfg: &rlspb.RouteLookupConfig{
|
||||
GrpcKeybuilders: []*rlspb.GrpcKeyBuilder{goodKeyBuilder1, goodKeyBuilder1},
|
||||
},
|
||||
wantErrPrefix: "rls: GrpcKeyBuilder in RouteLookupConfig contains repeated Name field",
|
||||
},
|
||||
{
|
||||
desc: "GrpcKeyBuilder with empty Service field",
|
||||
cfg: &rlspb.RouteLookupConfig{
|
||||
GrpcKeybuilders: []*rlspb.GrpcKeyBuilder{emptyServiceKeyBuilder, goodKeyBuilder1},
|
||||
},
|
||||
wantErrPrefix: "rls: GrpcKeyBuilder in RouteLookupConfig contains a Name field with no Service",
|
||||
},
|
||||
{
|
||||
desc: "GrpcKeyBuilder with no Name",
|
||||
cfg: &rlspb.RouteLookupConfig{
|
||||
GrpcKeybuilders: []*rlspb.GrpcKeyBuilder{{}, goodKeyBuilder1},
|
||||
},
|
||||
wantErrPrefix: "rls: GrpcKeyBuilder in RouteLookupConfig does not contain any Name",
|
||||
},
|
||||
{
|
||||
desc: "GrpcKeyBuilder with requiredMatch field set",
|
||||
cfg: &rlspb.RouteLookupConfig{
|
||||
GrpcKeybuilders: []*rlspb.GrpcKeyBuilder{requiredMatchKeyBuilder, goodKeyBuilder1},
|
||||
},
|
||||
wantErrPrefix: "rls: GrpcKeyBuilder in RouteLookupConfig has required_match field set",
|
||||
},
|
||||
{
|
||||
desc: "GrpcKeyBuilder two headers with same key",
|
||||
cfg: &rlspb.RouteLookupConfig{
|
||||
GrpcKeybuilders: []*rlspb.GrpcKeyBuilder{repeatedHeadersKeyBuilder, goodKeyBuilder1},
|
||||
},
|
||||
wantErrPrefix: "rls: GrpcKeyBuilder in RouteLookupConfig contains repeated Key field in headers",
|
||||
},
|
||||
{
|
||||
desc: "GrpcKeyBuilder with slash in method name",
|
||||
cfg: &rlspb.RouteLookupConfig{
|
||||
GrpcKeybuilders: []*rlspb.GrpcKeyBuilder{methodNameWithSlashKeyBuilder},
|
||||
},
|
||||
wantErrPrefix: "rls: GrpcKeyBuilder in RouteLookupConfig contains a method with a slash",
|
||||
},
|
||||
}
|
||||
|
||||
for _, test := range tests {
|
||||
t.Run(test.desc, func(t *testing.T) {
|
||||
builderMap, err := MakeBuilderMap(test.cfg)
|
||||
if builderMap != nil || !strings.HasPrefix(fmt.Sprint(err), test.wantErrPrefix) {
|
||||
t.Errorf("MakeBuilderMap(%+v) returned {%v, %v}, want: {nil, %v}", test.cfg, builderMap, err, test.wantErrPrefix)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestRLSKey(t *testing.T) {
|
||||
bm, err := MakeBuilderMap(&rlspb.RouteLookupConfig{
|
||||
GrpcKeybuilders: []*rlspb.GrpcKeyBuilder{goodKeyBuilder1, goodKeyBuilder2},
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("MakeBuilderMap() failed: %v", err)
|
||||
}
|
||||
|
||||
tests := []struct {
|
||||
desc string
|
||||
path string
|
||||
md metadata.MD
|
||||
wantKM KeyMap
|
||||
}{
|
||||
{
|
||||
// No keyBuilder is found for the provided service.
|
||||
desc: "service not found in key builder map",
|
||||
path: "/notFoundService/method",
|
||||
md: metadata.Pairs("n1", "v1", "n2", "v2"),
|
||||
wantKM: KeyMap{},
|
||||
},
|
||||
{
|
||||
// No keyBuilder is found for the provided method.
|
||||
desc: "method not found in key builder map",
|
||||
path: "/gBar/notFoundMethod",
|
||||
md: metadata.Pairs("n1", "v1", "n2", "v2"),
|
||||
wantKM: KeyMap{},
|
||||
},
|
||||
{
|
||||
// A keyBuilder is found, but none of the headers match.
|
||||
desc: "directPathMatch-NoMatchingKey",
|
||||
path: "/gBar/method1",
|
||||
md: metadata.Pairs("notMatchingKey", "v1"),
|
||||
wantKM: KeyMap{Map: map[string]string{}, Str: ""},
|
||||
},
|
||||
{
|
||||
// A keyBuilder is found, and a single headers matches.
|
||||
desc: "directPathMatch-SingleKey",
|
||||
path: "/gBar/method1",
|
||||
md: metadata.Pairs("n1", "v1"),
|
||||
wantKM: KeyMap{Map: map[string]string{"k1": "v1"}, Str: "k1=v1"},
|
||||
},
|
||||
{
|
||||
// A keyBuilder is found, and multiple headers match, but the first
|
||||
// match is chosen.
|
||||
desc: "directPathMatch-FirstMatchingKey",
|
||||
path: "/gBar/method1",
|
||||
md: metadata.Pairs("n2", "v2", "n1", "v1"),
|
||||
wantKM: KeyMap{Map: map[string]string{"k1": "v1"}, Str: "k1=v1"},
|
||||
},
|
||||
{
|
||||
// A keyBuilder is found as a wildcard match, but none of the
|
||||
// headers match.
|
||||
desc: "wildcardPathMatch-NoMatchingKey",
|
||||
path: "/gFoobar/method1",
|
||||
md: metadata.Pairs("notMatchingKey", "v1"),
|
||||
wantKM: KeyMap{Map: map[string]string{}, Str: ""},
|
||||
},
|
||||
{
|
||||
// A keyBuilder is found as a wildcard match, and a single headers
|
||||
// matches.
|
||||
desc: "wildcardPathMatch-SingleKey",
|
||||
path: "/gFoobar/method1",
|
||||
md: metadata.Pairs("n1", "v1"),
|
||||
wantKM: KeyMap{Map: map[string]string{"k1": "v1"}, Str: "k1=v1"},
|
||||
},
|
||||
{
|
||||
// A keyBuilder is found as a wildcard match, and multiple headers
|
||||
// match, but the first match is chosen.
|
||||
desc: "wildcardPathMatch-FirstMatchingKey",
|
||||
path: "/gFoobar/method1",
|
||||
md: metadata.Pairs("n2", "v2", "n1", "v1"),
|
||||
wantKM: KeyMap{Map: map[string]string{"k1": "v1"}, Str: "k1=v1"},
|
||||
},
|
||||
{
|
||||
// Multiple matchers find hits in the provided request headers.
|
||||
desc: "multipleMatchers",
|
||||
path: "/gFoo/method1",
|
||||
md: metadata.Pairs("n2", "v2", "n1", "v1"),
|
||||
wantKM: KeyMap{Map: map[string]string{"k1": "v1", "k2": "v1"}, Str: "k1=v1,k2=v1"},
|
||||
},
|
||||
{
|
||||
// A match is found for a header which is specified multiple times.
|
||||
// So, the values are joined with commas separating them.
|
||||
desc: "commaSeparated",
|
||||
path: "/gBar/method1",
|
||||
md: metadata.Pairs("n1", "v1", "n1", "v2", "n1", "v3"),
|
||||
wantKM: KeyMap{Map: map[string]string{"k1": "v1,v2,v3"}, Str: "k1=v1,v2,v3"},
|
||||
},
|
||||
}
|
||||
|
||||
for _, test := range tests {
|
||||
t.Run(test.desc, func(t *testing.T) {
|
||||
if gotKM := bm.RLSKey(test.md, test.path); !cmp.Equal(gotKM, test.wantKM) {
|
||||
t.Errorf("RLSKey(%+v, %s) = %+v, want %+v", test.md, test.path, gotKM, test.wantKM)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestMapToString(t *testing.T) {
|
||||
tests := []struct {
|
||||
desc string
|
||||
input map[string]string
|
||||
wantStr string
|
||||
}{
|
||||
{
|
||||
desc: "empty map",
|
||||
input: nil,
|
||||
wantStr: "",
|
||||
},
|
||||
{
|
||||
desc: "one key",
|
||||
input: map[string]string{
|
||||
"k1": "v1",
|
||||
},
|
||||
wantStr: "k1=v1",
|
||||
},
|
||||
{
|
||||
desc: "sorted keys",
|
||||
input: map[string]string{
|
||||
"k1": "v1",
|
||||
"k2": "v2",
|
||||
"k3": "v3",
|
||||
},
|
||||
wantStr: "k1=v1,k2=v2,k3=v3",
|
||||
},
|
||||
{
|
||||
desc: "unsorted keys",
|
||||
input: map[string]string{
|
||||
"k3": "v3",
|
||||
"k1": "v1",
|
||||
"k2": "v2",
|
||||
},
|
||||
wantStr: "k1=v1,k2=v2,k3=v3",
|
||||
},
|
||||
}
|
||||
|
||||
for _, test := range tests {
|
||||
t.Run(test.desc, func(t *testing.T) {
|
||||
if gotStr := mapToString(test.input); gotStr != test.wantStr {
|
||||
t.Errorf("mapToString(%v) = %s, want %s", test.input, gotStr, test.wantStr)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user