From fe3419b8984e2830776c990546ca9b23c6b79d9b Mon Sep 17 00:00:00 2001
From: Agost Biro <agostbiro@gmail.com>
Date: Wed, 5 Jun 2024 23:19:15 +0200
Subject: [PATCH] Add nekryptology compatibility tests

---
 vdf/go.mod      | 11 ++++++++++-
 vdf/go.sum      |  4 ++++
 vdf/vdf_test.go | 46 ++++++++++++++++++++++++++++++++++++++++++----
 3 files changed, 56 insertions(+), 5 deletions(-)
 create mode 100644 vdf/go.sum

diff --git a/vdf/go.mod b/vdf/go.mod
index 7369603..1f2fa5e 100644
--- a/vdf/go.mod
+++ b/vdf/go.mod
@@ -1,3 +1,12 @@
-module "source.quilibrium.com/quilibrium/cermonyclient/vdf"
+module source.quilibrium.com/quilibrium/cermonyclient/vdf
 
 go 1.20
+
+// A necessary hack until source.quilibrium.com is open to all
+replace source.quilibrium.com/quilibrium/monorepo/nekryptology => ../nekryptology
+
+require (
+	golang.org/x/crypto v0.24.0 // indirect
+	golang.org/x/sys v0.21.0 // indirect
+	source.quilibrium.com/quilibrium/monorepo/nekryptology v0.0.0-00010101000000-000000000000 // indirect
+)
diff --git a/vdf/go.sum b/vdf/go.sum
new file mode 100644
index 0000000..3983382
--- /dev/null
+++ b/vdf/go.sum
@@ -0,0 +1,4 @@
+golang.org/x/crypto v0.24.0 h1:mnl8DM0o513X8fdIkmyFE/5hTYxbwYOjDS/+rK6qpRI=
+golang.org/x/crypto v0.24.0/go.mod h1:Z1PMYSOR5nyMcyAVAIQSKCDwalqy85Aqn1x3Ws4L5DM=
+golang.org/x/sys v0.21.0 h1:rF+pYz3DAGSQAxAu1CbC7catZg4ebC4UIeIhKxBZvws=
+golang.org/x/sys v0.21.0/go.mod h1:/VUhepiaJMQUp4+oa/7Zr1D23ma6VTLIYjOOTFZPUcA=
diff --git a/vdf/vdf_test.go b/vdf/vdf_test.go
index a2e1585..e934327 100644
--- a/vdf/vdf_test.go
+++ b/vdf/vdf_test.go
@@ -1,17 +1,55 @@
 package vdf_test
 
 import (
+	"golang.org/x/crypto/sha3"
 	"source.quilibrium.com/quilibrium/cermonyclient/vdf"
+	nekrovdf "source.quilibrium.com/quilibrium/monorepo/nekryptology/pkg/vdf"
 	"testing"
 )
 
 const intSizeBits = uint16(2048)
-const difficulty = uint64(10000)
+
+func getChallenge(seed string) [32]byte {
+	return sha3.Sum256([]byte(seed))
+}
 
 func TestProveVerify(t *testing.T) {
-	solution := vdf.WesolowskiSolve(intSizeBits, []byte{0x01, 0x02, 0x03}, difficulty)
-	isOk := vdf.WesolowskiVerify(intSizeBits, []byte{0x01, 0x02, 0x03}, difficulty, solution)
+	difficulty := uint64(10000)
+	challenge := getChallenge("TestProveVerify")
+	solution := vdf.WesolowskiSolve(intSizeBits, challenge[:], difficulty)
+	isOk := vdf.WesolowskiVerify(intSizeBits, challenge[:], difficulty, solution)
 	if !isOk {
-		t.Errorf("WesolowskiVerify failed")
+		t.Fatalf("Verification failed")
+	}
+}
+
+func TestProveRustVerifyNekro(t *testing.T) {
+	difficulty := 100
+	challenge := getChallenge("TestProveRustVerifyNekro")
+
+	for i := 0; i < 100; i++ {
+		solution := vdf.WesolowskiSolve(intSizeBits, challenge[:], uint64(difficulty))
+		nekroVdf := nekrovdf.New(uint32(difficulty), challenge)
+		isOk := nekroVdf.Verify([516]byte(solution))
+		if !isOk {
+			t.Fatalf("Verification failed")
+		}
+		challenge = sha3.Sum256(solution)
+	}
+}
+
+func TestProveNekroVerifyRust(t *testing.T) {
+	difficulty := 100
+	challenge := getChallenge("TestProveNekroVerifyRust")
+
+	for i := 0; i < 100; i++ {
+		nekroVdf := nekrovdf.New(uint32(difficulty), challenge)
+		nekroVdf.Execute()
+		proof := nekroVdf.GetOutput()
+		isOk := vdf.WesolowskiVerify(intSizeBits, challenge[:], uint64(difficulty), proof[:])
+		if !isOk {
+			t.Fatalf("Verification failed")
+		}
+		challenge = sha3.Sum256(proof[:])
 	}
 }