From 168403073244f46bd1fffa5e3b28bb4175bb16a7 Mon Sep 17 00:00:00 2001
From: josh <josh@elixxir.io>
Date: Mon, 28 Sep 2020 13:12:34 -0700
Subject: [PATCH] Add rekey test to keyexchange

Go format, improve  existing keyexchange tests
---
 interfaces/params/keyExchange.go |  2 +-
 keyExchange/confirm_test.go      |  8 ++---
 keyExchange/exchange_test.go     | 43 ++++++++++++++++++++++-----
 keyExchange/rekey.go             | 14 +++++++--
 keyExchange/rekey_test.go        | 50 ++++++++++++++++++++++++++++++++
 keyExchange/trigger.go           |  2 --
 keyExchange/trigger_test.go      |  4 +--
 keyExchange/utils_test.go        |  4 ---
 storage/e2e/session.go           |  3 +-
 9 files changed, 104 insertions(+), 26 deletions(-)
 create mode 100644 keyExchange/rekey_test.go

diff --git a/interfaces/params/keyExchange.go b/interfaces/params/keyExchange.go
index 49ef1719b..712b16580 100644
--- a/interfaces/params/keyExchange.go
+++ b/interfaces/params/keyExchange.go
@@ -3,7 +3,7 @@ package params
 import "time"
 
 type Rekey struct {
-	RoundTimeout             time.Duration
+	RoundTimeout time.Duration
 }
 
 func GetDefaultRekey() Rekey {
diff --git a/keyExchange/confirm_test.go b/keyExchange/confirm_test.go
index 9088fc4f6..a9314aa05 100644
--- a/keyExchange/confirm_test.go
+++ b/keyExchange/confirm_test.go
@@ -44,10 +44,9 @@ func TestHandleConfirm(t *testing.T) {
 		SessionID: sessionID.Marshal(),
 	})
 
-
 	receiveMsg := message.Receive{
 		Payload:     rekey,
-		MessageType: message.NoType,
+		MessageType: message.KeyExchangeConfirm,
 		Sender:      bobID,
 		Timestamp:   time.Now(),
 		Encryption:  message.E2E,
@@ -62,10 +61,9 @@ func TestHandleConfirm(t *testing.T) {
 	// Check that the session is in the proper status
 	newSession := receivedManager.GetSendSession(sessionID)
 	if newSession.NegotiationStatus() != e2e.Confirmed {
-		t.Errorf("Session not in confirmed status!" +
-			"\n\tExpected: Confirmed" +
+		t.Errorf("Session not in confirmed status!"+
+			"\n\tExpected: Confirmed"+
 			"\n\tReceived: %s", confirmedSession.NegotiationStatus())
 	}
 
-
 }
diff --git a/keyExchange/exchange_test.go b/keyExchange/exchange_test.go
index 94e118772..71530f0c0 100644
--- a/keyExchange/exchange_test.go
+++ b/keyExchange/exchange_test.go
@@ -1,6 +1,7 @@
 package keyExchange
 
 import (
+	"fmt"
 	"github.com/golang/protobuf/proto"
 	"gitlab.com/elixxir/client/interfaces"
 	"gitlab.com/elixxir/client/interfaces/message"
@@ -8,6 +9,8 @@ import (
 	"gitlab.com/elixxir/client/storage"
 	"gitlab.com/elixxir/client/storage/e2e"
 	"gitlab.com/elixxir/client/switchboard"
+	"gitlab.com/elixxir/crypto/csprng"
+	dh "gitlab.com/elixxir/crypto/diffieHellman"
 	"gitlab.com/xx_network/primitives/id"
 	"testing"
 	"time"
@@ -31,6 +34,9 @@ func TestFullExchange(t *testing.T) {
 	alicePrivKey := aliceSession.E2e().GetDHPrivateKey()
 	alicePubKey := aliceSession.E2e().GetDHPublicKey()
 	bobPubKey := bobSession.E2e().GetDHPublicKey()
+	// Generate bob's new keypair
+	newBobPrivKey := dh.GeneratePrivateKey(dh.DefaultPrivateKeyLength, genericGroup, csprng.NewSystemRNG())
+	newBobPubKey := dh.GeneratePublicKey(newBobPrivKey, genericGroup)
 
 	// Add Alice and Bob as partners
 	aliceSession.E2e().AddPartner(exchangeBobId, bobPubKey,
@@ -40,32 +46,55 @@ func TestFullExchange(t *testing.T) {
 
 	// Start the listeners for alice and bob
 	rekeyParams := params.GetDefaultRekey()
-	rekeyParams.RoundTimeout = 5 * time.Second
+	rekeyParams.RoundTimeout = 1 * time.Second
 	Start(aliceSwitchboard, aliceSession, aliceManager, rekeyParams)
 	Start(bobSwitchboard, bobSession, bobManager, rekeyParams)
 
 	// Generate a session ID, bypassing some business logic here
-	sessionID := GeneratePartnerID(alicePrivKey, bobPubKey, genericGroup)
+	oldSessionID := GeneratePartnerID(alicePrivKey, bobPubKey, genericGroup)
 
 	// Generate the message
 	rekeyTrigger, _ := proto.Marshal(&RekeyTrigger{
-		SessionID: sessionID.Marshal(),
-		PublicKey: bobPubKey.Bytes(),
+		SessionID: oldSessionID.Marshal(),
+		PublicKey: newBobPubKey.Bytes(),
 	})
-	payload := make([]byte, 0)
-	payload = append(payload, rekeyTrigger...)
+
 	triggerMsg := message.Receive{
-		Payload:     payload,
+		Payload:     rekeyTrigger,
 		MessageType: message.KeyExchangeTrigger,
 		Sender:      exchangeBobId,
 		Timestamp:   time.Now(),
 		Encryption:  message.E2E,
 	}
 
+	// Get Alice's manager for reception from Bob
+	receivedManager, err := aliceSession.E2e().GetPartner(exchangeBobId)
+	if err != nil {
+		t.Errorf("Failed to get bob's manager: %v", err)
+	}
+
 	// Speak the message to Bob, triggers the SendE2E in utils_test
 	aliceSwitchboard.Speak(triggerMsg)
 
 	// Allow the test time to work it's goroutines
 	time.Sleep(1 * time.Second)
 
+	// Get Alice's session for Bob
+	confirmedSession := receivedManager.GetSendSession(oldSessionID)
+
+	// Generate the new session ID based off of Bob's new keys
+	baseKey := dh.GenerateSessionKey(alicePrivKey, newBobPubKey, genericGroup)
+	newSessionID := e2e.GetSessionIDFromBaseKeyForTesting(baseKey, t)
+
+	// Check that the Alice's session for Bob is in the proper status
+	newSession := receivedManager.GetReceiveSession(newSessionID)
+	fmt.Printf("newSession: %v\n", newSession)
+	if newSession == nil || newSession.NegotiationStatus() != e2e.Confirmed {
+		t.Errorf("Session not in confirmed status!"+
+			"\n\tExpected: Confirmed"+
+			"\n\tReceived: %s", confirmedSession.NegotiationStatus())
+	}
+
+	fmt.Printf("after status: %v\n", confirmedSession.NegotiationStatus())
+
 }
diff --git a/keyExchange/rekey.go b/keyExchange/rekey.go
index 342712fe8..923f6c9e4 100644
--- a/keyExchange/rekey.go
+++ b/keyExchange/rekey.go
@@ -7,6 +7,7 @@
 package keyExchange
 
 import (
+	"fmt"
 	"github.com/golang/protobuf/proto"
 	"github.com/pkg/errors"
 	jww "github.com/spf13/jwalterweatherman"
@@ -32,24 +33,31 @@ func CheckKeyExchanges(instance *network.Instance, sendE2E interfaces.SendE2E,
 }
 
 // There are two types of key negotiations that can be triggered, creating a new
-// session and negotiation, or resenting a negotiation for an already created
+// session and negotiation, or resetting a negotiation for an already created
 // session. They run the same negotiation, the former does it on a newly created
-// session while the latter on an extand
+// session while the latter on an extant session
 func trigger(instance *network.Instance, sendE2E interfaces.SendE2E,
 	sess *storage.Session, manager *e2e.Manager, session *e2e.Session,
 	sendTimeout time.Duration) {
 	var negotiatingSession *e2e.Session
+	fmt.Printf("session status: %v\n", session.NegotiationStatus())
 	switch session.NegotiationStatus() {
 	// If the passed session is triggering a negotiation on a new session to
 	// replace itself, then create the session
 	case e2e.NewSessionTriggered:
+		fmt.Printf("in new session triggered\n")
 		//create the session, pass a nil private key to generate a new one
 		negotiatingSession = manager.NewSendSession(nil,
 			e2e.GetDefaultSessionParams())
 		//move the state of the triggering session forward
 		session.SetNegotiationStatus(e2e.NewSessionCreated)
+		fmt.Printf("after setting session: %v\n", negotiatingSession.NegotiationStatus())
+
 	// If the session has not successfully negotiated, redo its negotiation
+	// Fixme: It doesn't seem possible to have an unconfirmed status in the codepath
 	case e2e.Unconfirmed:
+		fmt.Printf("in new Unconfirmed\n")
+
 		negotiatingSession = session
 	default:
 		jww.FATAL.Panicf("Session %s provided invalid e2e "+
@@ -133,7 +141,7 @@ func negotiate(instance *network.Instance, sendE2E interfaces.SendE2E,
 
 	// otherwise, the transmission is a success and this should be denoted
 	// in the session and the log
-	jww.INFO.Printf("Key Negotiation transmission for %s sucesfull",
+	jww.INFO.Printf("Key Negotiation transmission for %s sucesful",
 		session)
 	session.SetNegotiationStatus(e2e.Sent)
 
diff --git a/keyExchange/rekey_test.go b/keyExchange/rekey_test.go
new file mode 100644
index 000000000..8b9876ca7
--- /dev/null
+++ b/keyExchange/rekey_test.go
@@ -0,0 +1,50 @@
+package keyExchange
+
+import (
+	"gitlab.com/elixxir/client/storage/e2e"
+	"gitlab.com/xx_network/primitives/id"
+	"testing"
+	"time"
+)
+
+func TestRekey(t *testing.T) {
+	// Generate alice and bob's session
+	aliceSession, networkManager := InitTestingContextGeneric(t)
+	bobSession, _ := InitTestingContextGeneric(t)
+
+	// Pull the keys for Alice and Bob
+	alicePrivKey := aliceSession.E2e().GetDHPrivateKey()
+	bobPubKey := bobSession.E2e().GetDHPublicKey()
+
+	// Maintain an ID for bob
+	bobID := id.NewIdFromBytes([]byte("test"), t)
+
+	// Generate a session ID, bypassing some business logic here
+	sessionID := GeneratePartnerID(alicePrivKey, bobPubKey, genericGroup)
+
+	// Add bob as a partner
+	aliceSession.E2e().AddPartner(bobID, bobPubKey,
+		e2e.GetDefaultSessionParams(), e2e.GetDefaultSessionParams())
+
+	// Get Alice's manager for Bob
+	bobManager, err := aliceSession.E2e().GetPartner(bobID)
+	if err != nil {
+		t.Errorf("Bob is not recognized as Alice's partner: %v", err)
+	}
+	//// Trigger negotiations, so that negotiation statuses
+	//// can be transitioned
+	bobManager.TriggerNegotiations()
+
+	bobE2ESession := bobManager.GetSendSession(sessionID)
+
+	err = negotiate(networkManager.GetInstance(), networkManager.SendE2E, aliceSession, bobE2ESession, 1*time.Second)
+	if err != nil {
+		t.Errorf("Negotiate resulted in error: %v", err)
+	}
+
+	if bobE2ESession.NegotiationStatus() != e2e.Sent {
+		t.Errorf("Session not in expected state after negotiation."+
+			"\n\tExpected: %v"+
+			"\n\tReceived: %v", e2e.Sent, bobE2ESession.NegotiationStatus())
+	}
+}
diff --git a/keyExchange/trigger.go b/keyExchange/trigger.go
index 60fb18d2e..d7c654a46 100644
--- a/keyExchange/trigger.go
+++ b/keyExchange/trigger.go
@@ -125,13 +125,11 @@ func handleTrigger(sess *storage.Session, net interfaces.NetworkManager,
 			states.COMPLETED, states.FAILED)
 	}
 
-	fmt.Println("before tracking")
 	//Wait until the result tracking responds
 	success, numTimeOut, numRoundFail := utility.TrackResults(sendResults, len(rounds))
 	// If a single partition of the Key Negotiation request does not
 	// transmit, the partner will not be able to read the confirmation. If
 	// such a failure occurs
-	fmt.Println("after tracking")
 	if !success {
 		jww.ERROR.Printf("Key Negotiation for %s failed to "+
 			"transmit %v/%v paritions: %v round failures, %v timeouts",
diff --git a/keyExchange/trigger_test.go b/keyExchange/trigger_test.go
index 824e5f752..be1cdf8db 100644
--- a/keyExchange/trigger_test.go
+++ b/keyExchange/trigger_test.go
@@ -84,10 +84,8 @@ func TestHandleTrigger(t *testing.T) {
 	// Check that this session with unrecognized keys is not valid
 	badSession := receivedManager.GetReceiveSession(badSessionID)
 	if badSession != nil {
-		t.Errorf("Alice found a session from an unknown keypair. " +
+		t.Errorf("Alice found a session from an unknown keypair. "+
 			"\nSession: %v", badSession)
 	}
 
-
-
 }
diff --git a/keyExchange/utils_test.go b/keyExchange/utils_test.go
index da439420b..7c1f43f61 100644
--- a/keyExchange/utils_test.go
+++ b/keyExchange/utils_test.go
@@ -108,10 +108,6 @@ func InitTestingContextGeneric(i interface{}) (*storage.Session, interfaces.Netw
 		Manager: commsManager,
 	}
 
-	_, err := instanceComms.AddHost(&id.Permissioning, "0.0.0.0:420", []byte(pub), connect.GetDefaultHostParams())
-	if err != nil {
-		return nil, nil
-	}
 	thisInstance, err := network.NewInstanceTesting(instanceComms, def, def, nil, nil, i)
 	if err != nil {
 		return nil, nil
diff --git a/storage/e2e/session.go b/storage/e2e/session.go
index e2e23221d..b512368e7 100644
--- a/storage/e2e/session.go
+++ b/storage/e2e/session.go
@@ -234,7 +234,6 @@ func GetSessionIDFromBaseKeyForTesting(baseKey *cyclic.Int, i interface{}) Sessi
 	return getSessionIDFromBaseKey(baseKey)
 }
 
-
 //Blake2B hash of base key used for storage
 func (s *Session) GetID() SessionID {
 	return getSessionIDFromBaseKey(s.baseKey)
@@ -447,6 +446,8 @@ func (s *Session) triggerNegotiation() bool {
 			s.mux.Unlock()
 			return false
 		}
+		// fixme: Is this a bug? In rekey.go, it seems a session would never be unconfirmed
+		//  as it would be set to sending. Possible that this is wrong or the switch statement is
 	} else if s.negotiationStatus == Unconfirmed {
 		// retrigger this sessions negotiation
 		s.mux.RUnlock()
-- 
GitLab