Add tenant resolver (cortexproject/cortex#3486)

* Add tenant resolver package

This implements the multi tenant resolver as described by the [proposal]
for multi tenant query-federation.

By default it behaves like before, but it's implementation can be
swapped out.

[proposal]: cortexproject/cortex#3364

Signed-off-by: Christian Simon <simon@swine.de>

* Replace usages of `ExtractOrgID`

Use TenantID or UserID depending on which of the methods are meant to be
used.

Signed-off-by: Christian Simon <simon@swine.de>

* Replace usages of `ExtractOrgIDFromHTTPRequest`

This is replaced by ExtractTenantIDFromHTTPRequest, which makes sure
that exactly one tenant ID is set.

Signed-off-by: Christian Simon <simon@swine.de>

* Add methods to `tenant` package to use resolver directly

Signed-off-by: Christian Simon <simon@swine.de>

* Remove UserID method from Resolver interface

We need a better definition for what we are trying to achieve with
UserID before we can add it to the interface

Signed-off-by: Christian Simon <simon@swine.de>

* Update comment on the TenantID/TenantIDs

Signed-off-by: Christian Simon <simon@swine.de>

* Improve performance of NormalizeTenantIDs

- reduce allocations by reusing the input slice during de-duplication

Signed-off-by: Christian Simon <simon@swine.de>

tenant/resolver.go

@@ -0,0 +1,132 @@
+package tenant
+
+import (
+	"context"
+	"net/http"
+	"strings"
+
+	"github.com/weaveworks/common/user"
+)
+
+var defaultResolver Resolver = NewSingleResolver()
+
+// WithDefaultResolver updates the resolver used for the package methods.
+func WithDefaultResolver(r Resolver) {
+	defaultResolver = r
+}
+
+// TenantID returns exactly a single tenant ID from the context. It should be
+// used when a certain endpoint should only support exactly a single
+// tenant ID. It returns an error user.ErrNoOrgID if there is no tenant ID
+// supplied or user.ErrTooManyOrgIDs if there are multiple tenant IDs present.
+//
+// ignore stutter warning
+//nolint:golint
+func TenantID(ctx context.Context) (string, error) {
+	return defaultResolver.TenantID(ctx)
+}
+
+// TenantIDs returns all tenant IDs from the context. It should return
+// normalized list of ordered and distinct tenant IDs (as produced by
+// NormalizeTenantIDs).
+//
+// ignore stutter warning
+//nolint:golint
+func TenantIDs(ctx context.Context) ([]string, error) {
+	return defaultResolver.TenantIDs(ctx)
+}
+
+type Resolver interface {
+	// TenantID returns exactly a single tenant ID from the context. It should be
+	// used when a certain endpoint should only support exactly a single
+	// tenant ID. It returns an error user.ErrNoOrgID if there is no tenant ID
+	// supplied or user.ErrTooManyOrgIDs if there are multiple tenant IDs present.
+	TenantID(context.Context) (string, error)
+
+	// TenantIDs returns all tenant IDs from the context. It should return
+	// normalized list of ordered and distinct tenant IDs (as produced by
+	// NormalizeTenantIDs).
+	TenantIDs(context.Context) ([]string, error)
+}
+
+// NewSingleResolver creates a tenant resolver, which restricts all requests to
+// be using a single tenant only. This allows a wider set of characters to be
+// used within the tenant ID and should not impose a breaking change.
+func NewSingleResolver() *SingleResolver {
+	return &SingleResolver{}
+}
+
+type SingleResolver struct {
+}
+
+func (t *SingleResolver) TenantID(ctx context.Context) (string, error) {
+	//lint:ignore faillint wrapper around upstream method
+	return user.ExtractOrgID(ctx)
+}
+
+func (t *SingleResolver) TenantIDs(ctx context.Context) ([]string, error) {
+	//lint:ignore faillint wrapper around upstream method
+	orgID, err := user.ExtractOrgID(ctx)
+	if err != nil {
+		return nil, err
+	}
+	return []string{orgID}, err
+}
+
+type MultiResolver struct {
+}
+
+// NewMultiResolver creates a tenant resolver, which allows request to have
+// multiple tenant ids submitted separated by a '|' character. This enforces
+// further limits on the character set allowed within tenants as detailed here:
+// https://cortexmetrics.io/docs/guides/limitations/#tenant-id-naming)
+func NewMultiResolver() *MultiResolver {
+	return &MultiResolver{}
+}
+
+func (t *MultiResolver) TenantID(ctx context.Context) (string, error) {
+	orgIDs, err := t.TenantIDs(ctx)
+	if err != nil {
+		return "", err
+	}
+
+	if len(orgIDs) > 1 {
+		return "", user.ErrTooManyOrgIDs
+	}
+
+	return orgIDs[0], nil
+}
+
+func (t *MultiResolver) TenantIDs(ctx context.Context) ([]string, error) {
+	//lint:ignore faillint wrapper around upstream method
+	orgID, err := user.ExtractOrgID(ctx)
+	if err != nil {
+		return nil, err
+	}
+
+	orgIDs := strings.Split(orgID, tenantIDsLabelSeparator)
+	for _, orgID := range orgIDs {
+		if err := ValidTenantID(orgID); err != nil {
+			return nil, err
+		}
+	}
+
+	return NormalizeTenantIDs(orgIDs), nil
+}
+
+// ExtractTenantIDFromHTTPRequest extracts a single TenantID through a given
+// resolver directly from a HTTP request.
+func ExtractTenantIDFromHTTPRequest(req *http.Request) (string, context.Context, error) {
+	//lint:ignore faillint wrapper around upstream method
+	_, ctx, err := user.ExtractOrgIDFromHTTPRequest(req)
+	if err != nil {
+		return "", nil, err
+	}
+
+	tenantID, err := defaultResolver.TenantID(ctx)
+	if err != nil {
+		return "", nil, err
+	}
+
+	return tenantID, ctx, nil
+}

tenant/resolver_test.go

@@ -0,0 +1,107 @@
+package tenant
+
+import (
+	"context"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/weaveworks/common/user"
+)
+
+func strptr(s string) *string {
+	return &s
+}
+
+type resolverTestCase struct {
+	name         string
+	headerValue  *string
+	errTenantID  error
+	errTenantIDs error
+	tenantID     string
+	tenantIDs    []string
+}
+
+func (tc *resolverTestCase) test(r Resolver) func(t *testing.T) {
+	return func(t *testing.T) {
+
+		ctx := context.Background()
+		if tc.headerValue != nil {
+			ctx = user.InjectOrgID(ctx, *tc.headerValue)
+		}
+
+		tenantID, err := r.TenantID(ctx)
+		if tc.errTenantID != nil {
+			assert.Equal(t, tc.errTenantID, err)
+		} else {
+			assert.NoError(t, err)
+			assert.Equal(t, tc.tenantID, tenantID)
+		}
+
+		tenantIDs, err := r.TenantIDs(ctx)
+		if tc.errTenantIDs != nil {
+			assert.Equal(t, tc.errTenantIDs, err)
+		} else {
+			assert.NoError(t, err)
+			assert.Equal(t, tc.tenantIDs, tenantIDs)
+		}
+	}
+}
+
+var commonResolverTestCases = []resolverTestCase{
+	{
+		name:         "no-header",
+		errTenantID:  user.ErrNoOrgID,
+		errTenantIDs: user.ErrNoOrgID,
+	},
+	{
+		name:        "empty",
+		headerValue: strptr(""),
+		tenantIDs:   []string{""},
+	},
+	{
+		name:        "single-tenant",
+		headerValue: strptr("tenant-a"),
+		tenantID:    "tenant-a",
+		tenantIDs:   []string{"tenant-a"},
+	},
+}
+
+func TestSingleResolver(t *testing.T) {
+	r := NewSingleResolver()
+	for _, tc := range append(commonResolverTestCases, []resolverTestCase{
+		{
+			name:        "multi-tenant",
+			headerValue: strptr("tenant-a|tenant-b"),
+			tenantID:    "tenant-a|tenant-b",
+			tenantIDs:   []string{"tenant-a|tenant-b"},
+		},
+	}...) {
+		t.Run(tc.name, tc.test(r))
+	}
+}
+
+func TestMultiResolver(t *testing.T) {
+	r := NewMultiResolver()
+	for _, tc := range append(commonResolverTestCases, []resolverTestCase{
+		{
+			name:        "multi-tenant",
+			headerValue: strptr("tenant-a|tenant-b"),
+			errTenantID: user.ErrTooManyOrgIDs,
+			tenantIDs:   []string{"tenant-a", "tenant-b"},
+		},
+		{
+			name:        "multi-tenant-wrong-order",
+			headerValue: strptr("tenant-b|tenant-a"),
+			errTenantID: user.ErrTooManyOrgIDs,
+			tenantIDs:   []string{"tenant-a", "tenant-b"},
+		},
+		{
+			name:        "multi-tenant-duplicate-order",
+			headerValue: strptr("tenant-b|tenant-b|tenant-a"),
+			errTenantID: user.ErrTooManyOrgIDs,
+			tenantIDs:   []string{"tenant-a", "tenant-b"},
+		},
+	}...) {
+		t.Run(tc.name, tc.test(r))
+	}
+}

tenant/tenant.go

@@ -0,0 +1,89 @@
+package tenant
+
+import (
+	"errors"
+	"fmt"
+	"sort"
+)
+
+var (
+	errTenantIDTooLong = errors.New("tenant ID is too long: max 150 characters")
+)
+
+type errTenantIDUnsupportedCharacter struct {
+	pos      int
+	tenantID string
+}
+
+func (e *errTenantIDUnsupportedCharacter) Error() string {
+	return fmt.Sprintf(
+		"tenant ID '%s' contains unsupported character '%c'",
+		e.tenantID,
+		e.tenantID[e.pos],
+	)
+}
+
+const tenantIDsLabelSeparator = "|"
+
+// NormalizeTenantIDs is creating a normalized form by sortiing and de-duplicating the list of tenantIDs
+func NormalizeTenantIDs(tenantIDs []string) []string {
+	sort.Strings(tenantIDs)
+
+	count := len(tenantIDs)
+	if count <= 1 {
+		return tenantIDs
+	}
+
+	posOut := 1
+	for posIn := 1; posIn < count; posIn++ {
+		if tenantIDs[posIn] != tenantIDs[posIn-1] {
+			tenantIDs[posOut] = tenantIDs[posIn]
+			posOut++
+		}
+	}
+
+	return tenantIDs[0:posOut]
+}
+
+// ValidTenantID
+func ValidTenantID(s string) error {
+	// check if it contains invalid runes
+	for pos, r := range s {
+		if !isSupported(r) {
+			return &errTenantIDUnsupportedCharacter{
+				tenantID: s,
+				pos:      pos,
+			}
+		}
+	}
+
+	if len(s) > 150 {
+		return errTenantIDTooLong
+	}
+
+	return nil
+}
+
+// this checks if a rune is supported in tenant IDs (according to
+// https://cortexmetrics.io/docs/guides/limitations/#tenant-id-naming)
+func isSupported(c rune) bool {
+	// characters
+	if ('a' <= c && c <= 'z') || ('A' <= c && c <= 'Z') {
+		return true
+	}
+
+	// digits
+	if '0' <= c && c <= '9' {
+		return true
+	}
+
+	// special
+	return c == '!' ||
+		c == '-' ||
+		c == '_' ||
+		c == '.' ||
+		c == '*' ||
+		c == '\'' ||
+		c == '(' ||
+		c == ')'
+}

tenant/tenant_test.go

@@ -0,0 +1,42 @@
+package tenant
+
+import (
+	"strings"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+)
+
+func TestValidTenantIDs(t *testing.T) {
+	for _, tc := range []struct {
+		name string
+		err  *string
+	}{
+		{
+			name: "tenant-a",
+		},
+		{
+			name: "ABCDEFGHIJKLMNOPQRSTUVWXYZ-abcdefghijklmnopqrstuvwxyz_0987654321!.*'()",
+		},
+		{
+			name: "invalid|",
+			err:  strptr("tenant ID 'invalid|' contains unsupported character '|'"),
+		},
+		{
+			name: strings.Repeat("a", 150),
+		},
+		{
+			name: strings.Repeat("a", 151),
+			err:  strptr("tenant ID is too long: max 150 characters"),
+		},
+	} {
+		t.Run(tc.name, func(t *testing.T) {
+			err := ValidTenantID(tc.name)
+			if tc.err == nil {
+				assert.Nil(t, err)
+			} else {
+				assert.EqualError(t, err, *tc.err)
+			}
+		})
+	}
+}