keygen_test.go

// Copyright 2021 TNO
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//      http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package dpaillier

import (
	"math/big"
	"testing"
	"time"

	"encoding/json"
	"fmt"

	"github.com/stretchr/testify/assert"

	secret "github.com/TNO-MPC/shamir"
)

// Test whether p and q are successfully shared
func TestKeygenStage0(t *testing.T) {
	assert := assert.New(t)
	testConfig0 := KeyGenerationParameters{
		NumberOfParticipants:             3,
		ParticipantIndex:                 0,
		PaillierBitSize:                  64,
		SecretSharingDegree:              1,
		SecretSharingStatisticalSecurity: 20,
		SecretSharingModulus:             nil,
		BiprimalityCheckTimes:            20,
	}
	err := testConfig0.Validate()
	assert.NoError(err, "Validating the testing configuration")
	testConfig1 := testConfig0
	testConfig1.ParticipantIndex = 1
	testConfig2 := testConfig0
	testConfig2.ParticipantIndex = 2

	states := make([]*KeyGenerationStage1, 3)
	msgsfrom := make([][]*KeyGenerationMessage1, 3)

	states[0], msgsfrom[0], err = NewKeyGenerationProtocol(testConfig0)
	assert.NoError(err, "New keygen protocol for player 0")
	states[1], msgsfrom[1], err = NewKeyGenerationProtocol(testConfig1)
	assert.NoError(err, "New keygen protocol for player 1")
	states[2], msgsfrom[2], err = NewKeyGenerationProtocol(testConfig2)
	assert.NoError(err, "New keygen protocol for player 2")

	// Check that TermOfP and TermOfQ were correctly shared
	for i := range states {
		assert.Equal(2, len(msgsfrom[i]), "Player %d did not send 2 messages", i)
		shares := []secret.Share{
			states[i].MessageToSelf.ShareOfTermOfP,
			msgsfrom[i][0].ShareOfTermOfP,
			msgsfrom[i][1].ShareOfTermOfP,
		}
		recoveredP, err := secret.ShareCombine(shares)
		if assert.NoError(err, "Error combining shares of P from player %d! Shares are %#v", i, shares) {
			assert.Equal(states[i].TermOfP.String(), recoveredP.String(), "Wrong answer combining shares of P from player %d", i)
		}
		shares = []secret.Share{
			states[i].MessageToSelf.ShareOfTermOfQ,
			msgsfrom[i][0].ShareOfTermOfQ,
			msgsfrom[i][1].ShareOfTermOfQ,
		}
		recoveredQ, err := secret.ShareCombine(shares)
		if assert.NoError(err, "Error combining shares of Q from player %d! Shares are %#v", i, shares) {
			assert.Equal(states[i].TermOfQ.String(), recoveredQ.String(), "Wrong answer combining shares of Q from player %d", i)
		}

		// Also check modularity
		if i == 0 {
			assert.Equal(int64(3), big.NewInt(0).Mod(states[i].TermOfP, big.NewInt(4)).Int64(),
				"Player 0's P term should be 3 mod 4")
			assert.Equal(int64(3), big.NewInt(0).Mod(states[i].TermOfQ, big.NewInt(4)).Int64(),
				"Player 0's Q term should be 3 mod 4")
		} else {
			assert.Equal(int64(0), big.NewInt(0).Mod(states[i].TermOfP, big.NewInt(4)).Int64(),
				"Player %d's P term should be 0 mod 4", i)
			assert.Equal(int64(0), big.NewInt(0).Mod(states[i].TermOfQ, big.NewInt(4)).Int64(),
				"Player %d's Q term should be 0 mod 4", i)
		}
	}
}

// Test whether N is successfully constructed
func TestKeygenStage1(t *testing.T) {
	assert := assert.New(t)
	testConfig0 := KeyGenerationParameters{
		NumberOfParticipants:             3,
		ParticipantIndex:                 0,
		PaillierBitSize:                  512,
		SecretSharingDegree:              1,
		SecretSharingStatisticalSecurity: 20,
		SecretSharingModulus:             nil,
		BiprimalityCheckTimes:            20,
	}
	err := testConfig0.Validate()
	assert.NoError(err, "Validating the testing configuration")
	testConfig1 := testConfig0
	testConfig1.ParticipantIndex = 1
	testConfig2 := testConfig0
	testConfig2.ParticipantIndex = 2

	states1 := make([]*KeyGenerationStage1, 3)
	msgs1from := make([][]*KeyGenerationMessage1, 3)

	states1[0], msgs1from[0], err = NewKeyGenerationProtocol(testConfig0)
	assert.NoError(err, "New keygen protocol for player 0")
	states1[1], msgs1from[1], err = NewKeyGenerationProtocol(testConfig1)
	assert.NoError(err, "New keygen protocol for player 1")
	states1[2], msgs1from[2], err = NewKeyGenerationProtocol(testConfig2)
	assert.NoError(err, "New keygen protocol for player 2")

	// Reroute messages
	msgs1to := make([][]*KeyGenerationMessage1, 3)
	for i := range msgs1to {
		msgs1to[i] = make([]*KeyGenerationMessage1, 0)
		for j := range msgs1from {
			for k := range msgs1from[j] {
				if msgs1from[j][k].To == i {
					msgs1to[i] = append(msgs1to[i], msgs1from[j][k])
				}
			}
		}
	}

	states2 := make([]*KeyGenerationStage2, 3)
	msgs2from := make([][]*KeyGenerationMessage2, 3)

	states2[0], msgs2from[0], err = states1[0].Advance(msgs1to[0])
	assert.NoError(err, "Stage 1 advance to 2 for player 0")
	states2[1], msgs2from[1], err = states1[1].Advance(msgs1to[1])
	assert.NoError(err, "Stage 1 advance to 2 for player 1")
	states2[2], msgs2from[2], err = states1[2].Advance(msgs1to[2])
	assert.NoError(err, "Stage 1 advance to 2 for player 2")

	// Test if the shares of P and Q correspond to the shares of N
	totalP := big.NewInt(0)
	totalP.
		Add(totalP, states2[0].TermOfP).
		Add(totalP, states2[1].TermOfP).
		Add(totalP, states2[2].TermOfP)
	totalQ := big.NewInt(0)
	totalQ.
		Add(totalQ, states2[0].TermOfQ).
		Add(totalQ, states2[1].TermOfQ).
		Add(totalQ, states2[2].TermOfQ)

	computedN := big.NewInt(0).Mul(totalP, totalQ)
	recoveredN, err := secret.ShareCombine([]secret.Share{
		states2[0].ShareOfN,
		states2[1].ShareOfN,
		states2[2].ShareOfN,
	})
	if assert.NoError(err, "Error combining shares of N") {
		assert.Equal(computedN.String(), recoveredN.String(), "Wrong answer combining shares of N")
	}
}

// Test that the small prime test fails for p and q chosen so they are divisible by 3
func TestKeygenStage2FailSmallPrime(t *testing.T) {
	assert := assert.New(t)
	states2BytesWrong := []byte(`[{"Parameters":{"NumberOfParticipants":3,"ParticipantIndex":0,"PaillierBitSize":512,"SecretSharingDegree":1,"SecretSharingStatisticalSecurity":20,"SecretSharingModulus":2644661360355513958309655291853666463931958219241406377635074629142248030292025554190315393636252969949819450760875772323203997303838685181784576130412059800469654057299979726042241507848270140787201112605451700199487489997161403787870085759414117093021883584106919660272272706494256405606106151069655151877553,"NumProcessors":4,"BiprimalityCheckTimes":20},"TermOfP":4664563479187872757232163154071027585990812201075636200488486027352876950690062530911101175685438711140263151889058278017337724434269584558921149202388207,"TermOfQ":6448754197262872714356745108552160479533222235944327792670306030321276737969542044460602775561330868356525665190782534915495427098176666966426562383941487,"ShareOfN":{"FieldSize":2644661360355513958309655291853666463931958219241406377635074629142248030292025554190315393636252969949819450760875772323203997303838685181784576130412059800469654057299979726042241507848270140787201112605451700199487489997161403787870085759414117093021883584106919660272272706494256405606106151069655151877553,"Factor":null,"Degree":2,"X":1,"Y":307102069999111651616387222521238797055825865687245448960129340331715219802221492803174847768390292875800736114473252394277162815236359055074733271681802894972852641109458045626304893772894172335362526717841425552287168699345607477039386937193238377427764259360542349003714645198812414080450344799651586916124}},{"Parameters":{"NumberOfParticipants":3,"ParticipantIndex":1,"PaillierBitSize":512,"SecretSharingDegree":1,"SecretSharingStatisticalSecurity":20,"SecretSharingModulus":2644661360355513958309655291853666463931958219241406377635074629142248030292025554190315393636252969949819450760875772323203997303838685181784576130412059800469654057299979726042241507848270140787201112605451700199487489997161403787870085759414117093021883584106919660272272706494256405606106151069655151877553,"NumProcessors":4,"BiprimalityCheckTimes":20},"TermOfP":12721910721676378007098295473957304622952627641787121654688651069391185265726133609778724264496831107305644694831786247812013495135664885852224897224701000,"TermOfQ":7101205273045859110101365472542479292585690470692903796667704873420940149668449259914732733754194529168992896784603331246939795372818787519077764991930828,"ShareOfN":{"FieldSize":2644661360355513958309655291853666463931958219241406377635074629142248030292025554190315393636252969949819450760875772323203997303838685181784576130412059800469654057299979726042241507848270140787201112605451700199487489997161403787870085759414117093021883584106919660272272706494256405606106151069655151877553,"Factor":null,"Degree":2,"X":2,"Y":307102069999111651616387222521238797055825865687245448960129340331715219802221492803174847768390292875800736114473252394277162815236359055074733271681803000500900323736642497980419220443343437974896740861386158711139243311142710416954682997793259834462800729929861486339244991849358947736045616728916040935941}},{"Parameters":{"NumberOfParticipants":3,"ParticipantIndex":2,"PaillierBitSize":512,"SecretSharingDegree":1,"SecretSharingStatisticalSecurity":20,"SecretSharingModulus":2644661360355513958309655291853666463931958219241406377635074629142248030292025554190315393636252969949819450760875772323203997303838685181784576130412059800469654057299979726042241507848270140787201112605451700199487489997161403787870085759414117093021883584106919660272272706494256405606106151069655151877553,"NumProcessors":4,"BiprimalityCheckTimes":20},"TermOfP":1696844265263161136803336236431897176043248016092636225470088494722288787628862938068889012842777764934766288463568077575103396128302485103771418078323328,"TermOfQ":2542737957772917758526132663335280598107577505788555907734380862995364475963588048886150044811772031250663742552646707215327010675319455090221296270055080,"ShareOfN":{"FieldSize":2644661360355513958309655291853666463931958219241406377635074629142248030292025554190315393636252969949819450760875772323203997303838685181784576130412059800469654057299979726042241507848270140787201112605451700199487489997161403787870085759414117093021883584106919660272272706494256405606106151069655151877553,"Factor":null,"Degree":2,"X":3,"Y":307102069999111651616387222521238797055825865687245448960129340331715219802221492803174847768390292875800736114473252394277162815236359055074733271681803106028948006363826950334533547113792703614430955004930891869991317922939813356869979058393281291497837200499180623674775338499905481391640888658180494955776}}]`)
	var states2 []*KeyGenerationStage2
	json.Unmarshal(states2BytesWrong, &states2)

	msgs2toBytesWrong := []byte(`[[{"From":1,"To":0,"ShareOfN":{"FieldSize":2644661360355513958309655291853666463931958219241406377635074629142248030292025554190315393636252969949819450760875772323203997303838685181784576130412059800469654057299979726042241507848270140787201112605451700199487489997161403787870085759414117093021883584106919660272272706494256405606106151069655151877553,"Factor":null,"Degree":2,"X":2,"Y":307102069999111651616387222521238797055825865687245448960129340331715219802221492803174847768390292875800736114473252394277162815236359055074733271681803000500900323736642497980419220443343437974896740861386158711139243311142710416954682997793259834462800729929861486339244991849358947736045616728916040935941}},{"From":2,"To":0,"ShareOfN":{"FieldSize":2644661360355513958309655291853666463931958219241406377635074629142248030292025554190315393636252969949819450760875772323203997303838685181784576130412059800469654057299979726042241507848270140787201112605451700199487489997161403787870085759414117093021883584106919660272272706494256405606106151069655151877553,"Factor":null,"Degree":2,"X":3,"Y":307102069999111651616387222521238797055825865687245448960129340331715219802221492803174847768390292875800736114473252394277162815236359055074733271681803106028948006363826950334533547113792703614430955004930891869991317922939813356869979058393281291497837200499180623674775338499905481391640888658180494955776}}],[{"From":0,"To":1,"ShareOfN":{"FieldSize":2644661360355513958309655291853666463931958219241406377635074629142248030292025554190315393636252969949819450760875772323203997303838685181784576130412059800469654057299979726042241507848270140787201112605451700199487489997161403787870085759414117093021883584106919660272272706494256405606106151069655151877553,"Factor":null,"Degree":2,"X":1,"Y":307102069999111651616387222521238797055825865687245448960129340331715219802221492803174847768390292875800736114473252394277162815236359055074733271681802894972852641109458045626304893772894172335362526717841425552287168699345607477039386937193238377427764259360542349003714645198812414080450344799651586916124}},{"From":2,"To":1,"ShareOfN":{"FieldSize":2644661360355513958309655291853666463931958219241406377635074629142248030292025554190315393636252969949819450760875772323203997303838685181784576130412059800469654057299979726042241507848270140787201112605451700199487489997161403787870085759414117093021883584106919660272272706494256405606106151069655151877553,"Factor":null,"Degree":2,"X":3,"Y":307102069999111651616387222521238797055825865687245448960129340331715219802221492803174847768390292875800736114473252394277162815236359055074733271681803106028948006363826950334533547113792703614430955004930891869991317922939813356869979058393281291497837200499180623674775338499905481391640888658180494955776}}],[{"From":0,"To":2,"ShareOfN":{"FieldSize":2644661360355513958309655291853666463931958219241406377635074629142248030292025554190315393636252969949819450760875772323203997303838685181784576130412059800469654057299979726042241507848270140787201112605451700199487489997161403787870085759414117093021883584106919660272272706494256405606106151069655151877553,"Factor":null,"Degree":2,"X":1,"Y":307102069999111651616387222521238797055825865687245448960129340331715219802221492803174847768390292875800736114473252394277162815236359055074733271681802894972852641109458045626304893772894172335362526717841425552287168699345607477039386937193238377427764259360542349003714645198812414080450344799651586916124}},{"From":1,"To":2,"ShareOfN":{"FieldSize":2644661360355513958309655291853666463931958219241406377635074629142248030292025554190315393636252969949819450760875772323203997303838685181784576130412059800469654057299979726042241507848270140787201112605451700199487489997161403787870085759414117093021883584106919660272272706494256405606106151069655151877553,"Factor":null,"Degree":2,"X":2,"Y":307102069999111651616387222521238797055825865687245448960129340331715219802221492803174847768390292875800736114473252394277162815236359055074733271681803000500900323736642497980419220443343437974896740861386158711139243311142710416954682997793259834462800729929861486339244991849358947736045616728916040935941}}]]`)
	var msgs2to [][]*KeyGenerationMessage2
	json.Unmarshal(msgs2toBytesWrong, &msgs2to)

	states3 := make([]*KeyGenerationStage3, 3)
	msgs3from := make([][]*KeyGenerationMessage3, 3)
	var err error

	states3[0], msgs3from[0], err = states2[0].Advance(msgs2to[0])
	assert.NoError(err, "Stage 2 advance to 3 for player 0")
	states3[1], msgs3from[1], err = states2[1].Advance(msgs2to[1])
	assert.Equal(ErrorBiprimalityTestFailed, err, "Stage 2 advance to 3 for player 1 (should fail)")
	states3[2], msgs3from[2], err = states2[2].Advance(msgs2to[2])
	assert.NoError(err, "Stage 2 advance to 3 for player 2")
}

// Test that the small prime test succeeds for p and q chosen to be indivisible by them
func TestKeygenStage2Success(t *testing.T) {
	assert := assert.New(t)
	states2BytesWrong := []byte(`[{"Parameters":{"NumberOfParticipants":3,"ParticipantIndex":0,"PaillierBitSize":512,"SecretSharingDegree":1,"SecretSharingStatisticalSecurity":20,"SecretSharingModulus":2767358795579534626201522099944414057460669149746093712088535141119461552935518433598434092854122860825786069588973932199196606884533712870042425626025364973110966268760982999344839068940250371740097880566758540066569039367634494710949376247705939365037319266941187409397067481832747627771340381379225427948247,"NumProcessors":4,"BiprimalityCheckTimes":20},"TermOfP":561085336904163428030288783903473894165274999536382788870829764969055471703852407320244151609920041007617237089047021365131220346027082149949218938845075,"TermOfQ":6764972718620774731773858986688340963454732304038985296355982094395112443646477362082721237342012197433951917618307434090048492071004710761269168215877671,"ShareOfN":{"FieldSize":2767358795579534626201522099944414057460669149746093712088535141119461552935518433598434092854122860825786069588973932199196606884533712870042425626025364973110966268760982999344839068940250371740097880566758540066569039367634494710949376247705939365037319266941187409397067481832747627771340381379225427948247,"Factor":null,"Degree":2,"X":1,"Y":503866290021293091025774360186690101367021304358792448208020987385952147850154207221237269257513705402704973173050012162365377178881306849387994995786601117581193839554281348940864235842556061402767202680105727842927537284767540937228703083796664141686149198209853896452597138633433569389424507593659157249292}},{"Parameters":{"NumberOfParticipants":3,"ParticipantIndex":1,"PaillierBitSize":512,"SecretSharingDegree":1,"SecretSharingStatisticalSecurity":20,"SecretSharingModulus":2767358795579534626201522099944414057460669149746093712088535141119461552935518433598434092854122860825786069588973932199196606884533712870042425626025364973110966268760982999344839068940250371740097880566758540066569039367634494710949376247705939365037319266941187409397067481832747627771340381379225427948247,"NumProcessors":4,"BiprimalityCheckTimes":20},"TermOfP":7303533793414846968859713017182181989457047267302507350973976042908272174154327626526291142874493936114527768616783560166962889569809667822011007831362180,"TermOfQ":7445343306052853947619947872103083880330535974928572135639658208499860375899346151112232308083492841399392938183705600339307615205463182024789821460613072,"ShareOfN":{"FieldSize":2767358795579534626201522099944414057460669149746093712088535141119461552935518433598434092854122860825786069588973932199196606884533712870042425626025364973110966268760982999344839068940250371740097880566758540066569039367634494710949376247705939365037319266941187409397067481832747627771340381379225427948247,"Factor":null,"Degree":2,"X":2,"Y":503866290021293091025774360186690101367021304358792448208020987385952147850154207221237269257513705402704973173050012162365377178881306849387994995786601253964452863875278765213723885834083623015852701527232487815600343570584029897230285946693434083525420071549794536335177136314552226906349097282960416771165}},{"Parameters":{"NumberOfParticipants":3,"ParticipantIndex":2,"PaillierBitSize":512,"SecretSharingDegree":1,"SecretSharingStatisticalSecurity":20,"SecretSharingModulus":2767358795579534626201522099944414057460669149746093712088535141119461552935518433598434092854122860825786069588973932199196606884533712870042425626025364973110966268760982999344839068940250371740097880566758540066569039367634494710949376247705939365037319266941187409397067481832747627771340381379225427948247,"NumProcessors":4,"BiprimalityCheckTimes":20},"TermOfP":11286636008588111415419425071646018238958325238168969023166344400967791153480327172899428185808683122943870300714232717224130355617077014071135720190849424,"TermOfQ":12099515177859581980387719485140743554171779382306958991650767090355180544102336474346715230928677618058419817991217860146979800076457318034074830448959860,"ShareOfN":{"FieldSize":2767358795579534626201522099944414057460669149746093712088535141119461552935518433598434092854122860825786069588973932199196606884533712870042425626025364973110966268760982999344839068940250371740097880566758540066569039367634494710949376247705939365037319266941187409397067481832747627771340381379225427948247,"Factor":null,"Degree":2,"X":3,"Y":503866290021293091025774360186690101367021304358792448208020987385952147850154207221237269257513705402704973173050012162365377178881306849387994995786601390347711888196276181486583535825611184628938200374359247788273149856400518857231868809590204025364690944889735176217757133995670884423273686972261676293056}}]`)
	var states2 []*KeyGenerationStage2
	json.Unmarshal(states2BytesWrong, &states2)

	msgs2toBytesWrong := []byte(`[[{"From":1,"To":0,"ShareOfN":{"FieldSize":2767358795579534626201522099944414057460669149746093712088535141119461552935518433598434092854122860825786069588973932199196606884533712870042425626025364973110966268760982999344839068940250371740097880566758540066569039367634494710949376247705939365037319266941187409397067481832747627771340381379225427948247,"Factor":null,"Degree":2,"X":2,"Y":503866290021293091025774360186690101367021304358792448208020987385952147850154207221237269257513705402704973173050012162365377178881306849387994995786601253964452863875278765213723885834083623015852701527232487815600343570584029897230285946693434083525420071549794536335177136314552226906349097282960416771165}},{"From":2,"To":0,"ShareOfN":{"FieldSize":2767358795579534626201522099944414057460669149746093712088535141119461552935518433598434092854122860825786069588973932199196606884533712870042425626025364973110966268760982999344839068940250371740097880566758540066569039367634494710949376247705939365037319266941187409397067481832747627771340381379225427948247,"Factor":null,"Degree":2,"X":3,"Y":503866290021293091025774360186690101367021304358792448208020987385952147850154207221237269257513705402704973173050012162365377178881306849387994995786601390347711888196276181486583535825611184628938200374359247788273149856400518857231868809590204025364690944889735176217757133995670884423273686972261676293056}}],[{"From":0,"To":1,"ShareOfN":{"FieldSize":2767358795579534626201522099944414057460669149746093712088535141119461552935518433598434092854122860825786069588973932199196606884533712870042425626025364973110966268760982999344839068940250371740097880566758540066569039367634494710949376247705939365037319266941187409397067481832747627771340381379225427948247,"Factor":null,"Degree":2,"X":1,"Y":503866290021293091025774360186690101367021304358792448208020987385952147850154207221237269257513705402704973173050012162365377178881306849387994995786601117581193839554281348940864235842556061402767202680105727842927537284767540937228703083796664141686149198209853896452597138633433569389424507593659157249292}},{"From":2,"To":1,"ShareOfN":{"FieldSize":2767358795579534626201522099944414057460669149746093712088535141119461552935518433598434092854122860825786069588973932199196606884533712870042425626025364973110966268760982999344839068940250371740097880566758540066569039367634494710949376247705939365037319266941187409397067481832747627771340381379225427948247,"Factor":null,"Degree":2,"X":3,"Y":503866290021293091025774360186690101367021304358792448208020987385952147850154207221237269257513705402704973173050012162365377178881306849387994995786601390347711888196276181486583535825611184628938200374359247788273149856400518857231868809590204025364690944889735176217757133995670884423273686972261676293056}}],[{"From":0,"To":2,"ShareOfN":{"FieldSize":2767358795579534626201522099944414057460669149746093712088535141119461552935518433598434092854122860825786069588973932199196606884533712870042425626025364973110966268760982999344839068940250371740097880566758540066569039367634494710949376247705939365037319266941187409397067481832747627771340381379225427948247,"Factor":null,"Degree":2,"X":1,"Y":503866290021293091025774360186690101367021304358792448208020987385952147850154207221237269257513705402704973173050012162365377178881306849387994995786601117581193839554281348940864235842556061402767202680105727842927537284767540937228703083796664141686149198209853896452597138633433569389424507593659157249292}},{"From":1,"To":2,"ShareOfN":{"FieldSize":2767358795579534626201522099944414057460669149746093712088535141119461552935518433598434092854122860825786069588973932199196606884533712870042425626025364973110966268760982999344839068940250371740097880566758540066569039367634494710949376247705939365037319266941187409397067481832747627771340381379225427948247,"Factor":null,"Degree":2,"X":2,"Y":503866290021293091025774360186690101367021304358792448208020987385952147850154207221237269257513705402704973173050012162365377178881306849387994995786601253964452863875278765213723885834083623015852701527232487815600343570584029897230285946693434083525420071549794536335177136314552226906349097282960416771165}}]]`)
	var msgs2to [][]*KeyGenerationMessage2
	json.Unmarshal(msgs2toBytesWrong, &msgs2to)

	states3 := make([]*KeyGenerationStage3, 3)
	msgs3from := make([][]*KeyGenerationMessage3, 3)
	var err error

	states3[0], msgs3from[0], err = states2[0].Advance(msgs2to[0])
	assert.NoError(err, "Stage 2 advance to 3 for player 0")
	states3[1], msgs3from[1], err = states2[1].Advance(msgs2to[1])
	assert.NoError(err, "Stage 2 advance to 3 for player 1")
	states3[2], msgs3from[2], err = states2[2].Advance(msgs2to[2])
	assert.NoError(err, "Stage 2 advance to 3 for player 2")
}

// Test that the biprimality test passes for p and q chosen prime
func TestKeygenStage4BiprimeSuccess(t *testing.T) {
	assert := assert.New(t)
	testConfig0 := KeyGenerationParameters{
		NumberOfParticipants:             3,
		ParticipantIndex:                 0,
		PaillierBitSize:                  512,
		SecretSharingDegree:              1,
		SecretSharingStatisticalSecurity: 20,
		SecretSharingModulus:             nil,
		BiprimalityCheckTimes:            20,
	}
	err := testConfig0.Validate()
	assert.NoError(err, "Validating the testing configuration")
	testConfig1 := testConfig0
	testConfig1.ParticipantIndex = 1
	testConfig2 := testConfig0
	testConfig2.ParticipantIndex = 2

	// Simulate the first part of a keygen with two primes off-line
	// These are two primes, we split them so that each party has a "random term of P and Q"
	// they came up with, and do the secret share/add/multiply from there
	P, _ := big.NewInt(0).SetString("11120808620651904624786310689531044605752497671295053269308320157004124862913675467444400323814946379905247813626954435646456303989193301523714415415979603", 10)
	termP12 := big.NewInt(0).Div(P, big.NewInt(3))
	termP12.Mod(termP12, big.NewInt(4))
	termP0 := big.NewInt(0).Mul(termP12, big.NewInt(2))
	termP0.Sub(P, termP0)
	Q, _ := big.NewInt(0).SetString("10767776531951936098666435704607942082088824478217191538526941796775210438590741454374000946451448374862778186755376529814773752858068353699218714973392223", 10)
	termQ12 := big.NewInt(0).Div(Q, big.NewInt(3))
	termQ12.Mod(termQ12, big.NewInt(4))
	termQ0 := big.NewInt(0).Mul(termQ12, big.NewInt(2))
	termQ0.Sub(Q, termQ0)

	// Secret share "horizontally"
	shareP0 := secret.ShareFiniteField(termP0, testConfig0.SecretSharingModulus, testConfig0.SecretSharingDegree, testConfig0.NumberOfParticipants)
	shareP1 := secret.ShareFiniteField(termP12, testConfig0.SecretSharingModulus, testConfig0.SecretSharingDegree, testConfig0.NumberOfParticipants)
	shareP2 := secret.ShareFiniteField(termP12, testConfig0.SecretSharingModulus, testConfig0.SecretSharingDegree, testConfig0.NumberOfParticipants)
	shareQ0 := secret.ShareFiniteField(termQ0, testConfig0.SecretSharingModulus, testConfig0.SecretSharingDegree, testConfig0.NumberOfParticipants)
	shareQ1 := secret.ShareFiniteField(termQ12, testConfig0.SecretSharingModulus, testConfig0.SecretSharingDegree, testConfig0.NumberOfParticipants)
	shareQ2 := secret.ShareFiniteField(termQ12, testConfig0.SecretSharingModulus, testConfig0.SecretSharingDegree, testConfig0.NumberOfParticipants)

	// Add shares "vertically"
	addP0, _ := secret.ShareAdd([]secret.Share{shareP0[0], shareP1[0], shareP2[0]})
	addP1, _ := secret.ShareAdd([]secret.Share{shareP0[1], shareP1[1], shareP2[1]})
	addP2, _ := secret.ShareAdd([]secret.Share{shareP0[2], shareP1[2], shareP2[2]})
	addQ0, _ := secret.ShareAdd([]secret.Share{shareQ0[0], shareQ1[0], shareQ2[0]})
	addQ1, _ := secret.ShareAdd([]secret.Share{shareQ0[1], shareQ1[1], shareQ2[1]})
	addQ2, _ := secret.ShareAdd([]secret.Share{shareQ0[2], shareQ1[2], shareQ2[2]})

	// Multiply shares of P and Q
	shareN0, _ := secret.ShareMul([]secret.Share{addP0, addQ0})
	shareN1, _ := secret.ShareMul([]secret.Share{addP1, addQ1})
	shareN2, _ := secret.ShareMul([]secret.Share{addP2, addQ2})

	states2 := []*KeyGenerationStage2{{
		Parameters: &testConfig0,
		TermOfP:    termP0,
		TermOfQ:    termQ0,
		ShareOfN:   shareN0,
	}, {
		Parameters: &testConfig1,
		TermOfP:    termP12,
		TermOfQ:    termQ12,
		ShareOfN:   shareN1,
	}, {
		Parameters: &testConfig2,
		TermOfP:    termP12,
		TermOfQ:    termQ12,
		ShareOfN:   shareN2,
	}}

	msgs2to := [][]*KeyGenerationMessage2{{
		{From: 1, To: 0, ShareOfN: shareN1},
		{From: 2, To: 0, ShareOfN: shareN2},
	}, {
		{From: 0, To: 1, ShareOfN: shareN0},
		{From: 2, To: 1, ShareOfN: shareN2},
	}, {
		{From: 0, To: 2, ShareOfN: shareN0},
		{From: 1, To: 2, ShareOfN: shareN1},
	}}

	states3 := make([]*KeyGenerationStage3, 3)
	msgs3from := make([][]*KeyGenerationMessage3, 3)
	msgs3to := make([][]*KeyGenerationMessage3, 3)
	states4 := make([]*KeyGenerationStage4, 3)
	msgs4from := make([][]*KeyGenerationMessage4, 3)
	msgs4to := make([][]*KeyGenerationMessage4, 3)
	states5 := make([]*KeyGenerationStage5, 3)
	msgs5from := make([][]*KeyGenerationMessage5, 3)
	var err0, err1, err2 error

	states3[0], msgs3from[0], err0 = states2[0].Advance(msgs2to[0])
	assert.NoError(err0, "Stage 2 advance to 3 for player 0")
	states3[1], msgs3from[1], err1 = states2[1].Advance(msgs2to[1])
	assert.NoError(err1, "Stage 2 advance to 3 for player 1")
	states3[2], msgs3from[2], err2 = states2[2].Advance(msgs2to[2])
	assert.NoError(err2, "Stage 2 advance to 3 for player 2")

	// Reroute messages
	for i := range msgs3to {
		msgs3to[i] = make([]*KeyGenerationMessage3, 0)
		for j := range msgs3from {
			for k := range msgs3from[j] {
				if msgs3from[j][k].To == i {
					msgs3to[i] = append(msgs3to[i], msgs3from[j][k])
				}
			}
		}
	}

	states4[0], msgs4from[0], err0 = states3[0].Advance(msgs3to[0])
	assert.NoError(err0, "Stage 3 advance to 4 for player 0")
	states4[1], msgs4from[1], err1 = states3[1].Advance(msgs3to[1])
	assert.NoError(err1, "Stage 3 advance to 4 for player 1")
	states4[2], msgs4from[2], err2 = states3[2].Advance(msgs3to[2])
	assert.NoError(err2, "Stage 3 advance to 4 for player 2")

	// Reroute messages
	for i := range msgs4to {
		msgs4to[i] = make([]*KeyGenerationMessage4, 0)
		for j := range msgs4from {
			for k := range msgs4from[j] {
				if msgs4from[j][k].To == i {
					msgs4to[i] = append(msgs4to[i], msgs4from[j][k])
				}
			}
		}
	}

	states5[0], msgs5from[0], err0 = states4[0].Advance(msgs4to[0])
	assert.NoError(err0, "Stage 4 advance to 5 for player 0 (biprimality test)")
	states5[1], msgs5from[1], err1 = states4[1].Advance(msgs4to[1])
	assert.NoError(err1, "Stage 4 advance to 5 for player 1")
	states5[2], msgs5from[2], err2 = states4[2].Advance(msgs4to[2])
	assert.NoError(err2, "Stage 3 advance to 5 for player 2")
}

// Start from a state corresponding to a known good N, and test private keygen
func TestKeygenStage3Lambda(t *testing.T) {
	assert := assert.New(t)

	states3Bytes := []byte(`[{"Parameters":{"NumberOfParticipants":3,"ParticipantIndex":0,"PaillierBitSize":512,"SecretSharingDegree":1,"SecretSharingStatisticalSecurity":20,"SecretSharingModulus":2623222045352356732683025662305979215904046499376168877804324390266714056732042608675318325182279767188274546065280153959994987994681446800680132187027204726065259578691071312441425971149525208889491572233926587610577246558357978325575132606555085287066466813663551180648008762652591343459455276595418684019307,"NumProcessors":4,"BiprimalityCheckTimes":0},"TermOfP":12803860855266792015707421000341147329683001397665078863910739077342083201481596055578862246681369117726908521649592661613193660246129309533812235817002191,"TermOfQ":11691112965269741642710310239951428148824624216599551209455640938228901100844293068758078772654047418960574152297371646251736056716005176266470887289828403,"N":923588794041132910322504663489230159340645205888572410576204511494217137733476860130052233334326157027202966831612503418564291553442062779657809148246697677036362133578704288207973851468378736925970832206358261932321121488532519774503731055510154940781895101061565530550463014148475619594070354394490660675257,"TestElements":[]},{"Parameters":{"NumberOfParticipants":3,"ParticipantIndex":1,"PaillierBitSize":512,"SecretSharingDegree":1,"SecretSharingStatisticalSecurity":20,"SecretSharingModulus":2623222045352356732683025662305979215904046499376168877804324390266714056732042608675318325182279767188274546065280153959994987994681446800680132187027204726065259578691071312441425971149525208889491572233926587610577246558357978325575132606555085287066466813663551180648008762652591343459455276595418684019307,"NumProcessors":4,"BiprimalityCheckTimes":0},"TermOfP":12742261193987197851739933987716245988838680824380936634725103123644406499270331129522530647555789834741507862566885591202016390227120658237090490735408364,"TermOfQ":12815694744575221653712069934546572631152108091809604593011938294893015277970105750337439679951270254388366237340082013692604676279446925211665544973352996,"N":923588794041132910322504663489230159340645205888572410576204511494217137733476860130052233334326157027202966831612503418564291553442062779657809148246697677036362133578704288207973851468378736925970832206358261932321121488532519774503731055510154940781895101061565530550463014148475619594070354394490660675257,"TestElements":null},{"Parameters":{"NumberOfParticipants":3,"ParticipantIndex":2,"PaillierBitSize":512,"SecretSharingDegree":1,"SecretSharingStatisticalSecurity":20,"SecretSharingModulus":2623222045352356732683025662305979215904046499376168877804324390266714056732042608675318325182279767188274546065280153959994987994681446800680132187027204726065259578691071312441425971149525208889491572233926587610577246558357978325575132606555085287066466813663551180648008762652591343459455276595418684019307,"NumProcessors":4,"BiprimalityCheckTimes":0},"TermOfP":6217943986514694411881041065789078064812396459918119733985362339745527932350177614027381967435170629672583128985069267901496323799690751211860671895297372,"TermOfQ":4569721504050503264187619909243612614725010821916737084593198343059221888312372731164124724443166873831160338029914725359337473325781173503756176179234392,"N":923588794041132910322504663489230159340645205888572410576204511494217137733476860130052233334326157027202966831612503418564291553442062779657809148246697677036362133578704288207973851468378736925970832206358261932321121488532519774503731055510154940781895101061565530550463014148475619594070354394490660675257,"TestElements":null}]`)
	var states3 []*KeyGenerationStage3
	json.Unmarshal(states3Bytes, &states3)

	msgs3toBytes := []byte(`[[{"From":1,"To":0,"N":923588794041132910322504663489230159340645205888572410576204511494217137733476860130052233334326157027202966831612503418564291553442062779657809148246697677036362133578704288207973851468378736925970832206358261932321121488532519774503731055510154940781895101061565530550463014148475619594070354394490660675257,"TestElements":[]},{"From":2,"To":0,"N":923588794041132910322504663489230159340645205888572410576204511494217137733476860130052233334326157027202966831612503418564291553442062779657809148246697677036362133578704288207973851468378736925970832206358261932321121488532519774503731055510154940781895101061565530550463014148475619594070354394490660675257,"TestElements":[]}],[{"From":0,"To":1,"N":923588794041132910322504663489230159340645205888572410576204511494217137733476860130052233334326157027202966831612503418564291553442062779657809148246697677036362133578704288207973851468378736925970832206358261932321121488532519774503731055510154940781895101061565530550463014148475619594070354394490660675257,"TestElements":[]},{"From":2,"To":1,"N":923588794041132910322504663489230159340645205888572410576204511494217137733476860130052233334326157027202966831612503418564291553442062779657809148246697677036362133578704288207973851468378736925970832206358261932321121488532519774503731055510154940781895101061565530550463014148475619594070354394490660675257,"TestElements":[]}],[{"From":0,"To":2,"N":923588794041132910322504663489230159340645205888572410576204511494217137733476860130052233334326157027202966831612503418564291553442062779657809148246697677036362133578704288207973851468378736925970832206358261932321121488532519774503731055510154940781895101061565530550463014148475619594070354394490660675257,"TestElements":[]},{"From":1,"To":2,"N":923588794041132910322504663489230159340645205888572410576204511494217137733476860130052233334326157027202966831612503418564291553442062779657809148246697677036362133578704288207973851468378736925970832206358261932321121488532519774503731055510154940781895101061565530550463014148475619594070354394490660675257,"TestElements":[]}]]`)
	var msgs3to [][]*KeyGenerationMessage3
	json.Unmarshal(msgs3toBytes, &msgs3to)

	// Compute state 4 from state 3
	states4 := make([]*KeyGenerationStage4, 3)
	msgs4from := make([][]*KeyGenerationMessage4, 3)
	var err error

	states4[0], msgs4from[0], err = states3[0].Advance(msgs3to[0])
	assert.NoError(err, "Stage 3 advance to 4 for player 0")
	states4[1], msgs4from[1], err = states3[1].Advance(msgs3to[1])
	assert.NoError(err, "Stage 3 advance to 4 for player 1")
	states4[2], msgs4from[2], err = states3[2].Advance(msgs3to[2])
	assert.NoError(err, "Stage 3 advance to 4 for player 2")

	// The terms of lambda computed here are N - p1 - q1 + 1 and pi + qi.
	// We can test here if (lambda1 - sum_i lambda_i =) N - p - q + 1 == (p-1)(q-1)
	lhs := big.NewInt(0).Set(states4[0].TermOfLambda)
	for i := 1; i != len(states4); i++ {
		lhs.Sub(lhs, states4[i].TermOfLambda)
	}
	rhsP, rhsQ := big.NewInt(-1), big.NewInt(-1)
	for i := range states3 {
		rhsP.Add(rhsP, states3[i].TermOfP)
		rhsQ.Add(rhsQ, states3[i].TermOfQ)
	}
	rhs := big.NewInt(0).Mul(rhsP, rhsQ)
	assert.Equal(lhs.String(), rhs.String(), "Reconstruction of phi(N) from lambda terms failed")

	// We would expect that exp(r, lambda * n, n^2) == 1 for random r in 0..n, say 2
	random := big.NewInt(2)
	random.Exp(random, big.NewInt(0).Mul(rhs, states4[0].N), big.NewInt(0).Mul(states4[0].N, states4[0].N))
	assert.Equal("1", random.String(), "r^(lambda N) is not 1 mod N^2")
}

// Start from a state corresponding to a known good N, and test lambda key sharing
func TestKeygenStage5(t *testing.T) {
	assert := assert.New(t)

	states4Bytes := []byte(`[{"Parameters":{"NumberOfParticipants":3,"ParticipantIndex":0,"PaillierBitSize":512,"SecretSharingDegree":1,"SecretSharingStatisticalSecurity":20,"SecretSharingModulus":2623222045352356732683025662305979215904046499376168877804324390266714056732042608675318325182279767188274546065280153959994987994681446800680132187027204726065259578691071312441425971149525208889491572233926587610577246558357978325575132606555085287066466813663551180648008762652591343459455276595418684019307,"NumProcessors":4,"BiprimalityCheckTimes":0},"N":923588794041132910322504663489230159340645205888572410576204511494217137733476860130052233334326157027202966831612503418564291553442062779657809148246697677036362133578704288207973851468378736925970832206358261932321121488532519774503731055510154940781895101061565530550463014148475619594070354394490660675257,"TermOfLambda":923588794041132910322504663489230159340645205888572410576204511494217137733476860130052233334326157027202966831612503418564291553442062779657809148246697652541388313042170629790242611175803258418345217941728188565941105917548217448614606718569135605365358413578891583586155149218758657459584554111367553844664,"TestResponses":[]},{"Parameters":{"NumberOfParticipants":3,"ParticipantIndex":1,"PaillierBitSize":512,"SecretSharingDegree":1,"SecretSharingStatisticalSecurity":20,"SecretSharingModulus":2623222045352356732683025662305979215904046499376168877804324390266714056732042608675318325182279767188274546065280153959994987994681446800680132187027204726065259578691071312441425971149525208889491572233926587610577246558357978325575132606555085287066466813663551180648008762652591343459455276595418684019307,"NumProcessors":4,"BiprimalityCheckTimes":0},"N":923588794041132910322504663489230159340645205888572410576204511494217137733476860130052233334326157027202966831612503418564291553442062779657809148246697677036362133578704288207973851468378736925970832206358261932321121488532519774503731055510154940781895101061565530550463014148475619594070354394490660675257,"TermOfLambda":25557955938562419505452003922262818619990788916190541227737041418537421777240436879859970327507060089129874099906967604894621066506567583448756035708761360,"TestResponses":null},{"Parameters":{"NumberOfParticipants":3,"ParticipantIndex":2,"PaillierBitSize":512,"SecretSharingDegree":1,"SecretSharingStatisticalSecurity":20,"SecretSharingModulus":2623222045352356732683025662305979215904046499376168877804324390266714056732042608675318325182279767188274546065280153959994987994681446800680132187027204726065259578691071312441425971149525208889491572233926587610577246558357978325575132606555085287066466813663551180648008762652591343459455276595418684019307,"NumProcessors":4,"BiprimalityCheckTimes":0},"N":923588794041132910322504663489230159340645205888572410576204511494217137733476860130052233334326157027202966831612503418564291553442062779657809148246697677036362133578704288207973851468378736925970832206358261932321121488532519774503731055510154940781895101061565530550463014148475619594070354394490660675257,"TermOfLambda":10787665490565197676068660975032690679537407281834856818578560682804749820662550345191506691878337503503743467014983993260833797125471924715616848074531764,"TestResponses":null}]`)
	var states4 []*KeyGenerationStage4
	json.Unmarshal(states4Bytes, &states4)

	msgs4toBytes := []byte(`[[{"From":1,"To":0,"TestResponses":[]},{"From":2,"To":0,"TestResponses":[]}],[{"From":0,"To":1,"TestResponses":[]},{"From":2,"To":1,"TestResponses":[]}],[{"From":0,"To":2,"TestResponses":[]},{"From":1,"To":2,"TestResponses":[]}]]`)
	var msgs4to [][]*KeyGenerationMessage4
	json.Unmarshal(msgs4toBytes, &msgs4to)

	// Compute state 5 from state 4
	states5 := make([]*KeyGenerationStage5, 3)
	msgs5from := make([][]*KeyGenerationMessage5, 3)
	var err error

	states5[0], msgs5from[0], err = states4[0].Advance(msgs4to[0])
	assert.NoError(err, "Stage 4 advance to 5 for player 0")
	states5[1], msgs5from[1], err = states4[1].Advance(msgs4to[1])
	assert.NoError(err, "Stage 4 advance to 5 for player 1")
	states5[2], msgs5from[2], err = states4[2].Advance(msgs4to[2])
	assert.NoError(err, "Stage 4 advance to 5 for player 2")

	// Check that TermOfL and TermOfB were correctly shared
	for i := range states5 {
		assert.Equal(2, len(msgs5from[i]), "Player %d did not send 2 messages", i)
		recoveredL, err := secret.ShareCombine([]secret.Share{
			states5[i].MessageToSelf.ShareOfTermOfL,
			msgs5from[i][0].ShareOfTermOfL,
			msgs5from[i][1].ShareOfTermOfL,
		})
		if assert.NoError(err, "Error combining shares of L from player %d", i) {
			assert.Equal(states5[i].TermOfL.String(), recoveredL.String(), "Wrong answer combining shares of L from player %d", i)
		}
		recoveredB, err := secret.ShareCombine([]secret.Share{
			states5[i].MessageToSelf.ShareOfTermOfB,
			msgs5from[i][0].ShareOfTermOfB,
			msgs5from[i][1].ShareOfTermOfB,
		})
		if assert.NoError(err, "Error combining shares of B from player %d", i) {
			assert.Equal(states5[i].TermOfB.String(), recoveredB.String(), "Wrong answer combining shares of B from player %d", i)
		}
	}
}

// Starting from a known good N, test theta reconstruction from shares
func TestKeygenStage6(t *testing.T) {
	assert := assert.New(t)

	states4Bytes := []byte(`[{"Parameters":{"NumberOfParticipants":3,"ParticipantIndex":0,"PaillierBitSize":512,"SecretSharingDegree":1,"SecretSharingStatisticalSecurity":20,"SecretSharingModulus":2623222045352356732683025662305979215904046499376168877804324390266714056732042608675318325182279767188274546065280153959994987994681446800680132187027204726065259578691071312441425971149525208889491572233926587610577246558357978325575132606555085287066466813663551180648008762652591343459455276595418684019307,"NumProcessors":4,"BiprimalityCheckTimes":0},"N":923588794041132910322504663489230159340645205888572410576204511494217137733476860130052233334326157027202966831612503418564291553442062779657809148246697677036362133578704288207973851468378736925970832206358261932321121488532519774503731055510154940781895101061565530550463014148475619594070354394490660675257,"TermOfLambda":923588794041132910322504663489230159340645205888572410576204511494217137733476860130052233334326157027202966831612503418564291553442062779657809148246697652541388313042170629790242611175803258418345217941728188565941105917548217448614606718569135605365358413578891583586155149218758657459584554111367553844664,"TestResponses":[]},{"Parameters":{"NumberOfParticipants":3,"ParticipantIndex":1,"PaillierBitSize":512,"SecretSharingDegree":1,"SecretSharingStatisticalSecurity":20,"SecretSharingModulus":2623222045352356732683025662305979215904046499376168877804324390266714056732042608675318325182279767188274546065280153959994987994681446800680132187027204726065259578691071312441425971149525208889491572233926587610577246558357978325575132606555085287066466813663551180648008762652591343459455276595418684019307,"NumProcessors":4,"BiprimalityCheckTimes":0},"N":923588794041132910322504663489230159340645205888572410576204511494217137733476860130052233334326157027202966831612503418564291553442062779657809148246697677036362133578704288207973851468378736925970832206358261932321121488532519774503731055510154940781895101061565530550463014148475619594070354394490660675257,"TermOfLambda":25557955938562419505452003922262818619990788916190541227737041418537421777240436879859970327507060089129874099906967604894621066506567583448756035708761360,"TestResponses":null},{"Parameters":{"NumberOfParticipants":3,"ParticipantIndex":2,"PaillierBitSize":512,"SecretSharingDegree":1,"SecretSharingStatisticalSecurity":20,"SecretSharingModulus":2623222045352356732683025662305979215904046499376168877804324390266714056732042608675318325182279767188274546065280153959994987994681446800680132187027204726065259578691071312441425971149525208889491572233926587610577246558357978325575132606555085287066466813663551180648008762652591343459455276595418684019307,"NumProcessors":4,"BiprimalityCheckTimes":0},"N":923588794041132910322504663489230159340645205888572410576204511494217137733476860130052233334326157027202966831612503418564291553442062779657809148246697677036362133578704288207973851468378736925970832206358261932321121488532519774503731055510154940781895101061565530550463014148475619594070354394490660675257,"TermOfLambda":10787665490565197676068660975032690679537407281834856818578560682804749820662550345191506691878337503503743467014983993260833797125471924715616848074531764,"TestResponses":null}]`)
	var states4 []*KeyGenerationStage4
	json.Unmarshal(states4Bytes, &states4)

	msgs4toBytes := []byte(`[[{"From":1,"To":0,"TestResponses":[]},{"From":2,"To":0,"TestResponses":[]}],[{"From":0,"To":1,"TestResponses":[]},{"From":2,"To":1,"TestResponses":[]}],[{"From":0,"To":2,"TestResponses":[]},{"From":1,"To":2,"TestResponses":[]}]]`)
	var msgs4to [][]*KeyGenerationMessage4
	json.Unmarshal(msgs4toBytes, &msgs4to)

	// Compute state 5 from state 4
	states5 := make([]*KeyGenerationStage5, 3)
	msgs5from := make([][]*KeyGenerationMessage5, 3)
	var err error

	states5[0], msgs5from[0], err = states4[0].Advance(msgs4to[0])
	assert.NoError(err, "Stage 4 advance to 5 for player 0")
	states5[1], msgs5from[1], err = states4[1].Advance(msgs4to[1])
	assert.NoError(err, "Stage 4 advance to 5 for player 1")
	states5[2], msgs5from[2], err = states4[2].Advance(msgs4to[2])
	assert.NoError(err, "Stage 4 advance to 5 for player 2")

	// Reroute messages
	msgs5to := make([][]*KeyGenerationMessage5, 3)
	for i := range msgs5to {
		msgs5to[i] = make([]*KeyGenerationMessage5, 0)
		for j := range msgs5from {
			for k := range msgs5from[j] {
				if msgs5from[j][k].To == i {
					msgs5to[i] = append(msgs5to[i], msgs5from[j][k])
				}
			}
		}
	}

	// Compute state 6 from state 5
	states6 := make([]*KeyGenerationStage6, 3)
	msgs6from := make([][]*KeyGenerationMessage6, 3)

	states6[0], msgs6from[0], err = states5[0].Advance(msgs5to[0])
	assert.NoError(err, "Stage 5 advance to 6 for player 0")
	states6[1], msgs6from[1], err = states5[1].Advance(msgs5to[1])
	assert.NoError(err, "Stage 5 advance to 6 for player 1")
	states6[2], msgs6from[2], err = states5[2].Advance(msgs5to[2])
	assert.NoError(err, "Stage 5 advance to 6 for player 2")

	// Reroute messages
	msgs6to := make([][]*KeyGenerationMessage6, 3)
	for i := range msgs6to {
		msgs6to[i] = make([]*KeyGenerationMessage6, 0)
		for j := range msgs6from {
			for k := range msgs6from[j] {
				if msgs6from[j][k].To == i {
					msgs6to[i] = append(msgs6to[i], msgs6from[j][k])
				}
			}
		}
	}

	// Here we can test theta and h(i) as computed.
	// (i, h(i)) is a finite field share of lambda.beta (it autocorrects
	// for the extra factor of n! caused by the multiplication),
	// and theta is an integer share of lambda.beta.nu.nu, nu = n!.

	// First, reconstruct lambda and beta.
	lambda, beta := big.NewInt(0), big.NewInt(0)
	for i := range states5 {
		lambda.Add(lambda, states5[i].TermOfL)
		beta.Add(beta, states5[i].TermOfB)
	}
	nu := big.NewInt(6)
	nunu := big.NewInt(0).Mul(nu, nu)
	lb := big.NewInt(0).Mul(lambda, beta)

	// Reconstruct h from the h(i) shares and check that it's a secret sharing
	// of lambda.beta.
	h, err := secret.ShareCombine([]secret.Share{
		{Factor: nunu, Degree: 2, X: 1, Y: states6[0].Hi},
		{Factor: nunu, Degree: 2, X: 2, Y: states6[1].Hi},
		{Factor: nunu, Degree: 2, X: 3, Y: states6[2].Hi},
	})
	if assert.NoError(err, "Combining shares of h(i)") {
		assert.Equal(h.String(), lb.String(), "Combined H (expected) is not equal to lambda.beta (actual)")
	}
	h2, err := secret.ShareCombine([]secret.Share{
		{Factor: bigOne, Degree: 2, X: 1, Y: states6[0].Hi},
		{Factor: bigOne, Degree: 2, X: 2, Y: states6[1].Hi},
		{Factor: bigOne, Degree: 2, X: 3, Y: states6[2].Hi},
	})
	if assert.NoError(err, "Combining shares of h(i) (2)") {
		assert.Equal("0", big.NewInt(0).Mod(h2, lambda).String(), "Combined nu^2 lambda beta is not divisible by lambda")
	}

	theta0, err := secret.ShareCombine([]secret.Share{states6[0].MessageToSelf.ShareOfTheta, msgs6to[0][0].ShareOfTheta, msgs6to[0][1].ShareOfTheta})
	assert.NoError(err, "Combining shares of theta, player 0")
	theta1, err := secret.ShareCombine([]secret.Share{states6[1].MessageToSelf.ShareOfTheta, msgs6to[1][0].ShareOfTheta, msgs6to[1][1].ShareOfTheta})
	assert.NoError(err, "Combining shares of theta, player 1")
	theta2, err := secret.ShareCombine([]secret.Share{states6[2].MessageToSelf.ShareOfTheta, msgs6to[2][0].ShareOfTheta, msgs6to[2][1].ShareOfTheta})
	assert.NoError(err, "Combining shares of theta, player 2")

	assert.Equal(theta0.String(), theta1.String(), "Theta 0 and 1 differ")
	assert.Equal(theta0.String(), theta2.String(), "Theta 0 and 2 differ")

	if assert.NoError(err, "Combining shares of theta") {
		lb.Mul(lambda, beta).Mul(lb, nu).Mul(lb, nu).Mod(lb, states6[0].N)
		assert.Equal(theta0.String(), lb.String(), "Combined theta (expected) is not equiv (mod N) to lambda.beta.nu.nu (actual)")
	}
}

// Starting from a known good stage 6, test SK construction
func TestKeygenGenSK(t *testing.T) {
	assert := assert.New(t)

	states6Bytes := []byte(`[{"Parameters":{"NumberOfParticipants":3,"ParticipantIndex":0,"PaillierBitSize":512,"SecretSharingDegree":1,"SecretSharingStatisticalSecurity":20,"SecretSharingModulus":2623222045352356732683025662305979215904046499376168877804324390266714056732042608675318325182279767188274546065280153959994987994681446800680132187027204726065259578691071312441425971149525208889491572233926587610577246558357978325575132606555085287066466813663551180648008762652591343459455276595418684019307,"NumProcessors":4,"BiprimalityCheckTimes":0},"N":923588794041132910322504663489230159340645205888572410576204511494217137733476860130052233334326157027202966831612503418564291553442062779657809148246697677036362133578704288207973851468378736925970832206358261932321121488532519774503731055510154940781895101061565530550463014148475619594070354394490660675257,"I":1,"Hi":48938735228481605671506662510450651917225420709088722835831527327036543243711930821915805062356760485474508158912086145271911791041508609452706680749881605333082010585140898510660891757634755899441752457010570897966950967837887383087006028149229701618178026661292554946288902347203397401854638216743948189137442329908420367343706946673826242701381055223695141496310771597082998002652128978175203744487185810158252761746852357550692459267604898195079114769915771876490203867421540523010624150412030944174018880443891846140580169404351465610172768916571430784704146912769822075188808669369608285139381970859914135197778795,"MessageToSelf":{"From":0,"To":0,"ShareOfTheta":{"FieldSize":923588794041132910322504663489230159340645205888572410576204511494217137733476860130052233334326157027202966831612503418564291553442062779657809148246697677036362133578704288207973851468378736925970832206358261932321121488532519774503731055510154940781895101061565530550463014148475619594070354394490660675257,"Factor":null,"Degree":2,"X":1,"Y":69143941215688860464831230475269941825691340466071453383958233143863230246202903871468696661982579624617862499193110909207974006452758767359577033300026680587808301266980391282058768227665757538066324870173060324339894268637656133417342305848839031371961048497624866714522592439363659834836441051624715021753}}},{"Parameters":{"NumberOfParticipants":3,"ParticipantIndex":1,"PaillierBitSize":512,"SecretSharingDegree":1,"SecretSharingStatisticalSecurity":20,"SecretSharingModulus":2623222045352356732683025662305979215904046499376168877804324390266714056732042608675318325182279767188274546065280153959994987994681446800680132187027204726065259578691071312441425971149525208889491572233926587610577246558357978325575132606555085287066466813663551180648008762652591343459455276595418684019307,"NumProcessors":4,"BiprimalityCheckTimes":0},"N":923588794041132910322504663489230159340645205888572410576204511494217137733476860130052233334326157027202966831612503418564291553442062779657809148246697677036362133578704288207973851468378736925970832206358261932321121488532519774503731055510154940781895101061565530550463014148475619594070354394490660675257,"I":2,"Hi":48938735228481605671506662510450651917225420709088722835831527327036543243711930821915805062356760485474508158912086145271911791041508609452706680749881605333082010585140898510660891757634755899441752457010570897966950967837887383087006028149229701618178026661292554946288902347203397401854638216743948189137485448296776514499498081636331574247242734510498473234612884907078921183673112640075390810663175950953375673715009350130951317796393954035799500614761857263374926605380905050967953053023851627037874707923651368741931907289460729988952012611626013906553182932655866806990727945926452154637810872628659404425344728,"MessageToSelf":{"From":1,"To":1,"ShareOfTheta":{"FieldSize":923588794041132910322504663489230159340645205888572410576204511494217137733476860130052233334326157027202966831612503418564291553442062779657809148246697677036362133578704288207973851468378736925970832206358261932321121488532519774503731055510154940781895101061565530550463014148475619594070354394490660675257,"Factor":null,"Degree":2,"X":2,"Y":702447771470730776764579215302228473835298672923478868991860700333055915489929238076132139273120151496193356402010533914109091337173711623486201060037320421637888116011111091044164212502926722799235523137215612789306190905406125285489408806964833604440806285710342263312500498453354656936501907650323551525864}}},{"Parameters":{"NumberOfParticipants":3,"ParticipantIndex":2,"PaillierBitSize":512,"SecretSharingDegree":1,"SecretSharingStatisticalSecurity":20,"SecretSharingModulus":2623222045352356732683025662305979215904046499376168877804324390266714056732042608675318325182279767188274546065280153959994987994681446800680132187027204726065259578691071312441425971149525208889491572233926587610577246558357978325575132606555085287066466813663551180648008762652591343459455276595418684019307,"NumProcessors":4,"BiprimalityCheckTimes":0},"N":923588794041132910322504663489230159340645205888572410576204511494217137733476860130052233334326157027202966831612503418564291553442062779657809148246697677036362133578704288207973851468378736925970832206358261932321121488532519774503731055510154940781895101061565530550463014148475619594070354394490660675257,"I":3,"Hi":48938735228481605671506662510450651917225420709088722835831527327036543243711930821915805062356760485474508158912086145271911791041508609452706680749881605333082010585140898510660891757634755899441752457010570897966950967837887383087006028149229701618178026661292554946288902347203397401854638216743948189137528566685132661655289216598836905793104413797301804972914998217074844364694096301975577876839166091748498585683166342711210176325183009876519886459607942650259649343340269578925281955635672309901730535403410891343283645174569994367731256306680597028402218952541911538792647222483296024136239774397404673652910679,"MessageToSelf":{"From":2,"To":2,"ShareOfTheta":{"FieldSize":923588794041132910322504663489230159340645205888572410576204511494217137733476860130052233334326157027202966831612503418564291553442062779657809148246697677036362133578704288207973851468378736925970832206358261932321121488532519774503731055510154940781895101061565530550463014148475619594070354394490660675257,"Factor":null,"Degree":2,"X":3,"Y":412162807684639782741822536639956846504260799492313874023558656028031463000178712150743348549931566340565883473215453500445917114452601699955015938527916485651605797176537502598295805309808951134433889197899903321951366053642074663057744252570673236727756421861494129360015390318870034444097019854531727354736}}}]`)
	var states6 []*KeyGenerationStage6
	json.Unmarshal(states6Bytes, &states6)

	msgs6toBytes := []byte(`[[{"From":1,"To":0,"ShareOfTheta":{"FieldSize":923588794041132910322504663489230159340645205888572410576204511494217137733476860130052233334326157027202966831612503418564291553442062779657809148246697677036362133578704288207973851468378736925970832206358261932321121488532519774503731055510154940781895101061565530550463014148475619594070354394490660675257,"Factor":null,"Degree":2,"X":2,"Y":702447771470730776764579215302228473835298672923478868991860700333055915489929238076132139273120151496193356402010533914109091337173711623486201060037320421637888116011111091044164212502926722799235523137215612789306190905406125285489408806964833604440806285710342263312500498453354656936501907650323551525864}},{"From":2,"To":0,"ShareOfTheta":{"FieldSize":923588794041132910322504663489230159340645205888572410576204511494217137733476860130052233334326157027202966831612503418564291553442062779657809148246697677036362133578704288207973851468378736925970832206358261932321121488532519774503731055510154940781895101061565530550463014148475619594070354394490660675257,"Factor":null,"Degree":2,"X":3,"Y":412162807684639782741822536639956846504260799492313874023558656028031463000178712150743348549931566340565883473215453500445917114452601699955015938527916485651605797176537502598295805309808951134433889197899903321951366053642074663057744252570673236727756421861494129360015390318870034444097019854531727354736}}],[{"From":0,"To":1,"ShareOfTheta":{"FieldSize":923588794041132910322504663489230159340645205888572410576204511494217137733476860130052233334326157027202966831612503418564291553442062779657809148246697677036362133578704288207973851468378736925970832206358261932321121488532519774503731055510154940781895101061565530550463014148475619594070354394490660675257,"Factor":null,"Degree":2,"X":1,"Y":69143941215688860464831230475269941825691340466071453383958233143863230246202903871468696661982579624617862499193110909207974006452758767359577033300026680587808301266980391282058768227665757538066324870173060324339894268637656133417342305848839031371961048497624866714522592439363659834836441051624715021753}},{"From":2,"To":1,"ShareOfTheta":{"FieldSize":923588794041132910322504663489230159340645205888572410576204511494217137733476860130052233334326157027202966831612503418564291553442062779657809148246697677036362133578704288207973851468378736925970832206358261932321121488532519774503731055510154940781895101061565530550463014148475619594070354394490660675257,"Factor":null,"Degree":2,"X":3,"Y":412162807684639782741822536639956846504260799492313874023558656028031463000178712150743348549931566340565883473215453500445917114452601699955015938527916485651605797176537502598295805309808951134433889197899903321951366053642074663057744252570673236727756421861494129360015390318870034444097019854531727354736}}],[{"From":0,"To":2,"ShareOfTheta":{"FieldSize":923588794041132910322504663489230159340645205888572410576204511494217137733476860130052233334326157027202966831612503418564291553442062779657809148246697677036362133578704288207973851468378736925970832206358261932321121488532519774503731055510154940781895101061565530550463014148475619594070354394490660675257,"Factor":null,"Degree":2,"X":1,"Y":69143941215688860464831230475269941825691340466071453383958233143863230246202903871468696661982579624617862499193110909207974006452758767359577033300026680587808301266980391282058768227665757538066324870173060324339894268637656133417342305848839031371961048497624866714522592439363659834836441051624715021753}},{"From":1,"To":2,"ShareOfTheta":{"FieldSize":923588794041132910322504663489230159340645205888572410576204511494217137733476860130052233334326157027202966831612503418564291553442062779657809148246697677036362133578704288207973851468378736925970832206358261932321121488532519774503731055510154940781895101061565530550463014148475619594070354394490660675257,"Factor":null,"Degree":2,"X":2,"Y":702447771470730776764579215302228473835298672923478868991860700333055915489929238076132139273120151496193356402010533914109091337173711623486201060037320421637888116011111091044164212502926722799235523137215612789306190905406125285489408806964833604440806285710342263312500498453354656936501907650323551525864}}]]`)
	var msgs6to [][]*KeyGenerationMessage6
	json.Unmarshal(msgs6toBytes, &msgs6to)

	// Compute SK from state 6
	sks := make([]*PrivateKeyShare, 3)
	var err error

	sks[0], err = states6[0].Advance(msgs6to[0])
	assert.NoError(err, "Stage 6 advance to SK for player 0")
	sks[1], err = states6[1].Advance(msgs6to[1])
	assert.NoError(err, "Stage 6 advance to SK for player 1")
	sks[2], err = states6[2].Advance(msgs6to[2])
	assert.NoError(err, "Stage 6 advance to SK for player 2")

	testKeygenGenSKCorrectness(t, states6, msgs6to, sks)
}

func testKeygenGenSKCorrectness(t testing.TB, states6 []*KeyGenerationStage6, msgs6to [][]*KeyGenerationMessage6, sks []*PrivateKeyShare) {
	assert := assert.New(t)

	// Check sanity of the SKs
	for i := range sks {
		assert.Equal(states6[i].N.String(), sks[i].PublicKey.N.String(), "Public key N incorrect for player %d", i)
		assert.Equal(big.NewInt(0).Mul(states6[i].N, states6[i].N).String(), sks[i].PublicKey.N2.String(), "Public key N^2 incorrect for player %d", i)
		assert.Equal(big.NewInt(0).Add(states6[i].N, bigOne).String(), sks[i].PublicKey.Nplus1.String(), "Public key N+1 incorrect for player %d", i)

		assert.Equal(i+1, sks[i].ParticipantIndex, "Secret key participant index incorrect for player %d", i)
		assert.Equal(states6[i].Hi.String(), sks[i].Hi.String(), "Secret key h(i) incorrect for player %d", i)
		assert.Equal("6", sks[i].FactorialOfNPart.String(), "Secret key n! incorrect for player %d", i)

		theta, err := secret.ShareCombine([]secret.Share{states6[i].MessageToSelf.ShareOfTheta, msgs6to[i][0].ShareOfTheta, msgs6to[i][1].ShareOfTheta})
		if assert.NoError(err, "Can't combine shares of theta for player %d", i) {
			thetaInv := big.NewInt(0).ModInverse(theta, states6[i].N)
			assert.Equal(thetaInv.String(), sks[i].ThetaInv.String(), "Secret key theta^-1 incorrect for player %d", i)
		}
	}
}

func BenchmarkKeygenSmall(t *testing.B) {
	benchmarkKeygen(t, KeyGenerationParameters{
		NumberOfParticipants:             3,
		PaillierBitSize:                  512,
		SecretSharingDegree:              1,
		SecretSharingStatisticalSecurity: 20,
		BiprimalityCheckTimes:            10,
	})
}

func BenchmarkKeygenMedium(t *testing.B) {
	benchmarkKeygen(t, KeyGenerationParameters{
		NumberOfParticipants:             3,
		PaillierBitSize:                  1024,
		SecretSharingDegree:              1,
		SecretSharingStatisticalSecurity: 20,
		BiprimalityCheckTimes:            20,
	})
}

func BenchmarkKeygenBig(t *testing.B) {
	benchmarkKeygen(t, KeyGenerationParameters{
		NumberOfParticipants:             3,
		PaillierBitSize:                  2048,
		SecretSharingDegree:              1,
		SecretSharingStatisticalSecurity: 20,
		BiprimalityCheckTimes:            20,
	})
}

func benchmarkKeygen(t *testing.B, testConfig0 KeyGenerationParameters) {
	assert := assert.New(t)

	err := testConfig0.Validate()
	assert.NoError(err, "Validating the testing configuration")
	testConfig1 := testConfig0
	testConfig1.ParticipantIndex = 1
	testConfig2 := testConfig0
	testConfig2.ParticipantIndex = 2

	states1 := make([]*KeyGenerationStage1, 3)
	msgs1from := make([][]*KeyGenerationMessage1, 3)
	msgs1to := make([][]*KeyGenerationMessage1, 3)
	states2 := make([]*KeyGenerationStage2, 3)
	msgs2from := make([][]*KeyGenerationMessage2, 3)
	msgs2to := make([][]*KeyGenerationMessage2, 3)
	states3 := make([]*KeyGenerationStage3, 3)
	msgs3from := make([][]*KeyGenerationMessage3, 3)
	msgs3to := make([][]*KeyGenerationMessage3, 3)
	states4 := make([]*KeyGenerationStage4, 3)
	msgs4from := make([][]*KeyGenerationMessage4, 3)
	msgs4to := make([][]*KeyGenerationMessage4, 3)
	states5 := make([]*KeyGenerationStage5, 3)
	msgs5from := make([][]*KeyGenerationMessage5, 3)
	msgs5to := make([][]*KeyGenerationMessage5, 3)
	states6 := make([]*KeyGenerationStage6, 3)
	msgs6from := make([][]*KeyGenerationMessage6, 3)
	msgs6to := make([][]*KeyGenerationMessage6, 3)
	sks := make([]*PrivateKeyShare, 3)
	var err0, err1, err2 error

	tries1, tries2, tries3 := 0, 0, 0
	defer func() {
		fmt.Printf("Restarted because of the small prime test:  %d\n"+
			"Restarted because of the biprimality test:  %d\n"+
			"Restarted because theta was a divisor of N: %d\n",
			tries1, tries2, tries3,
		)
	}()
	for {
		states1[0], msgs1from[0], err = NewKeyGenerationProtocol(testConfig0)
		assert.NoError(err, "New keygen protocol for player 0")
		states1[1], msgs1from[1], err = NewKeyGenerationProtocol(testConfig1)
		assert.NoError(err, "New keygen protocol for player 1")
		states1[2], msgs1from[2], err = NewKeyGenerationProtocol(testConfig2)
		assert.NoError(err, "New keygen protocol for player 2")

		// Reroute messages
		for i := range msgs1to {
			msgs1to[i] = make([]*KeyGenerationMessage1, 0)
			for j := range msgs1from {
				for k := range msgs1from[j] {
					if msgs1from[j][k].To == i {
						msgs1to[i] = append(msgs1to[i], msgs1from[j][k])
					}
				}
			}
		}

		states2[0], msgs2from[0], err = states1[0].Advance(msgs1to[0])
		assert.NoError(err, "Stage 1 advance to 2 for player 0")
		states2[1], msgs2from[1], err = states1[1].Advance(msgs1to[1])
		assert.NoError(err, "Stage 1 advance to 2 for player 1")
		states2[2], msgs2from[2], err = states1[2].Advance(msgs1to[2])
		assert.NoError(err, "Stage 1 advance to 2 for player 2")

		// Reroute messages
		for i := range msgs2to {
			msgs2to[i] = make([]*KeyGenerationMessage2, 0)
			for j := range msgs2from {
				for k := range msgs2from[j] {
					if msgs2from[j][k].To == i {
						msgs2to[i] = append(msgs2to[i], msgs2from[j][k])
					}
				}
			}
		}

		states3[0], msgs3from[0], err = states2[0].Advance(msgs2to[0])
		assert.NoError(err, "Stage 2 advance to 3 for player 0")
		states3[1], msgs3from[1], err = states2[1].Advance(msgs2to[1])
		if err == ErrorBiprimalityTestFailed {
			// This output normally disabled to save on terminal printing time
			//fmt.Printf("Small prime test failed %v\n", time.Now())
			tries1++
			continue
		}
		assert.NoError(err, "Stage 2 advance to 3 for player 1")
		states3[2], msgs3from[2], err = states2[2].Advance(msgs2to[2])
		assert.NoError(err, "Stage 2 advance to 3 for player 2")

		// Reroute messages
		for i := range msgs3to {
			msgs3to[i] = make([]*KeyGenerationMessage3, 0)
			for j := range msgs3from {
				for k := range msgs3from[j] {
					if msgs3from[j][k].To == i {
						msgs3to[i] = append(msgs3to[i], msgs3from[j][k])
					}
				}
			}
		}

		states4[0], msgs4from[0], err0 = states3[0].Advance(msgs3to[0])
		assert.NoError(err0, "Stage 3 advance to 4 for player 0")
		states4[1], msgs4from[1], err1 = states3[1].Advance(msgs3to[1])
		assert.NoError(err1, "Stage 3 advance to 4 for player 1")
		states4[2], msgs4from[2], err2 = states3[2].Advance(msgs3to[2])
		assert.NoError(err2, "Stage 3 advance to 4 for player 2")

		// Reroute messages
		for i := range msgs4to {
			msgs4to[i] = make([]*KeyGenerationMessage4, 0)
			for j := range msgs4from {
				for k := range msgs4from[j] {
					if msgs4from[j][k].To == i {
						msgs4to[i] = append(msgs4to[i], msgs4from[j][k])
					}
				}
			}
		}

		states5[0], msgs5from[0], err0 = states4[0].Advance(msgs4to[0])
		if err0 == ErrorBiprimalityTestFailed {
			fmt.Printf("Biprimality test failed %v\n", time.Now())
			tries2++
			continue
		}
		assert.NoError(err0, "Stage 4 advance to 5 for player 0")
		states5[1], msgs5from[1], err1 = states4[1].Advance(msgs4to[1])
		assert.NoError(err1, "Stage 4 advance to 5 for player 1")
		states5[2], msgs5from[2], err2 = states4[2].Advance(msgs4to[2])
		assert.NoError(err2, "Stage 4 advance to 5 for player 2")

		// Reroute messages
		for i := range msgs5to {
			msgs5to[i] = make([]*KeyGenerationMessage5, 0)
			for j := range msgs5from {
				for k := range msgs5from[j] {
					if msgs5from[j][k].To == i {
						msgs5to[i] = append(msgs5to[i], msgs5from[j][k])
					}
				}
			}
		}

		states6[0], msgs6from[0], err0 = states5[0].Advance(msgs5to[0])
		assert.NoError(err0, "Stage 5 advance to 6 for player 0")
		states6[1], msgs6from[1], err1 = states5[1].Advance(msgs5to[1])
		assert.NoError(err1, "Stage 5 advance to 6 for player 1")
		states6[2], msgs6from[2], err2 = states5[2].Advance(msgs5to[2])
		assert.NoError(err2, "Stage 5 advance to 6 for player 2")

		// Reroute messages
		for i := range msgs6to {
			msgs6to[i] = make([]*KeyGenerationMessage6, 0)
			for j := range msgs6from {
				for k := range msgs6from[j] {
					if msgs6from[j][k].To == i {
						msgs6to[i] = append(msgs6to[i], msgs6from[j][k])
					}
				}
			}
		}

		sks[0], err0 = states6[0].Advance(msgs6to[0])
		if err0 == ErrorBiprimalityTestFailed {
			fmt.Printf("Theta was a divisor %v\n", time.Now())
			tries3++
			continue
		}
		assert.NoError(err0, "Stage 6 advance to SK for player 0")
		sks[1], err1 = states6[1].Advance(msgs6to[1])
		assert.NoError(err1, "Stage 6 advance to SK for player 1")
		sks[2], err2 = states6[2].Advance(msgs6to[2])
		assert.NoError(err2, "Stage 6 advance to SK for player 2")

		// We reached the end! Hurray!
		break
	}

	// Test if the secret keys are actually correct
	testKeygenGenSKCorrectness(t, states6, msgs6to, sks)

	// Print the secret keys
	b, _ := json.Marshal(sks)
	fmt.Println(string(b))

	t.ReportMetric(float64(tries1), "small_prime_restarts/op")
	t.ReportMetric(float64(tries2), "biprimality_restarts/op")
	t.ReportMetric(float64(tries3), "theta_divisor_restarts/op")
}