From 675358ce5c7b867afb97eeb1b494dc17b29bdfa1 Mon Sep 17 00:00:00 2001
From: braginini <bangvalo@gmail.com>
Date: Wed, 14 Apr 2021 14:54:02 +0200
Subject: [PATCH] fix: engine flow on ASNWER/OFFER

---
 engine/agent.go  | 40 +++++++++++------------------
 engine/engine.go | 65 ++++++++++++++++++++++++++++++++++--------------
 2 files changed, 62 insertions(+), 43 deletions(-)

diff --git a/engine/agent.go b/engine/agent.go
index d8646a8a8..5f30c1c7e 100644
--- a/engine/agent.go
+++ b/engine/agent.go
@@ -151,14 +151,6 @@ func (pa *PeerAgent) OpenConnection(initiator bool) (net.Conn, error) {
 	}
 	pa.wgConn = wgConn
 
-	go func() {
-		pa.proxyToRemotePeer()
-	}()
-
-	go func() {
-		pa.proxyToLocalWireguard()
-	}()
-
 	return wgConn, nil
 }
 
@@ -201,7 +193,6 @@ func (pa *PeerAgent) onCandidate() error {
 
 // onConnectionStateChange listens on ice.Agent connection state change events and once connected checks a Candidate pair
 // the ice.Conn was established with
-// Mostly used for debugging purposes (e.g. connection time, etc)
 func (pa *PeerAgent) onConnectionStateChange() error {
 	return pa.iceAgent.OnConnectionStateChange(func(state ice.ConnectionState) {
 		log.Debugf("ICE Connection State has changed: %s", state.String())
@@ -213,36 +204,35 @@ func (pa *PeerAgent) onConnectionStateChange() error {
 				return
 			}
 			log.Debugf("connected to peer %s via selected candidate pair %s", pa.RemoteKey, pair)
+
+			// start proxying data between local Wireguard and remote peer
+			go func() {
+				pa.proxyToRemotePeer()
+			}()
+
+			go func() {
+				pa.proxyToLocalWireguard()
+			}()
 		}
 	})
 }
 
 // authenticate sets the signal.Credential of the remote peer
-// and sends local signal.Credential to teh remote peer via signal server
-func (pa *PeerAgent) Authenticate(credential *signal.Credential) error {
+// and returns local Credentials
+func (pa *PeerAgent) Authenticate(credential *signal.Credential) (*signal.Credential, error) {
 
 	err := pa.iceAgent.SetRemoteCredentials(credential.UFrag, credential.Pwd)
 	if err != nil {
-		return err
+		return nil, err
 	}
 
 	localUFrag, localPwd, err := pa.iceAgent.GetLocalUserCredentials()
 	if err != nil {
-		return err
+		return nil, err
 	}
 
-	// notify the remote peer about our credentials
-	answer := signal.MarshalCredential(pa.LocalKey, pa.RemoteKey, &signal.Credential{
+	return &signal.Credential{
 		UFrag: localUFrag,
-		Pwd:   localPwd,
-	}, sProto.Message_ANSWER)
-
-	//notify the remote peer of our credentials
-	err = pa.signal.Send(answer)
-	if err != nil {
-		return err
-	}
-
-	return nil
+		Pwd:   localPwd}, nil
 
 }
diff --git a/engine/engine.go b/engine/engine.go
index eedb6be13..5e10ab96e 100644
--- a/engine/engine.go
+++ b/engine/engine.go
@@ -88,29 +88,30 @@ func (e *Engine) receiveSignal(localKey string) {
 
 		switch msg.Type {
 		case sProto.Message_OFFER:
+
+			cred, err := e.handle(msg, peerAgent, initiator)
+			if err != nil {
+				return err
+			}
+			// notify the remote peer about our credentials
+			answer := signal.MarshalCredential(peerAgent.LocalKey, peerAgent.RemoteKey, &signal.Credential{
+				UFrag: cred.UFrag,
+				Pwd:   cred.Pwd,
+			}, sProto.Message_ANSWER)
+
+			//notify the remote peer of our credentials
+			err = peerAgent.signal.Send(answer)
+			if err != nil {
+				return err
+			}
+
+			return nil
 		case sProto.Message_ANSWER:
-			remoteCred, err := signal.UnMarshalCredential(msg)
+			_, err := e.handle(msg, peerAgent, initiator)
 			if err != nil {
 				return err
 			}
 
-			err = peerAgent.Authenticate(remoteCred)
-			if err != nil {
-				log.Errorf("error authenticating remote peer %s", msg.Key)
-				return err
-			}
-
-			conn, err := peerAgent.OpenConnection(initiator)
-			if err != nil {
-				log.Errorf("error opening connection ot remote peer %s", msg.Key)
-				return err
-			}
-
-			err = iface.UpdatePeer(e.wgIface, peerAgent.RemoteKey, "0.0.0.0/0", 15*time.Second, conn.LocalAddr().String())
-			if err != nil {
-				log.Errorf("error while configuring Wireguard peer [%s] %s", peerAgent.RemoteKey, err.Error())
-				return err
-			}
 		case sProto.Message_CANDIDATE:
 			err := peerAgent.OnRemoteCandidate(msg)
 			if err != nil {
@@ -124,3 +125,31 @@ func (e *Engine) receiveSignal(localKey string) {
 
 	e.signal.WaitConnected()
 }
+
+func (e *Engine) handle(msg *sProto.Message, peerAgent *PeerAgent, initiator bool) (*signal.Credential, error) {
+	remoteCred, err := signal.UnMarshalCredential(msg)
+	if err != nil {
+		return nil, err
+	}
+
+	cred, err := peerAgent.Authenticate(remoteCred)
+	if err != nil {
+		log.Errorf("error authenticating remote peer %s", msg.Key)
+		return nil, err
+	}
+
+	go func() {
+
+		conn, err := peerAgent.OpenConnection(initiator)
+		if err != nil {
+			log.Errorf("error opening connection ot remote peer %s", msg.Key)
+		}
+
+		err = iface.UpdatePeer(e.wgIface, peerAgent.RemoteKey, "0.0.0.0/0", 15*time.Second, conn.LocalAddr().String())
+		if err != nil {
+			log.Errorf("error while configuring Wireguard peer [%s] %s", peerAgent.RemoteKey, err.Error())
+		}
+	}()
+
+	return cred, nil
+}