OpenComputers WebSocket library

stub.lua

---@class BufferedFrame
---@field public fin boolean @true if frame is finalized, false if frame is still being building
---@field public type string @text if frame is text, binary if frame is binary
---@field public payload string @payload of (potentially fragmented) frame
local BufferedFrame = {}


---@class Options
---@field public headers table @name-to-value mapping of headers
---@field public protocols table @list of subprotocols (Sec-WebSocket-Protocol)
local Options = {}

websocket.lua

local ws = {}

---@class WebSocket
---@field private _closed boolean
---@field private _closing boolean
---@field private _frames BufferedFrame[]
---@field private _close_code number
---@field private _close_reason string|nil
---@field private _sending_fragments boolean
local prototype = {}

local stringChar = string.char

local function maskPayload(key, payload)
    local masked = ""
    for i = 1, #payload do
        masked = masked .. stringChar(payload:byte(i) ~ key:byte(((i - 1) % 4) + 1))
    end
    return masked
end

local function randomBytes(n)
    local bytes = "" -- n < 100 so performance impact is insignificant
    for _ = 1, n do
        bytes = bytes .. stringChar(math.random(0, 255))
    end
    return bytes
end

local function intToBytes(i, n)
    local bytes = ""
    while i > 0 do
        bytes = stringChar(i & 255) .. bytes
        i = i >> 8
    end
    if n ~= nil then
        while #bytes < n do
            bytes = stringChar(0) .. bytes
        end
    end
    return bytes
end

--- Connect to remote server using HTTP Upgrade
--- See https://www.rfc-editor.org/rfc/rfc6455#section-1.3
--- You should call WebSocket#update method either as fast as you can or in response to OpenComputers's internet_ready event
---@param connect_function function(host:string,port:number):userdata function to open <b>non-blocking raw</b> TCP socket (usually component.internet.connect)
---@param encode64_function function(in:string):string base64 encode function (usually component.data.encode64)
---@param host string remote host
---@param port number remote port
---@param path string remote path, should start with /
---@param options Options additional options
---@return WebSocket An initialized WebSocket instance
---@overload fun(connect_function:function(host: string, port: number), datacard_component:table, pull_function:function(timeout: number), host:string, port:number, path:string)
function ws.connect(connect_function, encode64_function,
                    host, port, path, options)
    local obj = setmetatable({}, { __index = prototype })
    obj:_constructor(connect_function, encode64_function, host, port, path, options.headers, options.protocols)
    return obj
end


-- https://www.rfc-editor.org/rfc/rfc6455#section-4
function prototype:_constructor(connect_function, encode64_function, host, port, path, headers, protocols)
    self._socket = connect_function(host, port)
    self._closed = false
    self._closing = false
    self._close_code = -1
    self._close_reason = nil
    self._frames = {}
    self._sending_fragments = false

    local s = self._socket

    while true do
        local success, reason = s.finishConnect()
        if reason then
            error(reason)
        end
        if success then
            break
        end
    end
    local headersString = ""
    if protocols then
        headers["Sec-WebSocket-Protocol"] = table.concat(protocols, ",")
    end
    if headers then
        for key, value in pairs(headers) do
            headersString = headersString .. ("\r\n%s:%s"):format(key, value)
        end
    end
    
    local key = encode64_function(randomBytes(16))
    s.write(
            ("GET %s HTTP/1.1\r\nHost:%s\r\nUpgrade:websocket\r\nConnection:upgrade\r\nSec-WebSocket-Key:%s\r\nSec-WebSocket-Version:13%s\r\n\r\n")
                    :format(path, host, key, headersString))
    local buffer = ""
    while true do
        buffer = buffer .. assert(s.read())
        if buffer:find("\r\n\r\n") ~= nil then
            break
        end
        if #buffer > 4 and buffer:sub(1, 4) ~= "HTTP" then
            s.close()
            error("Got invalid response")
        end
    end
    local httpStatusLineEnd = buffer:find("\r\n")
    local httpStatusLine = buffer:sub(1, httpStatusLineEnd - 1)
    buffer = buffer:sub(httpStatusLineEnd + 2)
    local statusCode, message = httpStatusLine:match(
            "HTTP/%d+%.%d+%s+(%d+)%s+(.+)")
    if statusCode ~= "101" then
        s.close()
        error(("ERROR_INVALID_RESPONSE_CODE_%s_%s"):format(statusCode, message:upper():gsub(" ", "_")))
    end
    local responseHeaders = {}
    for header, value in buffer:gmatch("%s*(%S+)%s*:%s*(%S+)%s-\r\n") do
        responseHeaders[header:lower()] = value
    end

    if (responseHeaders["connection"] or ""):lower() ~= "upgrade"
            or (responseHeaders["upgrade"] or ""):lower() ~= "websocket"
    -- TODO implement sha1
    -- or responseHeaders["sec-websocket-accept"] ~= datacard_component.encode64(datacard_component.sha1(key+"258EAFA5-E914-47DA-95CA-C5AB0DC85B11"))
    then
        s.close()
        error("ERROR_UPGRADE_FAILED")
    end
    local secWebSocketProtocolHeader = responseHeaders["sec-websocket-protocol"]
    if secWebSocketProtocolHeader ~= nil then
        if headers["Sec-WebSocket-Protocol"]:match("%f[^,\0]" + secWebSocketProtocolHeader + "%f[,\0]") ==
                nil then
            s.close()
            error("ERROR_INVALID_SUBPROTOCOL")
        end
    elseif protocols ~= nil then
        s.close()
        error("ERROR_SUBPROTOCOL_HEADER_ABSENT")
    end
end

--- _Start the WebSocket Closing Handshake_. You still should repeatedly call WebSocket#update method until it returns false
--- See https://www.rfc-editor.org/rfc/rfc6455#section-7.1.2 and https://www.rfc-editor.org/rfc/rfc6455#section-1.4
---@param closeCode number status code of close frame
---@param closeReason string data to include with close frame
---@overload fun(closeCode: number) Send frame without close reason
---@overload fun() Send frame without payload
---@see WebSocket#update
function prototype:startClosingHandshake(closeCode, closeReason)
    if self._closing or self._closed then
        return
    end
    local payload = ""
    if closeCode then
        payload = payload .. intToBytes(closeCode, 2)
    end
    if closeReason then
        payload = payload .. closeReason
    end
    self:_sendFrame(8, 1, payload)
    self._closing = true
    self._close_code = closeCode
    self._close_reason = closeReason
end

---@private
function prototype:_fail(closeReason)
    self._close_code = 1002
    self._close_reason = closeReason
    self:_close()
end

---@private
function prototype:_close()
    self._closed = true
    self._closing = false
    self._socket.close()
end

--- https://www.rfc-editor.org/rfc/rfc6455#section-5
---@returns boolean true if connection is still alive (false if, for example, server sent a close frame)
function prototype:update()
    local selfSocketRead = self._socket.read
    if self._closed then
        return false
    end
    local header, reason = selfSocketRead(2)
    if not header then
        if self._close_code == -1 then -- Abnormal closure
            self._close_code = 1006
            self._close_reason = reason
        end
        return false
    end
    if #header == 0 then
        return true
    end
    while #header < 2 do
        header = header .. assert(selfSocketRead(2 - #header))
    end
    local firstByte, secondByte = header:byte(1, 2)
    local fin, rsv1, rsv2, rsv3, opcode, mask, payloadLen = (firstByte & 128) >> 7, (firstByte & 64) >> 6, (firstByte & 32) >> 5, (firstByte & 16) >> 4, (firstByte & 15), (secondByte & 128) >> 7, secondByte & 127
    --[[
      MUST be 0 unless an extension is negotiated that defines meanings
      for non-zero values.  If a nonzero value is received and none of
      the negotiated extensions defines the meaning of such a nonzero
      value, the receiving endpoint MUST _Fail the WebSocket
      Connection_.
    ]]
    if (rsv1 + rsv2 + rsv3) > 0 then
        self:_fail("ERROR_INVALID_FRAME_RSV_USED")
        return false
    end

    if mask == 1 then
        -- A server MUST NOT mask any frames that it sends to the client
        -- A client MUST close a connection if it detects a masked frame.
        self:_close()
        return false
    end

    local headerLength = 2 + payloadLen > 125 and (payloadLen == 126 and 2 or 8) or 0
    while #header < headerLength do
        header = header .. assert(selfSocketRead(headerLength - #header))
    end
    -- The length of the "Payload data", in bytes: if 0-125, that is the payload length.
    local finalPayloadDataLength = payloadLen
    if payloadLen == 126 then
        -- If 126, the following 2 bytes interpreted as a 16-bit unsigned integer are the payload length.
        local a, b = header:byte(2, 3)
        finalPayloadDataLength = (a << 8) + b --  header:byte(2) << 8 + header:byte(3)
    elseif payloadLen == 127 then
        -- If 127, the following 8 bytes interpreted as a 64-bit unsigned integer (the most significant bit MUST be 0)
        local a, b, c, d, e, f, g, h = header:byte(2, 9)
        finalPayloadDataLength = (a << 56) + (b << 48) + (c << 40) + (d << 32) + (e << 24) + (f << 16) + (g << 8) + h
        -- What to do if the most significant bit is set ?
        -- Let's assume we're in ideal world
    end
    local payload = ""
    while #payload < finalPayloadDataLength do
        payload = payload .. assert(selfSocketRead(finalPayloadDataLength - #payload))
    end
    if opcode & 8 > 0 then
        -- Control frame
        if opcode == 8 then
            -- OPCODE 0b1000 CLOSE
            local closeCode
            local closeReason
            if #payload > 1 then
                closeCode = (payload:byte(1) << 8) + payload:byte(2)
            end
            if #payload > 2 then
                closeReason = payload:sub(3)
            end
            if not self._closing then
                self:_sendFrame(8, 1, payload) -- Echo
                self._close_code = closeCode
                self._close_reason = closeReason
            end
            -- https://www.rfc-editor.org/rfc/rfc6455#section-7.1.1
            -- but can't get TCP/FIN packet and OpenComputers don't send TCP/FIN automatically (thus leaving socket opened)
            -- so violating RFC and closing socket right here
            self:_close()
            return false
        elseif opcode == 9 then
            --OPCODE 0b1001 PING
            self:_sendFrame(10, 1, payload)
        end -- OPCODE 0b1010 PONG ignored
    else
        local selfFrames = self._frames
        local selfFramesLen = #selfFrames
        if selfFramesLen > 0 and not selfFrames[selfFramesLen].fin then
            if opcode ~= 0 then
                -- TODO find in RFC what do we need to do
                self:_fail("ERROR_FRAGMENTATION_ABORTED")
                return false
            end
            local frame = selfFrames[selfFramesLen]
            frame.payload = frame.payload + payload
            frame.fin = fin > 0
        else
            selfFrames[selfFramesLen + 1] = {
                fin = (fin > 0),
                type = (opcode == 1 and "text" or "binary"),
                payload = payload
            }
        end
    end
    return true
end

---@private
function prototype:_sendFrame(opcode, fin, payload)
    local payloadLen = #payload
    local extendedPayloadLen = ""
    if payloadLen > 65535 then
        payloadLen = 127
        extendedPayloadLen = intToBytes(#payload)
    elseif payloadLen > 125 then
        payloadLen = 126
        extendedPayloadLen = intToBytes(#payload)
    end
    local maskingKey = randomBytes(4)
    local frame = stringChar((fin << 7) + opcode)
            .. stringChar(128 + payloadLen) -- 1 << 7 + payloadLen
            .. extendedPayloadLen
            .. maskingKey
            .. maskPayload(maskingKey, payload)
    local written = 0
    while written < #frame do
        written = written + assert(self._socket.write(frame:sub(written + 1)))
    end
end

--- Send frame (potentially fragmented)
--- The fragments of one message MUST NOT be interleaved between the fragments of another message.
---@param initialOpcode number starting opcode of message
---@param payload string data to send
---@param fin boolean true if this fragment is final. If this parameter is false, you MUST send more fragments later. true by default
---@overload fun(initialOpcode:number, payload: string)
function prototype:sendFrame(initialOpcode, payload, fin)
    if self._closing or self._closed then
        return
    end
    if not fin then
        fin = true
    end
    local opcode = self._sending_fragments and 0 or (initialOpcode & 15)
    self:_sendFrame(opcode, fin and 1 or 0, payload)
    self._sending_fragments = not fin
end

---Send text frame (potentially fragmented) 
---The fragments of one message MUST NOT be interleaved between the fragments of another message.
---@param text string text to send
---@param fin boolean true if this fragment is final. If this parameter is false, you MUST send more fragments later. 1 by default
---@overload fun(text: string)
---@see WebSocket#sendFrame
function prototype:sendText(text, fin)
    return self:sendFrame(1, text, fin)
end

--- Send binary frame (potentially fragmented)
--- The fragments of one message MUST NOT be interleaved between the fragments of another message.
---@param data string data to send
---@param fin boolean true if this fragment is final. If this parameter is false, you MUST send more fragments later. 1 by default
---@overload fun(data: string)
---@see WebSocket#sendFrame
function prototype:sendBinary(data, fin)
    return self:sendFrame(2, data, fin)
end

--- Returns frames read previously in update() method (in receiving order)
---@see WebSocket#update
---@returns BufferedFrame|nil
function prototype:getMessage(canBeFragmented)
    if #self._frames == 0 then
        return nil
    end
    local first = self._frames[1]
    if first.fin then
        return table.remove(self._frames, 1)
    elseif canBeFragmented then
        -- Copy table as it will be mutated
        return {
            fin = first.fin,
            type = first.type,
            payload = first.payload
        }
    end
    return nil
end

--- See https://www.rfc-editor.org/rfc/rfc6455#section-7.1.4
---@return boolean true if _The WebSocket Connection is Closed_
function prototype:isClosed()
    return self._closed
end

--- See https://www.rfc-editor.org/rfc/rfc6455#section-7.1.3
---@return boolean true if _The WebSocket Closing Handshake is Started_
function prototype:isClosing()
    return self._closing
end

--- See https://www.rfc-editor.org/rfc/rfc6455#section-7.4
---@return number close code
function prototype:getCloseCode()
    return self._close_code
end

--- See https://www.rfc-editor.org/rfc/rfc6455#section-7.1.6
---@return string Close reason
function prototype:getCloseReason()
    return self._close_reason
end

---@return string socket id, to use with OpenComputers's internet_ready event
function prototype:getSocketId()
    return self._socket.id()
end

return ws