diff --git a/balancer/rls/internal/keys/builder.go b/balancer/rls/internal/keys/builder.go new file mode 100644 index 00000000..5b165892 --- /dev/null +++ b/balancer/rls/internal/keys/builder.go @@ -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() +} diff --git a/balancer/rls/internal/keys/builder_test.go b/balancer/rls/internal/keys/builder_test.go new file mode 100644 index 00000000..a894a393 --- /dev/null +++ b/balancer/rls/internal/keys/builder_test.go @@ -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) + } + }) + } +}