From 74322cb13f575992a56e3e5c4c78451032f8fdda Mon Sep 17 00:00:00 2001
From: stenolog <158459002+stenolog@users.noreply.github.com>
Date: Mon, 4 Mar 2024 19:09:42 +0200
Subject: [PATCH] squash

---
 .../nodeView/state/ErgoStateContextSpec.scala | 23 ++------------
 .../utils/generators/ErgoGenerators.scala     | 30 +++++++++++++++++++
 2 files changed, 32 insertions(+), 21 deletions(-)

diff --git a/src/test/scala/org/ergoplatform/nodeView/state/ErgoStateContextSpec.scala b/src/test/scala/org/ergoplatform/nodeView/state/ErgoStateContextSpec.scala
index 9ee5c1cd74..486a8d907b 100644
--- a/src/test/scala/org/ergoplatform/nodeView/state/ErgoStateContextSpec.scala
+++ b/src/test/scala/org/ergoplatform/nodeView/state/ErgoStateContextSpec.scala
@@ -63,17 +63,7 @@ class ErgoStateContextSpec extends HistoryTestHelpers {
     val imvKey = extensionKvGen(Extension.FieldKeySize - 1, Extension.FieldValueMaxSize).sample.get
     sc.appendFullBlock(fbWithFields(imvKey +: oldFields)) shouldBe 'failure
 
-    // https://github.com/ergoplatform/ergo/issues/2114
-    // this fails sporadically, when `(imvValue._1.head == 0)`, because less value bytes will be generated
-    // by extensionKvGen(). Workaround is to just generate again while `(imvValue._1.head == 0)`
-    // TODO: document, and possibly rewrite/replace extensionKvGen after the above issues are clarified/solved
-    //       https://github.com/ergoplatform/ergo/issues/2118
-
-    // validation of field value sizes
-    var imvValue = extensionKvGen(Extension.FieldKeySize, Extension.FieldValueMaxSize + 1).sample.get
-    while (imvValue._1.head == 0) {
-        imvValue = extensionKvGen(Extension.FieldKeySize, Extension.FieldValueMaxSize + 1).sample.get
-    }
+    val imvValue = extensionKvGenImvValue(Extension.FieldKeySize, Extension.FieldValueMaxSize + 1).sample.get
     sc.appendFullBlock(fbWithFields(imvValue +: oldFields)) shouldBe 'failure
 
     // validation of incorrect interlinks
@@ -82,16 +72,7 @@ class ErgoStateContextSpec extends HistoryTestHelpers {
     ).fields
     sc.appendFullBlock(fbWithFields(invalidInterlinks ++ oldFields)) shouldBe 'failure
 
-    // https://github.com/ergoplatform/ergo/issues/2114
-    // if validMKV._1.head is 1, appendFullBlock within "valid application of correct extension" will fail,
-    // because with "key.head == 1", improperly packed interlink would be generated.
-    // As a workaround, just generate new values until (validMKV._1.head != 1)
-    // TODO: investigate and provide a full fix (followup issue)
-    //       https://github.com/ergoplatform/ergo/issues/2117
-    var validMKV = extensionKvGen(Extension.FieldKeySize, Extension.FieldValueMaxSize).sample.get
-    while (validMKV._1.head == 1) {
-      validMKV = extensionKvGen(Extension.FieldKeySize, Extension.FieldValueMaxSize).sample.get
-    }
+    val validMKV = extensionKvGenValidMKV(Extension.FieldKeySize, Extension.FieldValueMaxSize).sample.get
     // validation of key duplicates in fields
     sc.appendFullBlock(fbWithFields(Seq(validMKV, validMKV) ++ oldFields)) shouldBe 'failure
 
diff --git a/src/test/scala/org/ergoplatform/utils/generators/ErgoGenerators.scala b/src/test/scala/org/ergoplatform/utils/generators/ErgoGenerators.scala
index 311ce9d00b..4acb06c64f 100644
--- a/src/test/scala/org/ergoplatform/utils/generators/ErgoGenerators.scala
+++ b/src/test/scala/org/ergoplatform/utils/generators/ErgoGenerators.scala
@@ -92,11 +92,41 @@ trait ErgoGenerators extends CoreGenerators with ChainGenerator with Generators
       scorex.util.Random.randomBytes
     }
 
+  /**
+   * Main KV generator: contains special handling for key.head == 0
+   */
   def extensionKvGen(keySize: Int, valuesSize: Int): Gen[(Array[Byte], Array[Byte])] = for {
     key <- genSecureBoundedBytes(keySize, keySize)
     value <- if (key.head == 0) genSecureBoundedBytes(4, 4) else genSecureBoundedBytes(valuesSize, valuesSize)
   } yield (key, value)
 
+  /**
+   * Random KV generator: all random key/values
+   */
+  def extensionKvGenImvValue(keySize: Int, valuesSize: Int): Gen[(Array[Byte], Array[Byte])] = for {
+    key <- genSecureBoundedBytes(keySize, keySize)
+    value <- genSecureBoundedBytes(valuesSize, valuesSize)
+  } yield (key, value)
+
+  /**
+   * Special KV generator: does not generate key.head == 1, because improperly packed interlink
+   * would be generated, leading to failure.
+   */
+  def extensionKvGenValidMKV(keySize: Int, valueSize: Int): Gen[(Array[Byte], Array[Byte])] = {
+    def genSecureBoundedBytesWrapper(minSize: Int, maxSize: Int): Gen[Array[Byte]]= {
+      genSecureBoundedBytes(minSize, maxSize)
+    }
+
+    def keyGen: Gen[Array[Byte]] = genSecureBoundedBytesWrapper(keySize, keySize).flatMap { key =>
+      if (key.headOption.contains(1.toByte)) keyGen else Gen.const(key)
+    }
+
+    for {
+      key <- keyGen
+      value <- genSecureBoundedBytesWrapper(valueSize, valueSize)
+    } yield (key, value)
+  }
+
   lazy val extensionGen: Gen[Extension] = for {
     headerId <- modifierIdGen
     mandatoryElements <- Gen.mapOf(extensionKvGen(Extension.FieldKeySize, Extension.FieldValueMaxSize))