From e489df1cab1a8bcdf31e5edb98036ad6da7235a9 Mon Sep 17 00:00:00 2001 From: Matthias <5011972+fasmat@users.noreply.github.com> Date: Wed, 3 Jul 2024 15:56:47 +0000 Subject: [PATCH 1/3] Fix GenerateProof returning an invalid proof on a short deadline --- prover/prover.go | 45 ++++++++++++++++++++++++++++++++++++++++++--- 1 file changed, 42 insertions(+), 3 deletions(-) diff --git a/prover/prover.go b/prover/prover.go index 461368b5..465aa019 100644 --- a/prover/prover.go +++ b/prover/prover.go @@ -62,7 +62,17 @@ func GenerateProof( } defer treeCache.Close() - return generateProof(ctx, leavesCounter, labelHashFunc, tree, treeCache, deadline, 0, securityParam, persist) + return generateProof( + ctx, + leavesCounter, + labelHashFunc, + tree, + treeCache, + deadline, + 0, + securityParam, + persist, + ) } // GenerateProofRecovery recovers proof generation, from a given 'nextLeafID'. @@ -241,7 +251,7 @@ func makeRecoveryProofTree( defer layerReader.Close() memCachedParkedNodes, readCache, err := recoverMemCachedParkedNodes(layerReader, merkleHashFunc) if err != nil { - return nil, nil, fmt.Errorf("recoveing parked nodes from top layer of disk-cache: %w", err) + return nil, nil, fmt.Errorf("recovering parked nodes from top layer of disk-cache: %w", err) } parkedNodes = append(parkedNodes, memCachedParkedNodes...) @@ -294,6 +304,7 @@ func sequentialWork( treeCache *cache.Writer, end time.Time, nextLeafID uint64, + securityParam uint8, persist persistFunc, ) (uint64, error) { var parkedNodes [][]byte @@ -302,6 +313,8 @@ func sequentialWork( finished := time.NewTimer(time.Until(end)) defer finished.Stop() + stop := make(chan struct{}) + leavesCounter.Add(float64(nextLeafID)) for { @@ -321,6 +334,22 @@ func sequentialWork( } return nextLeafID, ctx.Err() case <-finished.C: + close(stop) + if nextLeafID < uint64(securityParam) { + // we reached deadline, but we didn't generate enough leaves to generate a valid proof, so continue + // generating leaves until we have enough. + continue + } + if err := persist(ctx, treeCache, nextLeafID); err != nil { + return 0, fmt.Errorf("persisting execution state: %w", err) + } + return nextLeafID, nil + case <-stop: + if nextLeafID < uint64(securityParam) { + // we reached deadline, but we didn't generate enough leaves to generate a valid proof, so continue + // generating leaves until we have enough. + continue + } if err := persist(ctx, treeCache, nextLeafID); err != nil { return 0, fmt.Errorf("persisting execution state: %w", err) } @@ -350,7 +379,17 @@ func generateProof( logger := logging.FromContext(ctx) logger.Info("generating proof", zap.Time("end", end), zap.Uint64("nextLeafID", nextLeafID)) - leaves, err := sequentialWork(ctx, leavesCounter, labelHashFunc, tree, treeCache, end, nextLeafID, persist) + leaves, err := sequentialWork( + ctx, + leavesCounter, + labelHashFunc, + tree, + treeCache, + end, + nextLeafID, + securityParam, + persist, + ) if err != nil { return leaves, nil, err } From 06970d3505b9210e34658c6369b87165b623b72c Mon Sep 17 00:00:00 2001 From: Matthias <5011972+fasmat@users.noreply.github.com> Date: Wed, 3 Jul 2024 16:37:12 +0000 Subject: [PATCH 2/3] Update tests --- prover/prover.go | 21 ++++++--------------- prover/prover_test.go | 8 +++----- 2 files changed, 9 insertions(+), 20 deletions(-) diff --git a/prover/prover.go b/prover/prover.go index 465aa019..f55692a9 100644 --- a/prover/prover.go +++ b/prover/prover.go @@ -310,11 +310,13 @@ func sequentialWork( var parkedNodes [][]byte makeLabel := shared.MakeLabelFunc() - finished := time.NewTimer(time.Until(end)) - defer finished.Stop() - stop := make(chan struct{}) - + finished := time.AfterFunc(time.Until(end), func() { + // instead of using a timer that only fires once we want a flag that indicates the timer stopped + // we use a dedicated channel for this purpose that is closed when the timer would fire + close(stop) + }) + defer finished.Stop() leavesCounter.Add(float64(nextLeafID)) for { @@ -333,17 +335,6 @@ func sequentialWork( return 0, fmt.Errorf("persisting execution state: %w", err) } return nextLeafID, ctx.Err() - case <-finished.C: - close(stop) - if nextLeafID < uint64(securityParam) { - // we reached deadline, but we didn't generate enough leaves to generate a valid proof, so continue - // generating leaves until we have enough. - continue - } - if err := persist(ctx, treeCache, nextLeafID); err != nil { - return 0, fmt.Errorf("persisting execution state: %w", err) - } - return nextLeafID, nil case <-stop: if nextLeafID < uint64(securityParam) { // we reached deadline, but we didn't generate enough leaves to generate a valid proof, so continue diff --git a/prover/prover_test.go b/prover/prover_test.go index cd170113..2d1eb4be 100644 --- a/prover/prover_test.go +++ b/prover/prover_test.go @@ -62,11 +62,9 @@ func TestRecoverParkedNodes(t *testing.T) { challenge := []byte("challenge this") leavesCounter := prometheus.NewCounter(prometheus.CounterOpts{}) - ctx, cancel := context.WithDeadline(context.Background(), time.Now().Add(time.Second*2)) + ctx, cancel := context.WithDeadline(context.Background(), time.Now()) defer cancel() - limit := time.Now().Add(time.Second * 5) - persist := func(ctx context.Context, treeCache *cache.Writer, _ uint64) error { // Call GetReader() so that the cache would flush and validate structure. _, err := treeCache.GetReader() @@ -81,7 +79,7 @@ func TestRecoverParkedNodes(t *testing.T) { treeCfg, hash.GenLabelHashFunc(challenge), hash.GenMerkleHashFunc(challenge), - limit, + time.Now().Add(5*time.Second), 150, persist, ) @@ -94,7 +92,7 @@ func TestRecoverParkedNodes(t *testing.T) { treeCfg, hash.GenLabelHashFunc(challenge), hash.GenMerkleHashFunc(challenge), - limit, + time.Now(), 150, leaves, persist, From 50d0d2bd2b911b95dc3ee62b52e6b4d95dcebb1a Mon Sep 17 00:00:00 2001 From: Matthias <5011972+fasmat@users.noreply.github.com> Date: Wed, 3 Jul 2024 16:49:52 +0000 Subject: [PATCH 3/3] Update tests --- poetcore_test.go | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/poetcore_test.go b/poetcore_test.go index 0f069f62..4398b18f 100644 --- a/poetcore_test.go +++ b/poetcore_test.go @@ -31,7 +31,7 @@ func BenchmarkProverAndVerifierBig(b *testing.B) { prover.TreeConfig{Datadir: b.TempDir()}, hash.GenLabelHashFunc(challenge), hash.GenMerkleHashFunc(challenge), - time.Now().Add(time.Second), + time.Now(), securityParam, ) r.NoError(err, "Failed to generate proof") @@ -63,7 +63,7 @@ func TestNip(t *testing.T) { prover.TreeConfig{Datadir: t.TempDir()}, hash.GenLabelHashFunc(challenge), hash.GenMerkleHashFunc(challenge), - time.Now().Add(1*time.Second), + time.Now(), securityParam, ) require.NoError(t, err) @@ -93,7 +93,7 @@ func BenchmarkProofEx(t *testing.B) { prover.TreeConfig{Datadir: t.TempDir()}, hash.GenLabelHashFunc(challenge), hash.GenMerkleHashFunc(challenge), - time.Now().Add(time.Second), + time.Now(), securityParam, ) require.NoError(t, err)