lua-debug/extension/script/common/socket.lua at a38713466c8d9c6b9733517bbd7e20c826dbeb63 · actboy168/lua-debug · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
local net = require 'common.net'
local proto = require 'common.protocol'

local function parseAddress(param)
    local mode, address = param:match "^([a-z]+):(.*)"
    if address:sub(1,1) == '@' then
        local fs = require "bee.filesystem"
        address = address:gsub("%$tmp", (fs.temp_directory_path():string():gsub("([/\\])$", "")))
        return {
            protocol = 'unix',
            address = address:sub(2),
            mode = mode,
        }
    end
    local ipv4, port = address:match("(%d+%.%d+%.%d+%.%d+):(%d+)")
    if ipv4 then
        return {
            protocol = 'tcp',
            address = ipv4,
            port = tonumber(port),
            mode = mode,
        }
    end
    local ipv6, port = address:match("%[([%d:a-fA-F]+)%]:(%d+)")
    if ipv6 then
        return {
            protocol = 'tcp6',
            address = ipv6,
            port = tonumber(port),
            mode = mode,
        }
    end
    error "Invalid address."
end

return function (param)
    local t = parseAddress(param)
    local stat = {}
    local readbuf = ''
    local writebuf = ''
    local server = nil
    local session = nil
    local closefunc = function () end
    local m = {}
    local function cleanup()
        session = nil
        readbuf = ''
        writebuf = ''
        stat = { debug = stat.debug }
        closefunc()
    end
    local function init_session(s)
        if session then
            return
        end
        session = s
        function session:on_data(data)
            readbuf = readbuf .. data
        end
        function session:on_close()
            cleanup()
        end
        return true
    end
    if t.mode == "connect" then
        local function try_connect()
            local s = net.connect(t.protocol, t.address, t.port)
            if s then
                function s:on_connected()
                    local ok = init_session(s)
                    assert(ok)
                    if writebuf ~= '' then
                        s:write(writebuf)
                        writebuf = ''
                    end
                end
                function s:on_error()
                    try_connect()
                end
            else
                net.async(try_connect)
            end
        end
        try_connect()
    elseif t.mode == "listen" then
        server = assert(net.listen(t.protocol, t.address, t.port))
        function server:on_accepted(new_s)
            local ok = init_session(new_s)
            assert(ok)
            if writebuf ~= '' then
                new_s:write(writebuf)
                writebuf = ''
            end
            return ok
        end
    else
        return
    end
    function m.event_close(f)
        closefunc = f
    end
    function m.debug(v)
        stat.debug = v
    end
    function m.sendmsg(pkg)
        local data = proto.send(pkg, stat)
        if session == nil then
            writebuf = writebuf .. data
            return
        end
        session:write(data)
    end
    function m.recvmsg()
        local data = readbuf
        readbuf = ''
        return proto.recv(data, stat)
    end
    function m.update(timeout)
        net.update(timeout)
    end
    function m.closeall()
        local fds = {}
        if t.mode == "connect" then
            if session == nil then
                return
            end
            session:close()
            fds[1] = session
        elseif t.mode == "listen" then
			fds[1] = server
			if session ~= nil then
				session:close()
				fds[2] = session
			end
			server:close()
        end
        local function is_finish()
            for _, fd in ipairs(fds) do
                if not fd:is_closed() then
                    return false
                end
            end
            return true
        end
        while not is_finish() do
            net.update(0)
        end
    end
    return m
end