Merge pull request #1896 from lightninglabs/wip/re-org/ReplaceProofInFiles-refactor-test-coverage

proof: improve re-org proof replacement, strengthen error handling, expand test coverage

Author: ffranr

proof/archive.go

@@ -45,8 +45,7 @@ var (
 	// is valid.
 	emptyKey btcec.PublicKey
 
-	// ErrProofNotFound is returned when a user attempts to look up a proof
-	// based on a Locator, but we can't find it on disk.
+	// ErrProofNotFound reports that no proof satisfies the lookup criteria.
 	ErrProofNotFound = fmt.Errorf("unable to find proof")
 
 	// ErrInvalidLocatorID is returned when a specified has an invalid
@@ -1090,61 +1089,62 @@ func (m *MultiArchiver) RemoveSubscriber(
 // NotifyArchiver interface.
 var _ NotifyArchiver = (*MultiArchiver)(nil)
 
-// ReplaceProofInBlob attempts to replace a proof in all proof files we have for
-// assets of the same ID. This is useful when we want to update the proof with a
-// new one after a re-org.
-func ReplaceProofInBlob(ctx context.Context, p *Proof, archive Archiver,
-	vCtx VerifierCtx) error {
-
-	// This is a bit of a hacky part. If we have a chain of transactions
-	// that were re-organized, we can't verify the whole chain until all of
-	// the transactions were confirmed and all proofs were updated with the
-	// new blocks and merkle roots. So we'll skip the verification here
-	// since we don't know if the whole chain has been updated yet (the
-	// confirmations might come in out of order).
-	// TODO(guggero): Find a better way to do this.
-	vCtx.HeaderVerifier = func(wire.BlockHeader, uint32) error {
-		return nil
-	}
+// ReplaceProofInFiles attempts to replace a proof in all provided proof files
+// for assets of the same ID and returns the updated proof blobs. This is useful
+// when we want to update the proof with a new one after a re-org.
+func ReplaceProofInFiles(newProof *Proof,
+	proofFiles []*AnnotatedProof) ([]*AnnotatedProof, error) {
 
-	assetID := p.Asset.ID()
-	scriptPubKeyOfUpdate := p.Asset.ScriptKey.PubKey
+	var (
+		updatedProofs []*AnnotatedProof
 
-	// We now fetch all proofs of that same asset ID and filter out those
-	// that need updating.
-	proofs, err := archive.FetchProofs(ctx, assetID)
-	if err != nil {
-		return fmt.Errorf("unable to fetch all proofs for asset ID "+
-			"%x: %w", assetID[:], err)
-	}
+		scriptPubKeyOfUpdate = newProof.Asset.ScriptKey.PubKey
+		targetAssetID        = newProof.Asset.ID()
+	)
 
-	for idx := range proofs {
-		existingProof := proofs[idx]
+	for idx := range proofFiles {
+		proofFile := proofFiles[idx]
+
+		// Sanity check: make sure the asset ID of the proof file
+		// matches the target asset ID.
+		if *proofFile.Locator.AssetID != targetAssetID {
+			return nil, fmt.Errorf("mismatched asset ID: "+
+				"expected %s, got %s",
+				proofFile.Locator.AssetID.String(),
+				targetAssetID.String())
+		}
 
 		f := &File{}
-		err := f.Decode(bytes.NewReader(existingProof.Blob))
+		err := f.Decode(bytes.NewReader(proofFile.Blob))
 		if err != nil {
-			return fmt.Errorf("unable to decode current proof: %w",
-				err)
+			return nil, fmt.Errorf("unable to decode current "+
+				"proof: %w", err)
 		}
 
 		// We only need to update proofs that contain this asset in the
 		// chain and haven't been updated yet (i.e. the block hash of
 		// the proof is different from the block hash of the proof we
 		// want to update).
 		_, indexToUpdate, err := f.LocateProof(func(fp *Proof) bool {
-			fileScriptKey := fp.Asset.ScriptKey.PubKey
-			fileTxHash := fp.AnchorTx.TxHash()
-			fileBlockHash := fp.BlockHeader.BlockHash()
-			return fileScriptKey.IsEqual(scriptPubKeyOfUpdate) &&
-				fileTxHash == p.AnchorTx.TxHash() &&
-				fileBlockHash != p.BlockHeader.BlockHash()
+			fScriptKey := fp.Asset.ScriptKey.PubKey
+			fTxHash := fp.AnchorTx.TxHash()
+			fBlockHash := fp.BlockHeader.BlockHash()
+
+			return fScriptKey.IsEqual(scriptPubKeyOfUpdate) &&
+				fTxHash == newProof.AnchorTx.TxHash() &&
+				fBlockHash != newProof.BlockHeader.BlockHash()
 		})
-		if err != nil {
-			// Either we failed to decode the proof for some reason,
-			// or we didn't find a proof that needs updating. In
-			// either case, we can skip this file.
+		switch {
+		// This proof file doesn't contain a proof that needs to be
+		// updated, so we can skip it.
+		case errors.Is(err, ErrProofNotFound):
 			continue
+
+		// Any other error during proof location is fatal and should be
+		// returned.
+		case err != nil:
+			return nil, fmt.Errorf("unable to locate "+
+				"proof to update: %w", err)
 		}
 
 		log.Debugf("Updating descendant proof at index %d "+
@@ -1154,29 +1154,25 @@ func ReplaceProofInBlob(ctx context.Context, p *Proof, archive Archiver,
 
 		// All good, we can now replace the proof in the file with the
 		// new one.
-		err = f.ReplaceProofAt(indexToUpdate, *p)
+		err = f.ReplaceProofAt(indexToUpdate, *newProof)
 		if err != nil {
-			return fmt.Errorf("unable to replace proof at index "+
-				"%d with updated one: %w", indexToUpdate, err)
+			return nil, fmt.Errorf("unable to replace proof at "+
+				"index %d with updated one: %w", indexToUpdate,
+				err)
 		}
 
 		var buf bytes.Buffer
 		if err := f.Encode(&buf); err != nil {
-			return fmt.Errorf("unable to encode updated proof: %w",
-				err)
+			return nil, fmt.Errorf("unable to encode updated "+
+				"proof: %w", err)
 		}
 
-		// We now update this direct proof in the archive.
-		directProof := &AnnotatedProof{
-			Locator: existingProof.Locator,
+		updatedProof := &AnnotatedProof{
+			Locator: proofFile.Locator,
 			Blob:    buf.Bytes(),
 		}
-		err = archive.ImportProofs(ctx, vCtx, true, directProof)
-		if err != nil {
-			return fmt.Errorf("unable to import updated proof: %w",
-				err)
-		}
+		updatedProofs = append(updatedProofs, updatedProof)
 	}
 
-	return nil
+	return updatedProofs, nil
 }

proof/archive_test.go

@@ -333,3 +333,129 @@ func TestMigrateOldFileNames(t *testing.T) {
 	require.NoError(t, err)
 	assertProofAtNewName(proof6)
 }
+
+// TestReplaceProofInFilesUpdatesProof ensures that when a matching proof is
+// found, it is replaced with the updated version and returned.
+func TestReplaceProofInFilesUpdatesProof(t *testing.T) {
+	t.Parallel()
+
+	encodeProof := func(p Proof) []byte {
+		file, err := NewFile(V0, p)
+		require.NoError(t, err)
+
+		var buf bytes.Buffer
+		require.NoError(t, file.Encode(&buf))
+
+		return buf.Bytes()
+	}
+
+	testBlocks := readTestData(t)
+
+	genesis := asset.RandGenesis(t, asset.Collectible)
+	scriptKey := test.RandPubKey(t)
+
+	oldProof := RandProof(t, genesis, scriptKey, testBlocks[0], 0, 1)
+	locator := Locator{
+		AssetID:   fn.Ptr(oldProof.Asset.ID()),
+		ScriptKey: *oldProof.Asset.ScriptKey.PubKey,
+		OutPoint:  fn.Ptr(oldProof.OutPoint()),
+	}
+	oldProofAnnotated := &AnnotatedProof{
+		Locator: locator,
+		Blob:    encodeProof(oldProof),
+	}
+
+	newProof := oldProof
+	newProof.BlockHeader = testBlocks[1].Header
+	newProof.BlockHeight = oldProof.BlockHeight + 1
+
+	updated, err := ReplaceProofInFiles(
+		&newProof, []*AnnotatedProof{oldProofAnnotated},
+	)
+	require.NoError(t, err)
+	require.Len(t, updated, 1)
+	require.Equal(t, locator, updated[0].Locator)
+
+	var updatedFile File
+	err = updatedFile.Decode(bytes.NewReader(updated[0].Blob))
+	require.NoError(t, err)
+
+	updatedProof, err := updatedFile.ProofAt(0)
+	require.NoError(t, err)
+
+	require.Equal(
+		t, newProof.BlockHeader.BlockHash(),
+		updatedProof.BlockHeader.BlockHash(),
+	)
+	require.Equal(
+		t, newProof.AnchorTx.TxHash(), updatedProof.AnchorTx.TxHash(),
+	)
+}
+
+// TestReplaceProofInFilesSkipsNonMatchingProof ensures that proofs that do not
+// match the update predicate are left untouched and not returned.
+func TestReplaceProofInFilesSkipsNonMatchingProof(t *testing.T) {
+	t.Parallel()
+
+	testBlocks := readTestData(t)
+	genesis := asset.RandGenesis(t, asset.Collectible)
+	scriptKey := test.RandPubKey(t)
+
+	oldProof := RandProof(t, genesis, scriptKey, testBlocks[0], 0, 1)
+	newProof := RandProof(t, genesis, scriptKey, testBlocks[0], 1, 1)
+
+	oldProofFile, err := NewFile(V0, oldProof)
+	require.NoError(t, err)
+
+	var buf bytes.Buffer
+	require.NoError(t, oldProofFile.Encode(&buf))
+
+	locator := Locator{
+		AssetID:   fn.Ptr(oldProof.Asset.ID()),
+		ScriptKey: *oldProof.Asset.ScriptKey.PubKey,
+		OutPoint:  fn.Ptr(oldProof.OutPoint()),
+	}
+	oldProofAnnotated := &AnnotatedProof{
+		Locator: locator,
+		Blob:    buf.Bytes(),
+	}
+
+	updated, err := ReplaceProofInFiles(
+		&newProof, []*AnnotatedProof{oldProofAnnotated},
+	)
+	require.NoError(t, err)
+	require.Empty(t, updated)
+}
+
+// TestReplaceProofInFilesMismatchedAsset ensures we error if a proof file has
+// a different asset ID than the proof being updated.
+func TestReplaceProofInFilesMismatchedAsset(t *testing.T) {
+	t.Parallel()
+
+	testBlocks := readTestData(t)
+	genesis := asset.RandGenesis(t, asset.Collectible)
+	scriptKey := test.RandPubKey(t)
+
+	proof := RandProof(t, genesis, scriptKey, testBlocks[0], 0, 1)
+
+	file, err := NewFile(V0, proof)
+	require.NoError(t, err)
+
+	var buf bytes.Buffer
+	require.NoError(t, file.Encode(&buf))
+
+	mismatchedLocator := Locator{
+		AssetID:   fn.Ptr(asset.RandGenesis(t, asset.Collectible).ID()),
+		ScriptKey: *proof.Asset.ScriptKey.PubKey,
+		OutPoint:  fn.Ptr(proof.OutPoint()),
+	}
+
+	_, err = ReplaceProofInFiles(
+		&proof, []*AnnotatedProof{{
+			Locator: mismatchedLocator,
+			Blob:    buf.Bytes(),
+		}},
+	)
+	require.Error(t, err)
+	require.Contains(t, err.Error(), "mismatched asset ID")
+}