// 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.
@[has_globals]
module http

import net
import net.ssl
import sync
import time

// This file implements connection reuse for the HTTP client: a Transport owns
// a pool of keep-alive connections, keyed by origin and TLS configuration, and
// fetch()/Request.do() route through a process-global default Transport unless
// the request opts out (`disable_connection_reuse`) or uses a proxy. The first
// cut pools HTTP/1.1 connections (plain TCP and TLS); HTTP/2 requests run on
// the existing one-shot driver until the multiplexed H2 connection lands.

// transport_err_unsafe_retry tags a failed exchange on a transport connection
// for a non-idempotent method: the request bytes may have (partially) reached
// the server — the write helpers cannot prove otherwise — so it must not be
// replayed, neither by round_trip's inner loop nor the caller's outer retry
// loop. is_no_need_retry_error recognizes this code so the outer loop honors it.
const transport_err_unsafe_retry = -20013

// H1PooledConn is one keep-alive HTTP/1.1 connection in a Transport pool:
// either a plain TCP connection or a TLS one (exactly one of the two is set).
@[heap]
struct H1PooledConn {
mut:
    key        string
    tcp        &net.TcpConn = unsafe { nil }
    ssl        &ssl.SSLConn = unsafe { nil }
    idle_since time.Time
}

// close_conn shuts the underlying connection down and clears it.
fn (mut c H1PooledConn) close_conn() {
    if c.ssl != unsafe { nil } {
        c.ssl.shutdown() or {}
        c.ssl = unsafe { nil }
    }
    if c.tcp != unsafe { nil } {
        c.tcp.close() or {}
        c.tcp = unsafe { nil }
    }
}

// refresh_timeouts applies the current request's timeouts to a pooled
// connection, which may have been dialled by a request with different ones.
fn (mut c H1PooledConn) refresh_timeouts(req &Request) {
    if c.tcp != unsafe { nil } {
        c.tcp.set_read_timeout(req.read_timeout)
        c.tcp.set_write_timeout(req.write_timeout)
    } else if c.ssl != unsafe { nil } {
        if req.read_timeout > 0 {
            c.ssl.set_read_timeout(req.read_timeout)
        }
    }
}

// exchange runs one request/response over the pooled connection, leaving it
// open. The bool result reports whether the response was precisely framed, so
// the connection can safely carry another request.
fn (mut c H1PooledConn) exchange(req &Request, raw string) !(Response, bool) {
    if c.ssl != unsafe { nil } {
        return req.h1_exchange_ssl(mut c.ssl, raw)
    }
    return req.h1_exchange_tcp(mut c.tcp, raw)
}

// Transport holds the connection pools and reuse policy for HTTP requests. It
// is safe for concurrent use; fetch()/Request.do() share one process-global
// instance (see default_transport).
//
// Note: idle_timeout is enforced lazily on checkout of the same pool key; there
// is no background reaper, so a connection to an origin that is never revisited
// is reclaimed by the global max_idle_conns cap (or close_idle_connections())
// rather than on a timer. A time-based reaper is a planned follow-up.
@[heap]
pub struct Transport {
pub mut:
    // max_idle_conns_per_host caps the idle keep-alive connections retained per
    // pool key (origin + TLS configuration).
    max_idle_conns_per_host int = 4
    // max_idle_conns caps the total idle keep-alive connections kept across all
    // pool keys, bounding file-descriptor use when many distinct origins are
    // fetched. The least-recently-used idle connection is evicted on overflow.
    // 0 disables the global cap.
    max_idle_conns int = 100
    // idle_timeout is how long an idle connection may sit in the pool before it
    // is discarded instead of reused.
    idle_timeout time.Duration = 90 * time.second
mut:
    mu      &sync.Mutex = sync.new_mutex()
    h1_idle map[string][]&H1PooledConn
}

// new_transport creates an empty Transport with default limits.
pub fn new_transport() &Transport {
    return &Transport{}
}

// Initialized eagerly in _vinit() (before main() and before any thread is
// spawned), so the shared default Transport is fully constructed without a
// concurrent-startup race. Lazy init via sync.Once is unsafe here: Once sets its
// done flag before running the callback, so a second caller can observe "done"
// and read this global while it is still nil (vlang/v#27456).
__global http_transport_default = new_transport()

// default_transport returns the process-global Transport that fetch() and
// Request.do() route through.
pub fn default_transport() &Transport {
    return http_transport_default
}

// close_idle_connections closes every idle pooled connection held by the
// default Transport. Useful before fork()-style process handoffs, or in tests.
pub fn close_idle_connections() {
    mut t := default_transport()
    t.close_idle()
}

// close_idle closes every idle pooled connection held by this Transport.
// In-flight requests are unaffected.
pub fn (mut t Transport) close_idle() {
    mut all := []&H1PooledConn{}
    t.mu.lock()
    for _, list in t.h1_idle {
        for c in list {
            all << c
        }
    }
    t.h1_idle.clear()
    t.mu.unlock()
    for mut c in all {
        c.close_conn()
    }
}

// transport_pool_key builds the pool key for a request: connections are only
// shared between requests whose origin and TLS-relevant settings all match.
// enable_http2 is part of the key because it changes what the connection
// advertised via ALPN at dial time: a forced-HTTP/1.1 connection (no ALPN)
// must not satisfy an HTTP/2-enabled request, which would otherwise silently
// never negotiate h2 against that origin while the pooled connection lives.
fn transport_pool_key(req &Request, scheme string, host string, port int) string {
    // enable_http2 only affects https dials (plain http ignores it), so keep
    // the plain-http pool unsplit.
    h2 := scheme == 'https' && req.enable_http2
    // Length-prefix the free-form string fields (host and the TLS paths/PEM
    // blobs) so a value containing the '|' separator cannot collide with a
    // different field split — e.g. cert='a|b',cert_key='c' vs cert='a',
    // cert_key='b|c' — which would let a request reuse a connection dialed with
    // the wrong CA or client certificate. Bools and the int port cannot contain
    // '|', and scheme is a fixed literal, so they need no prefixing.
    return '${scheme}|${pk_part(host)}|${port}|${h2}|${req.validate}|${pk_part(req.verify)}|${pk_part(req.cert)}|${pk_part(req.cert_key)}|${req.in_memory_verification}'
}

// pk_part length-prefixes a string (`len:value`) so concatenated pool-key
// components stay unambiguous regardless of the bytes they contain.
fn pk_part(s string) string {
    return '${s.len}:${s}'
}

// transport_is_idempotent reports whether a request with this method can be
// replayed safely after a failure on a reused connection (RFC 7231 4.2.2).
fn transport_is_idempotent(method Method) bool {
    return method in [.get, .head, .options, .trace, .put, .delete]
}

// response_allows_reuse reports whether the server's response permits keeping
// the connection open for another request.
fn response_allows_reuse(resp &Response) bool {
    if resp.status_code < 200 {
        // 1xx handling stops at the headers; the connection state past them is
        // not tracked, so never reuse.
        return false
    }
    // A server may split Connection across repeated header lines; get() returns
    // only the first, so join all values before scanning for the close token.
    conn_tokens := resp.header.values(.connection).join(',').to_lower()
    if conn_tokens.contains('close') {
        return false
    }
    ver := resp.version()
    if ver == .v1_0 {
        // HTTP/1.0 defaults to close; reuse only with an explicit keep-alive.
        return conn_tokens.contains('keep-alive')
    }
    return ver == .v1_1
}

// checkout pops the most recently used idle connection for `key`, discarding
// any that sat idle past the timeout. Returns nil when the pool has none.
fn (mut t Transport) checkout(key string) &H1PooledConn {
    mut expired := []&H1PooledConn{}
    mut found := &H1PooledConn(unsafe { nil })
    t.mu.lock()
    for {
        mut list := t.h1_idle[key] or { break }
        if list.len == 0 {
            break
        }
        mut c := list.pop()
        t.h1_idle[key] = list
        if time.now() - c.idle_since > t.idle_timeout {
            expired << c
            continue
        }
        found = c
        break
    }
    t.mu.unlock()
    for mut c in expired {
        c.close_conn()
    }
    return found
}

// checkin returns a healthy connection to the idle pool, or closes it when the
// per-host pool is already at capacity. Adding it may push the total idle count
// over max_idle_conns, in which case the least-recently-used idle connection
// across all pools is evicted and closed.
fn (mut t Transport) checkin(mut conn H1PooledConn) {
    conn.idle_since = time.now()
    mut evicted := &H1PooledConn(unsafe { nil })
    t.mu.lock()
    mut list := t.h1_idle[conn.key] or { []&H1PooledConn{} }
    if list.len >= t.max_idle_conns_per_host {
        t.mu.unlock()
        conn.close_conn()
        return
    }
    list << conn
    t.h1_idle[conn.key] = list
    if t.max_idle_conns > 0 && t.total_idle_locked() > t.max_idle_conns {
        evicted = t.evict_oldest_locked()
    }
    t.mu.unlock()
    if evicted != unsafe { nil } {
        evicted.close_conn()
    }
}

// total_idle_locked sums the idle connections across all pool keys. The caller
// must hold t.mu.
fn (t &Transport) total_idle_locked() int {
    mut n := 0
    for _, list in t.h1_idle {
        n += list.len
    }
    return n
}

// evict_oldest_locked removes and returns the least-recently-used idle
// connection across all pool keys (nil if the pool is empty), pruning an
// emptied key. The caller must hold t.mu and close the returned connection.
fn (mut t Transport) evict_oldest_locked() &H1PooledConn {
    mut oldest_key := ''
    mut oldest_idx := -1
    mut oldest_since := time.now()
    for k, list in t.h1_idle {
        for i, c in list {
            if oldest_idx == -1 || c.idle_since < oldest_since {
                oldest_key = k
                oldest_idx = i
                oldest_since = c.idle_since
            }
        }
    }
    if oldest_idx == -1 {
        return &H1PooledConn(unsafe { nil })
    }
    mut list := t.h1_idle[oldest_key] or { return &H1PooledConn(unsafe { nil }) }
    victim := list[oldest_idx]
    list.delete(oldest_idx)
    if list.len == 0 {
        t.h1_idle.delete(oldest_key)
    } else {
        t.h1_idle[oldest_key] = list
    }
    return victim
}

// maybe_checkin pools the connection when both the read framing and the
// response (and request) headers allow reuse; otherwise it closes it.
fn (mut t Transport) maybe_checkin(mut conn H1PooledConn, header Header, reusable bool, resp &Response) {
    if reusable && !header.contains(.connection) && response_allows_reuse(resp) {
        t.checkin(mut conn)
        return
    }
    conn.close_conn()
}

// round_trip performs one HTTP request through the connection pool: it reuses
// an idle connection for the request's pool key when one exists (transparently
// retrying once on a connection that turned out to be stale), dials otherwise,
// and returns healthy connections to the pool afterwards.
fn (mut t Transport) round_trip(req &Request, method Method, scheme string, host string, port int, path string, data string, header Header) !Response {
    $if windows && !no_vschannel ? {
        if scheme == 'https' {
            // The SChannel backend keeps its proven one-shot path until SChannel
            // pooling lands; plain-http pooling below already works on Windows.
            return req.ssl_do(port, method, host, path, data, header)
        }
    }
    raw := req.build_request_headers_opts(method, host, port, path, data, header, false)
    $if trace_http_request ? {
        eprint('> ')
        eprint(raw)
        eprintln('')
    }
    key := transport_pool_key(req, scheme, host, port)
    // A stale pooled connection (closed by the server while idle) fails the
    // exchange; drain through the pool, then dial fresh.
    for _ in 0 .. t.max_idle_conns_per_host + 1 {
        mut conn := t.checkout(key)
        mut reused := true
        if conn == unsafe { nil } {
            reused = false
            if scheme == 'https' {
                return t.tls_fresh_round_trip(req, key, raw, method, host, port, path, data, header)
            }
            conn = t.dial_h1_tcp(req, key, host, port)!
        } else {
            conn.refresh_timeouts(req)
        }
        resp, reusable := conn.exchange(req, raw) or {
            conn.close_conn()
            // The write helpers do partial writes and cannot report how many
            // bytes reached the server, so a failed exchange may have delivered
            // part of the request. Only idempotent methods are safe to replay;
            // a non-idempotent one is tagged so neither this loop nor the
            // caller's outer retry loop re-sends it (avoiding duplicate side
            // effects). A pre-write dial failure, by contrast, propagates before
            // this point and stays freely retryable.
            if !transport_is_idempotent(method) {
                return error_with_code(err.msg(), transport_err_unsafe_retry)
            }
            if reused {
                // Drain stale pooled connections, then fall through to a fresh dial.
                continue
            }
            return err
        }
        t.maybe_checkin(mut conn, header, reusable, resp)
        return resp
    }
    return error('http.transport: request failed after retrying on a fresh connection')
}

// dial_h1_tcp opens a fresh plain-TCP connection for `key`.
fn (mut t Transport) dial_h1_tcp(req &Request, key string, host string, port int) !&H1PooledConn {
    mut client := net.dial_tcp('${host}:${port}')!
    client.set_read_timeout(req.read_timeout)
    client.set_write_timeout(req.write_timeout)
    return &H1PooledConn{
        key: key
        tcp: client
    }
}

// tls_fresh_round_trip dials a fresh TLS connection (advertising ALPN h2 when
// HTTP/2 is enabled) and completes the request on it. When the server selects
// h2, the request runs on the existing one-shot HTTP/2 driver and the
// connection is not pooled yet; an http/1.1 connection is pooled afterwards
// like any other.
fn (mut t Transport) tls_fresh_round_trip(req &Request, key string, raw string, method Method, host string, port int, path string, data string, header Header) !Response {
    alpn := if req.enable_http2 { ['h2', 'http/1.1'] } else { []string{} }
    mut ssl_conn := ssl.new_ssl_conn(
        verify:                 req.verify
        cert:                   req.cert
        cert_key:               req.cert_key
        validate:               req.validate
        in_memory_verification: req.in_memory_verification
        alpn_protocols:         alpn
    )!
    ssl_conn.dial(host, port)!
    if req.read_timeout > 0 {
        ssl_conn.set_read_timeout(req.read_timeout)
    }
    if req.enable_http2 && ssl_conn.negotiated_alpn() == 'h2' {
        return req.h2_do(mut ssl_conn, method, host, port, path, data, header)
    }
    mut conn := &H1PooledConn{
        key: key
        ssl: ssl_conn
    }
    resp, reusable := conn.exchange(req, raw) or {
        conn.close_conn()
        // Past the TLS handshake the request bytes may have been (partially)
        // written; a non-idempotent method must not be replayed by the outer
        // retry loop. (A dial/handshake failure above propagates before this and
        // stays retryable, since no request byte was sent.)
        if !transport_is_idempotent(method) {
            return error_with_code(err.msg(), transport_err_unsafe_retry)
        }
        return err
    }
    t.maybe_checkin(mut conn, header, reusable, resp)
    return resp
}