From 801ce0730d7d7e7b289d26d437aa87a746232951 Mon Sep 17 00:00:00 2001
From: braginini <bangvalo@gmail.com>
Date: Sun, 18 Apr 2021 15:54:10 +0200
Subject: [PATCH] feat: add retry on connection restart

---
 cmd/config.go             |  3 +--
 connection/connection.go  | 13 ++++++++++---
 signal/client.go          |  3 ++-
 {signal => util}/retry.go |  2 +-
 4 files changed, 14 insertions(+), 7 deletions(-)
 rename {signal => util}/retry.go (97%)

diff --git a/cmd/config.go b/cmd/config.go
index c7ea6e7d8..ad7640afe 100644
--- a/cmd/config.go
+++ b/cmd/config.go
@@ -10,8 +10,7 @@ import (
 
 type Config struct {
 	// Wireguard private key of local peer
-	PrivateKey string
-	// configured remote peers (Wireguard public keys)
+	PrivateKey   string
 	Peers        []connection.Peer
 	StunTurnURLs []*ice.URL
 	// host:port of the signal server
diff --git a/connection/connection.go b/connection/connection.go
index e18a87260..2b0e1a68d 100644
--- a/connection/connection.go
+++ b/connection/connection.go
@@ -5,6 +5,7 @@ import (
 	"github.com/pion/ice/v2"
 	log "github.com/sirupsen/logrus"
 	"github.com/wiretrustee/wiretrustee/iface"
+	"github.com/wiretrustee/wiretrustee/util"
 	"golang.zx2c4.com/wireguard/wgctrl/wgtypes"
 	"net"
 	"time"
@@ -156,6 +157,8 @@ func (conn *Connection) OnAnswer(remoteAuth IceCredentials) error {
 
 func (conn *Connection) OnOffer(remoteAuth IceCredentials) error {
 
+	conn.remoteAuthChannel <- remoteAuth
+
 	uFrag, pwd, err := conn.agent.GetLocalUserCredentials()
 	if err != nil {
 		return err
@@ -166,8 +169,6 @@ func (conn *Connection) OnOffer(remoteAuth IceCredentials) error {
 		return err
 	}
 
-	conn.remoteAuthChannel <- remoteAuth
-
 	return nil
 }
 
@@ -250,8 +251,14 @@ func (conn *Connection) listenOnConnectionStateChanges() error {
 			}
 			log.Debugf("connected to peer %s via selected candidate pair %s", conn.Config.RemoteWgKey.String(), pair)
 		} else if state == ice.ConnectionStateDisconnected || state == ice.ConnectionStateFailed {
-			err := conn.Restart()
+			err := util.Retry(15, time.Second, func() error {
+				return conn.Restart()
+			}, func(err error) {
+				log.Warnf("failed restarting connection, retrying ... %s", err)
+			})
+
 			if err != nil {
+				log.Errorf("failed restarting connection %s", err)
 				return
 			}
 		}
diff --git a/signal/client.go b/signal/client.go
index 9fd17d3cc..1d237c572 100644
--- a/signal/client.go
+++ b/signal/client.go
@@ -5,6 +5,7 @@ import (
 	"fmt"
 	log "github.com/sirupsen/logrus"
 	"github.com/wiretrustee/wiretrustee/signal/proto"
+	"github.com/wiretrustee/wiretrustee/util"
 	"google.golang.org/grpc"
 	"google.golang.org/grpc/codes"
 	"google.golang.org/grpc/keepalive"
@@ -64,7 +65,7 @@ func NewClient(addr string, ctx context.Context) (*Client, error) {
 func (client *Client) Receive(key string, msgHandler func(msg *proto.Message) error) {
 	client.connWg.Add(1)
 	go func() {
-		err := Retry(15, time.Second, func() error {
+		err := util.Retry(15, time.Second, func() error {
 			return client.connect(key, msgHandler)
 		}, func(err error) {
 			log.Warnf("disconnected from the Signal Exchange due to an error %s. Retrying ... ", err)
diff --git a/signal/retry.go b/util/retry.go
similarity index 97%
rename from signal/retry.go
rename to util/retry.go
index 06a42cd6f..b04854b24 100644
--- a/signal/retry.go
+++ b/util/retry.go
@@ -1,4 +1,4 @@
-package signal
+package util
 
 import (
 	"math/rand"