From a9f8090bb455c95f038fa5969d5713dd403431af Mon Sep 17 00:00:00 2001
From: Ben Bierens <39762930+benbierens@users.noreply.github.com>
Date: Mon, 11 Dec 2023 11:29:15 +0100
Subject: [PATCH] Add challenge to prove callback (#649)

* Sets up passing proof challenge to onProve callback

* Implements tests for proof challenge

* Removes unused import
---
 codex/contracts/market.nim                    |  3 +++
 codex/market.nim                              |  4 ++++
 codex/node.nim                                |  2 +-
 codex/sales/salescontext.nim                  |  2 +-
 codex/sales/states/initialproving.nim         |  5 ++++-
 codex/sales/states/proving.nim                |  8 ++++---
 codex/sales/states/provingsimulated.nim       |  4 ++--
 tests/codex/examples.nim                      |  1 -
 tests/codex/helpers/mockmarket.nim            |  5 ++++-
 .../codex/sales/states/testinitialproving.nim | 22 +++++++++++++++----
 tests/codex/sales/states/testproving.nim      | 12 +++++++++-
 .../sales/states/testsimulatedproving.nim     |  2 +-
 tests/codex/sales/testsales.nim               | 10 ++++-----
 13 files changed, 59 insertions(+), 21 deletions(-)

diff --git a/codex/contracts/market.nim b/codex/contracts/market.nim
index c8ed9ffb..413262cc 100644
--- a/codex/contracts/market.nim
+++ b/codex/contracts/market.nim
@@ -143,6 +143,9 @@ method willProofBeRequired*(market: OnChainMarket,
       return false
     raise e
 
+method getChallenge*(market: OnChainMarket, id: SlotId): Future[ProofChallenge] {.async.} =
+  return await market.contract.getChallenge(id)
+
 method submitProof*(market: OnChainMarket,
                     id: SlotId,
                     proof: seq[byte]) {.async.} =
diff --git a/codex/market.nim b/codex/market.nim
index be0d06fc..b9a7651a 100644
--- a/codex/market.nim
+++ b/codex/market.nim
@@ -28,6 +28,7 @@ type
     requestId*: RequestId
     ask*: StorageAsk
     expiry*: UInt256
+  ProofChallenge* = array[32, byte]
 
 method getSigner*(market: Market): Future[Address] {.base, async.} =
   raiseAssert("not implemented")
@@ -103,6 +104,9 @@ method willProofBeRequired*(market: Market,
                             id: SlotId): Future[bool] {.base, async.} =
   raiseAssert("not implemented")
 
+method getChallenge*(market: Market, id: SlotId): Future[ProofChallenge] {.base, async.} =
+  raiseAssert("not implemented")
+
 method submitProof*(market: Market,
                     id: SlotId,
                     proof: seq[byte]) {.base, async.} =
diff --git a/codex/node.nim b/codex/node.nim
index cd1af002..60606c36 100644
--- a/codex/node.nim
+++ b/codex/node.nim
@@ -458,7 +458,7 @@ proc start*(node: CodexNodeRef) {.async.} =
       # TODO: remove data from local storage
       discard
 
-    hostContracts.sales.onProve = proc(slot: Slot): Future[seq[byte]] {.async.} =
+    hostContracts.sales.onProve = proc(slot: Slot, challenge: ProofChallenge): Future[seq[byte]] {.async.} =
       # TODO: generate proof
       return @[42'u8]
 
diff --git a/codex/sales/salescontext.nim b/codex/sales/salescontext.nim
index a7924ce0..4afa5f23 100644
--- a/codex/sales/salescontext.nim
+++ b/codex/sales/salescontext.nim
@@ -25,7 +25,7 @@ type
   OnStore* = proc(request: StorageRequest,
                   slot: UInt256,
                   onBatch: BatchProc): Future[?!void] {.gcsafe, upraises: [].}
-  OnProve* = proc(slot: Slot): Future[seq[byte]] {.gcsafe, upraises: [].}
+  OnProve* = proc(slot: Slot, challenge: ProofChallenge): Future[seq[byte]] {.gcsafe, upraises: [].}
   OnExpiryUpdate* = proc(rootCid: string, expiry: SecondsSince1970): Future[?!void] {.gcsafe, upraises: [].}
   OnClear* = proc(request: StorageRequest,
                   slotIndex: UInt256) {.gcsafe, upraises: [].}
diff --git a/codex/sales/states/initialproving.nim b/codex/sales/states/initialproving.nim
index 0cbbb7b7..d44b058b 100644
--- a/codex/sales/states/initialproving.nim
+++ b/codex/sales/states/initialproving.nim
@@ -31,7 +31,10 @@ method run*(state: SaleInitialProving, machine: Machine): Future[?State] {.async
     raiseAssert "onProve callback not set"
 
   debug "Generating initial proof", requestId = $data.requestId
-  let proof = await onProve(Slot(request: request, slotIndex: data.slotIndex))
+  let
+    slot = Slot(request: request, slotIndex: data.slotIndex)
+    challenge = await context.market.getChallenge(slot.id)
+    proof = await onProve(slot, challenge)
   debug "Finished proof calculation", requestId = $data.requestId
 
   return some State(SaleFilling(proof: proof))
diff --git a/codex/sales/states/proving.nim b/codex/sales/states/proving.nim
index e4059f6d..367eaa9e 100644
--- a/codex/sales/states/proving.nim
+++ b/codex/sales/states/proving.nim
@@ -21,12 +21,13 @@ type
 method prove*(
   state: SaleProving,
   slot: Slot,
+  challenge: ProofChallenge,
   onProve: OnProve,
   market: Market,
   currentPeriod: Period
 ) {.base, async.} =
   try:
-    let proof = await onProve(slot)
+    let proof = await onProve(slot, challenge)
     debug "Submitting proof", currentPeriod = currentPeriod, slotId = $slot.id
     await market.submitProof(slot.id, proof)
   except CatchableError as e:
@@ -71,8 +72,9 @@ proc proveLoop(
     debug "Proving for new period", period = currentPeriod
 
     if (await market.isProofRequired(slotId)) or (await market.willProofBeRequired(slotId)):
-      debug "Proof is required", period = currentPeriod
-      await state.prove(slot, onProve, market, currentPeriod)
+      let challenge = await market.getChallenge(slotId)
+      debug "Proof is required", period = currentPeriod, challenge = challenge
+      await state.prove(slot, challenge, onProve, market, currentPeriod)
 
     await waitUntilPeriod(currentPeriod + 1)
 
diff --git a/codex/sales/states/provingsimulated.nim b/codex/sales/states/provingsimulated.nim
index 4ed9bba1..e3f5b2c2 100644
--- a/codex/sales/states/provingsimulated.nim
+++ b/codex/sales/states/provingsimulated.nim
@@ -22,7 +22,7 @@ when codex_enable_proof_failures:
   proc onSubmitProofError(error: ref CatchableError, period: UInt256, slotId: SlotId) =
     error "Submitting invalid proof failed", period = period, slotId = $slotId, msg = error.msg
 
-  method prove*(state: SaleProvingSimulated, slot: Slot, onProve: OnProve, market: Market, currentPeriod: Period) {.async.} =
+  method prove*(state: SaleProvingSimulated, slot: Slot, challenge: ProofChallenge, onProve: OnProve, market: Market, currentPeriod: Period) {.async.} =
     trace "Processing proving in simulated mode"
     state.proofCount += 1
     if state.failEveryNProofs > 0 and
@@ -38,4 +38,4 @@ when codex_enable_proof_failures:
       except CatchableError as e:
         onSubmitProofError(e, currentPeriod, slot.id)
     else:
-      await procCall SaleProving(state).prove(slot, onProve, market, currentPeriod)
+      await procCall SaleProving(state).prove(slot, challenge, onProve, market, currentPeriod)
diff --git a/tests/codex/examples.nim b/tests/codex/examples.nim
index 87397464..9b75e239 100644
--- a/tests/codex/examples.nim
+++ b/tests/codex/examples.nim
@@ -71,4 +71,3 @@ proc example*(_: type Reservation): Reservation =
     size = uint16.example.u256,
     slotId = SlotId.example
   )
-
diff --git a/tests/codex/helpers/mockmarket.nim b/tests/codex/helpers/mockmarket.nim
index 867a3ef5..22d62e1d 100644
--- a/tests/codex/helpers/mockmarket.nim
+++ b/tests/codex/helpers/mockmarket.nim
@@ -29,6 +29,7 @@ type
     withdrawn*: seq[RequestId]
     proofsRequired: HashSet[SlotId]
     proofsToBeRequired: HashSet[SlotId]
+    proofChallenge*: ProofChallenge
     proofEnds: Table[SlotId, UInt256]
     signer: Address
     subscriptions: Subscriptions
@@ -235,7 +236,6 @@ method freeSlot*(market: MockMarket, slotId: SlotId) {.async.} =
       break
   market.slotState[slotId] = SlotState.Free
 
-
 method withdrawFunds*(market: MockMarket,
                       requestId: RequestId) {.async.} =
   market.withdrawn.add(requestId)
@@ -261,6 +261,9 @@ method willProofBeRequired*(mock: MockMarket,
                             id: SlotId): Future[bool] {.async.} =
   return mock.proofsToBeRequired.contains(id)
 
+method getChallenge*(mock: MockMarket, id: SlotId): Future[ProofChallenge] {.async.} =
+  return mock.proofChallenge
+
 proc setProofEnd*(mock: MockMarket, id: SlotId, proofEnd: UInt256) =
   mock.proofEnds[id] = proofEnd
 
diff --git a/tests/codex/sales/states/testinitialproving.nim b/tests/codex/sales/states/testinitialproving.nim
index f6d8df4b..dcc4f406 100644
--- a/tests/codex/sales/states/testinitialproving.nim
+++ b/tests/codex/sales/states/testinitialproving.nim
@@ -8,22 +8,29 @@ import pkg/codex/sales/states/failed
 import pkg/codex/sales/states/filling
 import pkg/codex/sales/salesagent
 import pkg/codex/sales/salescontext
+import pkg/codex/market
 import ../../examples
 import ../../helpers
+import ../../helpers/mockmarket
 
 asyncchecksuite "sales state 'initialproving'":
-
   let proof = exampleProof()
   let request = StorageRequest.example
   let slotIndex = (request.ask.slots div 2).u256
+  let market = MockMarket.new()
 
   var state: SaleInitialProving
   var agent: SalesAgent
+  var receivedChallenge: ProofChallenge
 
   setup:
-    let onProve = proc (slot: Slot): Future[seq[byte]] {.async.} =
-                            return proof
-    let context = SalesContext(onProve: onProve.some)
+    let onProve = proc (slot: Slot, challenge: ProofChallenge): Future[seq[byte]] {.async.} =
+                          receivedChallenge = challenge
+                          return proof
+    let context = SalesContext(
+      onProve: onProve.some,
+      market: market
+    )
     agent = newSalesAgent(context,
                           request.id,
                           slotIndex,
@@ -42,3 +49,10 @@ asyncchecksuite "sales state 'initialproving'":
     let next = await state.run(agent)
     check !next of SaleFilling
     check SaleFilling(!next).proof == proof
+
+  test "onProve callback provides proof challenge":
+    market.proofChallenge = ProofChallenge.example
+
+    let future = state.run(agent)
+
+    check receivedChallenge == market.proofChallenge
diff --git a/tests/codex/sales/states/testproving.nim b/tests/codex/sales/states/testproving.nim
index fc3b4b08..f6473882 100644
--- a/tests/codex/sales/states/testproving.nim
+++ b/tests/codex/sales/states/testproving.nim
@@ -23,11 +23,13 @@ asyncchecksuite "sales state 'proving'":
   var clock: MockClock
   var agent: SalesAgent
   var state: SaleProving
+  var receivedChallenge: ProofChallenge
 
   setup:
     clock = MockClock.new()
     market = MockMarket.new()
-    let onProve = proc (slot: Slot): Future[seq[byte]] {.async.} =
+    let onProve = proc (slot: Slot, challenge: ProofChallenge): Future[seq[byte]] {.async.} =
+                        receivedChallenge = challenge
                         return proof
     let context = SalesContext(market: market, clock: clock, onProve: onProve.some)
     agent = newSalesAgent(context,
@@ -80,3 +82,11 @@ asyncchecksuite "sales state 'proving'":
     check eventually future.finished
     check !(future.read()) of SalePayout
 
+  test "onProve callback provides proof challenge":
+    market.proofChallenge = ProofChallenge.example
+    market.slotState[slot.id] = SlotState.Filled
+    market.setProofRequired(slot.id, true)
+
+    let future = state.run(agent)
+
+    check receivedChallenge == market.proofChallenge
diff --git a/tests/codex/sales/states/testsimulatedproving.nim b/tests/codex/sales/states/testsimulatedproving.nim
index ddd509ee..359af769 100644
--- a/tests/codex/sales/states/testsimulatedproving.nim
+++ b/tests/codex/sales/states/testsimulatedproving.nim
@@ -44,7 +44,7 @@ asyncchecksuite "sales state 'simulated-proving'":
     market.setProofRequired(slot.id, true)
     subscription = await market.subscribeProofSubmission(onProofSubmission)
 
-    let onProve = proc (slot: Slot): Future[seq[byte]] {.async.} =
+    let onProve = proc (slot: Slot, challenge: ProofChallenge): Future[seq[byte]] {.async.} =
                         return proof
     let context = SalesContext(market: market, clock: clock, onProve: onProve.some)
     agent = newSalesAgent(context,
diff --git a/tests/codex/sales/testsales.nim b/tests/codex/sales/testsales.nim
index c49dc613..7b47f375 100644
--- a/tests/codex/sales/testsales.nim
+++ b/tests/codex/sales/testsales.nim
@@ -64,7 +64,7 @@ asyncchecksuite "Sales - start":
       return success()
 
     queue = sales.context.slotQueue
-    sales.onProve = proc(slot: Slot): Future[seq[byte]] {.async.} =
+    sales.onProve = proc(slot: Slot, challenge: ProofChallenge): Future[seq[byte]] {.async.} =
       return proof
     itemsProcessed = @[]
     request.expiry = (clock.now() + 42).u256
@@ -166,7 +166,7 @@ asyncchecksuite "Sales":
       return success()
 
     queue = sales.context.slotQueue
-    sales.onProve = proc(slot: Slot): Future[seq[byte]] {.async.} =
+    sales.onProve = proc(slot: Slot, challenge: ProofChallenge): Future[seq[byte]] {.async.} =
       return proof
     await sales.start()
     request.expiry = (clock.now() + 42).u256
@@ -369,7 +369,7 @@ asyncchecksuite "Sales":
 
   test "handles errors during state run":
     var saleFailed = false
-    sales.onProve = proc(slot: Slot): Future[seq[byte]] {.async.} =
+    sales.onProve = proc(slot: Slot, challenge: ProofChallenge): Future[seq[byte]] {.async.} =
       # raise exception so machine.onError is called
       raise newException(ValueError, "some error")
 
@@ -394,7 +394,7 @@ asyncchecksuite "Sales":
   test "generates proof of storage":
     var provingRequest: StorageRequest
     var provingSlot: UInt256
-    sales.onProve = proc(slot: Slot): Future[seq[byte]] {.async.} =
+    sales.onProve = proc(slot: Slot, challenge: ProofChallenge): Future[seq[byte]] {.async.} =
       provingRequest = slot.request
       provingSlot = slot.slotIndex
     createAvailability()
@@ -426,7 +426,7 @@ asyncchecksuite "Sales":
   test "calls onClear when storage becomes available again":
     # fail the proof intentionally to trigger `agent.finish(success=false)`,
     # which then calls the onClear callback
-    sales.onProve = proc(slot: Slot): Future[seq[byte]] {.async.} =
+    sales.onProve = proc(slot: Slot, challenge: ProofChallenge): Future[seq[byte]] {.async.} =
       raise newException(IOError, "proof failed")
     var clearedRequest: StorageRequest
     var clearedSlotIndex: UInt256