From aa5534b7cf77ac58266cefa8837d04a1c83f0760 Mon Sep 17 00:00:00 2001
From: Jono Wenger <jono@elixxir.io>
Date: Tue, 12 Oct 2021 20:50:27 +0000
Subject: [PATCH] Replace UD host and contact in memory with looking up from
 the NDF every time

---
 .gitlab-ci.yml         |   1 +
 ud/addFact.go          |  20 ++++--
 ud/addFact_test.go     |  27 ++++----
 ud/confirmFact.go      |  13 +++-
 ud/confirmFact_test.go |  14 ++--
 ud/lookup.go           |  10 ++-
 ud/lookup_test.go      |  14 +++-
 ud/manager.go          | 103 +++++++++++++++++++----------
 ud/register.go         |   8 ++-
 ud/register_test.go    |  14 ++--
 ud/remove.go           |  30 ++++++---
 ud/remove_test.go      |  48 ++++++++------
 ud/search.go           |  10 ++-
 ud/search_test.go      |  28 ++++++--
 ud/utils_test.go       | 146 +++++++++++++++++++++++++++++++++++++++++
 15 files changed, 371 insertions(+), 115 deletions(-)
 create mode 100644 ud/utils_test.go

diff --git a/.gitlab-ci.yml b/.gitlab-ci.yml
index fa85eb1c1..11649fc9a 100644
--- a/.gitlab-ci.yml
+++ b/.gitlab-ci.yml
@@ -8,6 +8,7 @@ before_script:
   - mkdir -p ~/.ssh
   - chmod 700 ~/.ssh
   - ssh-keyscan -t rsa $GITLAB_SERVER > ~/.ssh/known_hosts
+  - rm -rf ~/.gitconfig
   - git config --global url."git@$GITLAB_SERVER:".insteadOf "https://gitlab.com/"
   - git config --global url."git@$GITLAB_SERVER:".insteadOf "https://git.xx.network/" --add
   - export PATH=$HOME/go/bin:$PATH
diff --git a/ud/addFact.go b/ud/addFact.go
index a32668f71..1fc1a15f0 100644
--- a/ud/addFact.go
+++ b/ud/addFact.go
@@ -17,9 +17,9 @@ type addFactComms interface {
 	SendRegisterFact(host *connect.Host, message *pb.FactRegisterRequest) (*pb.FactRegisterResponse, error)
 }
 
-// Adds a fact for the user to user discovery. Will only succeed if the
-// user is already registered and the system does not have the fact currently
-// registered for any user.
+// SendRegisterFact adds a fact for the user to user discovery. Will only
+// succeed if the user is already registered and the system does not have the
+// fact currently registered for any user.
 // This does not complete the fact registration process, it returns a
 // confirmation id instead. Over the communications system the fact is
 // associated with, a code will be sent. This confirmation ID needs to be
@@ -43,10 +43,10 @@ func (m *Manager) addFact(inFact fact.Fact, uid *id.ID, aFC addFactComms) (strin
 	}
 
 	// Create a hash of our fact
-	fhash := factID.Fingerprint(f)
+	fHash := factID.Fingerprint(f)
 
 	// Sign our inFact for putting into the request
-	fsig, err := rsa.Sign(rand.Reader, m.privKey, hash.CMixHash, fhash, nil)
+	fSig, err := rsa.Sign(rand.Reader, m.privKey, hash.CMixHash, fHash, nil)
 	if err != nil {
 		return "", err
 	}
@@ -58,11 +58,17 @@ func (m *Manager) addFact(inFact fact.Fact, uid *id.ID, aFC addFactComms) (strin
 			Fact:     inFact.Fact,
 			FactType: uint32(inFact.T),
 		},
-		FactSig: fsig,
+		FactSig: fSig,
+	}
+
+	// Get UD host
+	host, err := m.getHost()
+	if err != nil {
+		return "", err
 	}
 
 	// Send the message
-	response, err := aFC.SendRegisterFact(m.host, &remFactMsg)
+	response, err := aFC.SendRegisterFact(host, &remFactMsg)
 
 	confirmationID := ""
 	if response != nil {
diff --git a/ud/addFact_test.go b/ud/addFact_test.go
index 3e1ce5ad3..eebb92a1d 100644
--- a/ud/addFact_test.go
+++ b/ud/addFact_test.go
@@ -1,6 +1,7 @@
 package ud
 
 import (
+	"gitlab.com/elixxir/comms/client"
 	pb "gitlab.com/elixxir/comms/mixmessages"
 	"gitlab.com/elixxir/primitives/fact"
 	"gitlab.com/xx_network/comms/connect"
@@ -12,20 +13,16 @@ import (
 
 type testAFC struct{}
 
-// Dummy implementation of SendRegisterFact so we don't need
-// to run our own UDB server
-func (rFC *testAFC) SendRegisterFact(host *connect.Host, message *pb.FactRegisterRequest) (*pb.FactRegisterResponse, error) {
+// Dummy implementation of SendRegisterFact so that we don't need to run our own
+// UDB server.
+func (rFC *testAFC) SendRegisterFact(*connect.Host, *pb.FactRegisterRequest) (
+	*pb.FactRegisterResponse, error) {
 	return &pb.FactRegisterResponse{}, nil
 }
 
 // Test that the addFact function completes successfully
 func TestAddFact(t *testing.T) {
 	isReg := uint32(1)
-	// Add our host, addFact uses it to get the ID of the user
-	h, err := connect.NewHost(&id.DummyUser, "address", nil, connect.GetDefaultHostParams())
-	if err != nil {
-		t.Fatal(err)
-	}
 
 	// Create a new Private Key to use for signing the Fact
 	rng := csprng.NewSystemRNG()
@@ -34,9 +31,15 @@ func TestAddFact(t *testing.T) {
 		t.Fatal(err)
 	}
 
+	comms, err := client.NewClientComms(nil, nil, nil, nil)
+	if err != nil {
+		t.Errorf("Failed to start client comms: %+v", err)
+	}
+
 	// Create our Manager object
 	m := Manager{
-		host:       h,
+		comms:      comms,
+		net:        newTestNetworkManager(t),
 		privKey:    cpk,
 		registered: &isReg,
 	}
@@ -49,13 +52,13 @@ func TestAddFact(t *testing.T) {
 		T:    2,
 	}
 
-	// Setup a dummy comms that implements SendRegisterFact
+	// Set up a dummy comms that implements SendRegisterFact
 	// This way we don't need to run UDB just to check that this
 	// function works.
-	tafc := testAFC{}
+	tAFC := testAFC{}
 	uid := &id.ID{}
 	// Run addFact and see if it returns without an error!
-	_, err = m.addFact(f, uid, &tafc)
+	_, err = m.addFact(f, uid, &tAFC)
 	if err != nil {
 		t.Fatal(err)
 	}
diff --git a/ud/confirmFact.go b/ud/confirmFact.go
index a625e7e63..b7c4294e2 100644
--- a/ud/confirmFact.go
+++ b/ud/confirmFact.go
@@ -12,8 +12,9 @@ type confirmFactComm interface {
 	SendConfirmFact(host *connect.Host, message *pb.FactConfirmRequest) (*messages.Ack, error)
 }
 
-// Confirms a fact first registered via AddFact. The confirmation ID comes from
-// AddFact while the code will come over the associated communications system
+// SendConfirmFact confirms a fact first registered via AddFact. The
+// confirmation ID comes from AddFact while the code will come over the
+// associated communications system.
 func (m *Manager) SendConfirmFact(confirmationID, code string) error {
 	jww.INFO.Printf("ud.SendConfirmFact(%s, %s)", confirmationID, code)
 	if err := m.confirmFact(confirmationID, code, m.comms); err != nil {
@@ -28,10 +29,16 @@ func (m *Manager) confirmFact(confirmationID, code string, comm confirmFactComm)
 			"client is not registered")
 	}
 
+	// Get UD host
+	host, err := m.getHost()
+	if err != nil {
+		return err
+	}
+
 	msg := &pb.FactConfirmRequest{
 		ConfirmationID: confirmationID,
 		Code:           code,
 	}
-	_, err := comm.SendConfirmFact(m.host, msg)
+	_, err = comm.SendConfirmFact(host, msg)
 	return err
 }
diff --git a/ud/confirmFact_test.go b/ud/confirmFact_test.go
index 2e060c098..9fe29b0a4 100644
--- a/ud/confirmFact_test.go
+++ b/ud/confirmFact_test.go
@@ -1,10 +1,10 @@
 package ud
 
 import (
+	"gitlab.com/elixxir/comms/client"
 	pb "gitlab.com/elixxir/comms/mixmessages"
 	"gitlab.com/xx_network/comms/connect"
 	"gitlab.com/xx_network/comms/messages"
-	"gitlab.com/xx_network/primitives/id"
 	"reflect"
 	"testing"
 )
@@ -20,17 +20,17 @@ func (t *testComm) SendConfirmFact(_ *connect.Host, message *pb.FactConfirmReque
 
 // Happy path.
 func TestManager_confirmFact(t *testing.T) {
-	// Create new host
-	host, err := connect.NewHost(&id.UDB, "0.0.0.0", nil, connect.GetDefaultHostParams())
+	isReg := uint32(1)
+
+	comms, err := client.NewClientComms(nil, nil, nil, nil)
 	if err != nil {
-		t.Fatalf("Could not create a new host: %+v", err)
+		t.Errorf("Failed to start client comms: %+v", err)
 	}
 
-	isReg := uint32(1)
-
 	// Set up manager
 	m := &Manager{
-		host:       host,
+		comms:      comms,
+		net:        newTestNetworkManager(t),
 		registered: &isReg,
 	}
 
diff --git a/ud/lookup.go b/ud/lookup.go
index e7bdb0da0..2adc57a4f 100644
--- a/ud/lookup.go
+++ b/ud/lookup.go
@@ -1,7 +1,6 @@
 package ud
 
 import (
-	"fmt"
 	"github.com/golang/protobuf/proto"
 	"github.com/pkg/errors"
 	jww "github.com/spf13/jwalterweatherman"
@@ -38,7 +37,13 @@ func (m *Manager) Lookup(uid *id.ID, callback lookupCallback, timeout time.Durat
 		m.lookupResponseProcess(uid, callback, payload, err)
 	}
 
-	err = m.single.TransmitSingleUse(m.udContact, requestMarshaled, LookupTag,
+	// Get UD contact
+	c, err := m.getContact()
+	if err != nil {
+		return err
+	}
+
+	err = m.single.TransmitSingleUse(c, requestMarshaled, LookupTag,
 		maxLookupMessages, f, timeout)
 	if err != nil {
 		return errors.WithMessage(err, "Failed to transmit lookup request.")
@@ -67,7 +72,6 @@ func (m *Manager) lookupResponseProcess(uid *id.ID, callback lookupCallback,
 		return
 	}
 
-	fmt.Printf("pubKey: %+v\n", lookupResponse.PubKey)
 	c := contact.Contact{
 		ID:       uid,
 		DhPubKey: m.grp.NewIntFromBytes(lookupResponse.PubKey),
diff --git a/ud/lookup_test.go b/ud/lookup_test.go
index 02162faaf..63cb7ab60 100644
--- a/ud/lookup_test.go
+++ b/ud/lookup_test.go
@@ -5,6 +5,8 @@ import (
 	"github.com/pkg/errors"
 	"gitlab.com/elixxir/client/single"
 	"gitlab.com/elixxir/client/stoppable"
+	"gitlab.com/elixxir/client/storage"
+	"gitlab.com/elixxir/comms/client"
 	"gitlab.com/elixxir/crypto/contact"
 	"gitlab.com/elixxir/crypto/cyclic"
 	"gitlab.com/xx_network/crypto/large"
@@ -20,9 +22,17 @@ import (
 func TestManager_Lookup(t *testing.T) {
 	// Set up manager
 	isReg := uint32(1)
+
+	comms, err := client.NewClientComms(nil, nil, nil, nil)
+	if err != nil {
+		t.Errorf("Failed to start client comms: %+v", err)
+	}
+
 	m := &Manager{
+		comms:      comms,
+		storage:    storage.InitTestingSession(t),
+		net:        newTestNetworkManager(t),
 		grp:        cyclic.NewGroup(large.NewInt(107), large.NewInt(2)),
-		udContact:  contact.Contact{ID: &id.UDB},
 		single:     &mockSingleLookup{},
 		registered: &isReg,
 	}
@@ -41,7 +51,7 @@ func TestManager_Lookup(t *testing.T) {
 	uid := id.NewIdFromUInt(0x500000000000000, id.User, t)
 
 	// Run the lookup
-	err := m.Lookup(uid, callback, 10*time.Millisecond)
+	err = m.Lookup(uid, callback, 10*time.Millisecond)
 	if err != nil {
 		t.Errorf("Lookup() returned an error: %+v", err)
 	}
diff --git a/ud/manager.go b/ud/manager.go
index 9342419bd..ccf9e3311 100644
--- a/ud/manager.go
+++ b/ud/manager.go
@@ -35,59 +35,43 @@ type Manager struct {
 	net     interfaces.NetworkManager
 
 	// Loaded from external access
-	udContact contact.Contact
-	privKey   *rsa.PrivateKey
-	grp       *cyclic.Group
+	privKey *rsa.PrivateKey
+	grp     *cyclic.Group
 
 	// internal structures
-	host   *connect.Host
 	single SingleInterface
 	myID   *id.ID
 
 	registered *uint32
 }
 
-// New manager builds a new user discovery manager. It requires that an
-// updated NDF is available and will error if one is not.
+// NewManager builds a new user discovery manager. It requires that an updated
+// NDF is available and will error if one is not.
 func NewManager(client *api.Client, single *single.Manager) (*Manager, error) {
 	jww.INFO.Println("ud.NewManager()")
 	if client.NetworkFollowerStatus() != api.Running {
-		return nil, errors.New("cannot start UD Manager when network follower is not " +
-			"running.")
+		return nil, errors.New(
+			"cannot start UD Manager when network follower is not running.")
 	}
 
 	m := &Manager{
-		client:    client,
-		comms:     client.GetComms(),
-		rng:       client.GetRng(),
-		sw:        client.GetSwitchboard(),
-		storage:   client.GetStorage(),
-		net:       client.GetNetworkInterface(),
-		udContact: contact.Contact{},
-		single:    single,
+		client:  client,
+		comms:   client.GetComms(),
+		rng:     client.GetRng(),
+		sw:      client.GetSwitchboard(),
+		storage: client.GetStorage(),
+		net:     client.GetNetworkInterface(),
+		single:  single,
 	}
 
-	var err error
-
-	// check that user discovery is available in the ndf
+	// check that user discovery is available in the NDF
 	def := m.net.GetInstance().GetPartialNdf().Get()
-	if m.udContact.ID, err = id.Unmarshal(def.UDB.ID); err != nil {
-		return nil, errors.WithMessage(err, "NDF does not have User Discovery "+
-			"information; is there network access?: ID could not be "+
-			"unmarshaled.")
-	}
 
 	if def.UDB.Cert == "" {
 		return nil, errors.New("NDF does not have User Discovery information, " +
 			"is there network access?: Cert not present.")
 	}
 
-	// Unmarshal UD DH public key
-	m.udContact.DhPubKey = m.storage.E2e().GetGroup().NewInt(1)
-	if err = m.udContact.DhPubKey.UnmarshalJSON(def.UDB.DhPubKey); err != nil {
-		return nil, errors.WithMessage(err, "Failed to unmarshal UD DH public key.")
-	}
-
 	// Create the user discovery host object
 	hp := connect.GetDefaultHostParams()
 	// Client will not send KeepAlive packets
@@ -95,11 +79,6 @@ func NewManager(client *api.Client, single *single.Manager) (*Manager, error) {
 	hp.MaxRetries = 3
 	hp.SendTimeout = 3 * time.Second
 	hp.AuthEnabled = false
-	m.host, err = m.comms.AddHost(&id.UDB, def.UDB.Address, []byte(def.UDB.Cert), hp)
-	if err != nil {
-		return nil, errors.WithMessage(err, "User Discovery host object could "+
-			"not be constructed.")
-	}
 
 	m.myID = m.storage.User().GetCryptographicIdentity().GetReceptionID()
 
@@ -114,3 +93,57 @@ func NewManager(client *api.Client, single *single.Manager) (*Manager, error) {
 
 	return m, nil
 }
+
+// getHost returns the current UD host for the UD ID found in the NDF. If the
+// host does not exist, then it is added and returned
+func (m *Manager) getHost() (*connect.Host, error) {
+	netDef := m.net.GetInstance().GetFullNdf().Get()
+
+	// Unmarshal UD ID from the NDF
+	udID, err := id.Unmarshal(netDef.UDB.ID)
+	if err != nil {
+		return nil, errors.Errorf("failed to unmarshal UD ID from NDF: %+v", err)
+	}
+
+	// Return the host, if it exists
+	host, exists := m.comms.GetHost(udID)
+	if exists {
+		return host, nil
+	}
+
+	// Add a new host and return it if it does not already exist
+	host, err = m.comms.AddHost(udID, netDef.UDB.Address,
+		[]byte(netDef.UDB.Cert), connect.GetDefaultHostParams())
+	if err != nil {
+		return nil, errors.WithMessage(err, "User Discovery host object could "+
+			"not be constructed.")
+	}
+
+	return host, nil
+}
+
+// getContact returns the contact for UD as retrieved from the NDF.
+func (m *Manager) getContact() (contact.Contact, error) {
+	netDef := m.net.GetInstance().GetFullNdf().Get()
+
+	// Unmarshal UD ID from the NDF
+	udID, err := id.Unmarshal(netDef.UDB.ID)
+	if err != nil {
+		return contact.Contact{},
+			errors.Errorf("failed to unmarshal UD ID from NDF: %+v", err)
+	}
+
+	// Unmarshal UD DH public key
+	dhPubKey := m.storage.E2e().GetGroup().NewInt(1)
+	if err = dhPubKey.UnmarshalJSON(netDef.UDB.DhPubKey); err != nil {
+		return contact.Contact{},
+			errors.WithMessage(err, "Failed to unmarshal UD DH public key.")
+	}
+
+	return contact.Contact{
+		ID:             udID,
+		DhPubKey:       dhPubKey,
+		OwnershipProof: nil,
+		Facts:          nil,
+	}, nil
+}
diff --git a/ud/register.go b/ud/register.go
index 8a25bafc0..f6b205a06 100644
--- a/ud/register.go
+++ b/ud/register.go
@@ -81,8 +81,14 @@ func (m *Manager) register(username string, comm registerUserComms) error {
 		FactSig: signedFact,
 	}
 
+	// Get UD host
+	host, err := m.getHost()
+	if err != nil {
+		return err
+	}
+
 	// Register user with user discovery
-	_, err = comm.SendRegisterUser(m.host, msg)
+	_, err = comm.SendRegisterUser(host, msg)
 
 	if err == nil {
 		err = m.setRegistered()
diff --git a/ud/register_test.go b/ud/register_test.go
index 7cfd66f1d..0bdd3ef54 100644
--- a/ud/register_test.go
+++ b/ud/register_test.go
@@ -3,6 +3,7 @@ package ud
 import (
 	"bytes"
 	"gitlab.com/elixxir/client/storage"
+	"gitlab.com/elixxir/comms/client"
 	pb "gitlab.com/elixxir/comms/mixmessages"
 	"gitlab.com/elixxir/crypto/factID"
 	"gitlab.com/elixxir/crypto/fastRNG"
@@ -12,7 +13,6 @@ import (
 	"gitlab.com/xx_network/comms/messages"
 	"gitlab.com/xx_network/crypto/csprng"
 	"gitlab.com/xx_network/crypto/signature/rsa"
-	"gitlab.com/xx_network/primitives/id"
 	"reflect"
 	"testing"
 )
@@ -28,17 +28,17 @@ func (t *testRegisterComm) SendRegisterUser(_ *connect.Host, msg *pb.UDBUserRegi
 
 // Happy path.
 func TestManager_register(t *testing.T) {
-	// Create new host
-	host, err := connect.NewHost(&id.UDB, "0.0.0.0", nil, connect.GetDefaultHostParams())
+	isReg := uint32(0)
+
+	comms, err := client.NewClientComms(nil, nil, nil, nil)
 	if err != nil {
-		t.Fatalf("Could not create a new host: %+v", err)
+		t.Errorf("Failed to start client comms: %+v", err)
 	}
 
-	isReg := uint32(0)
-
 	// Set up manager
 	m := &Manager{
-		host:       host,
+		comms:      comms,
+		net:        newTestNetworkManager(t),
 		rng:        fastRNG.NewStreamGenerator(12, 3, csprng.NewSystemRNG),
 		storage:    storage.InitTestingSession(t),
 		registered: &isReg,
diff --git a/ud/remove.go b/ud/remove.go
index 99d9447f3..67f672177 100644
--- a/ud/remove.go
+++ b/ud/remove.go
@@ -17,7 +17,7 @@ type removeFactComms interface {
 	SendRemoveFact(host *connect.Host, message *mixmessages.FactRemovalRequest) (*messages.Ack, error)
 }
 
-// Removes a previously confirmed fact.  Will fail if the fact is not
+// RemoveFact removes a previously confirmed fact. Will fail if the fact is not
 // associated with this client.
 func (m *Manager) RemoveFact(fact fact.Fact) error {
 	jww.INFO.Printf("ud.RemoveFact(%s)", fact.Stringify())
@@ -38,10 +38,10 @@ func (m *Manager) removeFact(fact fact.Fact, rFC removeFactComms) error {
 	}
 
 	// Create a hash of our fact
-	fhash := factID.Fingerprint(fact)
+	fHash := factID.Fingerprint(fact)
 
 	// Sign our inFact for putting into the request
-	fsig, err := rsa.Sign(rand.Reader, m.privKey, hash.CMixHash, fhash, nil)
+	fSig, err := rsa.Sign(rand.Reader, m.privKey, hash.CMixHash, fHash, nil)
 	if err != nil {
 		return err
 	}
@@ -50,11 +50,17 @@ func (m *Manager) removeFact(fact fact.Fact, rFC removeFactComms) error {
 	remFactMsg := mixmessages.FactRemovalRequest{
 		UID:         m.myID.Marshal(),
 		RemovalData: &mmFact,
-		FactSig:     fsig,
+		FactSig:     fSig,
+	}
+
+	// Get UD host
+	host, err := m.getHost()
+	if err != nil {
+		return err
 	}
 
 	// Send the message
-	_, err = rFC.SendRemoveFact(m.host, &remFactMsg)
+	_, err = rFC.SendRemoveFact(host, &remFactMsg)
 
 	// Return the error
 	return err
@@ -64,7 +70,7 @@ type removeUserComms interface {
 	SendRemoveUser(host *connect.Host, message *mixmessages.FactRemovalRequest) (*messages.Ack, error)
 }
 
-// Removes a previously confirmed fact.  Will fail if the fact is not
+// RemoveUser removes a previously confirmed fact. Will fail if the fact is not
 // associated with this client.
 func (m *Manager) RemoveUser(fact fact.Fact) error {
 	jww.INFO.Printf("ud.RemoveUser(%s)", fact.Stringify())
@@ -85,10 +91,10 @@ func (m *Manager) removeUser(fact fact.Fact, rFC removeUserComms) error {
 	}
 
 	// Create a hash of our fact
-	fhash := factID.Fingerprint(fact)
+	fHash := factID.Fingerprint(fact)
 
 	// Sign our inFact for putting into the request
-	fsig, err := rsa.Sign(rand.Reader, m.privKey, hash.CMixHash, fhash, nil)
+	fsig, err := rsa.Sign(rand.Reader, m.privKey, hash.CMixHash, fHash, nil)
 	if err != nil {
 		return err
 	}
@@ -100,8 +106,14 @@ func (m *Manager) removeUser(fact fact.Fact, rFC removeUserComms) error {
 		FactSig:     fsig,
 	}
 
+	// Get UD host
+	host, err := m.getHost()
+	if err != nil {
+		return err
+	}
+
 	// Send the message
-	_, err = rFC.SendRemoveUser(m.host, &remFactMsg)
+	_, err = rFC.SendRemoveUser(host, &remFactMsg)
 
 	// Return the error
 	return err
diff --git a/ud/remove_test.go b/ud/remove_test.go
index d65cb0e8c..5705a843a 100644
--- a/ud/remove_test.go
+++ b/ud/remove_test.go
@@ -1,6 +1,7 @@
 package ud
 
 import (
+	"gitlab.com/elixxir/comms/client"
 	pb "gitlab.com/elixxir/comms/mixmessages"
 	"gitlab.com/elixxir/primitives/fact"
 	"gitlab.com/xx_network/comms/connect"
@@ -13,16 +14,12 @@ import (
 
 type testRFC struct{}
 
-func (rFC *testRFC) SendRemoveFact(host *connect.Host, message *pb.FactRemovalRequest) (*messages.Ack, error) {
+func (rFC *testRFC) SendRemoveFact(*connect.Host, *pb.FactRemovalRequest) (
+	*messages.Ack, error) {
 	return &messages.Ack{}, nil
 }
 
 func TestRemoveFact(t *testing.T) {
-	h, err := connect.NewHost(&id.DummyUser, "address", nil, connect.GetDefaultHostParams())
-	if err != nil {
-		t.Fatal(err)
-	}
-
 	rng := csprng.NewSystemRNG()
 	cpk, err := rsa.GenerateKey(rng, 2048)
 	if err != nil {
@@ -31,9 +28,15 @@ func TestRemoveFact(t *testing.T) {
 
 	isReg := uint32(1)
 
-	m := Manager{
-		comms:      nil,
-		host:       h,
+	comms, err := client.NewClientComms(nil, nil, nil, nil)
+	if err != nil {
+		t.Errorf("Failed to start client comms: %+v", err)
+	}
+
+	// Set up manager
+	m := &Manager{
+		comms:      comms,
+		net:        newTestNetworkManager(t),
 		privKey:    cpk,
 		registered: &isReg,
 		myID:       &id.ID{},
@@ -44,23 +47,20 @@ func TestRemoveFact(t *testing.T) {
 		T:    2,
 	}
 
-	trfc := testRFC{}
+	tRFC := testRFC{}
 
-	err = m.removeFact(f, &trfc)
+	err = m.removeFact(f, &tRFC)
 	if err != nil {
 		t.Fatal(err)
 	}
 }
 
-func (rFC *testRFC) SendRemoveUser(host *connect.Host, message *pb.FactRemovalRequest) (*messages.Ack, error) {
+func (rFC *testRFC) SendRemoveUser(*connect.Host, *pb.FactRemovalRequest) (
+	*messages.Ack, error) {
 	return &messages.Ack{}, nil
 }
 
 func TestRemoveUser(t *testing.T) {
-	h, err := connect.NewHost(&id.DummyUser, "address", nil, connect.GetDefaultHostParams())
-	if err != nil {
-		t.Fatal(err)
-	}
 
 	rng := csprng.NewSystemRNG()
 	cpk, err := rsa.GenerateKey(rng, 2048)
@@ -70,9 +70,15 @@ func TestRemoveUser(t *testing.T) {
 
 	isReg := uint32(1)
 
-	m := Manager{
-		comms:      nil,
-		host:       h,
+	comms, err := client.NewClientComms(nil, nil, nil, nil)
+	if err != nil {
+		t.Errorf("Failed to start client comms: %+v", err)
+	}
+
+	// Set up manager
+	m := &Manager{
+		comms:      comms,
+		net:        newTestNetworkManager(t),
 		privKey:    cpk,
 		registered: &isReg,
 		myID:       &id.ID{},
@@ -83,9 +89,9 @@ func TestRemoveUser(t *testing.T) {
 		T:    2,
 	}
 
-	trfc := testRFC{}
+	tRFC := testRFC{}
 
-	err = m.removeUser(f, &trfc)
+	err = m.removeUser(f, &tRFC)
 	if err != nil {
 		t.Fatal(err)
 	}
diff --git a/ud/search.go b/ud/search.go
index 410012348..86c51c5b5 100644
--- a/ud/search.go
+++ b/ud/search.go
@@ -45,7 +45,13 @@ func (m *Manager) Search(list fact.FactList, callback searchCallback, timeout ti
 		m.searchResponseHandler(factMap, callback, payload, err)
 	}
 
-	err = m.single.TransmitSingleUse(m.udContact, requestMarshaled, SearchTag,
+	// Get UD contact
+	c, err := m.getContact()
+	if err != nil {
+		return err
+	}
+
+	err = m.single.TransmitSingleUse(c, requestMarshaled, SearchTag,
 		maxSearchMessages, f, timeout)
 	if err != nil {
 		return errors.WithMessage(err, "Failed to transmit search request.")
@@ -86,7 +92,7 @@ func (m *Manager) searchResponseHandler(factMap map[string]fact.Fact,
 		return
 	}
 
-	//return an error if no facts are found
+	// return an error if no facts are found
 	if len(searchResponse.Contacts) == 0 {
 		go callback(nil, errors.New("No contacts found in search"))
 	}
diff --git a/ud/search_test.go b/ud/search_test.go
index c333cf08c..d0a1edfb9 100644
--- a/ud/search_test.go
+++ b/ud/search_test.go
@@ -3,9 +3,11 @@ package ud
 import (
 	"fmt"
 	"github.com/golang/protobuf/proto"
-	errors "github.com/pkg/errors"
+	"github.com/pkg/errors"
 	"gitlab.com/elixxir/client/single"
 	"gitlab.com/elixxir/client/stoppable"
+	"gitlab.com/elixxir/client/storage"
+	"gitlab.com/elixxir/comms/client"
 	"gitlab.com/elixxir/crypto/contact"
 	"gitlab.com/elixxir/crypto/cyclic"
 	"gitlab.com/elixxir/crypto/factID"
@@ -23,10 +25,19 @@ import (
 func TestManager_Search(t *testing.T) {
 	// Set up manager
 	isReg := uint32(1)
-	grp := cyclic.NewGroup(large.NewInt(107), large.NewInt(2))
+
+	comms, err := client.NewClientComms(nil, nil, nil, nil)
+	if err != nil {
+		t.Errorf("Failed to start client comms: %+v", err)
+	}
+
+	store := storage.InitTestingSession(t)
+
 	m := &Manager{
-		grp:        grp,
-		udContact:  contact.Contact{ID: &id.UDB, DhPubKey: grp.NewInt(42)},
+		comms:      comms,
+		storage:    store,
+		net:        newTestNetworkManager(t),
+		grp:        store.E2e().GetGroup(),
 		single:     &mockSingleSearch{},
 		registered: &isReg,
 	}
@@ -62,7 +73,7 @@ func TestManager_Search(t *testing.T) {
 		})
 	}
 
-	err := m.Search(factList, callback, 10*time.Millisecond)
+	err = m.Search(factList, callback, 10*time.Millisecond)
 	if err != nil {
 		t.Errorf("Search() returned an error: %+v", err)
 	}
@@ -74,7 +85,12 @@ func TestManager_Search(t *testing.T) {
 			t.Errorf("Callback returned an error: %+v", cb.err)
 		}
 
-		expectedContacts := []contact.Contact{m.udContact}
+		c, err := m.getContact()
+		if err != nil {
+			t.Errorf("Failed to get UD contact: %+v", err)
+		}
+
+		expectedContacts := []contact.Contact{c}
 		if !contact.Equal(expectedContacts[0], cb.c[0]) {
 			t.Errorf("Failed to get expected Contacts."+
 				"\n\texpected: %+v\n\treceived: %+v", expectedContacts, cb.c)
diff --git a/ud/utils_test.go b/ud/utils_test.go
new file mode 100644
index 000000000..c04b9cb15
--- /dev/null
+++ b/ud/utils_test.go
@@ -0,0 +1,146 @@
+////////////////////////////////////////////////////////////////////////////////
+// Copyright © 2020 xx network SEZC                                           //
+//                                                                            //
+// Use of this source code is governed by a license that can be found in the  //
+// LICENSE file                                                               //
+////////////////////////////////////////////////////////////////////////////////
+
+////////////////////////////////////////////////////////////////////////////////
+// Copyright © 2020 xx network SEZC                                           //
+//                                                                            //
+// Use of this source code is governed by a license that can be found in the  //
+// LICENSE file                                                               //
+////////////////////////////////////////////////////////////////////////////////
+
+package ud
+
+import (
+	"gitlab.com/elixxir/client/interfaces"
+	"gitlab.com/elixxir/client/interfaces/message"
+	"gitlab.com/elixxir/client/interfaces/params"
+	"gitlab.com/elixxir/client/network/gateway"
+	"gitlab.com/elixxir/client/stoppable"
+	"gitlab.com/elixxir/comms/network"
+	"gitlab.com/elixxir/crypto/e2e"
+	"gitlab.com/elixxir/primitives/format"
+	"gitlab.com/xx_network/comms/connect"
+	"gitlab.com/xx_network/primitives/id"
+	"gitlab.com/xx_network/primitives/id/ephemeral"
+	"gitlab.com/xx_network/primitives/ndf"
+	"testing"
+	"time"
+)
+
+func newTestNetworkManager(t *testing.T) interfaces.NetworkManager {
+	instanceComms := &connect.ProtoComms{
+		Manager: connect.NewManagerTesting(t),
+	}
+
+	thisInstance, err := network.NewInstanceTesting(instanceComms, getNDF(),
+		getNDF(), nil, nil, t)
+	if err != nil {
+		t.Fatalf("Failed to create new test instance: %v", err)
+	}
+
+	return &testNetworkManager{
+		instance: thisInstance,
+	}
+}
+
+// testNetworkManager is a test implementation of NetworkManager interface.
+type testNetworkManager struct {
+	instance *network.Instance
+}
+
+func (tnm *testNetworkManager) SendE2E(message.Send, params.E2E, *stoppable.Single) ([]id.Round, e2e.MessageID, time.Time, error) {
+	return nil, e2e.MessageID{}, time.Time{}, nil
+}
+
+func (tnm *testNetworkManager) SendUnsafe(message.Send, params.Unsafe) ([]id.Round, error) {
+	return nil, nil
+}
+
+func (tnm *testNetworkManager) GetVerboseRounds() string {
+	return ""
+}
+
+func (tnm *testNetworkManager) SendCMIX(format.Message, *id.ID, params.CMIX) (id.Round, ephemeral.Id, error) {
+	return 0, ephemeral.Id{}, nil
+}
+
+func (tnm *testNetworkManager) SendManyCMIX(map[id.ID]format.Message, params.CMIX) (id.Round, []ephemeral.Id, error) {
+	return 0, nil, nil
+}
+
+type dummyEventMgr struct{}
+
+func (d *dummyEventMgr) Report(int, string, string, string) {}
+func (tnm *testNetworkManager) GetEventManager() interfaces.EventManager {
+	return &dummyEventMgr{}
+}
+
+func (tnm *testNetworkManager) GetInstance() *network.Instance             { return tnm.instance }
+func (tnm *testNetworkManager) GetHealthTracker() interfaces.HealthTracker { return nil }
+func (tnm *testNetworkManager) Follow(interfaces.ClientErrorReport) (stoppable.Stoppable, error) {
+	return nil, nil
+}
+func (tnm *testNetworkManager) CheckGarbledMessages()        {}
+func (tnm *testNetworkManager) InProgressRegistrations() int { return 0 }
+func (tnm *testNetworkManager) GetSender() *gateway.Sender   { return nil }
+func (tnm *testNetworkManager) GetAddressSize() uint8        { return 0 }
+func (tnm *testNetworkManager) RegisterAddressSizeNotification(string) (chan uint8, error) {
+	return nil, nil
+}
+func (tnm *testNetworkManager) UnregisterAddressSizeNotification(string) {}
+func (tnm *testNetworkManager) SetPoolFilter(gateway.Filter)             {}
+
+func getNDF() *ndf.NetworkDefinition {
+	return &ndf.NetworkDefinition{
+		UDB: ndf.UDB{
+			ID:      id.DummyUser.Bytes(),
+			Cert:    "",
+			Address: "address",
+			DhPubKey: []byte{123, 34, 86, 97, 108, 117, 101, 34, 58, 49, 44, 34,
+				70, 105, 110, 103, 101, 114, 112, 114, 105, 110, 116, 34, 58,
+				51, 49, 54, 49, 50, 55, 48, 53, 56, 49, 51, 52, 50, 49, 54, 54,
+				57, 52, 55, 125},
+		},
+		E2E: ndf.Group{
+			Prime: "E2EE983D031DC1DB6F1A7A67DF0E9A8E5561DB8E8D49413394C049B7A" +
+				"8ACCEDC298708F121951D9CF920EC5D146727AA4AE535B0922C688B55B3D" +
+				"D2AEDF6C01C94764DAB937935AA83BE36E67760713AB44A6337C20E78615" +
+				"75E745D31F8B9E9AD8412118C62A3E2E29DF46B0864D0C951C394A5CBBDC" +
+				"6ADC718DD2A3E041023DBB5AB23EBB4742DE9C1687B5B34FA48C3521632C" +
+				"4A530E8FFB1BC51DADDF453B0B2717C2BC6669ED76B4BDD5C9FF558E88F2" +
+				"6E5785302BEDBCA23EAC5ACE92096EE8A60642FB61E8F3D24990B8CB12EE" +
+				"448EEF78E184C7242DD161C7738F32BF29A841698978825B4111B4BC3E1E" +
+				"198455095958333D776D8B2BEEED3A1A1A221A6E37E664A64B83981C46FF" +
+				"DDC1A45E3D5211AAF8BFBC072768C4F50D7D7803D2D4F278DE8014A47323" +
+				"631D7E064DE81C0C6BFA43EF0E6998860F1390B5D3FEACAF1696015CB79C" +
+				"3F9C2D93D961120CD0E5F12CBB687EAB045241F96789C38E89D796138E63" +
+				"19BE62E35D87B1048CA28BE389B575E994DCA755471584A09EC723742DC3" +
+				"5873847AEF49F66E43873",
+			Generator: "2",
+		},
+		CMIX: ndf.Group{
+			Prime: "9DB6FB5951B66BB6FE1E140F1D2CE5502374161FD6538DF1648218642" +
+				"F0B5C48C8F7A41AADFA187324B87674FA1822B00F1ECF8136943D7C55757" +
+				"264E5A1A44FFE012E9936E00C1D3E9310B01C7D179805D3058B2A9F4BB6F" +
+				"9716BFE6117C6B5B3CC4D9BE341104AD4A80AD6C94E005F4B993E14F091E" +
+				"B51743BF33050C38DE235567E1B34C3D6A5C0CEAA1A0F368213C3D19843D" +
+				"0B4B09DCB9FC72D39C8DE41F1BF14D4BB4563CA28371621CAD3324B6A2D3" +
+				"92145BEBFAC748805236F5CA2FE92B871CD8F9C36D3292B5509CA8CAA77A" +
+				"2ADFC7BFD77DDA6F71125A7456FEA153E433256A2261C6A06ED3693797E7" +
+				"995FAD5AABBCFBE3EDA2741E375404AE25B",
+			Generator: "5C7FF6B06F8F143FE8288433493E4769C4D988ACE5BE25A0E2480" +
+				"9670716C613D7B0CEE6932F8FAA7C44D2CB24523DA53FBE4F6EC3595892D" +
+				"1AA58C4328A06C46A15662E7EAA703A1DECF8BBB2D05DBE2EB956C142A33" +
+				"8661D10461C0D135472085057F3494309FFA73C611F78B32ADBB5740C361" +
+				"C9F35BE90997DB2014E2EF5AA61782F52ABEB8BD6432C4DD097BC5423B28" +
+				"5DAFB60DC364E8161F4A2A35ACA3A10B1C4D203CC76A470A33AFDCBDD929" +
+				"59859ABD8B56E1725252D78EAC66E71BA9AE3F1DD2487199874393CD4D83" +
+				"2186800654760E1E34C09E4D155179F9EC0DC4473F996BDCE6EED1CABED8" +
+				"B6F116F7AD9CF505DF0F998E34AB27514B0FFE7",
+		},
+	}
+}
-- 
GitLab