Properly repopulate UD store after backup

api/client.go

@@ -167,17 +167,19 @@ func NewVanityClient(ndfJSON, storageDir string, password []byte,
 	return nil
 }
 
-// NewClientFromBackup constructs a new Client from an encrypted backup. The backup
-// is decrypted using the backupPassphrase. On success a successful client creation,
-// the function will return a JSON encoded list of the E2E partners
-// contained in the backup and a json-encoded string containing parameters stored in the backup
+// NewClientFromBackup constructs a new Client from an encrypted backup.
+// The backup is decrypted using the backupPassphrase. On success a
+// successful client creation, the function will return a JSON encoded
+// list of the E2E partners contained in the backup and a json-encoded
+//string containing parameters stored in the backup.
 func NewClientFromBackup(ndfJSON, storageDir string, sessionPassword,
 	backupPassphrase []byte, backupFileContents []byte) ([]*id.ID, string, error) {
 
 	backUp := &backup.Backup{}
 	err := backUp.Decrypt(string(backupPassphrase), backupFileContents)
 	if err != nil {
-		return nil, "", errors.WithMessage(err, "Failed to unmarshal decrypted client contents.")
+		return nil, "", errors.WithMessage(err, "Failed to "+
+			"unmarshal decrypted client contents.")
 	}
 
 	usr := user.NewUserFromBackup(backUp)
@@ -195,20 +197,31 @@ func NewClientFromBackup(ndfJSON, storageDir string, sessionPassword,
 
 	// Create storage object.
 	// Note we do not need registration
-	storageSess, err := checkVersionAndSetupStorage(def, storageDir, []byte(sessionPassword), usr,
-		cmixGrp, e2eGrp, rngStreamGen, false, backUp.RegistrationCode)
+	storageSess, err := checkVersionAndSetupStorage(def, storageDir,
+		[]byte(sessionPassword), usr, cmixGrp, e2eGrp, rngStreamGen,
+		false, backUp.RegistrationCode)
 
 	// Set registration values in storage
-	storageSess.User().SetReceptionRegistrationValidationSignature(backUp.ReceptionIdentity.RegistrarSignature)
-	storageSess.User().SetTransmissionRegistrationValidationSignature(backUp.TransmissionIdentity.RegistrarSignature)
+	storageSess.User().SetReceptionRegistrationValidationSignature(backUp.
+		ReceptionIdentity.RegistrarSignature)
+	storageSess.User().SetTransmissionRegistrationValidationSignature(backUp.
+		TransmissionIdentity.RegistrarSignature)
 	storageSess.User().SetRegistrationTimestamp(backUp.RegistrationTimestamp)
 
-	//move the registration state to indicate registered with registration on proto client
+	//move the registration state to indicate registered with registration
+	//on proto client
 	err = storageSess.ForwardRegistrationStatus(storage.PermissioningComplete)
 	if err != nil {
 		return nil, "", err
 	}
 
+	err = storageSess.GetUd().RestoreFromBackUp(backUp.
+		UserDiscoveryRegistration.FactList)
+	if err != nil {
+		return nil, "", errors.WithMessage(err, "Could not restore user "+
+			"discover storage")
+	}
+
 	return backUp.Contacts.Identities, backUp.JSONParams, nil
 }
 

storage/ud/facts.go

@@ -38,7 +38,6 @@ type Store struct {
 // NewStore creates a new, empty Store object.
 func NewStore(kv *versioned.KV) (*Store, error) {
 	kv = kv.Prefix(prefix)
-
 	s := &Store{
 		confirmedFacts:   make(map[fact.Fact]struct{}, 0),
 		unconfirmedFacts: make(map[string]fact.Fact, 0),
@@ -48,6 +47,24 @@ func NewStore(kv *versioned.KV) (*Store, error) {
 	return s, s.save()
 }
 
+// RestoreFromBackUp initializes the confirmedFacts map
+// with the backed up fact data. This will error if
+// the store is already stateful.
+func (s *Store) RestoreFromBackUp(backupData fact.FactList) error {
+	s.mux.Lock()
+	defer s.mux.Unlock()
+
+	if len(s.confirmedFacts) != 0 || len(s.unconfirmedFacts) != 0 {
+		return errors.New("cannot overwrite ud store with existing data")
+	}
+
+	for _, f := range backupData {
+		s.confirmedFacts[f] = struct{}{}
+	}
+
+	return nil
+}
+
 // StoreUnconfirmedFact stores a fact that has been added to UD but has not been
 // confirmed by the user. It is keyed on the confirmation ID given by UD.
 func (s *Store) StoreUnconfirmedFact(confirmationId string, f fact.Fact) error {