support connecting by socket with `--socket=PORT`

8 files changed, 175 insertions, 11 deletions

diff --git a/changelog.md b/changelog.md
index 0a01507b..50328c9c 100644
--- a/changelog.md
+++ b/changelog.md
@@ -1,5 +1,8 @@
 # changelog
 
+## 3.6.20
+* `NEW` support connecting by socket with `--socket=PORT`
+
 ## 3.6.19
 `2023-4-26`
 * `FIX` commandline parameter `checklevel` may not work
diff --git a/script/brave/work.lua b/script/brave/work.lua
index 3a565e08..3b82b452 100644
--- a/script/brave/work.lua
+++ b/script/brave/work.lua
@@ -1,6 +1,6 @@
 local brave   = require 'brave.brave'
 
-brave.on('loadProto', function ()
+brave.on('loadProtoByStdio', function ()
     local jsonrpc = require 'jsonrpc'
     while true do
         local proto, err = jsonrpc.decode(io.read)
@@ -13,6 +13,37 @@ brave.on('loadProto', function ()
     end
 end)
 
+brave.on('loadProtoBySocket', function (fdHandle)
+    local jsonrpc = require 'jsonrpc'
+    local socket  = require 'bee.socket'
+    local thread  = require 'bee.thread'
+    local fd = socket.fd(fdHandle)
+    local buf = ''
+    while true do
+        local proto, err = jsonrpc.decode(function (len)
+            while true do
+                if #buf >= len then
+                    local res = buf:sub(1, len)
+                    buf = buf:sub(len + 1)
+                    return res
+                end
+                local data = fd:recv()
+                if data then
+                    buf = buf .. data
+                else
+                    thread.sleep(0.01)
+                end
+            end
+        end)
+        --log.debug('loaded proto', proto.method)
+        if not proto then
+            brave.push('protoerror', err)
+            return
+        end
+        brave.push('proto', proto)
+    end
+end)
+
 brave.on('timer', function (time)
     local thread = require 'bee.thread'
     while true do
diff --git a/script/global.d.lua b/script/global.d.lua
index c120c153..f84ff0e4 100644
--- a/script/global.d.lua
+++ b/script/global.d.lua
@@ -73,3 +73,7 @@ COMPILECORES = 0
 -- TODO: delete this after new config
 ---@diagnostic disable-next-line: lowercase-global
 jit = false
+
+-- connect to client by socket
+---@type integer
+SOCKET = 0
diff --git a/script/jsonrpc.lua b/script/jsonrpc.lua
index 7411fee8..6de6da5b 100644
--- a/script/jsonrpc.lua
+++ b/script/jsonrpc.lua
@@ -14,17 +14,23 @@ function m.encode(pack)
     return buf
 end
 
+---@param reader fun(arg: integer):string
 local function readProtoHead(reader)
     local head = {}
+    local line = ''
     while true do
-        local line = reader 'L'
-        if line == nil then
+        local char = reader(1)
+        if char == nil then
             -- 说明管道已经关闭了
             return nil, 'Disconnected!'
         end
+        line = line .. char
         if line == '\r\n' then
             break
         end
+        if line:sub(-2) ~= '\r\n' then
+            goto continue
+        end
         local k, v = line:match '^([^:]+)%s*%:%s*(.+)\r\n$'
         if not k then
             return nil, 'Proto header error: ' .. line
@@ -33,10 +39,13 @@ local function readProtoHead(reader)
             v = tonumber(v)
         end
         head[k] = v
+        line = ''
+        ::continue::
     end
     return head
 end
 
+---@param reader fun(arg: integer):string
 function m.decode(reader)
     local head, err = readProtoHead(reader)
     if not head then
diff --git a/script/meta/bee/socket.lua b/script/meta/bee/socket.lua
new file mode 100644
index 00000000..b77c498b
--- /dev/null
+++ b/script/meta/bee/socket.lua
@@ -0,0 +1,62 @@
+---@meta
+
+---@alias bee.socket.protocol
+---| 'tcp'
+---| 'udp'
+---| 'unix'
+---| 'tcp6'
+---| 'udp6'
+
+---@class bee.socket
+---@overload fun(protocol: bee.socket.protocol): bee.socket.fd?, string?
+local socket = {}
+
+---@param readfds? bee.socket.fd[]
+---@param writefds? bee.socket.fd[]
+---@param timeout number
+---@return bee.socket.fd[] # readfds
+---@return bee.socket.fd[] # writefds
+function socket.select(readfds, writefds, timeout) end
+
+---@param handle lightuserdata
+---@return bee.socket.fd
+function socket.fd(handle) end
+
+---@class bee.socket.fd
+local fd = {}
+
+---@param addr string
+---@param port? integer
+---@return boolean
+---@return string?
+function fd:bind(addr, port) end
+
+function fd:close() end
+
+---@return boolean
+---@return string?
+function fd:listen() end
+
+---@param addr string
+---@param port integer
+---@return boolean
+---@return string?
+function fd:connect(addr, port) end
+
+---@param len? integer
+---@return string | false
+function fd:recv(len) end
+
+---@param content string
+function fd:send(content) end
+
+---@return lightuserdata
+function fd:handle() end
+
+---@return boolean
+function fd:status() end
+
+---@return bee.socket.fd
+function fd:accept() end
+
+return socket
diff --git a/script/meta/bee/thread.lua b/script/meta/bee/thread.lua
new file mode 100644
index 00000000..6b4323a4
--- /dev/null
+++ b/script/meta/bee/thread.lua
@@ -0,0 +1,33 @@
+---@meta
+
+---@class bee.thread
+local thread = {}
+
+---@param time number
+function thread.sleep(time) end
+
+---@param name string
+function thread.newchannel(name) end
+
+---@param name string
+---@return bee.thread.channel
+function thread.channel(name) end
+
+---@param script string
+---@return bee.thread.thread
+function thread.thread(script) end
+
+---@class bee.thread.channel
+local channel = {}
+
+function channel:push(...) end
+
+---@return ...
+function channel:pop() end
+
+---@return ...
+function channel:bpop() end
+
+---@class bee.thread.thread
+
+return thread
diff --git a/script/proto/proto.lua b/script/proto/proto.lua
index 7875b919..fb623106 100644
--- a/script/proto/proto.lua
+++ b/script/proto/proto.lua
@@ -1,4 +1,5 @@
 local subprocess = require 'bee.subprocess'
+local socket     = require 'bee.socket'
 local util       = require 'utility'
 local await      = require 'await'
 local pub        = require 'pub'
@@ -6,6 +7,7 @@ local jsonrpc    = require 'jsonrpc'
 local define     = require 'proto.define'
 local json       = require 'json'
 local inspect    = require 'inspect'
+local thread     = require 'bee.thread'
 
 local reqCounter = util.counter()
 
@@ -29,6 +31,9 @@ local m = {}
 m.ability = {}
 m.waiting = {}
 m.holdon  = {}
+m.mode    = 'stdio'
+---@type bee.socket.fd
+m.fd      = nil
 
 function m.getMethodName(proto)
     if proto.method:sub(1, 2) == '$/' then
@@ -46,7 +51,11 @@ end
 function m.send(data)
     local buf = jsonrpc.encode(data)
     logSend(buf)
-    io.write(buf)
+    if m.mode == 'stdio' then
+        io.write(buf)
+    elseif m.mode == 'socket' then
+        m.fd:send(buf)
+    end
 end
 
 function m.response(id, res)
@@ -219,12 +228,20 @@ function m.doResponse(proto)
     waiting.resume(proto.result)
 end
 
-function m.listen()
-    subprocess.filemode(io.stdin,  'b')
-    subprocess.filemode(io.stdout, 'b')
-    io.stdin:setvbuf  'no'
-    io.stdout:setvbuf 'no'
-    pub.task('loadProto')
+function m.listen(mode, socketPort)
+    m.mode = mode
+    if mode == 'stdio' then
+        subprocess.filemode(io.stdin,  'b')
+        subprocess.filemode(io.stdout, 'b')
+        io.stdin:setvbuf  'no'
+        io.stdout:setvbuf 'no'
+        pub.task('loadProtoByStdio')
+    elseif mode == 'socket' then
+        local fd = assert(socket('tcp'))
+        fd:connect('127.0.0.1', socketPort)
+        m.fd = fd
+        pub.task('loadProtoBySocket', fd:handle())
+    end
 end
 
 return m
diff --git a/script/service/service.lua b/script/service/service.lua
index cd83dd54..7011ec4f 100644
--- a/script/service/service.lua
+++ b/script/service/service.lua
@@ -267,7 +267,12 @@ function m.start()
     if COMPILECORES and COMPILECORES > 0 then
         pub.recruitBraves(COMPILECORES, 'compile')
     end
-    proto.listen()
+    if SOCKET then
+        assert(math.tointeger(SOCKET), '`socket` must be integer')
+        proto.listen('socket', SOCKET)
+    else
+        proto.listen('stdio')
+    end
     m.report()
     m.testVersion()
     m.lockCache()