From f5e52eb1d9f93569901953da3f5a1115b5fd3b39 Mon Sep 17 00:00:00 2001
From: braginini <bangvalo@gmail.com>
Date: Mon, 15 Nov 2021 15:51:54 +0100
Subject: [PATCH] fix: minor refactor to consider signal status

---
 signal/client/grpc.go      | 18 ++++++++----------
 signal/client/websocket.go | 24 ++++++++++++++++++------
 2 files changed, 26 insertions(+), 16 deletions(-)

diff --git a/signal/client/grpc.go b/signal/client/grpc.go
index e44ae8e76..ac576ecfb 100644
--- a/signal/client/grpc.go
+++ b/signal/client/grpc.go
@@ -22,12 +22,6 @@ import (
 
 // A set of tools to exchange connection details (Wireguard endpoints) with the remote peer.
 
-// Status is the status of the client
-type Status string
-
-const streamConnected Status = "streamConnected"
-const streamDisconnected Status = "streamDisconnected"
-
 // GrpcClient Wraps the Signal Exchange Service gRpc client
 type GrpcClient struct {
 	key        wgtypes.Key
@@ -42,6 +36,10 @@ type GrpcClient struct {
 	status Status
 }
 
+func (c *GrpcClient) GetStatus() Status {
+	return c.status
+}
+
 // Close Closes underlying connections to the Signal Exchange
 func (c *GrpcClient) Close() error {
 	return c.signalConn.Close()
@@ -79,7 +77,7 @@ func NewClient(ctx context.Context, addr string, key wgtypes.Key, tlsEnabled boo
 		signalConn: conn,
 		key:        key,
 		mux:        sync.Mutex{},
-		status:     streamDisconnected,
+		status:     StreamDisconnected,
 	}, nil
 }
 
@@ -148,13 +146,13 @@ func (c *GrpcClient) Receive(msgHandler func(msg *proto.Message) error) error {
 func (c *GrpcClient) notifyStreamDisconnected() {
 	c.mux.Lock()
 	defer c.mux.Unlock()
-	c.status = streamDisconnected
+	c.status = StreamDisconnected
 }
 
 func (c *GrpcClient) notifyStreamConnected() {
 	c.mux.Lock()
 	defer c.mux.Unlock()
-	c.status = streamConnected
+	c.status = StreamConnected
 	if c.connectedCh != nil {
 		// there are goroutines waiting on this channel -> release them
 		close(c.connectedCh)
@@ -206,7 +204,7 @@ func (c *GrpcClient) ready() bool {
 // WaitStreamConnected waits until the client is connected to the Signal stream
 func (c *GrpcClient) WaitStreamConnected() {
 
-	if c.status == streamConnected {
+	if c.status == StreamConnected {
 		return
 	}
 
diff --git a/signal/client/websocket.go b/signal/client/websocket.go
index 2ec37a57c..19689a76c 100644
--- a/signal/client/websocket.go
+++ b/signal/client/websocket.go
@@ -8,14 +8,21 @@ import (
 	"github.com/wiretrustee/wiretrustee/signal/proto"
 	"golang.zx2c4.com/wireguard/wgctrl/wgtypes"
 	"nhooyr.io/websocket"
+	"sync"
 	"time"
 )
 
 //WebsocketClient is a Signal server websocket client (alternative to the original gRPC Client)
 type WebsocketClient struct {
-	key  wgtypes.Key
-	ctx  context.Context
-	conn *websocket.Conn
+	key    wgtypes.Key
+	ctx    context.Context
+	conn   *websocket.Conn
+	status Status
+	mu     sync.Mutex
+}
+
+func (c *WebsocketClient) GetStatus() Status {
+	return c.status
 }
 
 func NewWebsocketClient(ctx context.Context, endpoint string, wgPrivateKey wgtypes.Key) (*WebsocketClient, error) {
@@ -32,13 +39,18 @@ func NewWebsocketClient(ctx context.Context, endpoint string, wgPrivateKey wgtyp
 	}
 
 	return &WebsocketClient{
-		key:  wgPrivateKey,
-		ctx:  ctx,
-		conn: conn,
+		key:    wgPrivateKey,
+		ctx:    ctx,
+		conn:   conn,
+		status: StreamConnected,
+		mu:     sync.Mutex{},
 	}, nil
 }
 
 func (c *WebsocketClient) Close() error {
+	c.mu.Lock()
+	defer c.mu.Unlock()
+	c.status = StreamDisconnected
 	return c.conn.Close(websocket.StatusNormalClosure, "close")
 }