From 41db4bcea193f2c5eddd075f07ddf1ba874d1de1 Mon Sep 17 00:00:00 2001
From: Bernardo Cardoso <bernardo@elixxir.io>
Date: Wed, 1 May 2019 19:24:03 -0600
Subject: [PATCH] Add command line parameter to define E2E key generation
 parameters

---
 README.md            |  1 +
 api/client.go        |  4 ++++
 api/client_test.go   |  2 +-
 cmd/root.go          | 44 ++++++++++++++++++++++++++++++++++++++++++++
 keyStore/keyStore.go |  8 ++++----
 5 files changed, 54 insertions(+), 5 deletions(-)

diff --git a/README.md b/README.md
index d1db8870f..d546e8514 100644
--- a/README.md
+++ b/README.md
@@ -55,6 +55,7 @@ Optional args:
 |--registrationaddr|-a|Address:Port for connecting to the registration server|-a "localhost:11420"|
 |--dummyfrequency| |How often dummy messages should be sent per second. This flag is likely to be replaced when we implement better dummy message sending.|--dummyfrequency 0.5|
 |--end2end| |Send messages with E2E encryption to destination user|--end2end|
+|--keyParams| |Set E2E key generation parameters. Pass values in comma separated list, with the following order: MinKeys,MaxKeys,NumRekeys,TTLScalar,MinNumKeys|--keyParams 100,200,32,1.2,50|
 
 ##Project Structure
 
diff --git a/api/client.go b/api/client.go
index a60843f59..efbc52355 100644
--- a/api/client.go
+++ b/api/client.go
@@ -403,6 +403,10 @@ func (cl *Client) GetCurrentUser() *id.User {
 	return cl.sess.GetCurrentUser().User
 }
 
+func (cl *Client) GetKeyParams() *keyStore.KeyParams {
+	return cl.sess.GetKeyStore().GetKeyParams()
+}
+
 // Logout closes the connection to the server at this time and does
 // nothing with the user id. In the future this will release resources
 // and safely release any sensitive memory.
diff --git a/api/client_test.go b/api/client_test.go
index 1bc1d5e22..ec0af42b6 100644
--- a/api/client_test.go
+++ b/api/client_test.go
@@ -301,7 +301,7 @@ func TestRegisterUserE2E_CheckAllKeys(t *testing.T) {
 	testClient.registerUserE2E(partner, partnerPubKeyCyclic.Bytes())
 
 	// Generate all keys and confirm they all match
-	keyParams := session.GetKeyStore().GetKeyParams()
+	keyParams := testClient.GetKeyParams()
 	baseKey, _ := diffieHellman.CreateDHSessionKey(partnerPubKeyCyclic, myPrivKeyCyclic, grp)
 	keyTTL, numKeys := e2e.GenerateKeyTTL(baseKey.GetLargeInt(),
 		keyParams.MinKeys, keyParams.MaxKeys, keyParams.TTLParams)
diff --git a/cmd/root.go b/cmd/root.go
index a3fbc4ab9..fe5d64039 100644
--- a/cmd/root.go
+++ b/cmd/root.go
@@ -31,6 +31,7 @@ import (
 	"log"
 	"math/big"
 	"os"
+	"strconv"
 	"sync/atomic"
 	"time"
 )
@@ -52,6 +53,7 @@ var registrationAddr string
 var registrationCode string
 var userEmail string
 var end2end bool
+var keyParams []string
 var client *api.Client
 
 // Execute adds all child commands to the root command and sets flags
@@ -173,6 +175,40 @@ func sessionInitialization() *id.User {
 	return uid
 }
 
+func setKeyParams() {
+	minKeys, err := strconv.Atoi(keyParams[0])
+	if err != nil {
+		return
+	}
+
+	maxKeys, err := strconv.Atoi(keyParams[1])
+	if err != nil {
+		return
+	}
+
+	numRekeys, err := strconv.Atoi(keyParams[2])
+	if err != nil {
+		return
+	}
+
+	ttlScalar, err := strconv.ParseFloat(keyParams[3], 64)
+	if err != nil {
+		return
+	}
+
+	minNumKeys, err := strconv.Atoi(keyParams[4])
+	if err != nil {
+		return
+	}
+
+	params := client.GetKeyParams()
+	params.MinKeys = uint16(minKeys)
+	params.MaxKeys = uint16(maxKeys)
+	params.NumRekeys = uint16(numRekeys)
+	params.TTLScalar = ttlScalar
+	params.MinNumKeys = uint16(minNumKeys)
+}
+
 type FallbackListener struct {
 	messagesReceived int64
 }
@@ -274,6 +310,10 @@ var rootCmd = &cobra.Command{
 		SetCertPaths(gwCertPath, registrationCertPath)
 
 		userID := sessionInitialization()
+		// Set Key parameters if defined
+		if len(keyParams) == 5 {
+			setKeyParams()
+		}
 		// Set up the listeners for both of the types the client needs for
 		// the integration test
 		// Normal text messages
@@ -454,6 +494,10 @@ func init() {
 
 	rootCmd.PersistentFlags().BoolVarP(&end2end, "end2end", "", false,
 		"Send messages with E2E encryption to destination user")
+
+	rootCmd.PersistentFlags().StringArrayVarP(&keyParams, "keyParams", "",
+		make([]string, 0), "Define key generation parameters. Pass values in comma separated list"+
+			" in the following order: MinKeys,MaxKeys,NumRekeys,TTLScalar,MinNumKeys")
 }
 
 // Sets the cert paths in comms
diff --git a/keyStore/keyStore.go b/keyStore/keyStore.go
index 868a892fc..85392d0e3 100644
--- a/keyStore/keyStore.go
+++ b/keyStore/keyStore.go
@@ -110,17 +110,17 @@ type KeyStore struct {
 
 	// Reception Key Managers map
 	recvKeyManagers map[id.User]*KeyManager
-	lock sync.Mutex
+	lock            sync.Mutex
 }
 
 func NewStore() *KeyStore {
 	ks := new(KeyStore)
 	ks.params = &KeyParams{
-		MinKeys: minKeys,
-		MaxKeys: maxKeys,
+		MinKeys:   minKeys,
+		MaxKeys:   maxKeys,
 		NumRekeys: numReKeys,
 		TTLParams: e2e.TTLParams{
-			TTLScalar: ttlScalar,
+			TTLScalar:  ttlScalar,
 			MinNumKeys: threshold,
 		},
 	}
-- 
GitLab