From e6607b3bef1e19589b87305af1b0383376869f65 Mon Sep 17 00:00:00 2001
From: Liang Ding <dl88250@gmail.com>
Date: Thu, 20 Nov 2014 14:11:54 +0800
Subject: [PATCH] :fire: channel enhance

---
 editor/editors.go             |  3 ++-
 notification/notifications.go |  2 +-
 session/sessions.go           |  2 +-
 shell/shells.go               |  3 ++-
 util/websocket.go             | 15 +++++++++++++++
 5 files changed, 21 insertions(+), 4 deletions(-)

diff --git a/editor/editors.go b/editor/editors.go
index 5945f33..f4c75d3 100644
--- a/editor/editors.go
+++ b/editor/editors.go
@@ -55,7 +55,7 @@ func WSHandler(w http.ResponseWriter, r *http.Request) {
 
 	args := map[string]interface{}{}
 	for {
-		if err := session.EditorWS[sid].Conn.ReadJSON(&args); err != nil {
+		if err := session.EditorWS[sid].ReadJSON(&args); err != nil {
 			if err.Error() == "EOF" {
 				return
 			}
@@ -65,6 +65,7 @@ func WSHandler(w http.ResponseWriter, r *http.Request) {
 			}
 
 			glog.Error("Editor WS ERROR: " + err.Error())
+
 			return
 		}
 
diff --git a/notification/notifications.go b/notification/notifications.go
index 3c557ba..483c3d3 100644
--- a/notification/notifications.go
+++ b/notification/notifications.go
@@ -113,7 +113,7 @@ func WSHandler(w http.ResponseWriter, r *http.Request) {
 	input := map[string]interface{}{}
 
 	for {
-		if err := wsChan.Conn.ReadJSON(&input); err != nil {
+		if err := wsChan.ReadJSON(&input); err != nil {
 			if err.Error() == "EOF" {
 				return
 			}
diff --git a/session/sessions.go b/session/sessions.go
index 3c5d434..951c08a 100644
--- a/session/sessions.go
+++ b/session/sessions.go
@@ -139,7 +139,7 @@ func WSHandler(w http.ResponseWriter, r *http.Request) {
 	input := map[string]interface{}{}
 
 	for {
-		if err := wsChan.Conn.ReadJSON(&input); err != nil {
+		if err := wsChan.ReadJSON(&input); err != nil {
 			glog.V(3).Infof("[Session Channel] of session [%s] disconnected, releases all resources with it", sid)
 
 			WideSessions.Remove(sid)
diff --git a/shell/shells.go b/shell/shells.go
index 9615c06..ff4e0f0 100644
--- a/shell/shells.go
+++ b/shell/shells.go
@@ -102,7 +102,7 @@ func WSHandler(w http.ResponseWriter, r *http.Request) {
 	input := map[string]interface{}{}
 
 	for {
-		if err := wsChan.Conn.ReadJSON(&input); err != nil {
+		if err := wsChan.ReadJSON(&input); err != nil {
 			if err.Error() == "EOF" {
 				return
 			}
@@ -112,6 +112,7 @@ func WSHandler(w http.ResponseWriter, r *http.Request) {
 			}
 
 			glog.Error("Shell WS ERROR: " + err.Error())
+
 			return
 		}
 
diff --git a/util/websocket.go b/util/websocket.go
index 0781920..1e826e0 100644
--- a/util/websocket.go
+++ b/util/websocket.go
@@ -45,6 +45,21 @@ func (c *WSChannel) WriteJSON(v interface{}) (ret error) {
 	return c.Conn.WriteJSON(v)
 }
 
+// ReadJSON reads the next JSON-encoded message from the channel and stores it in the value pointed to by v.
+func (c *WSChannel) ReadJSON(v interface{}) (ret error) {
+	if nil == c.Conn {
+		return errors.New("connection is nil, channel has been closed")
+	}
+
+	defer func() {
+		if r := recover(); nil != r {
+			ret = errors.New("channel has been closed")
+		}
+	}()
+
+	return c.Conn.ReadJSON(v)
+}
+
 // Close closed the channel.
 func (c *WSChannel) Close() {
 	if nil != c.Conn {