simulation_test.go

package nmt_test

import (
	"crypto/sha256"
	"os"
	"testing"

	"github.com/celestiaorg/nmt"
	"github.com/celestiaorg/nmt/namespace"
	"github.com/tidwall/gjson"
)

const nidSize = 2

func getLeavesFromState(state gjson.Result, nidSize int, key string) [][]byte {
	itfLeaves := state.Get(key).Array()

	leaves := make([][]byte, len(itfLeaves))

	for idx, itfLeaf := range itfLeaves {
		stateNamespace := int(itfLeaf.Get("namespaceId").Int())
		dataString := itfLeaf.Get("value.#tup.1").String()
		pushData := namespace.PrefixedData(append(
			intToBytes(stateNamespace, nidSize),
			[]byte(dataString)...))
		leaves[idx] = pushData

	}
	return leaves
}

// checking if the proof generated by the execution matches the proof in the model
// hashes are not checked because the model is using a simplified hash function
func checkProof(modelProof gjson.Result, execProof nmt.Proof, nidSize int) bool {
	modelProofStart := int(modelProof.Get("start").Int())
	execProofStart := execProof.Start()

	modelProofEnd := int(modelProof.Get("end").Int())
	execProofEnd := execProof.End()
	if modelProofStart != execProofStart || modelProofEnd != execProofEnd {
		return false
	}
	supportingHashes := modelProof.Get("supporting_hashes").Array()
	if len(supportingHashes) != len(execProof.Nodes()) {
		return false
	}
	for idx, supportingHash := range supportingHashes {
		modelMinNS := int(supportingHash.Get("minNS").Int())
		execMinNs := bytesToInt(execProof.Nodes()[idx][:nidSize])
		if modelMinNS != execMinNs {
			return false
		}

		modelMaxNS := int(supportingHash.Get("maxNS").Int())
		execMaxNs := bytesToInt(execProof.Nodes()[idx][nidSize:(2 * nidSize)])
		if modelMaxNS != execMaxNs {
			return false
		}
	}

	return true
}

func intToBytes(n int, bytesSize int) []byte {
	b := make([]byte, bytesSize)
	for i := 0; i < bytesSize; i++ {
		b[i] = byte(n >> (8 * (bytesSize - i - 1)))
	}
	return b
}

func bytesToInt(bytes []byte) int {
	var result int
	for i, b := range bytes {
		shift := uint((len(bytes) - 1 - i) * 8)
		result |= int(b) << shift
	}
	return result
}

func TestFromITF(t *testing.T) {
	var tree *nmt.NamespacedMerkleTree

	itfFileName := "nmt_quint_model/ITF_files/runTest.itf.json"
	// itfFileName := "nmt_quint_model/ITF_files/panicIssue.itf.json"
	data, err := os.ReadFile(itfFileName)
	if err != nil {
		t.Errorf("Error opening file: %v", err)
	}
	states := gjson.GetBytes(data, "states").Array()

	// iterate over all states of the test run
	for stateIdx, state := range states {
		t.Logf("State %v\n", stateIdx)
		// because the test run consists of multiple steps to generate the tree and the proof, test
		// only the states marked "final"
		if state.Get("state_v").String() == "final" {
			modelNamespace := int(state.Get("namespace_v").Int())
			modelLeaves := getLeavesFromState(state, nidSize, "leaves_v")
			modelProof := state.Get("proof_v")
			corrupted := state.Get("corrupted").Bool()
			t.Logf("Obtained state data:\n\tState namespace: %v\n\tleaves: %v\n\tcorrupted: %v\n",
				modelNamespace, modelLeaves, corrupted)

			tree = nmt.New(sha256.New(), nmt.NamespaceIDSize(nidSize))
			for _, leaf := range modelLeaves {
				err := tree.Push(leaf)
				if err != nil {
					t.Errorf("Error on push: %v", err)
				}
			}

			namespaceBytes := intToBytes(modelNamespace, nidSize)

			proof, err := tree.ProveNamespace(namespaceBytes)
			if err != nil {
				t.Errorf("Error on prove: %v", err)
			}
			var newHashes [][]byte
			t.Logf("Proof: %v\n", proof)

			var corruptionType string
			if !corrupted {
				proofMatching := checkProof(modelProof, proof, nidSize)
				if !proofMatching {
					t.Errorf("Proofs do not match. Expected: %v, got: %v", state, proof)
				}
			} else {
				// CORRUPTED: now we are modifying the proof to be corrupted here as well
				corruptionType = state.Get("corruption_type").String()
				t.Log("Corruption type: ", corruptionType)
				if corruptionType == "supporting_hashes" {
					for idx := 0; idx < len(state.Get("corruption_diff.changed_indices").Array()); idx++ {
						newHashes = append(newHashes, proof.Nodes()[idx])
					}
				} else {
					newHashes = proof.Nodes()
				}
				// We could examine for different corruption types and then only user `modelProof` for
				// the corrupted parts of the proof and the rest from the `proof`. The way we do it now,
				// however, is not a problem since for non-corrupted parts, `modelProof` and `proof` are
				// the same.
				proof = nmt.NewInclusionProof(
					int(modelProof.Get("start").Int()),
					int(modelProof.Get("end").Int()),
					newHashes,
					true,
				)
				t.Logf("Corrupted proof: %v\n", proof)

			}

			root, _ := tree.Root()

			dataToVerify := modelLeaves[proof.Start():proof.End()]

			t.Logf("Verifying. Root: %v, Namespace: %x, Leaves: %v, Proof: %v\n", root, namespaceBytes, dataToVerify, proof)

			successVerification := proof.VerifyNamespace(sha256.New(), namespaceBytes, dataToVerify, root)
			t.Logf("Verification result: %v\n", successVerification)

			if successVerification == corrupted {
				successVerification := proof.VerifyNamespace(sha256.New(), namespaceBytes, dataToVerify, root)
				t.Errorf("Mismatch! Data corrupted = %v, verification result = %v", corrupted, successVerification)
			}

		} else { // these are the intermediate states, at which we are not checking anything
			continue
		}

	}
}

func TestFromScenario(t *testing.T) {
	itfFileName := "nmt_quint_model/ITF_files/scenario.json"
	data, err := os.ReadFile(itfFileName)
	if err != nil {
		t.Errorf("Error opening file: %v", err)
	}
	tests := gjson.GetBytes(data, "testCases").Array()

	for _, tt := range tests {
		states := tt.Get("states").Array()
		lastState := states[len(states)-1]
		namespaceID := int(lastState.Get("namespace_v").Int())

		itfLeaves := lastState.Get("leaves_namespace_idx_v").Array()

		tree := nmt.New(sha256.New(), nmt.NamespaceIDSize(nidSize))
		for _, itfLeaf := range itfLeaves {
			leafID := int(itfLeaf.Int())
			pushData := namespace.PrefixedData(append(
				intToBytes(leafID, nidSize),
				[]byte("dummy_data")...))
			err := tree.Push(pushData)
			if err != nil {
				t.Fatalf("invalid test case: %v, error on Push(): %v", tt.Get("name").String(), err)
			}
		}
		proveNID := intToBytes(namespaceID, nidSize)
		gotProof, err := tree.ProveNamespace(proveNID)
		if err != nil {
			t.Fatalf("ProveNamespace() unexpected error: %v", err)
		}

		gotFound := gotProof.IsNonEmptyRange() && len(gotProof.LeafHash()) == 0
		wantFound := lastState.Get("namespaceIdFound_v").Bool()

		if gotFound != wantFound {
			t.Errorf("Proof.ProveNamespace() gotFound = %v, wantFound = %v ", gotFound, wantFound)
		}

		// Verification round-trip should always pass:
		gotGetLeaves := tree.Get(proveNID)
		r, err := tree.Root()
		if err != nil {
			t.Fatalf("Root() unexpected error: %v", err)
		}
		gotChecksOut := gotProof.VerifyNamespace(sha256.New(), proveNID, gotGetLeaves, r)
		if !gotChecksOut {
			t.Errorf("Proof.VerifyNamespace() gotChecksOut: %v, want: true", gotChecksOut)
		}
	}
}