path: root/xmake/modules/private/service/stream.lua

--!A cross-platform build utility based on Lua
--
-- Licensed under the Apache License, Version 2.0 (the "License");
-- you may not use this file except in compliance with the License.
-- You may obtain a copy of the License at
--
--     http://www.apache.org/licenses/LICENSE-2.0
--
-- Unless required by applicable law or agreed to in writing, software
-- distributed under the License is distributed on an "AS IS" BASIS,
-- WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-- See the License for the specific language governing permissions and
-- limitations under the License.
--
-- Copyright (C) 2015-present, TBOOX Open Source Group.
--
-- @author      ruki
-- @file        stream.lua
--

-- imports
import("core.base.object")
import("core.base.socket")
import("core.base.bytes")
import("private.service.message")

-- define module
local stream = stream or object()

-- max data buffer size
local STREAM_STRING_MAXN = 1024 * 1024

-- init stream
function stream:init(sock)
    self._SOCK = sock
    self._BUFF = bytes(65536)
    self._RCACHE = bytes(8192)
    self._RCACHE_SIZE = 0
    self._WCACHE = bytes(8192)
    self._WCACHE_SIZE = 0
end

-- get socket
function stream:sock()
    return self._SOCK
end

-- flush data
function stream:flush()
    local cache = self._WCACHE
    local cache_size = self._WCACHE_SIZE
    if cache_size > 0 then
        local sock = self._SOCK
        local real = sock:send(cache, {block = true, last = cache_size})
        if real > 0 then
            self._WCACHE_SIZE = 0
            return true
        end
    else
        return true
    end
end

-- send the given bytes
function stream:send(data, start, last)
    start = start or 1
    last = last or data:size()
    local size = last + 1 - start
    assert(size <= data:size())

    -- write data to cache first
    local cache = self._WCACHE
    local cache_size = self._WCACHE_SIZE
    local cache_maxn = cache:size()
    local cache_left = cache_maxn - cache_size
    if size <= cache_left then
        cache:copy2(cache_size + 1, data, start, last)
        cache_size = cache_size + size
        self._WCACHE_SIZE = cache_size
        return true
    elseif cache_left > 0 then
        cache:copy2(cache_size + 1, data, start, start + cache_left - 1)
        cache_size = cache_size + cache_left
        start = start + cache_left
        size = last + 1 - start
    end
    assert(cache_size == cache_maxn)

    -- send data to socket
    local sock = self._SOCK
    local real = sock:send(cache, {block = true})
    if real > 0 then
        -- copy left data to cache
        assert(size <= cache_maxn)
        cache:copy2(1, data, start, last)
        self._WCACHE_SIZE = size
        return true
    end
end

-- send message
function stream:send_msg(msg)
    return self:send_object(msg:body())
end

-- send object
function stream:send_object(obj)
    local str, errors = string.serialize(obj, {strip = true, indent = false})
    if errors then
        raise(errors)
    end
    if str then
        return self:send_string(str)
    end
end

-- send string
function stream:send_string(str)
    local buff = self._BUFF
    local size = #str
    assert(size < STREAM_STRING_MAXN, "too large string size(%d)", size)
    buff:u32be_set(1, size)
    if self:send(buff, 1, 4) then
        local data = bytes(str)
        local send = 0
        local cache = self._WCACHE
        local cache_maxn = cache:size()
        while send < size do
            local left = math.min(cache_maxn, size - send)
            if self:send(data, send + 1, send + left) then
                send = send + left
            else
                break
            end
        end
        if send == size then
            return true
        end
    end
end

-- send file
function stream:send_file(filepath)

    -- send size
    local buff = self._BUFF
    local size = os.filesize(filepath)
    buff:u32be_set(1, size)
    if not self:send(buff, 1, 4) then
        return
    end

    -- flush cache data first
    if not self:flush() then
        return
    end

    -- send file
    local ok = false
    local sock = self._SOCK
    local file = io.open(filepath, 'rb')
    if file then
        local send = sock:sendfile(file, {block = true})
        if send > 0 then
            ok = true
        end
        file:close()
    end
    return ok
end

-- recv the given bytes
function stream:recv(buff, size)
    assert(size <= buff:size(), "too large size(%d)", size)

    -- read data from cache first
    local buffsize = 0
    local cache = self._RCACHE
    local cache_size = self._RCACHE_SIZE
    local cache_maxn = cache:size()
    if size <= cache_size then
        buff:copy(cache, 1, size)
        cache:move(size + 1, cache_size)
        cache_size = cache_size - size
        self._RCACHE_SIZE = cache_size
        return buff:slice(1, size)
    elseif cache_size > 0 then
        buff:copy(cache, 1, cache_size)
        buffsize = cache_size
        cache_size = 0
    end
    assert(cache_size == 0)

    -- recv data from socket
    local real = 0
    local data = nil
    local wait = false
    local sock = self._SOCK
    while buffsize < size do
        real, data = sock:recv(cache)
        if real > 0 then
            -- append data to buffer
            local leftsize = size - buffsize
            if real < leftsize then
                buff:copy2(buffsize + 1, data)
                buffsize = buffsize + real
            else
                buff:copy2(buffsize + 1, data, 1, leftsize)
                buffsize = buffsize + leftsize

                -- move left cache to head
                cache_size = real - leftsize
                if cache_size > 0 then
                    cache:move(leftsize + 1, real)
                end
                self._RCACHE_SIZE = cache_size
                return buff:slice(1, buffsize)
            end
            wait = false
        elseif real == 0 and not wait then
            if sock:wait(socket.EV_RECV, -1) == socket.EV_RECV then
                wait = true
            else
                break
            end
        else
            break
        end
    end
end

-- recv u32be
function stream:recv_u32be()
    local data = self:recv(self._BUFF, 4)
    if data then
        return data:u32be(1)
    end
end

-- recv message
function stream:recv_msg()
    local body = self:recv_object()
    if body then
        return message(body)
    end
end

-- recv object
function stream:recv_object()
    local str = self:recv_string()
    if str then
        local obj, errors = str:deserialize()
        if errors then
            raise(errors)
        end
        return obj
    end
end

-- recv string
function stream:recv_string()
    local size = self:recv_u32be()
    if size then
        local str
        local recv = 0
        local buff = self._BUFF
        assert(size < STREAM_STRING_MAXN, "too large string size(%d)", size)
        while recv < size do
            local data = self:recv(buff, math.min(buff:size(), size - recv))
            if data then
                if str then
                    str = str .. data:str()
                else
                    str = data:str()
                end
                recv = recv + data:size()
            else
                break
            end
        end
        if recv == size then
            return str
        end
    end
end

-- recv file
function stream:recv_file(filepath)
    local size = self:recv_u32be()
    if size then
        local buff = self._BUFF
        local recv = 0
        local file = io.open(filepath, "wb")
        while recv < size do
            local data = self:recv(buff, math.min(buff:size(), size - recv))
            if data then
                file:write(data)
                recv = recv + data:size()
            end
        end
        file:close()
        if recv == size then
            return true
        end
    end
end

function main(sock)
    local instance = stream()
    instance:init(sock)
    return instance
end