From 59ae872c2499ee7fd24e3b30337c2357087af817 Mon Sep 17 00:00:00 2001
From: Ojima Abraham <ojima.abraham@viam.com>
Date: Thu, 5 Oct 2023 14:40:03 -0400
Subject: [PATCH] APP-2592:  Goutils: implement auth handler that supports (id,
 key) pairs (#203)

---
 rpc/auth.go      | 20 ++++++++++++++++++++
 rpc/auth_test.go | 23 +++++++++++++++++++++++
 2 files changed, 43 insertions(+)

diff --git a/rpc/auth.go b/rpc/auth.go
index 032a9201..9f30a80d 100644
--- a/rpc/auth.go
+++ b/rpc/auth.go
@@ -200,6 +200,26 @@ func MakeSimpleMultiAuthHandler(forEntities, expectedPayloads []string) AuthHand
 	})
 }
 
+// MakeSimpleMultiAuthPairHandler works similarly to MakeSimpleMultiAuthHandler with the addition of
+// supporting a key, id pair used to ensure that a key that maps to the id matches the key passed
+// during the function call.
+func MakeSimpleMultiAuthPairHandler(expectedPayloads map[string]string) AuthHandler {
+	if len(expectedPayloads) == 0 {
+		panic("expected at least one payload")
+	}
+
+	return AuthHandlerFunc(func(ctx context.Context, entity, payload string) (map[string]string, error) {
+		if _, ok := expectedPayloads[entity]; !ok {
+			return nil, errInvalidCredentials
+		}
+
+		if subtle.ConstantTimeCompare([]byte(expectedPayloads[entity]), []byte(payload)) == 1 {
+			return map[string]string{}, nil
+		}
+		return nil, errInvalidCredentials
+	})
+}
+
 // MakeEntitiesChecker checks a list of entities against a given one for use in an auth handler.
 func MakeEntitiesChecker(forEntities []string) func(ctx context.Context, entities ...string) error {
 	return func(ctx context.Context, entities ...string) error {
diff --git a/rpc/auth_test.go b/rpc/auth_test.go
index 95241ec1..6149fd7a 100644
--- a/rpc/auth_test.go
+++ b/rpc/auth_test.go
@@ -137,6 +137,29 @@ func TestMakeSimpleMultiAuthHandler(t *testing.T) {
 	})
 }
 
+func TestMakeSimpleMultiAuthPairHandler(t *testing.T) {
+	test.That(t, func() {
+		MakeSimpleMultiAuthPairHandler(map[string]string{})
+	}, test.ShouldPanicWith, "expected at least one payload")
+
+	t.Run("should validate (keyID, key) mappings", func(t *testing.T) {
+		expectedKeysMap := map[string]string{"myKeyID": "someKey", "somethingElseKeyID": "someOtherKeyID"}
+		handler := MakeSimpleMultiAuthPairHandler(expectedKeysMap)
+
+		for key, value := range expectedKeysMap {
+			t.Run(key, func(t *testing.T) {
+				_, err := handler.Authenticate(context.Background(), key, value)
+				test.That(t, err, test.ShouldBeNil)
+				_, err = handler.Authenticate(context.Background(), key, value+"1")
+				test.That(t, err, test.ShouldEqual, errInvalidCredentials)
+
+				_, err = handler.Authenticate(context.Background(), "notent", key)
+				test.That(t, err, test.ShouldBeError, errInvalidCredentials)
+			})
+		}
+	})
+}
+
 func TestTokenVerificationKeyProviderFunc(t *testing.T) {
 	err1 := errors.New("whoops")
 	capCtx := make(chan struct{})