|
| 1 | +--- VM Service WebSocket client for Flutter/Dart debugging |
| 2 | +--- WebSocket framing per RFC 6455: |
| 3 | +--- Client frames must be masked with a 4-byte key |
| 4 | +--- Frame format: [FIN/opcode][mask/length][mask-key][payload] |
| 5 | +local uv = vim.uv or vim.loop |
| 6 | + |
| 7 | +local M = {} |
| 8 | + |
| 9 | +local OPCODE_TEXT = 1 |
| 10 | +local OPCODE_CLOSE = 8 |
| 11 | +local OPCODE_PING = 9 |
| 12 | + |
| 13 | +local tcp = nil |
| 14 | +local connected = false |
| 15 | +local handshake_complete = false |
| 16 | +local request_id = 0 |
| 17 | +local pending_requests = {} |
| 18 | +local event_handlers = {} |
| 19 | +local read_buffer = "" |
| 20 | + |
| 21 | +local function generate_handshake(host, port, path) |
| 22 | + local lines = { |
| 23 | + "GET " .. path .. " HTTP/1.1", |
| 24 | + "Host: " .. host .. ":" .. port, |
| 25 | + "Upgrade: websocket", |
| 26 | + "Connection: Upgrade", |
| 27 | + "Sec-WebSocket-Key: dGhlIHNhbXBsZSBub25jZQ==", |
| 28 | + "Sec-WebSocket-Version: 13", |
| 29 | + "", |
| 30 | + "", |
| 31 | + } |
| 32 | + return table.concat(lines, "\r\n") |
| 33 | +end |
| 34 | + |
| 35 | +local function generate_mask_key() |
| 36 | + return { |
| 37 | + math.random(0, 255), |
| 38 | + math.random(0, 255), |
| 39 | + math.random(0, 255), |
| 40 | + math.random(0, 255), |
| 41 | + } |
| 42 | +end |
| 43 | + |
| 44 | +local function mask_payload(payload, key) |
| 45 | + local masked = {} |
| 46 | + for i = 1, #payload do |
| 47 | + local byte = string.byte(payload, i, i) |
| 48 | + local mask_byte = key[((i - 1) % 4) + 1] |
| 49 | + table.insert(masked, string.char(bit.bxor(byte, mask_byte))) |
| 50 | + end |
| 51 | + return table.concat(masked, "") |
| 52 | +end |
| 53 | + |
| 54 | +local function create_frame(payload) |
| 55 | + local key = generate_mask_key() |
| 56 | + local len = #payload |
| 57 | + local frame = {} |
| 58 | + |
| 59 | + table.insert(frame, string.char(0x81)) |
| 60 | + |
| 61 | + if len < 126 then |
| 62 | + table.insert(frame, string.char(0x80 + len)) |
| 63 | + elseif len < 65536 then |
| 64 | + table.insert(frame, string.char(0x80 + 126)) |
| 65 | + table.insert(frame, string.char(bit.rshift(len, 8))) |
| 66 | + table.insert(frame, string.char(bit.band(len, 0xFF))) |
| 67 | + else |
| 68 | + table.insert(frame, string.char(0x80 + 127)) |
| 69 | + for i = 7, 0, -1 do |
| 70 | + table.insert(frame, string.char(bit.band(bit.rshift(len, i * 8), 0xFF))) |
| 71 | + end |
| 72 | + end |
| 73 | + |
| 74 | + for _, k in ipairs(key) do |
| 75 | + table.insert(frame, string.char(k)) |
| 76 | + end |
| 77 | + |
| 78 | + table.insert(frame, mask_payload(payload, key)) |
| 79 | + |
| 80 | + return table.concat(frame, "") |
| 81 | +end |
| 82 | + |
| 83 | +local function parse_frame(data) |
| 84 | + if #data < 2 then return nil, data end |
| 85 | + |
| 86 | + local b1 = string.byte(data, 1) |
| 87 | + local b2 = string.byte(data, 2) |
| 88 | + |
| 89 | + local opcode = bit.band(b1, 0x0F) |
| 90 | + local payload_len = bit.band(b2, 0x7F) |
| 91 | + |
| 92 | + local header_len = 2 |
| 93 | + if payload_len == 126 then |
| 94 | + if #data < 4 then return nil, data end |
| 95 | + payload_len = bit.lshift(string.byte(data, 3), 8) + string.byte(data, 4) |
| 96 | + header_len = 4 |
| 97 | + elseif payload_len == 127 then |
| 98 | + if #data < 10 then return nil, data end |
| 99 | + payload_len = 0 |
| 100 | + for i = 3, 10 do |
| 101 | + payload_len = bit.lshift(payload_len, 8) + string.byte(data, i) |
| 102 | + end |
| 103 | + header_len = 10 |
| 104 | + end |
| 105 | + |
| 106 | + local has_mask = bit.band(b2, 0x80) > 0 |
| 107 | + if has_mask then header_len = header_len + 4 end |
| 108 | + |
| 109 | + local total_len = header_len + payload_len |
| 110 | + if #data < total_len then return nil, data end |
| 111 | + |
| 112 | + local payload = data:sub(header_len + 1, total_len) |
| 113 | + local remaining = data:sub(total_len + 1) |
| 114 | + |
| 115 | + return { opcode = opcode, payload = payload }, remaining |
| 116 | +end |
| 117 | + |
| 118 | +local function handle_message(message) |
| 119 | + local ok, data = pcall(vim.json.decode, message) |
| 120 | + if not ok then return end |
| 121 | + |
| 122 | + if data.id and pending_requests[data.id] then |
| 123 | + local callback = pending_requests[data.id] |
| 124 | + pending_requests[data.id] = nil |
| 125 | + vim.schedule(function() callback(data.error, data.result) end) |
| 126 | + return |
| 127 | + end |
| 128 | + |
| 129 | + if data.method == "streamNotify" and data.params then |
| 130 | + local stream_id = data.params.streamId |
| 131 | + local event = data.params.event |
| 132 | + if event_handlers[stream_id] then |
| 133 | + vim.schedule(function() event_handlers[stream_id](event) end) |
| 134 | + end |
| 135 | + end |
| 136 | +end |
| 137 | + |
| 138 | +local function parse_uri(uri) |
| 139 | + local protocol, rest = uri:match("^(wss?)://(.+)$") |
| 140 | + if not protocol then |
| 141 | + protocol, rest = uri:match("^(https?)://(.+)$") |
| 142 | + end |
| 143 | + if not rest then return nil, nil, nil end |
| 144 | + |
| 145 | + local host_port, path = rest:match("^([^/]+)(/.*)$") |
| 146 | + if not host_port then |
| 147 | + host_port = rest |
| 148 | + path = "/" |
| 149 | + end |
| 150 | + |
| 151 | + local host, port = host_port:match("^([^:]+):(%d+)$") |
| 152 | + if not host then |
| 153 | + host = host_port |
| 154 | + port = (protocol == "wss" or protocol == "https") and 443 or 80 |
| 155 | + end |
| 156 | + |
| 157 | + if not path:match("/ws$") then path = path:gsub("/$", "") .. "/ws" end |
| 158 | + |
| 159 | + return host, tonumber(port), path |
| 160 | +end |
| 161 | + |
| 162 | +function M.connect(uri, on_connect, on_error) |
| 163 | + if connected then M.disconnect() end |
| 164 | + |
| 165 | + local host, port, path = parse_uri(uri) |
| 166 | + if not host or not port then |
| 167 | + if on_error then on_error("Invalid URI: " .. uri) end |
| 168 | + return |
| 169 | + end |
| 170 | + |
| 171 | + tcp = uv.new_tcp() |
| 172 | + handshake_complete = false |
| 173 | + read_buffer = "" |
| 174 | + |
| 175 | + tcp:connect(host, port, function(err) |
| 176 | + if err then |
| 177 | + vim.schedule(function() |
| 178 | + if on_error then on_error("Connection failed: " .. tostring(err)) end |
| 179 | + end) |
| 180 | + return |
| 181 | + end |
| 182 | + |
| 183 | + tcp:write(generate_handshake(host, port, path)) |
| 184 | + |
| 185 | + tcp:read_start(function(read_err, chunk) |
| 186 | + if read_err then |
| 187 | + vim.schedule(function() |
| 188 | + if on_error then on_error("Read error: " .. tostring(read_err)) end |
| 189 | + end) |
| 190 | + M.disconnect() |
| 191 | + return |
| 192 | + end |
| 193 | + |
| 194 | + if not chunk then |
| 195 | + M.disconnect() |
| 196 | + return |
| 197 | + end |
| 198 | + |
| 199 | + if not handshake_complete then |
| 200 | + if chunk:match("HTTP/1.1 101") then |
| 201 | + handshake_complete = true |
| 202 | + connected = true |
| 203 | + vim.schedule(function() |
| 204 | + if on_connect then on_connect() end |
| 205 | + end) |
| 206 | + end |
| 207 | + return |
| 208 | + end |
| 209 | + |
| 210 | + read_buffer = read_buffer .. chunk |
| 211 | + while true do |
| 212 | + local frame, remaining = parse_frame(read_buffer) |
| 213 | + if not frame then break end |
| 214 | + read_buffer = remaining |
| 215 | + |
| 216 | + if frame.opcode == OPCODE_TEXT then |
| 217 | + handle_message(frame.payload) |
| 218 | + elseif frame.opcode == OPCODE_PING then |
| 219 | + local pong = string.char(0x8A, 0x80 + #frame.payload) |
| 220 | + .. mask_payload(frame.payload, generate_mask_key()) |
| 221 | + tcp:write(pong) |
| 222 | + elseif frame.opcode == OPCODE_CLOSE then |
| 223 | + M.disconnect() |
| 224 | + return |
| 225 | + end |
| 226 | + end |
| 227 | + end) |
| 228 | + end) |
| 229 | +end |
| 230 | + |
| 231 | +function M.request(method, params, callback) |
| 232 | + if not connected or not tcp then |
| 233 | + if callback then callback("Not connected", nil) end |
| 234 | + return |
| 235 | + end |
| 236 | + |
| 237 | + request_id = request_id + 1 |
| 238 | + local id = tostring(request_id) |
| 239 | + |
| 240 | + local message = vim.json.encode({ |
| 241 | + jsonrpc = "2.0", |
| 242 | + id = id, |
| 243 | + method = method, |
| 244 | + params = params or {}, |
| 245 | + }) |
| 246 | + |
| 247 | + if callback then pending_requests[id] = callback end |
| 248 | + |
| 249 | + tcp:write(create_frame(message)) |
| 250 | +end |
| 251 | + |
| 252 | +function M.stream_listen(stream_id, handler, callback) |
| 253 | + event_handlers[stream_id] = handler |
| 254 | + M.request("streamListen", { streamId = stream_id }, callback) |
| 255 | +end |
| 256 | + |
| 257 | +function M.is_connected() return connected end |
| 258 | + |
| 259 | +function M.disconnect() |
| 260 | + connected = false |
| 261 | + handshake_complete = false |
| 262 | + |
| 263 | + for _, callback in pairs(pending_requests) do |
| 264 | + vim.schedule(function() callback("Service connection closed", nil) end) |
| 265 | + end |
| 266 | + pending_requests = {} |
| 267 | + event_handlers = {} |
| 268 | + read_buffer = "" |
| 269 | + |
| 270 | + if tcp then |
| 271 | + if not tcp:is_closing() then |
| 272 | + tcp:read_stop() |
| 273 | + tcp:close() |
| 274 | + end |
| 275 | + tcp = nil |
| 276 | + end |
| 277 | +end |
| 278 | + |
| 279 | +return M |
0 commit comments