stage2 now works; tests pass

1 files changed, 101 insertions, 8 deletions

diff --git a/nizk/stage2_test.go b/nizk/stage2_test.go
index 3a2d51c..3aeb894 100644
--- a/nizk/stage2_test.go
+++ b/nizk/stage2_test.go
@@ -12,7 +12,7 @@ func TestStage2Simple1(t *testing.T) {
 	for _, lost := range []bool{true, false} {
 		b1 := NewBit(id, !lost)
 		c1 := b1.StageCommit()
-		r1, _ := b1.RevealStage1()
+		r1, _ := b1.RevealStage1(c1.X)
 
 		// Because the first index is a junction, any subsequent
 		// combination of Bits must verify with 'lost' set to true
@@ -31,17 +31,17 @@ func TestStage2Simple1(t *testing.T) {
 			c3 := b3.StageCommit()
 			c4 := b4.StageCommit()
 
-			r2, p2 := b2.RevealStage2(lost, b1)
+			r2, p2 := b2.RevealStage2(lost, b1, c1.X, c2.X, c3.X)
 			if !b2.Commitment.VerifyStage2(c1, c2, r1, r2, p2) {
 				t.Fatalf("failed to verify b2: %t b3: %t bc2/b1", s[0], s[1])
 			}
 
-			r3, p3 := b3.RevealStage2(lost, b1)
+			r3, p3 := b3.RevealStage2(lost, b1, c1.X, c2.X, c3.X)
 			if !b3.Commitment.VerifyStage2(c1, c3, r1, r3, p3) {
 				t.Fatalf("failed to verify b1: %t b3: %t bc3/b1", s[0], s[1])
 			}
 
-			r4, p4 := b4.RevealStage2(lost, b1)
+			r4, p4 := b4.RevealStage2(lost, b1, c1.X, c2.X, c3.X, c4.X)
 			if !b4.Commitment.VerifyStage2(c1, c4, r1, r4, p4) {
 				t.Fatalf("failed to verify b1: %t b4: %t bc4/b1", s[0], s[1])
 			}
@@ -49,7 +49,7 @@ func TestStage2Simple1(t *testing.T) {
 	}
 }
 
-func bit2bit(bid uint) [4]*Bit {
+func uint2bits(bid int) [4]*Bit {
 	id := Curve.RandomScalar()
 
 	return [4]*Bit{
@@ -60,10 +60,103 @@ func bit2bit(bid uint) [4]*Bit {
 	}
 }
 
+var Id = Curve.Identity()
+
 func TestStage2Complex(t *testing.T) {
-	bits1 := 0b0101
-	bits2 := 0b0010
-	t.Logf("testing bits1: %b vs. bits2: %b", bits1, bits2)
+	bid1 := 0b0101
+	bid2 := 0b0010
+	t.Logf("testing bid1: %04b vs. bid2: %04b", bid1, bid2)
+
+	bits1 := uint2bits(bid1)
+	bits2 := uint2bits(bid2)
+
+	lost1 := false
+	lost2 := false
+
+	if len(bits1) != len(bits2) || len(bits1) != 4 {
+		t.Fatalf("oops")
+	}
+
+	instage1 := true
+	junction := -1
+	result := 0
+
+	for c := 0; c < len(bits1); c++ {
+		b1 := bits1[c]
+		b2 := bits2[c]
+
+		c1 := b1.StageCommit()
+		c2 := b2.StageCommit()
+
+		if instage1 {
+			t.Logf("Testing bit b1[%d] = %t vs b2[%d] = %t", c, b1.IsSet(), c, b2.IsSet())
+
+			r1, p1 := b1.RevealStage1(c1.X, c2.X)
+			r2, p2 := b2.RevealStage1(c1.X, c2.X)
+
+			if !b1.Commitment.VerifyStage1(c1, r1, p1) {
+				t.Fatalf("b1 commitment failed to verify in stage1")
+			}
+			if !b2.Commitment.VerifyStage1(c2, r2, p2) {
+				t.Fatalf("b2 commitment failed to verify in stage1")
+			}
+
+			Z := Curve.Product(r1.Z, r2.Z)
+			if !Id.Equal(Z) {
+				t.Logf("Aha! Z[%d] != Id, switch to stage2", c)
+				junction = c
+				instage1 = false
+
+				if !lost1 && !b1.IsSet() {
+					t.Logf("setting lost1 to true")
+					lost1 = true
+				}
+
+				if !lost2 && !b2.IsSet() {
+					t.Logf("setting lost2 to true")
+					lost2 = true
+				}
+				result |= 1 << (3 - c)
+			} else {
+				t.Logf("Z[%d] == Id, staying in stage1", c)
+			}
+		} else {
+			t.Logf("Testing bit b1[%d]∧lost1 = %t vs b2[%d]∧lost2 = %t", c, b1.IsSet() && !lost1, c, b2.IsSet() && !lost2)
+
+			bj1 := bits1[junction]
+			bj2 := bits2[junction]
+
+			r1, p1 := b1.RevealStage2(lost1, bj1, c1.X, c2.X)
+			r2, p2 := b2.RevealStage2(lost2, bj2, c1.X, c2.X)
+
+			if !b1.Commitment.VerifyStage2(bj1.StageCommitment, c1, bj1.StageReveal, r1, p1) {
+				t.Fatalf("b1 commitment failed to verify in stage1")
+			}
+			if !b2.Commitment.VerifyStage2(bj2.StageCommitment, c2, bj2.StageReveal, r2, p2) {
+				t.Fatalf("b2 commitment failed to verify in stage1")
+			}
+
+			Z := Curve.Product(r1.Z, r2.Z)
+			if !Id.Equal(Z) {
+				t.Logf("Aha! Z[%d] != Id, new junction!", c)
+				junction = c
+
+				if !lost1 && !b1.IsSet() {
+					t.Logf("setting lost1 to true")
+					lost1 = true
+				}
+
+				if !lost2 && !b2.IsSet() {
+					t.Logf("setting lost2 to true")
+					lost2 = true
+				}
+				result |= 1 << (3 - c)
+			}
+		}
+	}
+	if result != bid1 {
+		t.Fatalf("wrong result: %04b", result)
+	}
 }
 
 func TestFromPaper(t *testing.T) {