From 2a57f9182a025ca55956c86b71949ca507479707 Mon Sep 17 00:00:00 2001
From: Olaoluwa Osuntokun <laolu32@gmail.com>
Date: Sat, 16 Jul 2016 18:00:52 -0700
Subject: [PATCH] lndc+lnd: fix panic when connecting to multiple peers, plus
 duplicate conn detection

---
 lndc/lnadr.go | 46 ++++++++++++++++++++++++++++++++++------------
 peer.go       | 16 +++++++++++++---
 rpcserver.go  |  2 +-
 server.go     |  4 ++--
 4 files changed, 50 insertions(+), 18 deletions(-)

diff --git a/lndc/lnadr.go b/lndc/lnadr.go
index 1cab4756..8daa9f11 100644
--- a/lndc/lnadr.go
+++ b/lndc/lnadr.go
@@ -25,22 +25,30 @@ type LNAdr struct {
 	name        string // human readable name?  Not a thing yet.
 	host        string // internet host this ID is reachable at. also not a thing
 	endorsement []byte // a sig confirming the name?  Not implemented
+
+	net *chaincfg.Params
 }
 
-// String...
-func (l *LNAdr) String() string {
-	var encodedId []byte
-	if l.PubKey == nil {
-		encodedId = l.Base58Adr.ScriptAddress()
-	} else {
-		encodedId = l.PubKey.SerializeCompressed()
+// newLnAdr....
+func NewLnAdr(addr *net.TCPAddr, pubkey *btcec.PublicKey,
+	net *chaincfg.Params) (*LNAdr, error) {
+
+	hash160 := btcutil.Hash160(pubkey.SerializeCompressed())
+	pkh, err := btcutil.NewAddressPubKeyHash(hash160, net)
+	if err != nil {
+		return nil, err
 	}
 
-	return fmt.Sprintf("%v@%v", hex.EncodeToString(encodedId), l.NetAddr)
+	return &LNAdr{
+		PubKey:    pubkey,
+		Base58Adr: pkh,
+		NetAddr:   addr,
+		net:       net,
+	}, nil
 }
 
 // newLnAddr...
-func LnAddrFromString(encodedAddr string) (*LNAdr, error) {
+func LnAddrFromString(encodedAddr string, netParams *chaincfg.Params) (*LNAdr, error) {
 	// The format of an lnaddr is "<pubkey or pkh>@host"
 	idHost := strings.Split(encodedAddr, "@")
 	if len(idHost) != 2 {
@@ -54,7 +62,7 @@ func LnAddrFromString(encodedAddr string) (*LNAdr, error) {
 		return nil, err
 	}
 
-	addr := &LNAdr{NetAddr: ipAddr}
+	addr := &LNAdr{NetAddr: ipAddr, net: netParams}
 
 	idLen := len(idHost[0])
 	switch {
@@ -73,14 +81,14 @@ func LnAddrFromString(encodedAddr string) (*LNAdr, error) {
 		// got pubey, populate address from pubkey
 		pkh := btcutil.Hash160(addr.PubKey.SerializeCompressed())
 		addr.Base58Adr, err = btcutil.NewAddressPubKeyHash(pkh,
-			&chaincfg.TestNet3Params)
+			netParams)
 		if err != nil {
 			return nil, err
 		}
 	// Is the ID a string encoded bitcoin address?
 	case idLen > 33 && idLen < 37:
 		addr.Base58Adr, err = btcutil.DecodeAddress(idHost[0],
-			&chaincfg.TestNet3Params)
+			netParams)
 		if err != nil {
 			return nil, err
 		}
@@ -167,3 +175,17 @@ func (l *LNAdr) Deserialize(s []byte) error {
 
 	return nil
 }
+
+// String...
+func (l *LNAdr) String() string {
+	var encodedId []byte
+	if l.Base58Adr != nil {
+		encodedId = l.Base58Adr.ScriptAddress()
+	} else {
+		pubKey := l.PubKey.SerializeCompressed()
+		pkh, _ := btcutil.NewAddressPubKeyHash(pubKey, l.net)
+		encodedId = pkh.ScriptAddress()
+	}
+
+	return fmt.Sprintf("%v@%v", hex.EncodeToString(encodedId), l.NetAddr)
+}
diff --git a/peer.go b/peer.go
index 416517e8..73cb6549 100644
--- a/peer.go
+++ b/peer.go
@@ -142,14 +142,15 @@ type peer struct {
 
 // newPeer creates a new peer from an establish connection object, and a
 // pointer to the main server.
-func newPeer(conn net.Conn, server *server, net wire.BitcoinNet, inbound bool) (*peer, error) {
-	nodePub := conn.(*lndc.LNDConn).RemotePub
+func newPeer(conn net.Conn, server *server, btcNet wire.BitcoinNet, inbound bool) (*peer, error) {
+	lndcConn := conn.(*lndc.LNDConn)
+	nodePub := lndcConn.RemotePub
 
 	p := &peer{
 		conn:        conn,
 		lightningID: wire.ShaHash(fastsha256.Sum256(nodePub.SerializeCompressed())),
 		id:          atomic.AddInt32(&numNodes, 1),
-		chainNet:    net,
+		chainNet:    btcNet,
 		inbound:     inbound,
 
 		server: server,
@@ -174,6 +175,15 @@ func newPeer(conn net.Conn, server *server, net wire.BitcoinNet, inbound bool) (
 		quit:      make(chan struct{}),
 	}
 
+	// TODO(roasbeef): re-write after lnaddr revamp, shouldn't need to use
+	// type assertions
+	var err error
+	tcpAddr := lndcConn.Conn.(*net.TCPConn).RemoteAddr().(*net.TCPAddr)
+	p.lightningAddr, err = lndc.NewLnAdr(tcpAddr, nodePub, activeNetParams.Params)
+	if err != nil {
+		return nil, err
+	}
+
 	// Initiate the pending channel identifier properly depending on if this
 	// node is inbound or outbound. This value will be used in an increasing
 	// manner to track pending channels.
diff --git a/rpcserver.go b/rpcserver.go
index d66b578e..d2a49653 100644
--- a/rpcserver.go
+++ b/rpcserver.go
@@ -172,7 +172,7 @@ func (r *rpcServer) ConnectPeer(ctx context.Context,
 	idAtHost := fmt.Sprintf("%v@%v", in.Addr.PubKeyHash, in.Addr.Host)
 	rpcsLog.Debugf("[connectpeer] peer=%v", idAtHost)
 
-	peerAddr, err := lndc.LnAddrFromString(idAtHost)
+	peerAddr, err := lndc.LnAddrFromString(idAtHost, activeNetParams.Params)
 	if err != nil {
 		rpcsLog.Errorf("(connectpeer): error parsing ln addr: %v", err)
 		return nil, err
diff --git a/server.go b/server.go
index f920a2f7..27264a0e 100644
--- a/server.go
+++ b/server.go
@@ -283,13 +283,13 @@ func (s *server) handleConnectPeer(msg *connectPeerMsg) {
 	// Ensure we're not already connected to this
 	// peer.
 	for _, peer := range s.peers {
-		if peer.lightningAddr.String() ==
-			addr.String() {
+		if peer.lightningAddr.String() == addr.String() {
 			msg.err <- fmt.Errorf(
 				"already connected to peer: %v",
 				peer.lightningAddr,
 			)
 			msg.resp <- -1
+			return
 		}
 	}