优化代码逻辑与判断

websocket/protocol.lua

@@ -10,6 +10,8 @@ local str_char = string.char
 local rand = math.random
 local tostring = tostring
 local type = type
+local error = error
+local assert = assert
 
 local new_tab = function () return {} end
 
@@ -101,9 +103,6 @@ function _M.recv_frame(sock, max_payload_len, force_masking)
         end
     end
 
-    -- print("payload len: ", payload_len, ", max payload len: ",
-          -- max_payload_len)
-
     if payload_len > max_payload_len then
         return nil, nil, "exceeding max payload len"
     end
@@ -115,7 +114,6 @@ function _M.recv_frame(sock, max_payload_len, force_masking)
     else
         rest = payload_len
     end
-    -- print("rest: ", rest)
 
     local data
     if rest > 0 then
@@ -189,7 +187,7 @@ end
 
 
 local function build_frame(fin, opcode, payload_len, payload, masking)
-    -- XXX optimize this when we have string.buffer in LuaJIT 2.1
+
     local fst
     if fin then
         fst = 0x80 | opcode
@@ -245,33 +243,25 @@ _M.build_frame = build_frame
 
 function _M.send_frame(sock, fin, opcode, payload, max_payload_len, masking)
 
-    if not payload then
-        payload = ""
-
-    elseif type(payload) ~= "string" then
-        payload = tostring(payload)
-    end
+  assert(type(payload) == 'string' and #payload <= max_payload_len, "无效的数据类型或长度超出预期")
 
-    local payload_len = #payload
+  local payload_len = #payload
 
-    if payload_len > max_payload_len then
-        return nil, "payload too big"
+  if opcode & 0x8 ~= 0 then
+    if payload_len > 125 then
+      return error("控制帧的有效载荷长度太多")
     end
-
-    if opcode & 0x8 ~= 0 then
-        if payload_len > 125 then
-            return nil, "too much payload for control frame"
-        end
-        if not fin then
-            return nil, "fragmented control frame"
-        end
+    if not fin then
+        return error("畸形的控制帧")
     end
+  end
 
-    local frame, err = build_frame(fin, opcode, payload_len, payload, masking)
-    if not frame then
-        return nil, "failed to build frame: " .. err
-    end
-    return sock_send(sock, frame)
+  local frame, err = build_frame(fin, opcode, payload_len, payload, masking)
+  if not frame then
+    return error("错误的数据帧:"..err)
+  end
+
+  return sock_send(sock, frame)
 end
 
 return _M

websocket/server.lua

@@ -15,14 +15,15 @@ local _recv_frame = wbproto.recv_frame
 local _send_frame = wbproto.send_frame
 
 local type = type
-local char = string.char
 local pcall = pcall
+local assert = assert
 local tostring = tostring
 local tonumber = tonumber
+
+local char = string.char
 local match = string.match
 local lower = string.lower
 local concat = table.concat
-local co_self = coroutine.running
 
 local HTTP_CODE = {
   [101] = "HTTP/1.1 101 Switching Protocol",
@@ -81,8 +82,7 @@ end
 -- 握手
 local function do_handshak (self)
   local fd, auth = self.fd, false
-  skynet.fork(function (...)
-    skynet.sleep(self.timeout or 3 * 100)
+  skynet.timeout(self.timeout or 3 * 100, function (...)
     if not auth then
       self.fd = nil
       return sock_close(fd)
@@ -97,7 +97,7 @@ local function do_handshak (self)
     sock_send(fd, ERROR_RESPONSE(400))
     return nil, "1. 错误的HTTP Protocol Path 或者 methond."
   end
-  if method ~= "GET" or (self.path and self.path ~= path) or tonumber(version) ~= 1.1 then
+  if method ~= "GET" or (self.path and self.path ~= path) or version ~= '1.1' then
     sock_send(fd, ERROR_RESPONSE(405))
     return nil, "2. 不支持的HTTP Protocol 版本."
   end
@@ -207,9 +207,7 @@ function websocket:start()
       skynet.fork(on_message, cls, data, typ)
     end
     if typ == "ping" then
-      skynet.fork(function (...)
-        return _send_frame(self.fd, true, 0xA, data or '', max_payload_len, send_masked)
-      end)
+      skynet.fork(_send_frame, self.fd, true, 0xA, data or '', max_payload_len, send_masked)
     end
   end
   return self:close()