Consolidate remote peer checks

router/connection.go

@@ -283,10 +283,15 @@ func (conn *LocalConnection) run(actionChan <-chan ConnectionAction, finished ch
 
 	conn.TCPConn.SetLinger(0)
 
-	tcpReceiver, err := conn.handshake(acceptNewPeer)
+	tcpReceiver, remote, err := conn.handshake()
 	if err != nil {
 		return
 	}
+
+	if err = conn.registerRemote(remote, acceptNewPeer); err != nil {
+		return
+	}
+
 	conn.Log("completed handshake; using protocol version", conn.version)
 
 	// The ordering of the following is very important. [1]
@@ -351,6 +356,27 @@ func (conn *LocalConnection) run(actionChan <-chan ConnectionAction, finished ch
 // prevent that completely, since, for example, forwarder can only be
 // created when we know the remote UDP address, but it helps to try.
 
+func (conn *LocalConnection) registerRemote(remote *Peer, acceptNewPeer bool) error {
+	if acceptNewPeer {
+		conn.remote = conn.Router.Peers.FetchWithDefault(remote)
+	} else {
+		conn.remote = conn.Router.Peers.FetchAndAddRef(remote.Name)
+		if conn.remote == nil {
+			return fmt.Errorf("Found unknown remote name: %s at %s", remote.Name, conn.remoteTCPAddr)
+		}
+	}
+
+	if conn.remote.UID != remote.UID {
+		return fmt.Errorf("Connection appears to be with different version of a peer we already know of")
+	}
+
+	if conn.remote == conn.local {
+		return fmt.Errorf("Cannot connect to ourself")
+	}
+
+	return nil
+}
+
 func (conn *LocalConnection) initHeartbeats() error {
 	conn.heartbeatTCP = time.NewTicker(TCPHeartbeat)
 	conn.heartbeatTimeout = time.NewTimer(HeartbeatTimeout)

router/handshake.go

@@ -45,7 +45,7 @@ func (fv *FieldValidator) Err() error {
 	return fv.err
 }
 
-func (conn *LocalConnection) handshake(acceptNewPeer bool) (TCPReceiver, error) {
+func (conn *LocalConnection) handshake() (TCPReceiver, *Peer, error) {
 	// We do not need to worry about locking in here as at this point
 	// the connection is not reachable by any go-routine other than
 	// ourself. Only when we add this connection to the conn.local
@@ -54,7 +54,7 @@ func (conn *LocalConnection) handshake(acceptNewPeer bool) (TCPReceiver, error)
 	conn.TCPConn.SetDeadline(time.Now().Add(HeaderTimeout))
 	version, err := exchangeProtocolHeader(conn.TCPConn, conn.TCPConn)
 	if err != nil {
-		return nil, err
+		return nil, nil, err
 	}
 	conn.version = version
 	conn.TCPConn.SetDeadline(time.Time{})
@@ -66,7 +66,7 @@ func (conn *LocalConnection) handshake(acceptNewPeer bool) (TCPReceiver, error)
 	dec := gob.NewDecoder(conn.TCPConn)
 	fv, private, err := conn.exchangeHandshake(localConnID, usingPassword, enc, dec)
 	if err != nil {
-		return nil, err
+		return nil, nil, err
 	}
 
 	nameStr, _ := fv.Value("Name")
@@ -74,25 +74,20 @@ func (conn *LocalConnection) handshake(acceptNewPeer bool) (TCPReceiver, error)
 	uidStr, _ := fv.Value("UID")
 	remoteConnIDStr, _ := fv.Value("ConnID")
 	if err := fv.Err(); err != nil {
-		return nil, err
+		return nil, nil, err
 	}
 
 	name, err := PeerNameFromString(nameStr)
 	if err != nil {
-		return nil, err
-	}
-	if !acceptNewPeer {
-		if conn.Router.Peers.Fetch(name) != nil {
-			return nil, fmt.Errorf("Found unknown remote name: %s at %s", name, conn.remoteTCPAddr)
-		}
+		return nil, nil, err
 	}
 	uid, err := ParsePeerUID(uidStr)
 	if err != nil {
-		return nil, err
+		return nil, nil, err
 	}
 	remoteConnID, err := strconv.ParseUint(remoteConnIDStr, 10, 64)
 	if err != nil {
-		return nil, err
+		return nil, nil, err
 	}
 	conn.uid = localConnID ^ remoteConnID
 
@@ -100,11 +95,11 @@ func (conn *LocalConnection) handshake(acceptNewPeer bool) (TCPReceiver, error)
 	remotePublicStr, rpErr := fv.Value("PublicKey")
 	if usingPassword {
 		if rpErr != nil {
-			return nil, rpErr
+			return nil, nil, rpErr
 		}
 		remotePublicSlice, rpErr := hex.DecodeString(remotePublicStr)
 		if rpErr != nil {
-			return nil, rpErr
+			return nil, nil, rpErr
 		}
 		remotePublic := [32]byte{}
 		for idx, elem := range remotePublicSlice {
@@ -116,14 +111,14 @@ func (conn *LocalConnection) handshake(acceptNewPeer bool) (TCPReceiver, error)
 		conn.Decryptor = NewNaClDecryptor(conn.SessionKey, conn.outbound)
 	} else {
 		if rpErr == nil {
-			return nil, fmt.Errorf("Remote network is encrypted. Password required.")
+			return nil, nil, fmt.Errorf("Remote network is encrypted. Password required.")
 		}
 		conn.tcpSender = NewSimpleTCPSender(enc)
 		tcpReceiver = NewSimpleTCPReceiver(dec)
 		conn.Decryptor = NewNonDecryptor()
 	}
 
-	return tcpReceiver, conn.setRemote(NewPeer(name, nickNameStr, uid, 0))
+	return tcpReceiver, NewPeer(name, nickNameStr, uid, 0), nil
 }
 
 func (conn *LocalConnection) exchangeHandshake(localConnID uint64, usingPassword bool, enc *gob.Encoder, dec *gob.Decoder) (*FieldValidator, *[32]byte, error) {
@@ -151,17 +146,3 @@ func (conn *LocalConnection) exchangeHandshake(localConnID uint64, usingPassword
 	fv.CheckEqual("PeerNameFlavour", PeerNameFlavour)
 	return fv, private, nil
 }
-
-func (conn *LocalConnection) setRemote(toPeer *Peer) error {
-	toPeer = conn.Router.Peers.FetchWithDefault(toPeer)
-	switch toPeer {
-	case nil:
-		return fmt.Errorf("Connection appears to be with different version of a peer we already know of")
-	case conn.local:
-		conn.remote = toPeer // have to do assigment here to ensure Shutdown releases ref count
-		return fmt.Errorf("Cannot connect to ourself")
-	default:
-		conn.remote = toPeer
-		return nil
-	}
-}

router/peers.go

@@ -64,9 +64,6 @@ func (peers *Peers) FetchWithDefault(peer *Peer) *Peer {
 	peers.Lock()
 	defer peers.Unlock()
 	if existingPeer, found := peers.table[peer.Name]; found {
-		if existingPeer.UID != peer.UID {
-			return nil
-		}
 		existingPeer.localRefCount++
 		return existingPeer
 	}
@@ -81,6 +78,16 @@ func (peers *Peers) Fetch(name PeerName) *Peer {
 	return peers.table[name]
 }
 
+func (peers *Peers) FetchAndAddRef(name PeerName) *Peer {
+	peers.Lock()
+	defer peers.Unlock()
+	peer := peers.table[name]
+	if peer != nil {
+		peer.localRefCount++
+	}
+	return peer
+}
+
 func (peers *Peers) Dereference(peer *Peer) {
 	peers.Lock()
 	defer peers.Unlock()