From 51c6e0091f240dcbcff07ca83bcd24fc843a2449 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Miroslav=20=C5=A0ediv=C3=BD?= <sedivy.miro@gmail.com>
Date: Sun, 1 Nov 2020 21:03:25 +0100
Subject: [PATCH] Connected / Disconnected refactor.

---
 internal/session/manager.go   |  2 +-
 internal/session/session.go   | 24 ++++++++++++------------
 internal/types/session.go     |  3 +--
 internal/webrtc/webrtc.go     |  4 ++--
 internal/websocket/manager.go |  2 +-
 5 files changed, 17 insertions(+), 18 deletions(-)

diff --git a/internal/session/manager.go b/internal/session/manager.go
index c8ac91da..9af44f22 100644
--- a/internal/session/manager.go
+++ b/internal/session/manager.go
@@ -168,7 +168,7 @@ func (manager *SessionManagerCtx) OnConnected(listener func(session types.Sessio
 func (manager *SessionManagerCtx) OnDisconnected(listener func(session types.Session)) {
 	manager.emmiter.On("disconnected", func(payload ...interface{}) {
 		// Stop streaming, if everyone left
-		if manager.capture.Streaming() && len(manager.members) == 0 {
+		if manager.capture.Streaming() && len(manager.Members()) == 0 {
 			manager.capture.StopStream()
 		}
 
diff --git a/internal/session/session.go b/internal/session/session.go
index efcbd2e2..b1f7071c 100644
--- a/internal/session/session.go
+++ b/internal/session/session.go
@@ -69,18 +69,18 @@ func (session *SessionCtx) SetPeer(peer types.Peer) {
 	session.peer = peer
 }
 
-func (session *SessionCtx) SetConnected() {
-	session.connected = true
-	session.manager.emmiter.Emit("connected", session)
-}
+func (session *SessionCtx) SetConnected(connected bool) {
+	session.connected = connected
 
-func (session *SessionCtx) SetDisconnected() {
-	session.connected = false
-	session.manager.emmiter.Emit("disconnected", session)
-	session.socket = nil
-
-	// TODO: Refactor.
-	session.manager.Destroy(session.id)
+	if connected {
+		session.manager.emmiter.Emit("connected", session)
+	} else {
+		session.manager.emmiter.Emit("disconnected", session)
+		session.socket = nil
+	
+		// TODO: Refactor.
+		session.manager.Destroy(session.id)
+	}
 }
 
 func (session *SessionCtx) Disconnect(reason string) error {
@@ -96,7 +96,7 @@ func (session *SessionCtx) Disconnect(reason string) error {
 		return err
 	}
 
-	session.SetDisconnected()
+	session.SetConnected(false)
 	return nil
 }
 
diff --git a/internal/types/session.go b/internal/types/session.go
index 7c5715df..a80dd9a6 100644
--- a/internal/types/session.go
+++ b/internal/types/session.go
@@ -14,8 +14,7 @@ type Session interface {
 	SetName(name string)
 	SetSocket(socket WebSocket)
 	SetPeer(peer Peer)
-	SetConnected()
-	SetDisconnected()
+	SetConnected(connected bool)
 	Disconnect(reason string) error
 	Send(v interface{}) error
 	SignalAnswer(sdp string) error
diff --git a/internal/webrtc/webrtc.go b/internal/webrtc/webrtc.go
index 3b20ebea..a8b3dcc5 100644
--- a/internal/webrtc/webrtc.go
+++ b/internal/webrtc/webrtc.go
@@ -156,10 +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")
-			session.SetDisconnected()
+			session.SetConnected(false)
 		case webrtc.PeerConnectionStateConnected:
 			manager.logger.Info().Str("id", session.ID()).Msg("peer connected")
-			session.SetConnected()
+			session.SetConnected(true)
 		}
 	})
 
diff --git a/internal/websocket/manager.go b/internal/websocket/manager.go
index f8884ea6..0a727644 100644
--- a/internal/websocket/manager.go
+++ b/internal/websocket/manager.go
@@ -197,7 +197,7 @@ func (ws *WebSocketManagerCtx) handle(connection *websocket.Conn, session types.
 		defer func() {
 			ticker.Stop()
 			ws.logger.Debug().Str("address", connection.RemoteAddr().String()).Msg("handle socket ending")
-			session.SetDisconnected()
+			session.SetConnected(false)
 		}()
 
 		for {