From d76e57cb09ef2ee6bde18f2be21c04ad5bcd4ddd Mon Sep 17 00:00:00 2001
From: Easwar Swaminathan <easwars@google.com>
Date: Thu, 6 Feb 2020 20:57:36 -0800
Subject: [PATCH] rls: Implementation of RLS key builders. (#3344)

---
 balancer/rls/internal/keys/builder.go      | 174 +++++++++++
 balancer/rls/internal/keys/builder_test.go | 332 +++++++++++++++++++++
 2 files changed, 506 insertions(+)
 create mode 100644 balancer/rls/internal/keys/builder.go
 create mode 100644 balancer/rls/internal/keys/builder_test.go

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)
+			}
+		})
+	}
+}