From 652c248af0ef302591a4ee02ad106f98f1d6cdd3 Mon Sep 17 00:00:00 2001
From: lostlevels <lostlevels@users.noreply.github.com>
Date: Mon, 6 Feb 2023 13:25:31 -0800
Subject: [PATCH] Return 404 if keyId not found.

---
 appvalidate/repository.go                   | 1 +
 auth/service/api/v1/appvalidate.go          | 4 ++++
 auth/store/mongo/app_validate_repository.go | 5 ++++-
 3 files changed, 9 insertions(+), 1 deletion(-)

diff --git a/appvalidate/repository.go b/appvalidate/repository.go
index 00be9c8fb2..b64da845bb 100644
--- a/appvalidate/repository.go
+++ b/appvalidate/repository.go
@@ -8,6 +8,7 @@ import (
 
 var (
 	ErrDuplicateKeyId = errors.New("duplicate key id")
+	ErrKeyIdNotFound  = errors.New("key id not found")
 )
 
 type Filter struct {
diff --git a/auth/service/api/v1/appvalidate.go b/auth/service/api/v1/appvalidate.go
index 1871f9b890..ccdf05c802 100644
--- a/auth/service/api/v1/appvalidate.go
+++ b/auth/service/api/v1/appvalidate.go
@@ -95,6 +95,10 @@ func (r *Router) VerifyAttestation(res rest.ResponseWriter, req *rest.Request) {
 			"keyId":  attestVerify.KeyID,
 		}
 		log.LoggerFromContext(ctx).WithFields(fields).WithError(err).Error("unable to verify attestation")
+		if errors.Is(err, appvalidate.ErrKeyIdNotFound) {
+			responder.Error(http.StatusNotFound, err)
+			return
+		}
 		if errors.Is(err, appvalidate.ErrAttestationVerificationFailed) {
 			responder.Error(http.StatusBadRequest, err)
 			return
diff --git a/auth/store/mongo/app_validate_repository.go b/auth/store/mongo/app_validate_repository.go
index 108f58189e..67a2745063 100644
--- a/auth/store/mongo/app_validate_repository.go
+++ b/auth/store/mongo/app_validate_repository.go
@@ -67,6 +67,9 @@ func (r *AppValidateRepository) GetAttestationChallenge(ctx context.Context, f a
 		WithError(err).
 		Debug("GetAttestationChallenge")
 	if err != nil {
+		if err == mongo.ErrNoDocuments {
+			return "", appvalidate.ErrKeyIdNotFound
+		}
 		return "", err
 	}
 	return av.AttestationChallenge, nil
@@ -124,7 +127,7 @@ func (r *AppValidateRepository) UpdateAttestation(ctx context.Context, f appvali
 		return err
 	}
 	if res.ModifiedCount == 0 {
-		return errors.New("unable to find app validation object")
+		return appvalidate.ErrKeyIdNotFound
 	}
 	return nil
 }