From 69900cccc8d06367cbeeee0b15a0628404302546 Mon Sep 17 00:00:00 2001
From: Oliver Gugger <gugger@gmail.com>
Date: Wed, 23 Feb 2022 09:41:48 +0100
Subject: [PATCH] waddrmgr: refactor newWitnessScriptAddress to return
 interface

To make integration of the upcoming taproot script address type a bit
easier, we allow returning a ManagedScriptAddress interface type in
newWitnessScriptAddress so we can re-use that function for the taproot
script address too.
---
 waddrmgr/address.go        | 17 +++++++++++++++++
 waddrmgr/scoped_manager.go | 27 +++++++++------------------
 2 files changed, 26 insertions(+), 18 deletions(-)

diff --git a/waddrmgr/address.go b/waddrmgr/address.go
index cf58531e83..c19fcd16a7 100644
--- a/waddrmgr/address.go
+++ b/waddrmgr/address.go
@@ -536,6 +536,14 @@ func newManagedAddressFromExtKey(s *ScopedKeyManager,
 	return managedAddr, nil
 }
 
+// clearTextScriptSetter is a non-exported interface to identify script types
+// that allow their clear text script to be set.
+type clearTextScriptSetter interface {
+	// setClearText sets the unencrypted script on the struct after
+	// unlocking/decrypting it.
+	setClearTextScript([]byte)
+}
+
 // baseScriptAddress represents the common fields of a pay-to-script-hash and
 // a pay-to-witness-script-hash address.
 type baseScriptAddress struct {
@@ -547,6 +555,8 @@ type baseScriptAddress struct {
 	scriptMutex     sync.Mutex
 }
 
+var _ clearTextScriptSetter = (*baseScriptAddress)(nil)
+
 // unlock decrypts and stores the associated script.  It will fail if the key is
 // invalid or the encrypted script is not available.  The returned clear text
 // script will always be a copy that may be safely used by the caller without
@@ -605,6 +615,13 @@ func (a *baseScriptAddress) Internal() bool {
 	return false
 }
 
+// setClearText sets the unencrypted script on the struct after unlocking/
+// decrypting it.
+func (a *baseScriptAddress) setClearTextScript(script []byte) {
+	a.scriptClearText = make([]byte, len(script))
+	copy(a.scriptClearText, script)
+}
+
 // scriptAddress represents a pay-to-script-hash address.
 type scriptAddress struct {
 	baseScriptAddress
diff --git a/waddrmgr/scoped_manager.go b/waddrmgr/scoped_manager.go
index 6d7f435cde..ab3b67e9a2 100644
--- a/waddrmgr/scoped_manager.go
+++ b/waddrmgr/scoped_manager.go
@@ -2215,38 +2215,29 @@ func (s *ScopedKeyManager) importScriptAddress(ns walletdb.ReadWriteBucket,
 	// when not a watching-only address manager, make a copy of the script
 	// since it will be cleared on lock and the script the caller passed
 	// should not be cleared out from under the caller.
-	var (
-		managedAddr    ManagedScriptAddress
-		baseScriptAddr *baseScriptAddress
-	)
+	var managedAddr ManagedScriptAddress
 	switch addrType {
 	case WitnessScript:
-		witnessAddr, err := newWitnessScriptAddress(
+		managedAddr, err = newWitnessScriptAddress(
 			s, ImportedAddrAccount, scriptIdent, encryptedScript,
 			witnessVersion, isSecretScript,
 		)
-		if err != nil {
-			return nil, err
-		}
-		managedAddr = witnessAddr
-		baseScriptAddr = &witnessAddr.baseScriptAddress
 
 	default:
-		scriptAddr, err := newScriptAddress(
+		managedAddr, err = newScriptAddress(
 			s, ImportedAddrAccount, scriptIdent, encryptedScript,
 		)
-		if err != nil {
-			return nil, err
-		}
-		managedAddr = scriptAddr
-		baseScriptAddr = &scriptAddr.baseScriptAddress
+	}
+	if err != nil {
+		return nil, err
 	}
 
 	// Even if the script is secret, we are currently unlocked, so we keep a
 	// clear text copy of the script around to avoid decrypting it on each
 	// access.
-	baseScriptAddr.scriptClearText = make([]byte, len(script))
-	copy(baseScriptAddr.scriptClearText, script)
+	if cts, ok := managedAddr.(clearTextScriptSetter); ok {
+		cts.setClearTextScript(script)
+	}
 
 	// Add the new managed address to the cache of recent addresses and
 	// return it.