From ac871a530650eabc3167a54a29d23b938b75371b Mon Sep 17 00:00:00 2001
From: braginini <bangvalo@gmail.com>
Date: Mon, 19 Apr 2021 22:56:01 +0200
Subject: [PATCH] feat: add timeout while initiating conenction

---
 connection/connection.go | 19 +++++++++++--------
 connection/engine.go     |  7 +++++--
 2 files changed, 16 insertions(+), 10 deletions(-)

diff --git a/connection/connection.go b/connection/connection.go
index b21aaeef1..e3231e051 100644
--- a/connection/connection.go
+++ b/connection/connection.go
@@ -53,7 +53,7 @@ type Connection struct {
 	// remoteAuthChannel is a channel used to wait for remote credentials to proceed with the connection
 	remoteAuthChannel chan IceCredentials
 
-	closeChannel  chan bool
+	closeChannel  chan struct{}
 	closedChannel chan struct{}
 
 	// agent is an actual ice.Agent that is used to negotiate and maintain a connection to a remote peer
@@ -78,7 +78,7 @@ func NewConnection(config ConnConfig,
 		signalOffer:       signalOffer,
 		signalAnswer:      signalAnswer,
 		remoteAuthChannel: make(chan IceCredentials, 1),
-		closeChannel:      make(chan bool, 2),
+		closeChannel:      make(chan struct{}),
 		closedChannel:     make(chan struct{}),
 		agent:             nil,
 		isActive:          false,
@@ -108,8 +108,7 @@ func (conn *Connection) Close() error {
 
 	log.Debugf("closing connection to peer %s", conn.Config.RemoteWgKey.String())
 
-	conn.closeChannel <- true
-	conn.closeChannel <- true
+	close(conn.closeChannel)
 
 	log.Debugf("closed connection to peer %s", conn.Config.RemoteWgKey.String())
 
@@ -364,8 +363,10 @@ func (conn *Connection) proxyToRemotePeer(wgConn net.Conn, remoteConn *ice.Conn)
 			if err != nil {
 				log.Warnln("Error writing to remote peer: ", err.Error())
 			}
-		case <-conn.closeChannel:
-			log.Infof("stopped proxying to remote peer %s", conn.Config.RemoteWgKey.String())
+		case _, ok := <-conn.closeChannel:
+			if !ok {
+				log.Infof("stopped proxying to remote peer %s", conn.Config.RemoteWgKey.String())
+			}
 			return
 		}
 	}
@@ -388,8 +389,10 @@ func (conn *Connection) proxyToLocalWireguard(wgConn net.Conn, remoteConn *ice.C
 			if err != nil {
 				log.Errorf("failed writing to local Wireguard instance %s", err)
 			}
-		case <-conn.closeChannel:
-			log.Infof("stopped proxying from remote peer %s", conn.Config.RemoteWgKey.String())
+		case _, ok := <-conn.closeChannel:
+			if !ok {
+				log.Infof("stopped proxying from remote peer %s", conn.Config.RemoteWgKey.String())
+			}
 			return
 		}
 	}
diff --git a/connection/engine.go b/connection/engine.go
index b41e05027..c71ef15b7 100644
--- a/connection/engine.go
+++ b/connection/engine.go
@@ -77,14 +77,16 @@ func (e *Engine) Start(privateKey string, peers []Peer) error {
 		go func() {
 
 			operation := func() error {
-				_, closed, err := e.openConnection(*wgPort, myKey, peer)
+				_, closed, err := e.openPeerConnection(*wgPort, myKey, peer)
 				if err != nil {
+					e.conns[peer.WgPubKey] = nil
 					return err
 				}
 
 				select {
 				case _, ok := <-closed:
 					if !ok {
+						e.conns[peer.WgPubKey] = nil
 						return fmt.Errorf("connection to peer %s has been closed", peer.WgPubKey)
 					}
 					return nil
@@ -103,7 +105,8 @@ func (e *Engine) Start(privateKey string, peers []Peer) error {
 	return nil
 }
 
-func (e *Engine) openConnection(wgPort int, myKey wgtypes.Key, peer Peer) (*Connection, chan struct{}, error) {
+func (e *Engine) openPeerConnection(wgPort int, myKey wgtypes.Key, peer Peer) (*Connection, chan struct{}, error) {
+
 	remoteKey, _ := wgtypes.ParseKey(peer.WgPubKey)
 	connConfig := &ConnConfig{
 		WgListenAddr: fmt.Sprintf("127.0.0.1:%d", wgPort),