From 05e3a7f6c09949dcde3ca43f5ffefa46871f734c Mon Sep 17 00:00:00 2001
From: Conner Fromknecht <conner@lightning.engineering>
Date: Fri, 15 Mar 2019 02:34:00 -0700
Subject: [PATCH] watchtower/wtmock/peer: set local pubkey

---
 watchtower/wtmock/peer.go          |  7 +++++--
 watchtower/wtserver/server_test.go | 20 +++++++++++++-------
 2 files changed, 18 insertions(+), 9 deletions(-)

diff --git a/watchtower/wtmock/peer.go b/watchtower/wtmock/peer.go
index 16aeb0d8..fc1ff9af 100644
--- a/watchtower/wtmock/peer.go
+++ b/watchtower/wtmock/peer.go
@@ -27,14 +27,17 @@ type MockPeer struct {
 }
 
 // NewMockPeer returns a fresh MockPeer.
-func NewMockPeer(pk *btcec.PublicKey, addr net.Addr, bufferSize int) *MockPeer {
+func NewMockPeer(lpk, rpk *btcec.PublicKey, addr net.Addr,
+	bufferSize int) *MockPeer {
+
 	return &MockPeer{
-		remotePub:  pk,
+		remotePub:  rpk,
 		remoteAddr: addr,
 		localAddr: &net.TCPAddr{
 			IP:   net.IP{0x32, 0x31, 0x30, 0x29},
 			Port: 36723,
 		},
+		localPub:     lpk,
 		IncomingMsgs: make(chan []byte, bufferSize),
 		OutgoingMsgs: make(chan []byte, bufferSize),
 		Quit:         make(chan struct{}),
diff --git a/watchtower/wtserver/server_test.go b/watchtower/wtserver/server_test.go
index 6df7cd58..cdbaa281 100644
--- a/watchtower/wtserver/server_test.go
+++ b/watchtower/wtserver/server_test.go
@@ -87,10 +87,12 @@ func TestServerOnlyAcceptOnePeer(t *testing.T) {
 	s := initServer(t, nil, timeoutDuration)
 	defer s.Stop()
 
+	localPub := randPubKey(t)
+
 	// Create two peers using the same session id.
 	peerPub := randPubKey(t)
-	peer1 := wtmock.NewMockPeer(peerPub, nil, 0)
-	peer2 := wtmock.NewMockPeer(peerPub, nil, 0)
+	peer1 := wtmock.NewMockPeer(localPub, peerPub, nil, 0)
+	peer2 := wtmock.NewMockPeer(localPub, peerPub, nil, 0)
 
 	// Serialize a Init message to be sent by both peers.
 	init := wtwire.NewInitMessage(
@@ -219,9 +221,11 @@ func testServerCreateSession(t *testing.T, i int, test createSessionTestCase) {
 	s := initServer(t, nil, timeoutDuration)
 	defer s.Stop()
 
+	localPub := randPubKey(t)
+
 	// Create a new client and connect to server.
 	peerPub := randPubKey(t)
-	peer := wtmock.NewMockPeer(peerPub, nil, 0)
+	peer := wtmock.NewMockPeer(localPub, peerPub, nil, 0)
 	connect(t, i, s, peer, test.initMsg, timeoutDuration)
 
 	// Send the CreateSession message, and wait for a reply.
@@ -249,7 +253,7 @@ func testServerCreateSession(t *testing.T, i int, test createSessionTestCase) {
 
 	// Simulate a peer with the same session id connection to the server
 	// again.
-	peer = wtmock.NewMockPeer(peerPub, nil, 0)
+	peer = wtmock.NewMockPeer(localPub, peerPub, nil, 0)
 	connect(t, i, s, peer, test.initMsg, timeoutDuration)
 
 	// Send the _same_ CreateSession message as the first attempt.
@@ -559,9 +563,11 @@ func testServerStateUpdates(t *testing.T, i int, test stateUpdateTestCase) {
 	s := initServer(t, nil, timeoutDuration)
 	defer s.Stop()
 
+	localPub := randPubKey(t)
+
 	// Create a new client and connect to the server.
 	peerPub := randPubKey(t)
-	peer := wtmock.NewMockPeer(peerPub, nil, 0)
+	peer := wtmock.NewMockPeer(localPub, peerPub, nil, 0)
 	connect(t, i, s, peer, test.initMsg, timeoutDuration)
 
 	// Register a session for this client to use in the subsequent tests.
@@ -581,7 +587,7 @@ func testServerStateUpdates(t *testing.T, i int, test stateUpdateTestCase) {
 
 	// Now that the original connection has been closed, connect a new
 	// client with the same session id.
-	peer = wtmock.NewMockPeer(peerPub, nil, 0)
+	peer = wtmock.NewMockPeer(localPub, peerPub, nil, 0)
 	connect(t, i, s, peer, test.initMsg, timeoutDuration)
 
 	// Send the intended StateUpdate messages in series.
@@ -592,7 +598,7 @@ func testServerStateUpdates(t *testing.T, i int, test stateUpdateTestCase) {
 		if update == nil {
 			assertConnClosed(t, peer, 2*timeoutDuration)
 
-			peer = wtmock.NewMockPeer(peerPub, nil, 0)
+			peer = wtmock.NewMockPeer(localPub, peerPub, nil, 0)
 			connect(t, i, s, peer, test.initMsg, timeoutDuration)
 
 			continue