Fix GenerateProof returning an invalid proof on a short deadline

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)

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,14 +304,19 @@ func sequentialWork(
 	treeCache *cache.Writer,
 	end time.Time,
 	nextLeafID uint64,
+	securityParam uint8,
 	persist persistFunc,
 ) (uint64, error) {
 	var parkedNodes [][]byte
 	makeLabel := shared.MakeLabelFunc()
 
-	finished := time.NewTimer(time.Until(end))
+	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 {
@@ -320,7 +335,12 @@ func sequentialWork(
 				return 0, fmt.Errorf("persisting execution state: %w", err)
 			}
 			return nextLeafID, ctx.Err()
-		case <-finished.C:
+		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 +370,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
 	}

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,