From 83bf984ca266ad448095a508c7eddcb22ac4b420 Mon Sep 17 00:00:00 2001
From: Donald Adu-Poku <donald.adu@gmail.com>
Date: Tue, 13 Aug 2019 16:30:21 +0000
Subject: [PATCH] multi: store unencrypted scripts in address mgr.

This allows storing unencrypted script addresses
 in the address manager. This behaviour will become
default after the wallet is upgraded.
---
 wallet/udb/addressdb.go      | 66 +++++++++++++++++------------
 wallet/udb/addressmanager.go | 56 ++++++++++---------------
 wallet/udb/upgrades.go       | 80 +++++++++++++++++++++++++++++++++++-
 wallet/udb/upgrades_test.go  | 33 +++++++++++++++
 4 files changed, 174 insertions(+), 61 deletions(-)

diff --git a/wallet/udb/addressdb.go b/wallet/udb/addressdb.go
index da323b4b1..daf6d2bfe 100644
--- a/wallet/udb/addressdb.go
+++ b/wallet/udb/addressdb.go
@@ -108,8 +108,9 @@ type dbImportedAddressRow struct {
 // address in the database.
 type dbScriptAddressRow struct {
 	dbAddressRow
-	encryptedHash   []byte
-	encryptedScript []byte
+	hash      []byte
+	script    []byte
+	encrypted bool
 }
 
 // Key names for various database fields.
@@ -968,15 +969,16 @@ func serializeImportedAddress(encryptedPubKey, encryptedPrivKey []byte) []byte {
 // row as a script address.
 func deserializeScriptAddress(row *dbAddressRow) (*dbScriptAddressRow, error) {
 	// The serialized script address raw data format is:
-	//   <encscripthashlen><encscripthash><encscriptlen><encscript>
+	//   <scripthashlen><scripthash><scriptlen><script><encryptedbool>
 	//
-	// 4 bytes encrypted script hash len + encrypted script hash + 4 bytes
-	// encrypted script len + encrypted script
+	// 4 bytes script hash len + script hash + 4 bytes script len + script +
+	// 1 byte encrypted boolean
 
 	// Given the above, the length of the entry must be at a minimum
 	// the constant value sizes.
-	if len(row.rawData) < 8 {
-		return nil, errors.E(errors.IO, errors.Errorf("bad script address len %d", len(row.rawData)))
+	if len(row.rawData) < 9 {
+		return nil, errors.E(errors.IO,
+			errors.Errorf("bad script address len %d", len(row.rawData)))
 	}
 
 	retRow := dbScriptAddressRow{
@@ -984,35 +986,49 @@ func deserializeScriptAddress(row *dbAddressRow) (*dbScriptAddressRow, error) {
 	}
 
 	hashLen := binary.LittleEndian.Uint32(row.rawData[0:4])
-	retRow.encryptedHash = make([]byte, hashLen)
-	copy(retRow.encryptedHash, row.rawData[4:4+hashLen])
+	retRow.hash = make([]byte, hashLen)
+	copy(retRow.hash, row.rawData[4:4+hashLen])
 	offset := 4 + hashLen
 	scriptLen := binary.LittleEndian.Uint32(row.rawData[offset : offset+4])
 	offset += 4
-	retRow.encryptedScript = make([]byte, scriptLen)
-	copy(retRow.encryptedScript, row.rawData[offset:offset+scriptLen])
+	retRow.script = make([]byte, scriptLen)
+	copy(retRow.script, row.rawData[offset:offset+scriptLen])
+	offset += scriptLen
+
+	encrypted := row.rawData[offset]
+	if encrypted == 1 {
+		retRow.encrypted = true
+	}
 
 	return &retRow, nil
 }
 
 // serializeScriptAddress returns the serialization of the raw data field for
 // a script address.
-func serializeScriptAddress(encryptedHash, encryptedScript []byte) []byte {
+func serializeScriptAddress(hash, script []byte, encrypted bool) []byte {
 	// The serialized script address raw data format is:
-	//   <encscripthashlen><encscripthash><encscriptlen><encscript>
+	//   <scripthashlen><scripthash><scriptlen><script><encryptedbool>
 	//
-	// 4 bytes encrypted script hash len + encrypted script hash + 4 bytes
-	// encrypted script len + encrypted script
+	// 4 bytes script hash len + script hash + 4 bytes script len + script +
+	// 1 byte encrypted boolean
 
-	hashLen := uint32(len(encryptedHash))
-	scriptLen := uint32(len(encryptedScript))
-	rawData := make([]byte, 8+hashLen+scriptLen)
+	hashLen := uint32(len(hash))
+	scriptLen := uint32(len(script))
+	rawData := make([]byte, 8+hashLen+scriptLen+1)
 	binary.LittleEndian.PutUint32(rawData[0:4], hashLen)
-	copy(rawData[4:4+hashLen], encryptedHash)
+	copy(rawData[4:4+hashLen], hash)
 	offset := 4 + hashLen
 	binary.LittleEndian.PutUint32(rawData[offset:offset+4], scriptLen)
 	offset += 4
-	copy(rawData[offset:offset+scriptLen], encryptedScript)
+	copy(rawData[offset:offset+scriptLen], script)
+	offset += scriptLen
+
+	if encrypted {
+		rawData[offset] = byte(1)
+	} else {
+		rawData[offset] = byte(0)
+	}
+
 	return rawData
 }
 
@@ -1111,9 +1127,9 @@ func putImportedAddress(ns walletdb.ReadWriteBucket, addressID []byte, account u
 // putScriptAddress stores the provided script address information to the
 // database.
 func putScriptAddress(ns walletdb.ReadWriteBucket, addressID []byte, account uint32,
-	status syncStatus, encryptedHash, encryptedScript []byte) error {
+	status syncStatus, hash, script []byte, encrypted bool) error {
 
-	rawData := serializeScriptAddress(encryptedHash, encryptedScript)
+	rawData := serializeScriptAddress(hash, script, encrypted)
 	addrRow := dbAddressRow{
 		addrType:   adtScript,
 		account:    account,
@@ -1316,8 +1332,7 @@ func deletePrivateKeys(ns walletdb.ReadWriteBucket, dbVersion uint32) error {
 
 		// Reserialize the imported address without the private
 		// key and store it.
-		row.rawData = serializeImportedAddress(
-			irow.encryptedPubKey, nil)
+		row.rawData = serializeImportedAddress(irow.encryptedPubKey, nil)
 		err = bucket.Put([]byte(k), serializeAddressRow(row))
 		if err != nil {
 			return errors.E(errors.IO, err)
@@ -1332,8 +1347,7 @@ func deletePrivateKeys(ns walletdb.ReadWriteBucket, dbVersion uint32) error {
 
 		// Reserialize the script address without the script
 		// and store it.
-		row.rawData = serializeScriptAddress(srow.encryptedHash,
-			nil)
+		row.rawData = serializeScriptAddress(srow.hash, nil, srow.encrypted)
 		err = bucket.Put([]byte(k), serializeAddressRow(row))
 		if err != nil {
 			return errors.E(errors.IO, err)
diff --git a/wallet/udb/addressmanager.go b/wallet/udb/addressmanager.go
index e43a18f7f..dfec398c0 100644
--- a/wallet/udb/addressmanager.go
+++ b/wallet/udb/addressmanager.go
@@ -855,13 +855,17 @@ func (m *Manager) importedAddressRowToManaged(row *dbImportedAddressRow) (Manage
 // scriptAddressRowToManaged returns a new managed address based on script
 // address data loaded from the database.
 func (m *Manager) scriptAddressRowToManaged(row *dbScriptAddressRow) (ManagedAddress, error) {
-	// Use the crypto public key to decrypt the imported script hash.
-	scriptHash, err := m.cryptoKeyPub.Decrypt(row.encryptedHash)
-	if err != nil {
-		return nil, errors.E(errors.Crypto, errors.Errorf("decrypt imported P2SH address: %v", err))
+	if row.encrypted {
+		// Use the crypto public key to decrypt the imported script hash.
+		scriptHash, err := m.cryptoKeyPub.Decrypt(row.hash)
+		if err != nil {
+			return nil, errors.E(errors.Crypto, errors.Errorf("decrypt imported P2SH address: %v", err))
+		}
+
+		return newScriptAddress(m, row.account, scriptHash)
 	}
 
-	return newScriptAddress(m, row.account, scriptHash)
+	return newScriptAddress(m, row.account, row.hash)
 }
 
 // rowInterfaceToManaged returns a new managed address based on the given
@@ -1245,32 +1249,10 @@ func (m *Manager) ImportScript(ns walletdb.ReadWriteBucket, script []byte) (Mana
 		return nil, errors.E(errors.Exist, "script already exists")
 	}
 
-	// The manager must be unlocked to encrypt the imported script.
-	if !m.watchingOnly && m.locked {
-		return nil, errors.E(errors.Locked)
-	}
-
-	// Encrypt the script hash using the crypto public key so it is
-	// accessible when the address manager is locked or watching-only.
-	encryptedHash, err := m.cryptoKeyPub.Encrypt(scriptHash)
-	if err != nil {
-		return nil, errors.E(errors.Crypto, errors.Errorf("encrypt script hash: %v", err))
-	}
-
-	// Encrypt the script for storage in database using the crypto script
-	// key when not a watching-only address manager.
-	var encryptedScript []byte
-	if !m.watchingOnly {
-		encryptedScript, err = m.cryptoKeyScript.Encrypt(script)
-		if err != nil {
-			return nil, errors.E(errors.Crypto, errors.Errorf("encrypt script: %v", err))
-		}
-	}
-
 	// Save the new imported address to the db and update start block (if
 	// needed) in a single transaction.
-	err = putScriptAddress(ns, scriptHash, ImportedAddrAccount,
-		ssNone, encryptedHash, encryptedScript)
+	err := putScriptAddress(ns, scriptHash, ImportedAddrAccount, ssNone,
+		scriptHash, script, false)
 	if err != nil {
 		return nil, err
 	}
@@ -1994,16 +1976,22 @@ func (m *Manager) RedeemScript(ns walletdb.ReadBucket, addr dcrutil.Address) (sc
 	}
 	switch a := addrInterface.(type) {
 	case *dbScriptAddressRow:
-		script, err = m.cryptoKeyScript.Decrypt(a.encryptedScript)
-		if err != nil {
-			return nil, nil, errors.E(errors.Crypto, errors.Errorf("decrypt imported script: %v", err))
+		script = a.script
+		if a.encrypted {
+			script, err = m.cryptoKeyScript.Decrypt(a.script)
+			if err != nil {
+				return nil, nil, errors.E(errors.Crypto,
+					errors.Errorf("decrypt imported script: %v", err))
+			}
 		}
 
 	case *dbChainAddressRow, *dbImportedAddressRow:
-		return nil, nil, errors.E(errors.Invalid, "redeem script lookup requires P2SH address")
+		return nil, nil, errors.E(errors.Invalid,
+			"redeem script lookup requires P2SH address")
 
 	default:
-		return nil, nil, errors.E(errors.Invalid, errors.Errorf("address row type %T", addrInterface))
+		return nil, nil, errors.E(errors.Invalid,
+			errors.Errorf("address row type %T", addrInterface))
 	}
 
 	// Lock the RWMutex for reads for the caller and prepare to return the
diff --git a/wallet/udb/upgrades.go b/wallet/udb/upgrades.go
index 986aa8bda..27a0e53be 100644
--- a/wallet/udb/upgrades.go
+++ b/wallet/udb/upgrades.go
@@ -6,6 +6,7 @@ package udb
 
 import (
 	"crypto/sha256"
+	"encoding/binary"
 
 	"github.com/decred/dcrd/blockchain/stake/v2"
 	"github.com/decred/dcrd/chaincfg/chainhash"
@@ -110,10 +111,15 @@ const (
 	// accounting of total locked funds.
 	ticketCommitmentsVersion = 12
 
+	// scriptStorageVersion is the thirteenth version of the database. It
+	// updates the serialization semantics of the address row data and
+	// allows storing raw scripts in the address manager.
+	scriptStorageVersion = 13
+
 	// DBVersion is the latest version of the database that is understood by the
 	// program.  Databases with recorded versions higher than this will fail to
 	// open (meaning any upgrades prevent reverting to older software).
-	DBVersion = ticketCommitmentsVersion
+	DBVersion = scriptStorageVersion
 )
 
 // upgrades maps between old database versions and the upgrade function to
@@ -131,6 +137,7 @@ var upgrades = [...]func(walletdb.ReadWriteTx, []byte, *chaincfg.Params) error{
 	cfVersion - 1:                    cfUpgrade,
 	lastProcessedTxsBlockVersion - 1: lastProcessedTxsBlockUpgrade,
 	ticketCommitmentsVersion - 1:     ticketCommitmentsUpgrade,
+	scriptStorageVersion - 1:         scriptStorageUpgrade,
 }
 
 func lastUsedAddressIndexUpgrade(tx walletdb.ReadWriteTx, publicPassphrase []byte, params *chaincfg.Params) error {
@@ -994,6 +1001,77 @@ func ticketCommitmentsUpgrade(tx walletdb.ReadWriteTx, publicPassphrase []byte,
 	return unifiedDBMetadata{}.putVersion(metadataBucket, newVersion)
 }
 
+func scriptStorageUpgrade(tx walletdb.ReadWriteTx, publicPassphrase []byte, params *chaincfg.Params) error {
+	const oldVersion = 12
+	const newVersion = 13
+
+	metadataBucket := tx.ReadWriteBucket(unifiedDBMetadata{}.rootBucketKey())
+	addrmgrBucket := tx.ReadWriteBucket(waddrmgrBucketKey)
+
+	// Assert that this function is only called on version 12 databases.
+	dbVersion, err := unifiedDBMetadata{}.getVersion(metadataBucket)
+	if err != nil {
+		return err
+	}
+	if dbVersion != oldVersion {
+		return errors.E(errors.Invalid, "scriptStorageUpgrade inappropriately called")
+	}
+
+	addrbkt := addrmgrBucket.NestedReadBucket(addrBucketName)
+	toUpdate := make(map[string]*dbScriptAddressRow, 0)
+	cursor := addrbkt.ReadCursor()
+	for k, v := cursor.First(); k != nil; k, v = cursor.Next() {
+		row, err := deserializeAddressRow(v)
+		if err != nil {
+			return err
+		}
+
+		switch row.addrType {
+		case adtScript:
+			// deserialize the address row.
+			hashLen := binary.LittleEndian.Uint32(row.rawData[0:4])
+			encryptedHash := make([]byte, hashLen)
+			copy(encryptedHash, row.rawData[4:4+hashLen])
+			offset := 4 + hashLen
+			scriptLen := binary.LittleEndian.Uint32(row.rawData[offset : offset+4])
+			offset += 4
+			encryptedScript := make([]byte, scriptLen)
+			copy(encryptedScript, row.rawData[offset:offset+scriptLen])
+
+			// reconstruct the address row.
+			row := &dbScriptAddressRow{
+				dbAddressRow: *row,
+				hash:         encryptedHash,
+				script:       encryptedScript,
+				encrypted:    true,
+			}
+
+			toUpdate[string(k)] = row
+		}
+	}
+
+	// persist reserialized address rows.
+	abkt := addrmgrBucket.NestedReadWriteBucket(addrBucketName)
+	for k, row := range toUpdate {
+		rawData := serializeScriptAddress(row.hash, row.script, row.encrypted)
+		addrRow := &dbAddressRow{
+			addrType:   adtScript,
+			account:    row.account,
+			addTime:    row.addTime,
+			syncStatus: row.syncStatus,
+			rawData:    rawData,
+		}
+
+		err := abkt.Put([]byte(k), serializeAddressRow(addrRow))
+		if err != nil {
+			return errors.E(errors.IO, err)
+		}
+	}
+
+	// Write the new database version.
+	return unifiedDBMetadata{}.putVersion(metadataBucket, newVersion)
+}
+
 // Upgrade checks whether the any upgrades are necessary before the database is
 // ready for application usage.  If any are, they are performed.
 func Upgrade(db walletdb.DB, publicPassphrase []byte, params *chaincfg.Params) error {
diff --git a/wallet/udb/upgrades_test.go b/wallet/udb/upgrades_test.go
index af3eee6da..8f6eb4ce2 100644
--- a/wallet/udb/upgrades_test.go
+++ b/wallet/udb/upgrades_test.go
@@ -37,6 +37,7 @@ var dbUpgradeTests = [...]struct {
 	// No upgrade test for V9, it is a fix for V8 and the previous test still applies
 	// TODO: V10 upgrade test
 	{verifyV12Upgrade, "v11.db.gz"},
+	{verifyV13Upgrade, "v12.db.gz"},
 }
 
 var pubPass = []byte("public")
@@ -531,3 +532,35 @@ func verifyV12Upgrade(t *testing.T, db walletdb.DB) {
 		t.Error(err)
 	}
 }
+
+func verifyV13Upgrade(t *testing.T, db walletdb.DB) {
+	err := walletdb.View(db, func(tx walletdb.ReadTx) error {
+		amgrns := tx.ReadBucket(waddrmgrBucketKey)
+		addrbkt := amgrns.NestedReadBucket(addrBucketName)
+		err := addrbkt.ForEach(func(k []byte, v []byte) error {
+			row, err := deserializeAddressRow(v)
+			if err != nil {
+				return fmt.Errorf("unable to deserialize address row: %v", err)
+			}
+
+			switch row.addrType {
+			case adtScript:
+				scrAddr, err := deserializeScriptAddress(row)
+				if err != nil {
+					return fmt.Errorf("unable to deserialize script address: %v", err)
+				}
+
+				if !scrAddr.encrypted {
+					return fmt.Errorf("expected an encrypted script address for %v", scrAddr.hash)
+				}
+			}
+
+			return nil
+		})
+
+		return err
+	})
+	if err != nil {
+		t.Error(err)
+	}
+}