From bdec2abc0ddd4396152bcb7bf9f1a9fede901b65 Mon Sep 17 00:00:00 2001
From: DaevMithran <daevmithran1999@gmail.com>
Date: Fri, 11 Oct 2024 16:01:27 +0530
Subject: [PATCH] Add fragment validation & tests

---
 x/did/types/diddoc_diddoc.go      |  2 +-
 x/did/types/diddoc_diddoc_test.go | 32 +++++++++++++++++++++++++++++++
 x/did/types/validate.go           | 24 ++++++++++++++---------
 3 files changed, 48 insertions(+), 10 deletions(-)

diff --git a/x/did/types/diddoc_diddoc.go b/x/did/types/diddoc_diddoc.go
index 74ae9fb7a..ce1ff8202 100644
--- a/x/did/types/diddoc_diddoc.go
+++ b/x/did/types/diddoc_diddoc.go
@@ -96,7 +96,7 @@ func (didDoc DidDoc) Validate(allowedNamespaces []string) error {
 			IsUniqueStrList(), validation.Each(IsDIDUrl(allowedNamespaces, Empty, Empty, Required), HasPrefix(didDoc.Id)),
 		),
 		validation.Field(&didDoc.AssertionMethod,
-			IsUniqueStrList(), validation.Each(IsAssertionMethod(allowedNamespaces, didDoc.Id))),
+			IsUniqueStrList(), validation.Each(IsAssertionMethod(allowedNamespaces, didDoc))),
 		validation.Field(&didDoc.CapabilityInvocation,
 			IsUniqueStrList(), validation.Each(IsDIDUrl(allowedNamespaces, Empty, Empty, Required), HasPrefix(didDoc.Id)),
 		),
diff --git a/x/did/types/diddoc_diddoc_test.go b/x/did/types/diddoc_diddoc_test.go
index 39cbb76cf..2805abc19 100644
--- a/x/did/types/diddoc_diddoc_test.go
+++ b/x/did/types/diddoc_diddoc_test.go
@@ -263,4 +263,36 @@ var _ = DescribeTable("DIDDoc Validation tests", func(testCase DIDDocTestCase) {
 			isValid:  true,
 			errorMsg: "",
 		}),
+	Entry(
+		"Assertion method is has wrong fragment",
+		DIDDocTestCase{
+			didDoc: &DidDoc{
+				Id:         ValidTestDID,
+				Controller: []string{ValidTestDID},
+				VerificationMethod: []*VerificationMethod{
+					{
+						Id:                     fmt.Sprintf("%s#fragment", ValidTestDID),
+						VerificationMethodType: "Ed25519VerificationKey2020",
+						Controller:             ValidTestDID,
+						VerificationMaterial:   ValidEd25519VerificationKey2020VerificationMaterial,
+					},
+				},
+				AssertionMethod: []string{fmt.Sprintf("%s#fragment", ValidTestDID), func() string {
+					b, _ := json.Marshal(struct {
+						Id              string
+						Type            string
+						Controller      string
+						PublicKeyBase58 string
+					}{
+						Id:              fmt.Sprintf("%s#fragment-1", ValidTestDID),
+						Type:            "Ed25519VerificationKey2020",
+						Controller:      ValidTestDID,
+						PublicKeyBase58: "base58",
+					})
+					return strconv.Quote(string(b))
+				}()},
+			},
+			isValid:  false,
+			errorMsg: "assertionMethod should be a valid key reference within the DID document's verification method",
+		}),
 )
diff --git a/x/did/types/validate.go b/x/did/types/validate.go
index cc0beb1fd..dd151f664 100644
--- a/x/did/types/validate.go
+++ b/x/did/types/validate.go
@@ -113,15 +113,15 @@ type AssertionMethod struct {
 	Controller string
 }
 
-func IsAssertionMethod(allowedNamespaces []string, prefix string) *CustomErrorRule {
+func IsAssertionMethod(allowedNamespaces []string, didDoc DidDoc) *CustomErrorRule {
 	return NewCustomErrorRule(func(value interface{}) error {
-		err := validation.Validate(value, IsDIDUrl(allowedNamespaces, Empty, Empty, Required), HasPrefix(prefix))
-		if err != nil {
-			casted, ok := value.(string)
-			if !ok {
-				panic("IsAssertionMethod must be only applied on string properties")
-			}
+		err := validation.Validate(value, IsDIDUrl(allowedNamespaces, Empty, Empty, Required), HasPrefix(didDoc.Id))
+		casted, ok := value.(string)
+		if !ok {
+			panic("IsAssertionMethod must be only applied on string properties")
+		}
 
+		if err != nil {
 			unescapedJSON, err := strconv.Unquote(casted)
 			if err != nil {
 				return errors.New("assertionMethod should be a DIDUrl or an Escaped JSON string")
@@ -134,13 +134,19 @@ func IsAssertionMethod(allowedNamespaces []string, prefix string) *CustomErrorRu
 			}
 
 			return validation.ValidateStruct(&result,
-				validation.Field(&result.Id, validation.Required, IsDIDUrl(allowedNamespaces, Empty, Empty, Required), HasPrefix(prefix)),
+				validation.Field(&result.Id, validation.Required, IsAssertionMethod(allowedNamespaces, didDoc)),
 				validation.Field(&result.Controller, validation.Required, IsDID(allowedNamespaces)),
 				validation.Field(&result.Type, IsURI()),
 			)
 		}
 
-		return nil
+		for _, v := range didDoc.VerificationMethod {
+			if v.Id == casted {
+				return nil
+			}
+		}
+
+		return errors.New("assertionMethod should be a valid key reference within the DID document's verification method")
 	})
 }