// Copyright (c) 2019-2024 Alexander Medvednikov. All rights reserved.
// Use of this source code is governed by an MIT license
// that can be found in the LICENSE file.
module http

// This file implements a minimal HTTP/2 client connection (RFC 7540) on top of
// the framing and HPACK layers. It is intentionally synchronous and handles a
// single request at a time: it sends one request, then reads frames until that
// stream completes, servicing connection-level frames (SETTINGS, PING,
// WINDOW_UPDATE, GOAWAY) inline. Stream multiplexing and a background reader
// are a follow-up; this is the smallest useful, testable client.

// h2_client_preface is the fixed sequence a client sends to start an HTTP/2
// connection, immediately followed by a SETTINGS frame (RFC 7540 Section 3.5).
pub const h2_client_preface = 'PRI * HTTP/2.0\r\n\r\nSM\r\n\r\n'

// h2_default_initial_window is the initial flow-control window for both the
// connection and new streams (RFC 7540 Section 6.9.2).
pub const h2_default_initial_window = u32(65535)

// h2_conn_read_chunk is the size of each transport read.
const h2_conn_read_chunk = 16 * 1024

// H2Transport is the byte transport an H2Conn runs over: typically an
// ALPN-negotiated `h2` TLS connection, but any reader/writer works (which makes
// the connection testable without a socket). Its method signatures match
// net.ssl.SSLConn, so an SSLConn satisfies it directly.
pub interface H2Transport {
mut:
    read(mut buf []u8) !int
    write(buf []u8) !int
}

// H2PeerSettings holds the peer's SETTINGS, with HTTP/2 defaults.
struct H2PeerSettings {
mut:
    header_table_size      u32  = 4096
    enable_push            bool = true
    max_concurrent_streams u32  = max_u32
    initial_window_size    u32  = 65535
    max_frame_size         u32  = 16384
    max_header_list_size   u32  = max_u32
}

// H2DataFn is called for each DATA frame received on the response stream, with
// the chunk's bytes, the cumulative body bytes received (including this chunk),
// the body length from Content-Length if known (else 0), and the response
// status code.
pub type H2DataFn = fn (chunk []u8, body_so_far u64, body_expected u64, status int) !

// H2ClientRequest describes a single HTTP/2 request. Header names in `headers`
// must be lowercase (RFC 7540 Section 8.1.2); the pseudo-headers are filled in
// from the other fields.
pub struct H2ClientRequest {
pub:
    method    string = 'GET'
    scheme    string = 'https'
    authority string
    path      string = '/'
    headers   []H2HeaderField
    body      []u8
    // Optional response chunk callback, called after each DATA frame's payload
    // is received. The arguments are the chunk bytes (not yet copied into the
    // response body), the cumulative body bytes received so far (including this
    // chunk), the body length from Content-Length (0 when not present), and the
    // response status code.
    on_data H2DataFn = unsafe { nil }
    // stop_copying_limit, when >= 0, caps the cumulative body bytes copied into
    // the response body; further DATA chunks are dropped but the callback keeps
    // firing and the stream is drained to completion.
    stop_copying_limit i64 = -1
    // stop_receiving_limit, when >= 0, causes the response read loop to break
    // once that many body bytes have been received. The callback fires for the
    // final chunk; no further callbacks fire after that.
    stop_receiving_limit i64 = -1
}

// H2ClientResponse is the result of an HTTP/2 request.
pub struct H2ClientResponse {
pub mut:
    status  int
    headers []H2HeaderField
    body    []u8
}

// H2Conn is a client-side HTTP/2 connection.
pub struct H2Conn {
mut:
    transport          H2Transport
    encoder            H2HpackEncoder
    decoder            H2HpackDecoder
    peer               H2PeerSettings
    rbuf               []u8      // buffered bytes read from the transport, not yet consumed
    pending            []H2Frame // stream frames read early (while sending), to replay
    next_stream_id     u32 = 1 // clients use odd stream ids
    cur_stream_id      u32 // the stream currently being driven by do()
    send_window        i64 = 65535 // connection-level send flow-control window
    stream_send_window i64 // send window for cur_stream_id
    handshaked         bool
    goaway             bool
    // aborted is set when this connection terminated a stream early
    // (RST_STREAM sent without draining the remaining DATA). Subsequent
    // requests on the same connection must fail rather than risk being starved
    // by leftover DATA frames the peer had already sent for the cancelled
    // stream.
    aborted bool
}

// new_h2_conn creates a client connection over `transport`. The HTTP/2
// connection preface is sent lazily on the first request.
pub fn new_h2_conn(transport H2Transport) &H2Conn {
    return &H2Conn{
        transport: transport
    }
}

// do sends `req` and returns the response, after the request's stream closes.
pub fn (mut c H2Conn) do(req H2ClientRequest) !H2ClientResponse {
    c.handshake()!
    if c.goaway {
        return error('h2: connection is shutting down (GOAWAY)')
    }
    if c.aborted {
        return error('h2: connection is no longer usable after an early stream termination')
    }
    stream_id := c.next_stream_id
    c.next_stream_id += 2
    c.cur_stream_id = stream_id
    c.stream_send_window = i64(c.peer.initial_window_size)

    mut fields := [
        H2HeaderField{':method', req.method},
        H2HeaderField{':scheme', req.scheme},
        H2HeaderField{':authority', req.authority},
        H2HeaderField{':path', req.path},
    ]
    for h in req.headers {
        fields << h
    }
    block := c.encoder.encode(fields)
    has_body := req.body.len > 0
    c.send_header_block(stream_id, block, !has_body)!
    if has_body {
        c.send_body(stream_id, req.body)!
    }
    return c.read_response(stream_id, req)!
}

fn (mut c H2Conn) handshake() ! {
    if c.handshaked {
        return
    }
    c.write_all(h2_client_preface.bytes())!
    // Advertise our SETTINGS: refuse server push, and use the protocol defaults
    // otherwise. (Our decoder uses the default 4096-byte HPACK table.)
    c.send_frame(H2SettingsFrame{
        settings: [
            H2Setting{h2_settings_enable_push, 0},
            H2Setting{h2_settings_initial_window_size, h2_default_initial_window},
            H2Setting{h2_settings_max_frame_size, h2_default_max_frame_size},
        ]
    })!
    c.handshaked = true
}

// read_response reads frames until `stream_id` is closed, returning its
// response and servicing connection-level frames along the way. The streaming
// options on `req` (on_data callback and the two stop limits) are honored
// while reading DATA frames, matching the HTTP/1.1 streaming semantics.
fn (mut c H2Conn) read_response(stream_id u32, req H2ClientRequest) !H2ClientResponse {
    mut resp := H2ClientResponse{}
    mut got_headers := false
    mut body_so_far := u64(0)
    mut body_expected := u64(0)
    for {
        frame := c.next_frame()!
        if c.handle_conn_frame(frame)! {
            continue
        }
        match frame {
            H2HeadersFrame {
                if frame.stream_id != stream_id {
                    continue
                }
                fragment := c.collect_header_block(frame.fragment, frame.end_headers, stream_id)!
                for f in c.decoder.decode(fragment)! {
                    if f.name == ':status' {
                        resp.status = f.value.int()
                    } else if !f.name.starts_with(':') {
                        resp.headers << f
                        if f.name == 'content-length' {
                            body_expected = f.value.u64()
                        }
                    }
                }
                got_headers = true
                if frame.end_stream {
                    break
                }
            }
            H2DataFrame {
                if frame.stream_id != stream_id {
                    continue
                }
                if frame.data.len > 0 {
                    body_so_far += u64(frame.data.len)
                    // Append the chunk to the response body unless the copy
                    // limit has been reached; the callback still fires.
                    if req.stop_copying_limit < 0
                        || i64(body_so_far) - i64(frame.data.len) < req.stop_copying_limit {
                        if req.stop_copying_limit >= 0 && i64(body_so_far) > req.stop_copying_limit {
                            remaining := req.stop_copying_limit - (i64(body_so_far) - i64(frame.data.len))
                            if remaining > 0 {
                                resp.body << frame.data[..int(remaining)]
                            }
                        } else {
                            resp.body << frame.data
                        }
                    }
                    if req.on_data != unsafe { nil } {
                        req.on_data(frame.data, body_so_far, body_expected, resp.status)!
                    }
                    // Replenish flow control so the peer keeps sending.
                    c.send_window_update(0, u32(frame.data.len))!
                    c.send_window_update(stream_id, u32(frame.data.len))!
                }
                if frame.end_stream {
                    break
                }
                if req.stop_receiving_limit >= 0 && i64(body_so_far) >= req.stop_receiving_limit {
                    // Cancel the stream (RFC 7540 Section 8.1.4 / 5.4.2) so the
                    // peer stops sending more DATA, and mark the connection
                    // unusable: in-flight DATA frames that the peer has already
                    // sent for this stream would otherwise consume the
                    // connection-level receive window and block subsequent
                    // requests on the same H2Conn.
                    c.send_frame(H2RstStreamFrame{
                        stream_id:  stream_id
                        error_code: u32(H2ErrorCode.cancel)
                    })!
                    c.aborted = true
                    break
                }
            }
            H2RstStreamFrame {
                if frame.stream_id == stream_id {
                    return error('h2: stream reset by peer (${h2_error_code_name(frame.error_code)})')
                }
            }
            else {
                // PRIORITY / PUSH_PROMISE / stray CONTINUATION / unknown: ignore.
            }
        }
    }
    if !got_headers {
        return error('h2: stream closed without a response')
    }
    return resp
}

// collect_header_block returns the full HPACK block for a HEADERS frame,
// reading and concatenating CONTINUATION frames until END_HEADERS.
fn (mut c H2Conn) collect_header_block(first []u8, end_headers bool, stream_id u32) ![]u8 {
    if end_headers {
        return first
    }
    mut fragment := first.clone()
    for {
        frame := c.read_frame()!
        if frame is H2ContinuationFrame {
            if frame.stream_id != stream_id {
                return error('h2: CONTINUATION on the wrong stream')
            }
            fragment << frame.fragment
            if frame.end_headers {
                break
            }
        } else {
            return error('h2: expected a CONTINUATION frame')
        }
    }
    return fragment
}

// handle_conn_frame services a connection-level frame, returning true if the
// frame was fully handled here (so the caller should skip it).
fn (mut c H2Conn) handle_conn_frame(frame H2Frame) !bool {
    match frame {
        H2SettingsFrame {
            if !frame.ack {
                c.apply_settings(frame.settings)
                c.send_frame(H2SettingsFrame{
                    ack: true
                })!
            }
            return true
        }
        H2PingFrame {
            if !frame.ack {
                c.send_frame(H2PingFrame{
                    ack:  true
                    data: frame.data
                })!
            }
            return true
        }
        H2WindowUpdateFrame {
            if frame.stream_id == 0 {
                c.send_window += i64(frame.window_size_increment)
            } else if frame.stream_id == c.cur_stream_id {
                c.stream_send_window += i64(frame.window_size_increment)
            }
            return true
        }
        H2GoawayFrame {
            c.goaway = true
            return error('h2: GOAWAY received (${h2_error_code_name(frame.error_code)})')
        }
        else {
            return false
        }
    }
}

fn (mut c H2Conn) apply_settings(settings []H2Setting) {
    for s in settings {
        match s.id {
            h2_settings_header_table_size {
                c.peer.header_table_size = s.value
            }
            h2_settings_enable_push {
                c.peer.enable_push = s.value != 0
            }
            h2_settings_max_concurrent_streams {
                c.peer.max_concurrent_streams = s.value
            }
            h2_settings_initial_window_size {
                // RFC 7540 Section 6.9.2: a change to the initial window size
                // retroactively adjusts the active stream's send window by the
                // delta.
                delta := i64(s.value) - i64(c.peer.initial_window_size)
                c.peer.initial_window_size = s.value
                c.stream_send_window += delta
            }
            h2_settings_max_frame_size {
                c.peer.max_frame_size = s.value
            }
            h2_settings_max_header_list_size {
                c.peer.max_header_list_size = s.value
            }
            else {} // unknown settings are ignored (RFC 7540 Section 6.5.2)
        }
    }
}

// send_header_block sends an HPACK header block as a HEADERS frame, splitting
// it across CONTINUATION frames when it exceeds the peer's max frame size
// (RFC 7540 Section 4.3). END_STREAM, when set, goes on the HEADERS frame.
fn (mut c H2Conn) send_header_block(stream_id u32, block []u8, end_stream bool) ! {
    max := int(c.peer.max_frame_size)
    if block.len <= max {
        c.send_frame(H2HeadersFrame{
            stream_id:   stream_id
            fragment:    block
            end_headers: true
            end_stream:  end_stream
        })!
        return
    }
    c.send_frame(H2HeadersFrame{
        stream_id:   stream_id
        fragment:    block[..max]
        end_headers: false
        end_stream:  end_stream
    })!
    mut off := max
    for off < block.len {
        mut next := off + max
        if next > block.len {
            next = block.len
        }
        c.send_frame(H2ContinuationFrame{
            stream_id:   stream_id
            fragment:    block[off..next]
            end_headers: next == block.len
        })!
        off = next
    }
}

// send_body writes the request body as DATA frames, chunked to the peer's max
// frame size and bounded by both the connection and the per-stream send
// flow-control windows (RFC 7540 Section 6.9).
fn (mut c H2Conn) send_body(stream_id u32, body []u8) ! {
    max := int(c.peer.max_frame_size)
    mut off := 0
    for off < body.len {
        for c.send_window <= 0 || c.stream_send_window <= 0 {
            // Wait for the peer to grow a window. Connection-level frames are
            // handled; any stream frames are stashed for read_response.
            frame := c.next_frame()!
            if !c.handle_conn_frame(frame)! {
                c.pending << frame
            }
        }
        avail := if c.send_window < c.stream_send_window {
            c.send_window
        } else {
            c.stream_send_window
        }
        mut chunk := body.len - off
        if chunk > max {
            chunk = max
        }
        if i64(chunk) > avail {
            chunk = int(avail)
        }
        next := off + chunk
        c.send_frame(H2DataFrame{
            stream_id:  stream_id
            data:       body[off..next]
            end_stream: next == body.len
        })!
        c.send_window -= i64(chunk)
        c.stream_send_window -= i64(chunk)
        off = next
    }
}

fn (mut c H2Conn) send_window_update(stream_id u32, inc u32) ! {
    if inc == 0 {
        return
    }
    c.send_frame(H2WindowUpdateFrame{
        stream_id:             stream_id
        window_size_increment: inc
    })!
}

// next_frame returns the next frame, replaying any stashed stream frames first.
fn (mut c H2Conn) next_frame() !H2Frame {
    if c.pending.len > 0 {
        f := c.pending[0]
        c.pending.delete(0)
        return f
    }
    return c.read_frame()!
}

// read_frame reads and decodes one frame from the transport, enforcing our
// advertised max frame size.
fn (mut c H2Conn) read_frame() !H2Frame {
    c.fill_at_least(h2_frame_header_len)!
    header := h2_parse_frame_header(c.rbuf)!
    if header.length > h2_default_max_frame_size {
        return error('h2: frame larger than SETTINGS_MAX_FRAME_SIZE (${header.length})')
    }
    total := h2_frame_header_len + int(header.length)
    c.fill_at_least(total)!
    frame := h2_parse_frame(header, c.rbuf[h2_frame_header_len..total])!
    c.rbuf = c.rbuf[total..].clone()
    return frame
}

// fill_at_least reads from the transport until rbuf holds at least n bytes.
fn (mut c H2Conn) fill_at_least(n int) ! {
    for c.rbuf.len < n {
        mut tmp := []u8{len: h2_conn_read_chunk}
        got := c.transport.read(mut tmp)!
        if got <= 0 {
            return error('h2: connection closed by peer')
        }
        c.rbuf << tmp[..got]
    }
}

fn (mut c H2Conn) send_frame(f H2Frame) ! {
    c.write_all(f.encode())!
}

fn (mut c H2Conn) write_all(data []u8) ! {
    mut sent := 0
    for sent < data.len {
        n := c.transport.write(data[sent..])!
        if n <= 0 {
            return error('h2: transport write returned ${n}')
        }
        sent += n
    }
}

// h2_error_code_name renders an HTTP/2 error code, falling back to hex for
// codes outside the defined range.
fn h2_error_code_name(code u32) string {
    if code <= u32(H2ErrorCode.http_1_1_required) {
        return unsafe { H2ErrorCode(code) }.str()
    }
    return 'unknown(0x${code.hex()})'
}