OpenAPI.jl/src/client/chunk_readers.jl at 6d7c490b4da2a2964c98b36d14035640f89f9c23 · JuliaComputing/OpenAPI.jl · 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
struct LineChunkReader <: AbstractChunkReader
    buffered_input::Base.BufferStream
end

function Base.iterate(iter::LineChunkReader, _state=nothing)
    if eof(iter.buffered_input)
        return nothing
    else
        out = IOBuffer()
        while !eof(iter.buffered_input)
            byte = read(iter.buffered_input, UInt8)
            (byte == codepoint('\n')) && break
            write(out, byte)
        end
        return (take!(out), iter)
    end
end

struct JSONChunkReader <: AbstractChunkReader
    buffered_input::Base.BufferStream
end

function Base.iterate(iter::JSONChunkReader, _state=nothing)
    if eof(iter.buffered_input)
        return nothing
    else
        # read all whitespaces
        while !eof(iter.buffered_input)
            byte = peek(iter.buffered_input, UInt8)
            if isspace(Char(byte))
                read(iter.buffered_input, UInt8)
            else
                break
            end
        end
        eof(iter.buffered_input) && return nothing
        valid_json = JSON.parse(iter.buffered_input)
        bytes = convert(Vector{UInt8}, codeunits(JSON.json(valid_json)))
        return (bytes, iter)
    end
end

# Ref: https://www.rfc-editor.org/rfc/rfc7464.html
const RFC7464_RECORD_SEPARATOR = UInt8(0x1E)
struct RFC7464ChunkReader <: AbstractChunkReader
    buffered_input::Base.BufferStream
end

function Base.iterate(iter::RFC7464ChunkReader, _state=nothing)
    if eof(iter.buffered_input)
        return nothing
    else
        out = IOBuffer()
        while !eof(iter.buffered_input)
            byte = read(iter.buffered_input, UInt8)
            if byte == RFC7464_RECORD_SEPARATOR
                bytes = take!(out)
                if isnothing(_state) || !isempty(bytes)
                    return (bytes, iter)
                end
            else
                write(out, byte)
            end
        end
        bytes = take!(out)
        return (bytes, iter)
    end
end