Improve tests for JWK vs. KID revoke auth flow

The logic for both test cases is fairly similar, but with some
small differences. Made those clearer by means of some comments.
Also added some comments to the middleware logic that decided
whether to extract JWK or lookup by KID.

acme/api/middleware.go

@@ -262,11 +262,11 @@ func (h *Handler) extractJWK(next nextHTTP) nextHTTP {
 		// Store the JWK in the context.
 		ctx = context.WithValue(ctx, jwkContextKey, jwk)
 
-		// Get Account or continue to generate a new one.
+		// Get Account OR continue to generate a new one OR continue Revoke with certificate private key
 		acc, err := h.db.GetAccountByKeyID(ctx, jwk.KeyID)
 		switch {
 		case errors.Is(err, acme.ErrNotFound):
-			// For NewAccount requests ...
+			// For NewAccount and Revoke requests ...
 			break
 		case err != nil:
 			api.WriteError(w, err)
@@ -364,14 +364,15 @@ func (h *Handler) extractOrLookupJWK(next nextHTTP) nextHTTP {
 		}
 
 		// at this point the JWS has already been verified (if correctly configured in middleware),
-		// and it can be used to check if a JWK exists.
+		// and it can be used to check if a JWK exists. This flow is used when the ACME client
+		// signed the payload with a certificate private key.
 		if canExtractJWKFrom(jws) {
 			h.extractJWK(next)(w, r)
 			return
 		}
 
-		// default to looking up the JWK based on KeyID
-		// NOTE: this is a JWK signed with the certificate private key
+		// default to looking up the JWK based on KeyID. This flow is used when the ACME client
+		// signed the payload with an account private key.
 		h.lookupJWK(next)(w, r)
 	}
 }

acme/api/middleware_test.go

@@ -1557,7 +1557,7 @@ func TestHandler_extractOrLookupJWK(t *testing.T) {
 			pub := jwk.Public()
 			pub.KeyID = base64.RawURLEncoding.EncodeToString(kid)
 			so := new(jose.SignerOptions)
-			so.WithHeader("jwk", pub)
+			so.WithHeader("jwk", pub) // JWK for certificate private key flow
 			signer, err := jose.NewSigner(jose.SigningKey{
 				Algorithm: jose.SignatureAlgorithm(jwk.Algorithm),
 				Key:       jwk.Key,
@@ -1569,14 +1569,12 @@ func TestHandler_extractOrLookupJWK(t *testing.T) {
 			assert.FatalError(t, err)
 			parsedJWS, err := jose.ParseJWS(raw)
 			assert.FatalError(t, err)
-			assert.FatalError(t, err)
-			acc := &acme.Account{Status: "valid"}
 			return test{
 				linker: NewLinker("dns", "acme"),
 				db: &acme.MockDB{
 					MockGetAccountByKeyID: func(ctx context.Context, kid string) (*acme.Account, error) {
 						assert.Equals(t, kid, pub.KeyID)
-						return acc, nil
+						return nil, acme.ErrNotFound
 					},
 				},
 				ctx:        context.WithValue(context.Background(), jwsContextKey, parsedJWS),
@@ -1595,7 +1593,7 @@ func TestHandler_extractOrLookupJWK(t *testing.T) {
 			accID := "accID"
 			prefix := fmt.Sprintf("%s/acme/%s/account/", baseURL, provName)
 			so := new(jose.SignerOptions)
-			so.WithHeader("kid", fmt.Sprintf("%s%s", prefix, accID))
+			so.WithHeader("kid", fmt.Sprintf("%s%s", prefix, accID)) // KID for account private key flow
 			signer, err := jose.NewSigner(jose.SigningKey{
 				Algorithm: jose.SignatureAlgorithm(jwk.Algorithm),
 				Key:       jwk.Key,
@@ -1607,8 +1605,6 @@ func TestHandler_extractOrLookupJWK(t *testing.T) {
 			assert.FatalError(t, err)
 			parsedJWS, err := jose.ParseJWS(raw)
 			assert.FatalError(t, err)
-			assert.FatalError(t, err)
-			//acc := &acme.Account{Status: "valid", Key: jwk}
 			acc := &acme.Account{ID: "accID", Key: jwk, Status: "valid"}
 			ctx := context.WithValue(context.Background(), provisionerContextKey, prov)
 			ctx = context.WithValue(ctx, baseURLContextKey, baseURL)