// 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 the HTTP/2 binary framing layer (RFC 7540 Sections 4
// and 6). It is self-contained and does not touch the rest of net.http; the
// connection layer (added separately) drives it.

// HTTP/2 frame types (RFC 7540 Section 6).
pub const h2_frame_data = u8(0x0)
pub const h2_frame_headers = u8(0x1)
pub const h2_frame_priority = u8(0x2)
pub const h2_frame_rst_stream = u8(0x3)
pub const h2_frame_settings = u8(0x4)
pub const h2_frame_push_promise = u8(0x5)
pub const h2_frame_ping = u8(0x6)
pub const h2_frame_goaway = u8(0x7)
pub const h2_frame_window_update = u8(0x8)
pub const h2_frame_continuation = u8(0x9)

// HTTP/2 frame flags (RFC 7540 Section 6). The same bit is reused across
// frame types with different meanings.
pub const h2_flag_end_stream = u8(0x1)
pub const h2_flag_ack = u8(0x1)
pub const h2_flag_end_headers = u8(0x4)
pub const h2_flag_padded = u8(0x8)
pub const h2_flag_priority = u8(0x20)

// h2_frame_header_len is the fixed size of a frame header in bytes.
pub const h2_frame_header_len = 9

// h2_default_max_frame_size is the initial SETTINGS_MAX_FRAME_SIZE
// (RFC 7540 Section 6.5.2), and the smallest value a peer may set it to.
pub const h2_default_max_frame_size = u32(16384)

// HTTP/2 setting identifiers (RFC 7540 Section 6.5.2).
pub const h2_settings_header_table_size = u16(0x1)
pub const h2_settings_enable_push = u16(0x2)
pub const h2_settings_max_concurrent_streams = u16(0x3)
pub const h2_settings_initial_window_size = u16(0x4)
pub const h2_settings_max_frame_size = u16(0x5)
pub const h2_settings_max_header_list_size = u16(0x6)

// H2FrameHeader is the 9-byte header that precedes every HTTP/2 frame.
pub struct H2FrameHeader {
pub:
    length    u32 // 24-bit payload length
    typ       u8
    flags     u8
    stream_id u32 // 31-bit; the reserved bit is ignored on read
}

pub struct H2DataFrame {
pub:
    stream_id  u32
    data       []u8
    end_stream bool
}

pub struct H2HeadersFrame {
pub:
    stream_id   u32
    fragment    []u8 // HPACK header block fragment
    end_stream  bool
    end_headers bool
    // Priority information, present only when the PRIORITY flag is set.
    has_priority bool
    exclusive    bool
    stream_dep   u32
    weight       u8
}

pub struct H2PriorityFrame {
pub:
    stream_id  u32
    exclusive  bool
    stream_dep u32
    weight     u8
}

pub struct H2RstStreamFrame {
pub:
    stream_id  u32
    error_code u32
}

pub struct H2Setting {
pub:
    id    u16
    value u32
}

pub struct H2SettingsFrame {
pub:
    ack      bool
    settings []H2Setting
}

pub struct H2PushPromiseFrame {
pub:
    stream_id          u32
    promised_stream_id u32
    fragment           []u8
    end_headers        bool
}

pub struct H2PingFrame {
pub:
    ack  bool
    data []u8 // 8 opaque bytes
}

pub struct H2GoawayFrame {
pub:
    last_stream_id u32
    error_code     u32
    debug_data     []u8
}

pub struct H2WindowUpdateFrame {
pub:
    stream_id             u32
    window_size_increment u32
}

pub struct H2ContinuationFrame {
pub:
    stream_id   u32
    fragment    []u8
    end_headers bool
}

// H2UnknownFrame preserves a frame of an unrecognised type, which receivers
// must ignore (RFC 7540 Section 4.1) but may want to inspect or forward.
pub struct H2UnknownFrame {
pub:
    header  H2FrameHeader
    payload []u8
}

// H2Frame is any HTTP/2 frame.
pub type H2Frame = H2ContinuationFrame
    | H2DataFrame
    | H2GoawayFrame
    | H2HeadersFrame
    | H2PingFrame
    | H2PriorityFrame
    | H2PushPromiseFrame
    | H2RstStreamFrame
    | H2SettingsFrame
    | H2UnknownFrame
    | H2WindowUpdateFrame

// --- Big-endian helpers ---

fn h2_be_u16(b []u8, o int) u16 {
    return (u16(b[o]) << 8) | u16(b[o + 1])
}

fn h2_be_u24(b []u8, o int) u32 {
    return (u32(b[o]) << 16) | (u32(b[o + 1]) << 8) | u32(b[o + 2])
}

fn h2_be_u32(b []u8, o int) u32 {
    return (u32(b[o]) << 24) | (u32(b[o + 1]) << 16) | (u32(b[o + 2]) << 8) | u32(b[o + 3])
}

fn h2_put_u16(mut b []u8, v u16) {
    b << u8(v >> 8)
    b << u8(v)
}

fn h2_put_u24(mut b []u8, v u32) {
    b << u8(v >> 16)
    b << u8(v >> 8)
    b << u8(v)
}

fn h2_put_u32(mut b []u8, v u32) {
    b << u8(v >> 24)
    b << u8(v >> 16)
    b << u8(v >> 8)
    b << u8(v)
}

// --- Frame header ---

// h2_parse_frame_header parses the 9-byte frame header at the start of `buf`.
pub fn h2_parse_frame_header(buf []u8) !H2FrameHeader {
    if buf.len < h2_frame_header_len {
        return error('h2: frame header truncated')
    }
    return H2FrameHeader{
        length:    h2_be_u24(buf, 0)
        typ:       buf[3]
        flags:     buf[4]
        stream_id: h2_be_u32(buf, 5) & 0x7fff_ffff
    }
}

// --- Decoding ---

// h2_read_frame parses one frame (header + payload) from the start of `buf`,
// returning the frame and the number of bytes consumed. The caller is
// responsible for enforcing the negotiated SETTINGS_MAX_FRAME_SIZE.
pub fn h2_read_frame(buf []u8) !(H2Frame, int) {
    header := h2_parse_frame_header(buf)!
    total := h2_frame_header_len + int(header.length)
    if buf.len < total {
        return error('h2: frame payload truncated (need ${total}, have ${buf.len})')
    }
    payload := buf[h2_frame_header_len..total]
    frame := h2_parse_frame(header, payload)!
    return frame, total
}

// h2_strip_padding removes the optional pad-length prefix byte and the
// trailing padding from a frame payload (RFC 7540 Section 6.1).
fn h2_strip_padding(payload []u8, padded bool) ![]u8 {
    if !padded {
        return payload
    }
    if payload.len < 1 {
        return error('h2: padded frame missing pad length')
    }
    pad_len := int(payload[0])
    if 1 + pad_len > payload.len {
        return error('h2: pad length exceeds frame payload')
    }
    return payload[1..payload.len - pad_len]
}

// h2_parse_frame decodes a frame from an already-parsed header and its payload.
pub fn h2_parse_frame(header H2FrameHeader, payload []u8) !H2Frame {
    match header.typ {
        h2_frame_data {
            if header.stream_id == 0 {
                return error('h2: DATA frame on stream 0')
            }
            body := h2_strip_padding(payload, header.flags & h2_flag_padded != 0)!
            return H2DataFrame{
                stream_id:  header.stream_id
                data:       body.clone()
                end_stream: header.flags & h2_flag_end_stream != 0
            }
        }
        h2_frame_headers {
            if header.stream_id == 0 {
                return error('h2: HEADERS frame on stream 0')
            }
            mut body := h2_strip_padding(payload, header.flags & h2_flag_padded != 0)!
            mut has_priority := false
            mut exclusive := false
            mut stream_dep := u32(0)
            mut weight := u8(0)
            if header.flags & h2_flag_priority != 0 {
                if body.len < 5 {
                    return error('h2: HEADERS priority section truncated')
                }
                dep := h2_be_u32(body, 0)
                has_priority = true
                exclusive = dep & 0x8000_0000 != 0
                stream_dep = dep & 0x7fff_ffff
                weight = body[4]
                body = unsafe { body[5..] }
            }
            return H2HeadersFrame{
                stream_id:    header.stream_id
                fragment:     body.clone()
                end_stream:   header.flags & h2_flag_end_stream != 0
                end_headers:  header.flags & h2_flag_end_headers != 0
                has_priority: has_priority
                exclusive:    exclusive
                stream_dep:   stream_dep
                weight:       weight
            }
        }
        h2_frame_priority {
            if header.stream_id == 0 {
                return error('h2: PRIORITY frame on stream 0')
            }
            if payload.len != 5 {
                return error('h2: PRIORITY frame must be 5 bytes')
            }
            // Note: a stream depending on itself (stream_dep == stream_id) is a
            // stream error (RFC 7540 Section 5.3.1), and a zero stream
            // dependency is otherwise valid. These are semantic checks left to
            // the connection layer, which must respond with RST_STREAM on the
            // affected stream rather than tearing down the whole connection.
            dep := h2_be_u32(payload, 0)
            return H2PriorityFrame{
                stream_id:  header.stream_id
                exclusive:  dep & 0x8000_0000 != 0
                stream_dep: dep & 0x7fff_ffff
                weight:     payload[4]
            }
        }
        h2_frame_rst_stream {
            if header.stream_id == 0 {
                return error('h2: RST_STREAM frame on stream 0')
            }
            if payload.len != 4 {
                return error('h2: RST_STREAM frame must be 4 bytes')
            }
            return H2RstStreamFrame{
                stream_id:  header.stream_id
                error_code: h2_be_u32(payload, 0)
            }
        }
        h2_frame_settings {
            if header.stream_id != 0 {
                return error('h2: SETTINGS frame on non-zero stream')
            }
            ack := header.flags & h2_flag_ack != 0
            if ack {
                if payload.len != 0 {
                    return error('h2: SETTINGS ACK must have empty payload')
                }
                return H2SettingsFrame{
                    ack: true
                }
            }
            if payload.len % 6 != 0 {
                return error('h2: SETTINGS payload not a multiple of 6')
            }
            mut settings := []H2Setting{cap: payload.len / 6}
            for i := 0; i < payload.len; i += 6 {
                settings << H2Setting{
                    id:    h2_be_u16(payload, i)
                    value: h2_be_u32(payload, i + 2)
                }
            }
            return H2SettingsFrame{
                ack:      false
                settings: settings
            }
        }
        h2_frame_push_promise {
            if header.stream_id == 0 {
                return error('h2: PUSH_PROMISE frame on stream 0')
            }
            mut body := h2_strip_padding(payload, header.flags & h2_flag_padded != 0)!
            if body.len < 4 {
                return error('h2: PUSH_PROMISE missing promised stream id')
            }
            promised := h2_be_u32(body, 0) & 0x7fff_ffff
            return H2PushPromiseFrame{
                stream_id:          header.stream_id
                promised_stream_id: promised
                fragment:           body[4..].clone()
                end_headers:        header.flags & h2_flag_end_headers != 0
            }
        }
        h2_frame_ping {
            if header.stream_id != 0 {
                return error('h2: PING frame on non-zero stream')
            }
            if payload.len != 8 {
                return error('h2: PING frame must be 8 bytes')
            }
            return H2PingFrame{
                ack:  header.flags & h2_flag_ack != 0
                data: payload.clone()
            }
        }
        h2_frame_goaway {
            if header.stream_id != 0 {
                return error('h2: GOAWAY frame on non-zero stream')
            }
            if payload.len < 8 {
                return error('h2: GOAWAY frame too short')
            }
            return H2GoawayFrame{
                last_stream_id: h2_be_u32(payload, 0) & 0x7fff_ffff
                error_code:     h2_be_u32(payload, 4)
                debug_data:     payload[8..].clone()
            }
        }
        h2_frame_window_update {
            if payload.len != 4 {
                return error('h2: WINDOW_UPDATE frame must be 4 bytes')
            }
            // Note: a zero increment is an error (RFC 7540 Section 6.9) — a
            // stream error on a stream, but a connection error on stream 0.
            // That stream-vs-connection distinction is the connection layer's
            // responsibility, so it is not rejected here.
            return H2WindowUpdateFrame{
                stream_id:             header.stream_id
                window_size_increment: h2_be_u32(payload, 0) & 0x7fff_ffff
            }
        }
        h2_frame_continuation {
            if header.stream_id == 0 {
                return error('h2: CONTINUATION frame on stream 0')
            }
            return H2ContinuationFrame{
                stream_id:   header.stream_id
                fragment:    payload.clone()
                end_headers: header.flags & h2_flag_end_headers != 0
            }
        }
        else {
            // Unknown frame types must be ignored (RFC 7540 Section 4.1);
            // preserve them so the caller can decide.
            return H2UnknownFrame{
                header:  header
                payload: payload.clone()
            }
        }
    }
}

// --- Encoding ---

// h2_frame_bytes builds a complete frame from its parts.
fn h2_frame_bytes(typ u8, flags u8, stream_id u32, payload []u8) []u8 {
    mut b := []u8{cap: h2_frame_header_len + payload.len}
    h2_put_u24(mut b, u32(payload.len))
    b << typ
    b << flags
    h2_put_u32(mut b, stream_id & 0x7fff_ffff)
    b << payload
    return b
}

// encode serialises a frame to its on-the-wire bytes. The encoder never emits
// padding.
pub fn (f H2Frame) encode() []u8 {
    match f {
        H2DataFrame {
            flags := if f.end_stream { h2_flag_end_stream } else { u8(0) }
            return h2_frame_bytes(h2_frame_data, flags, f.stream_id, f.data)
        }
        H2HeadersFrame {
            mut flags := u8(0)
            if f.end_stream {
                flags |= h2_flag_end_stream
            }
            if f.end_headers {
                flags |= h2_flag_end_headers
            }
            mut payload := []u8{}
            if f.has_priority {
                flags |= h2_flag_priority
                mut dep := f.stream_dep & 0x7fff_ffff
                if f.exclusive {
                    dep |= 0x8000_0000
                }
                h2_put_u32(mut payload, dep)
                payload << f.weight
            }
            payload << f.fragment
            return h2_frame_bytes(h2_frame_headers, flags, f.stream_id, payload)
        }
        H2PriorityFrame {
            mut payload := []u8{cap: 5}
            mut dep := f.stream_dep & 0x7fff_ffff
            if f.exclusive {
                dep |= 0x8000_0000
            }
            h2_put_u32(mut payload, dep)
            payload << f.weight
            return h2_frame_bytes(h2_frame_priority, 0, f.stream_id, payload)
        }
        H2RstStreamFrame {
            mut payload := []u8{cap: 4}
            h2_put_u32(mut payload, f.error_code)
            return h2_frame_bytes(h2_frame_rst_stream, 0, f.stream_id, payload)
        }
        H2SettingsFrame {
            if f.ack {
                return h2_frame_bytes(h2_frame_settings, h2_flag_ack, 0, [])
            }
            mut payload := []u8{cap: f.settings.len * 6}
            for s in f.settings {
                h2_put_u16(mut payload, s.id)
                h2_put_u32(mut payload, s.value)
            }
            return h2_frame_bytes(h2_frame_settings, 0, 0, payload)
        }
        H2PushPromiseFrame {
            mut flags := u8(0)
            if f.end_headers {
                flags |= h2_flag_end_headers
            }
            mut payload := []u8{cap: 4 + f.fragment.len}
            h2_put_u32(mut payload, f.promised_stream_id & 0x7fff_ffff)
            payload << f.fragment
            return h2_frame_bytes(h2_frame_push_promise, flags, f.stream_id, payload)
        }
        H2PingFrame {
            flags := if f.ack { h2_flag_ack } else { u8(0) }
            mut payload := []u8{len: 8}
            for i in 0 .. 8 {
                payload[i] = if i < f.data.len { f.data[i] } else { u8(0) }
            }
            return h2_frame_bytes(h2_frame_ping, flags, 0, payload)
        }
        H2GoawayFrame {
            mut payload := []u8{cap: 8 + f.debug_data.len}
            h2_put_u32(mut payload, f.last_stream_id & 0x7fff_ffff)
            h2_put_u32(mut payload, f.error_code)
            payload << f.debug_data
            return h2_frame_bytes(h2_frame_goaway, 0, 0, payload)
        }
        H2WindowUpdateFrame {
            mut payload := []u8{cap: 4}
            h2_put_u32(mut payload, f.window_size_increment & 0x7fff_ffff)
            return h2_frame_bytes(h2_frame_window_update, 0, f.stream_id, payload)
        }
        H2ContinuationFrame {
            flags := if f.end_headers { h2_flag_end_headers } else { u8(0) }
            return h2_frame_bytes(h2_frame_continuation, flags, f.stream_id, f.fragment)
        }
        H2UnknownFrame {
            return h2_frame_bytes(f.header.typ, f.header.flags, f.header.stream_id, f.payload)
        }
    }
}