From 7c342cef42670a030324d5c2577d57096207236f Mon Sep 17 00:00:00 2001
From: Seth Vargo <seth@sethvargo.com>
Date: Fri, 11 Sep 2020 10:32:56 -0400
Subject: [PATCH] Still rate limit by realm

---
 cmd/adminapi/main.go                  |  2 +-
 cmd/apiserver/main.go                 |  2 +-
 pkg/ratelimit/limitware/middleware.go | 38 +++++++++++++++++++++------
 3 files changed, 32 insertions(+), 10 deletions(-)

diff --git a/cmd/adminapi/main.go b/cmd/adminapi/main.go
index 679473c85..e9dff66f7 100644
--- a/cmd/adminapi/main.go
+++ b/cmd/adminapi/main.go
@@ -115,7 +115,7 @@ func realMain(ctx context.Context) error {
 	defer limiterStore.Close()
 
 	httplimiter, err := limitware.NewMiddleware(ctx, limiterStore,
-		limitware.APIKeyFunc(ctx, "adminapi:ratelimit:", config.RateLimit.HMACKey),
+		limitware.APIKeyFunc(ctx, db, "adminapi:ratelimit:", config.RateLimit.HMACKey),
 		limitware.AllowOnError(false))
 	if err != nil {
 		return fmt.Errorf("failed to create limiter middleware: %w", err)
diff --git a/cmd/apiserver/main.go b/cmd/apiserver/main.go
index b81a0c16e..af54d7a1e 100644
--- a/cmd/apiserver/main.go
+++ b/cmd/apiserver/main.go
@@ -128,7 +128,7 @@ func realMain(ctx context.Context) error {
 	defer limiterStore.Close()
 
 	httplimiter, err := limitware.NewMiddleware(ctx, limiterStore,
-		limitware.APIKeyFunc(ctx, "apiserver:ratelimit:", config.RateLimit.HMACKey),
+		limitware.APIKeyFunc(ctx, db, "apiserver:ratelimit:", config.RateLimit.HMACKey),
 		limitware.AllowOnError(false))
 	if err != nil {
 		return fmt.Errorf("failed to create limiter middleware: %w", err)
diff --git a/pkg/ratelimit/limitware/middleware.go b/pkg/ratelimit/limitware/middleware.go
index 84adb0539..398592bf3 100644
--- a/pkg/ratelimit/limitware/middleware.go
+++ b/pkg/ratelimit/limitware/middleware.go
@@ -12,6 +12,7 @@ import (
 	"time"
 
 	"github.com/google/exposure-notifications-verification-server/pkg/controller"
+	"github.com/google/exposure-notifications-verification-server/pkg/database"
 	"github.com/google/exposure-notifications-verification-server/pkg/observability"
 
 	"github.com/google/exposure-notifications-server/pkg/logging"
@@ -196,9 +197,9 @@ func (m *Middleware) Handle(next http.Handler) http.Handler {
 }
 
 // APIKeyFunc returns a default key function for ratelimiting on our API key
-// header. Since APIKeys are assumed to be "public" at some point, they are
-// rate limited by [api-key,ip].
-func APIKeyFunc(ctx context.Context, scope string, hmacKey []byte) httplimit.KeyFunc {
+// header. Since APIKeys are assumed to be "public" at some point, they are rate
+// limited by [realm,ip], and API keys have a 1-1 mapping to a realm.
+func APIKeyFunc(ctx context.Context, db *database.Database, scope string, hmacKey []byte) httplimit.KeyFunc {
 	logger := logging.FromContext(ctx).Named(scope + ".ratelimit")
 	ipAddrLimit := IPAddressKeyFunc(ctx, scope, hmacKey)
 
@@ -206,12 +207,15 @@ func APIKeyFunc(ctx context.Context, scope string, hmacKey []byte) httplimit.Key
 		// Procss the API key
 		v := r.Header.Get("x-api-key")
 		if v != "" {
-			logger.Debugw("limiting by apikey")
-			dig, err := digest(fmt.Sprintf("%s:%s", v, remoteIP(r)), hmacKey)
-			if err != nil {
-				return "", fmt.Errorf("failed to digest api key: %w", err)
+			realmID := realmIDFromAPIKey(db, v)
+			if realmID != 0 {
+				logger.Debugw("limiting by realm from apikey")
+				dig, err := digest(fmt.Sprintf("%d:%s", realmID, remoteIP(r)), hmacKey)
+				if err != nil {
+					return "", fmt.Errorf("failed to digest api key: %w", err)
+				}
+				return fmt.Sprintf("%srealm:%s", scope, dig), nil
 			}
-			return fmt.Sprintf("%sapikey:%s", scope, dig), nil
 		}
 
 		return ipAddrLimit(r)
@@ -288,3 +292,21 @@ func digest(in string, key []byte) (string, error) {
 	dig := h.Sum(nil)
 	return fmt.Sprintf("%x", dig), nil
 }
+
+// realmIDFromAPIKey extracts the realmID from the provided API key, handling v1
+// and v2 API key formats.
+func realmIDFromAPIKey(db *database.Database, apiKey string) uint {
+	// v2 API keys encode in the realm to limit the db calls
+	_, realmID, err := db.VerifyAPIKeySignature(apiKey)
+	if err == nil {
+		return realmID
+	}
+
+	// v1 API keys are more expensive
+	app, err := db.FindAuthorizedAppByAPIKey(apiKey)
+	if err == nil {
+		return app.RealmID
+	}
+
+	return 0
+}