From f9b2460c3b3e21fc33d2580ba2856a2e2db28434 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Miroslav=20=C5=A0ediv=C3=BD?= <sedivy.miro@gmail.com>
Date: Wed, 25 Nov 2020 20:26:11 +0100
Subject: [PATCH] replace SetConnected function.

---
 internal/session/manager.go   |  1 -
 internal/session/session.go   | 38 ++++++++++++++++-------------------
 internal/types/session.go     |  2 --
 internal/webrtc/manager.go    |  6 ++----
 internal/websocket/manager.go | 18 ++++++++---------
 5 files changed, 27 insertions(+), 38 deletions(-)

diff --git a/internal/session/manager.go b/internal/session/manager.go
index ac2e6e7d..f5cad684 100644
--- a/internal/session/manager.go
+++ b/internal/session/manager.go
@@ -47,7 +47,6 @@ func (manager *SessionManagerCtx) New(id string, admin bool) types.Session {
 		profile:   MemberProfile{
 			is_admin: admin,
 		},
-		connected: false,
 	}
 
 	manager.members[id] = session
diff --git a/internal/session/session.go b/internal/session/session.go
index ebf4bf0c..d775fd85 100644
--- a/internal/session/session.go
+++ b/internal/session/session.go
@@ -27,8 +27,6 @@ type SessionCtx struct {
 	websocket_connected bool
 	webrtc_peer         types.WebRTCPeer
 	webrtc_connected    bool
-	// TODO: Refactor.
-	connected           bool
 }
 
 func (session *SessionCtx) ID() string {
@@ -48,7 +46,7 @@ func (session *SessionCtx) IsHost() bool {
 }
 
 func (session *SessionCtx) Connected() bool {
-	return session.connected
+	return session.websocket_connected && session.webrtc_connected
 }
 
 func (session *SessionCtx) SetName(name string) {
@@ -66,7 +64,15 @@ func (session *SessionCtx) SetWebSocketConnected(connected bool) {
 		session.manager.emmiter.Emit("websocket_connected", session)
 	} else {
 		session.websocket_connected = false
-		session.manager.emmiter.Emit("websocket_disconnected", session)
+
+		// TODO: Refactor.
+		//session.manager.emmiter.Emit("websocket_disconnected", session)
+		session.manager.emmiter.Emit("disconnected", session)
+
+		session.websocket_peer = nil
+		
+		// TODO: Refactor.
+		_ = session.manager.Destroy(session.id)
 	}
 }
 
@@ -77,31 +83,21 @@ func (session *SessionCtx) SetWebRTCPeer(webrtc_peer types.WebRTCPeer) {
 func (session *SessionCtx) SetWebRTCConnected(connected bool) {
 	if connected {
 		session.webrtc_connected = true
-		session.manager.emmiter.Emit("webrtc_connected", session)
+
+		// TODO: Refactor.
+		//session.manager.emmiter.Emit("webrtc_connected", session)
+		session.manager.emmiter.Emit("connected", session)
 	} else {
 		session.webrtc_connected = false
 		session.manager.emmiter.Emit("webrtc_disconnected", session)
-	}
-}
 
-// TODO: Refactor.
-func (session *SessionCtx) SetConnected(connected bool) {
-	session.connected = connected
-
-	if connected {
-		session.manager.emmiter.Emit("connected", session)
-	} else {
-		session.manager.emmiter.Emit("disconnected", session)
-		session.websocket_peer = nil
-	
-		// TODO: Refactor.
-		_ = session.manager.Destroy(session.id)
+		session.webrtc_peer = nil
 	}
 }
 
 func (session *SessionCtx) Disconnect(reason string) error {
-	// TODO: Refactor.
-	session.SetConnected(false)
+	// TODO: End WebSocket connection.
+	// TODO: End WebRTC connection.
 
 	return session.Send(
 		message.Disconnect{
diff --git a/internal/types/session.go b/internal/types/session.go
index 29507962..70b0b533 100644
--- a/internal/types/session.go
+++ b/internal/types/session.go
@@ -13,8 +13,6 @@ type Session interface {
 	SetWebSocketConnected(connected bool)
 	SetWebRTCPeer(webrtc_peer WebRTCPeer)
 	SetWebRTCConnected(connected bool)
-	// TODO: Refactor.
-	SetConnected(connected bool)
 	Disconnect(reason string) error
 	Send(v interface{}) error
 	SignalAnswer(sdp string) error
diff --git a/internal/webrtc/manager.go b/internal/webrtc/manager.go
index 9997ec03..1af845aa 100644
--- a/internal/webrtc/manager.go
+++ b/internal/webrtc/manager.go
@@ -156,12 +156,10 @@ func (manager *WebRTCManagerCtx) CreatePeer(session types.Session) (string, bool
 		case webrtc.PeerConnectionStateDisconnected:
 		case webrtc.PeerConnectionStateFailed:
 			manager.logger.Info().Str("id", session.ID()).Msg("peer disconnected")
-			// TODO: Change WebRTC connection state.
-			session.SetConnected(false)
+			session.SetWebRTCConnected(false)
 		case webrtc.PeerConnectionStateConnected:
 			manager.logger.Info().Str("id", session.ID()).Msg("peer connected")
-			// TODO: Change WebRTC connection state.
-			session.SetConnected(true)
+			session.SetWebRTCConnected(true)
 		}
 	})
 
diff --git a/internal/websocket/manager.go b/internal/websocket/manager.go
index b4a48f4b..89319386 100644
--- a/internal/websocket/manager.go
+++ b/internal/websocket/manager.go
@@ -169,12 +169,16 @@ func (ws *WebSocketManagerCtx) Upgrade(w http.ResponseWriter, r *http.Request) e
 		Str("address", connection.RemoteAddr().String()).
 		Msg("connection started")
 
+	session.SetWebSocketConnected(true)
+
 	defer func() {
 		ws.logger.
 			Debug().
 			Str("session", session.ID()).
 			Str("address", connection.RemoteAddr().String()).
 			Msg("connection ended")
+
+		session.SetWebSocketConnected(false)
 	}()
 
 	ws.handle(connection, session)
@@ -186,17 +190,11 @@ func (ws *WebSocketManagerCtx) handle(connection *websocket.Conn, session types.
 	cancel := make(chan struct{})
 	ticker := time.NewTicker(pingPeriod)
 
-	go func() {
-		// TODO: Change WebSocket connection state.
-		//session.SetConnected(true)
-		
-		defer func() {
-			ticker.Stop()
-			ws.logger.Debug().Str("address", connection.RemoteAddr().String()).Msg("handle socket ending")
-			// TODO: Change WebSocket connection state.
-			session.SetConnected(false)
-		}()
+	defer func() {
+		ticker.Stop()
+	}()
 
+	go func() {
 		for {
 			_, raw, err := connection.ReadMessage()
 			if err != nil {