websocket: add method to merge all continuous frames

The method is called fin() with single parameter: the last frame that
contains finished state.

Also, use the Frames.fin method on ClientManager.finFrames, to minimize
duplicate code.

2 files changed, 30 insertions, 15 deletions

diff --git a/lib/websocket/clientmanager.go b/lib/websocket/clientmanager.go
index 313248ce..54710337 100644
--- a/lib/websocket/clientmanager.go
+++ b/lib/websocket/clientmanager.go
@@ -73,28 +73,17 @@ func (cls *ClientManager) All() (conns []int) {
 // finFrames merge all continuous frames into single frame and clear the
 // stored frame and frames on behalf of connection.
 //
-func (cls *ClientManager) finFrames(conn int, fin *Frame) (f *Frame) {
+func (cls *ClientManager) finFrames(conn int, last *Frame) (f *Frame) {
 	cls.Lock()
+
 	frames, ok := cls.frames[conn]
 	if !ok {
 		cls.Unlock()
-		return fin
+		return last
 	}
 
-	f = frames.v[0]
-	for x := 1; x < len(frames.v); x++ {
-		if frames.v[x].opcode == OpcodeClose {
-			break
-		}
+	f = frames.fin(last)
 
-		// Ignore control PING or PONG frame.
-		if frames.v[x].opcode == OpcodePing || frames.v[x].opcode == OpcodePong {
-			continue
-		}
-
-		f.payload = append(f.payload, frames.v[x].payload...)
-	}
-	f.payload = append(f.payload, fin.payload...)
 	delete(cls.frames, conn)
 	delete(cls.frame, conn)
 
diff --git a/lib/websocket/frames.go b/lib/websocket/frames.go
index 6a71a9ce..2de567a2 100644
--- a/lib/websocket/frames.go
+++ b/lib/websocket/frames.go
@@ -63,6 +63,32 @@ func (frames *Frames) Append(f *Frame) {
 }
 
 //
+// fin merge all continuous frame with last frame and return it as single
+// frame.
+//
+func (frames *Frames) fin(last *Frame) (frame *Frame) {
+	frame = frames.v[0]
+
+	for x := 1; x < len(frames.v); x++ {
+		if frames.v[x].opcode == OpcodeClose {
+			break
+		}
+
+		// Ignore control PING or PONG frame.
+		if frames.v[x].opcode == OpcodePing || frames.v[x].opcode == OpcodePong {
+			continue
+		}
+
+		frame.payload = append(frame.payload, frames.v[x].payload...)
+	}
+	if last != nil {
+		frame.payload = append(frame.payload, last.payload...)
+	}
+
+	return frame
+}
+
+//
 // isClosed will return true if one of the frame is control CLOSE frame.
 //
 func (frames *Frames) isClosed() bool {