整理代码(by Particle_G)，ws关闭前会发关闭消息

Zframework/websocket.lua

@@ -11,29 +11,30 @@
 		res, opcode = client:read()
 		print(res)
 		client:close()
-]]
+]] -- local debug_print=print
 
--- local debug_print=print
 local socket = require "socket"
 local band, bor, bxor = bit.band, bit.bor, bit.bxor
 local shl, shr = bit.lshift, bit.rshift
 
-local OPCODES = {
+OPCODES = {
 	CONTINUE = 0,
 	TEXT = 1,
 	BINARY = 2,
 	CLOSE = 8,
 	PING = 9,
-	PONG    =10,
+	PONG = 10
 }
 
 local _M = {
-	OPCODES = OPCODES,
+	OPCODES = OPCODES
 }
 _M.__index = _M
 
 function _M.new()
-	local m = {socket = socket.tcp()}
+	local m = {
+		socket = socket.tcp()
+	}
 	setmetatable(m, _M)
 	return m
 end
@@ -42,17 +43,21 @@ local seckey = "osT3F7mvlojIvf3/8uIsJQ=="
 function _M:connect(host, port, path)
 	local SOCK = self.socket
 	local res, err = SOCK:connect(host, port)
-	if res~=1 then return res, err end
+	if res ~= 1 then
+		return res, err
+	end
 	-- debug_print("[handshake] connected")
 
 	-- WebSocket handshake
 	res, err = SOCK:send("GET " .. (path or "/") .. " HTTP/1.1\r\nHost: " .. host .. ":" .. port .."\r\nConnection: Upgrade\r\nUpgrade: websocket\r\nSec-WebSocket-Version: 13\r\nSec-WebSocket-Key: " ..seckey .. "\r\n\r\n")
-	repeat res = SOCK:receive("*l") until res==""
+	repeat
+		res = SOCK:receive("*l")
+	until res == ""
 	-- debug_print("[handshake] succeed")
 end
 
 local mask_key = {1, 14, 5, 14}
-local function send(SOCK, opcode, message)
+local function _send(SOCK, opcode, message)
 	-- message type
 	SOCK:send(string.char(bor(0x80, opcode)))
 
@@ -65,15 +70,9 @@ local function send(SOCK, opcode, message)
 	local length = #message
 	-- debug_print("[encode] message length: "..length)
 	if length > 65535 then
-		SOCK:send(string.char(bor(127, 0x80), 0, 0, 0, 0,
-			band(shr(length, 24), 0xff),
-			band(shr(length, 16), 0xff),
-			band(shr(length, 8), 0xff),
-			band(length, 0xff)))
+		SOCK:send(string.char(bor(127, 0x80), 0, 0, 0, 0, band(shr(length, 24), 0xff), band(shr(length, 16), 0xff),band(shr(length, 8), 0xff), band(length, 0xff)))
 	elseif length > 125 then
-		SOCK:send(string.char(bor(126, 0x80),
-			band(shr(length, 8), 0xff),
-			band(length, 0xff)))
+		SOCK:send(string.char(bor(126, 0x80), band(shr(length, 8), 0xff), band(length, 0xff)))
 	else
 		SOCK:send(string.char(bor(length, 0x80)))
 	end
@@ -87,23 +86,23 @@ local function send(SOCK, opcode, message)
 	-- debug_print("[encode] end")
 end
 
-function _M:send(message)
-	send(self.socket, OPCODES.BINARY, message)
+function _M:send(message, type)
+	local tempType = OPCODES.BINARY
+	for _, opcode in pairs(_M.OPCODES) do
+		if type == opcode then
+			tempType = type
 		end
-
-function _M:ping(message)
-	send(self.socket, OPCODES.PING, message)
 	end
-
-function _M:pong(message)
-	send(self.socket, OPCODES.PONG, message)
+	_send(self.socket, tempType, message)
 end
 
 function _M:read()
 	-- byte 0-1
 	local SOCK = self.socket
 	local res, err = SOCK:receive(2)
-	if not res then return res, err end
+	if not res then
+		return res, err
+	end
 
 	local OPCODE = band(res:byte(), 0x0f)
 	-- local flag_FIN = res:byte()>=0x80
@@ -126,18 +125,21 @@ function _M:read()
 
 	-- data
 	res = SOCK:receive(length)
+	local closeCode = nil
 	if OPCODE == OPCODES.PING then
 		self:pong(res)
 	elseif OPCODE == OPCODES.CLOSE then
+		closeCode = shl(res:byte(1), 8) + res:byte(2)
+		res = string.sub(res, 3, string.len(res) - 2)
 		self:close()
 	end
 	-- debug_print("[decode] string length: "..#res)
 	-- debug_print("[decode] end")
-	return res, OPCODE
+	return OPCODE, res, closeCode
 end
 
-function _M:close() self.socket:close() end
-
-function _M:settimeout(t) self.socket:settimeout(t) end
+function _M:settimeout(t)
+	self.socket:settimeout(t)
+end
 
 return _M