websocket: prevent adding closed client connection back to epoll

When server received unmasked frame or CLOSE control frame, the client
connection will be closed.  To prevent the connection to be added back
to epoll we add a flag to indicate that connection will be closed.

2 files changed, 25 insertions, 8 deletions

diff --git a/lib/websocket/client_test.go b/lib/websocket/client_test.go
index 5269facf..93abe735 100644
--- a/lib/websocket/client_test.go
+++ b/lib/websocket/client_test.go
@@ -92,7 +92,14 @@ func TestClientPing(t *testing.T) {
 
 	recvHandler := func(ctx context.Context, resp []byte) (err error) {
 		exp := ctx.Value(ctxKeyBytes).([]byte)
+
 		test.Assert(t, "resp", exp, resp, true)
+
+		frames := Unpack(resp)
+		if frames.IsClosed() {
+			testClient.SendClose(false)
+		}
+
 		return
 	}
 
diff --git a/lib/websocket/server.go b/lib/websocket/server.go
index fc2b1ed8..4bc6aae9 100644
--- a/lib/websocket/server.go
+++ b/lib/websocket/server.go
@@ -482,7 +482,8 @@ func (serv *Server) handlePing(conn int, req *Frame) {
 //
 func (serv *Server) reader() {
 	var (
-		events [128]unix.EpollEvent
+		events    [128]unix.EpollEvent
+		isClosing bool
 	)
 
 	for {
@@ -506,10 +507,13 @@ func (serv *Server) reader() {
 				serv.ClientRemove(conn)
 				continue
 			}
+
+			isClosing = false
 			for _, frame := range frames.v {
 				if frame.masked != frameIsMasked {
 					serv.handleBadRequest(conn)
-					continue
+					isClosing = true
+					break
 				}
 
 				switch frame.opcode {
@@ -521,20 +525,26 @@ func (serv *Server) reader() {
 					serv.handleFragment(conn, frame)
 				case opcodeClose:
 					serv.handleClose(conn, frame)
+					isClosing = true
 				case opcodePing:
 					serv.handlePing(conn, frame)
 				case opcodePong:
 					// Ignore pong from client.
 				}
+				if isClosing {
+					break
+				}
 			}
 
-			// See https://idea.popcount.org/2017-02-20-epoll-is-fundamentally-broken-12/
-			events[x].Events = unix.EPOLLIN | unix.EPOLLONESHOT
+			if !isClosing {
+				// See https://idea.popcount.org/2017-02-20-epoll-is-fundamentally-broken-12/
+				events[x].Events = unix.EPOLLIN | unix.EPOLLONESHOT
 
-			err = unix.EpollCtl(serv.epollRead, unix.EPOLL_CTL_MOD, conn, &events[x])
-			if err != nil {
-				log.Println("websocket: server.reader: unix.EpollCtl: " + err.Error())
-				continue
+				err = unix.EpollCtl(serv.epollRead, unix.EPOLL_CTL_MOD, conn, &events[x])
+				if err != nil {
+					log.Println("websocket: server.reader: unix.EpollCtl: " + err.Error())
+					continue
+				}
 			}
 		}
 	}