From 9ea0bf223c9178b694830daf6e8d2cff87fdcb9e Mon Sep 17 00:00:00 2001
From: giongto35 <giongto35@yahoo.com>
Date: Thu, 16 May 2019 01:04:54 +0800
Subject: [PATCH] Check peerconnection is in room before detach and reconn

---
 handler/browser.go   | 16 ++++++++++++----
 handler/handlers.go  | 12 ++++--------
 handler/room/room.go | 18 +++++++++++++-----
 webrtc/webrtc.go     | 12 ++++++------
 4 files changed, 35 insertions(+), 23 deletions(-)

diff --git a/handler/browser.go b/handler/browser.go
index 6c174f5f..b76e9a55 100644
--- a/handler/browser.go
+++ b/handler/browser.go
@@ -46,6 +46,9 @@ func (s *Session) RouteBrowser() {
 		req.Data = "ok"
 		if s.RoomID != "" {
 			room := s.handler.getRoom(s.RoomID)
+			if room == nil {
+				return
+			}
 			err := room.SaveGame()
 			if err != nil {
 				log.Println("[!] Cannot save game state: ", err)
@@ -95,17 +98,22 @@ func (s *Session) RouteBrowser() {
 			}
 		}
 
-		// Create new room
+		// Get Room in local server
 		// TODO: check if roomID is in the current server
 		room := s.handler.getRoom(s.RoomID)
 		log.Println("Got Room from local ", room, " ID: ", s.RoomID)
+		// If room is not running
 		if room == nil {
+			// Create new room
 			room = s.handler.createNewRoom(s.GameName, s.RoomID, s.PlayerIndex)
 		}
 
-		// Attach peerconnection to room
-		s.handler.detachPeerConn(s.peerconnection)
-		room.AddConnectionToRoom(s.peerconnection, s.PlayerIndex)
+		// Attach peerconnection to room. If PC is already in room, don't detach
+		log.Println("Is PC in room", room.IsPCInRoom(s.peerconnection))
+		if !room.IsPCInRoom(s.peerconnection) {
+			s.handler.detachPeerConn(s.peerconnection)
+			room.AddConnectionToRoom(s.peerconnection, s.PlayerIndex)
+		}
 		s.RoomID = room.ID
 
 		// Register room to overlord if we are connecting to overlord
diff --git a/handler/handlers.go b/handler/handlers.go
index fd54d663..e75483f0 100644
--- a/handler/handlers.go
+++ b/handler/handlers.go
@@ -112,18 +112,14 @@ func (h *Handler) WS(w http.ResponseWriter, r *http.Request) {
 		Data: gamelist.GetEncodedGameList(h.gamePath),
 	}, nil)
 
-	wssession.BrowserClient.Listen()
-
-	// TODO: Use callback
 	// If peerconnection is done (client.Done is signalled), we close peerconnection
 	go func() {
-		for {
-			<-client.Done
-			h.detachPeerConn(wssession.peerconnection)
-			return
-		}
+		<-client.Done
+		log.Println("Socket terminated, detach connection")
+		h.detachPeerConn(wssession.peerconnection)
 	}()
 
+	wssession.BrowserClient.Listen()
 }
 
 // Detach peerconnection detach/remove a peerconnection from current room
diff --git a/handler/room/room.go b/handler/room/room.go
index a5b6f5f0..60c341cc 100644
--- a/handler/room/room.go
+++ b/handler/room/room.go
@@ -1,7 +1,6 @@
 package room
 
 import (
-	"fmt"
 	"image"
 	"io/ioutil"
 	"log"
@@ -147,14 +146,13 @@ func (r *Room) CleanSession(peerconnection *webrtc.WebRTC) {
 }
 
 func (r *Room) removeSession(w *webrtc.WebRTC) {
-	fmt.Println("Cleaning session: ", w)
-	fmt.Println("Sessions list", r.rtcSessions)
+	log.Println("Cleaning session: ", w.ID)
 	// TODO: get list of r.rtcSessions in lock
 	for i, s := range r.rtcSessions {
-		fmt.Println("found session: ", s, w)
+		log.Println("found session: ", w.ID)
 		if s.ID == w.ID {
 			r.rtcSessions = append(r.rtcSessions[:i], r.rtcSessions[i+1:]...)
-			fmt.Println("found session: ", len(r.rtcSessions))
+			log.Println("Removed session ", s.ID, " from room: ", r.ID)
 
 			// If room has no sessions, close room
 			// Note: this logic cannot be brought outside of forloop because we only close room if room had at least one session
@@ -169,6 +167,16 @@ func (r *Room) removeSession(w *webrtc.WebRTC) {
 	}
 }
 
+// TODO: Reuse for remove Session
+func (r *Room) IsPCInRoom(w *webrtc.WebRTC) bool {
+	for _, s := range r.rtcSessions {
+		if s.ID == w.ID {
+			return true
+		}
+	}
+	return false
+}
+
 func (r *Room) Close() {
 	if r.Done {
 		return
diff --git a/webrtc/webrtc.go b/webrtc/webrtc.go
index f1b0b1cd..c7a1d98e 100644
--- a/webrtc/webrtc.go
+++ b/webrtc/webrtc.go
@@ -93,12 +93,12 @@ type WebRTC struct {
 
 // StartClient start webrtc
 func (w *WebRTC) StartClient(remoteSession string, width, height int) (string, error) {
-	defer func() {
-		if err := recover(); err != nil {
-			log.Println(err)
-			w.StopClient()
-		}
-	}()
+	//defer func() {
+	//if err := recover(); err != nil {
+	//log.Println(err)
+	//w.StopClient()
+	//}
+	//}()
 
 	// reset client
 	if w.isConnected {