// 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

import io
import net
import time
import net.mbedtls

const tls_accept_poll_timeout = 100 * time.millisecond

// tls_handshake_timeout is the default value for Server.tls_handshake_timeout,
// used as the fallback handshake budget when Server.accept_timeout is zero or
// net.infinite_timeout. The handshake now runs on a worker thread (the accept
// thread only does the raw TCP accept, polled at tls_accept_poll_timeout, so
// stop() is observed promptly regardless of handshakes); without a finite bound a
// client that completes the TCP connect and then stalls mid-handshake would tie
// up a worker indefinitely.
const tls_handshake_timeout = 30 * time.second

fn tls_accept_timeouts(accept_timeout time.Duration, handshake_fallback time.Duration) (time.Duration, time.Duration) {
	// A finite `accept_timeout` doubles as the handshake budget; when it is
	// zero or net.infinite_timeout (i64.max), fall back to `handshake_fallback`
	// so the handshake still times out and shutdown stays responsive.
	// net.infinite_timeout is positive (i64.max), so the > 0 check alone is
	// not enough — mbedtls's ssl_timeout_deadline treats it as an infinite
	// deadline exactly like 0 or negative values.
	is_finite := accept_timeout > 0 && accept_timeout != net.infinite_timeout
	handshake_timeout := if is_finite { accept_timeout } else { handshake_fallback }
	accept_poll_timeout := if is_finite && accept_timeout < tls_accept_poll_timeout {
		accept_timeout
	} else {
		tls_accept_poll_timeout
	}
	return accept_poll_timeout, handshake_timeout
}

// This file implements TLS termination for net.http.Server on top of the
// mbedtls SSL listener. It is gated to the default TLS backend; the matching
// `server_tls_d_use_openssl.v` provides a clear-error stub when the project is
// built with `-d use_openssl`.

// listen_and_serve_tls is the TLS counterpart of listen_and_serve. It is
// dispatched to by listen_and_serve when `s.cert` and `s.cert_key` are set.
fn (mut s Server) listen_and_serve_tls() {
	// Pick a default port that's distinct from the plain-HTTP default if the
	// user hasn't overridden it.
	addr := if s.addr == '' || s.addr == ':${default_server_port}' {
		':${default_https_server_port}'
	} else {
		s.addr
	}

	// When HTTP/2 is enabled, advertise ALPN `h2, http/1.1` on the listener.
	// Clients that select `h2` are dispatched to the HTTP/2 driver after the
	// handshake; clients that select `http/1.1` (or send no ALPN extension)
	// keep the existing HTTP/1.1 worker path.
	alpn := if s.enable_http2 { ['h2', 'http/1.1'] } else { []string{} }
	mut listener := mbedtls.new_ssl_listener(addr, mbedtls.SSLConnectConfig{
		cert:                   s.cert
		cert_key:               s.cert_key
		in_memory_verification: s.in_memory_verification
		validate:               false // accept any client; servers don't verify clients by default
		read_timeout:           s.read_timeout
		alpn_protocols:         alpn
	}) or {
		eprintln('Listening TLS on ${addr} failed, err: ${err}')
		return
	}
	defer {
		listener.shutdown() or {}
		if s.state == .stopped {
			s.state = .closed
			if s.on_closed != unsafe { nil } {
				s.on_closed(mut s)
			}
		}
	}
	s.addr = addr

	accept_poll_timeout, handshake_timeout := tls_accept_timeouts(s.accept_timeout,
		s.tls_handshake_timeout)
	ch := chan &mbedtls.SSLConn{cap: s.pool_channel_slots}
	mut idle_conns := &TlsIdleConnTracker{}
	mut ws := []thread{cap: s.worker_num}
	for wid in 0 .. s.worker_num {
		ws << new_tls_handler_worker(wid, ch, s.handler, s.max_keep_alive_requests,
			handshake_timeout, idle_conns)
	}

	if s.show_startup_message {
		println('Listening on https://${s.addr}/')
		flush_stdout()
	}

	time.sleep(20 * time.millisecond)
	s.state = .running
	if s.on_running != unsafe { nil } {
		s.on_running(mut s)
	}
	for s.state == .running {
		mut conn := listener.accept_raw_with_timeout(accept_poll_timeout) or {
			if s.state != .running {
				break
			}
			if err.code() == net.err_timed_out_code {
				continue
			}
			$if debug {
				eprintln('TLS accept failed: ${err}; skipping')
			}
			continue
		}
		// Hand the raw (not-yet-handshaked) conn to a worker. Don't use a bare
		// blocking `ch <- conn`: the accept loop, this send, and the post-loop
		// shutdown (ch.close + close_idle + ws.wait) all run on this one thread,
		// while stop()/close() only flip s.state from another thread. Under a slow-
		// handshake flood -- every worker blocked in complete_handshake and the
		// channel buffer full of untracked conns -- a blocking send would wedge the
		// accept thread so it never re-checks s.state nor reaches close_idle(),
		// delaying shutdown until a worker handshake times out. Poll s.state via a
		// select timeout so shutdown is still observed promptly.
		mut queued := false
		for s.state == .running && !queued {
			select {
				ch <- conn {
					queued = true
				}
				accept_poll_timeout {
					// channel full; loop re-checks s.state
				}
			}
		}
		if !queued {
			// Shutting down before this conn could be handed off. It was never
			// mark_idle'd, so close_idle() won't see it; close it here (exactly
			// once -- no worker ever received it) to avoid leaking the fd.
			conn.shutdown() or {}
		}
	}
	ch.close()
	idle_conns.close_idle()
	ws.wait()
}

// TlsHandlerWorker serves HTTP/1.1 requests on TLS-wrapped connections.
struct TlsHandlerWorker {
	id                      int
	ch                      chan &mbedtls.SSLConn
	max_keep_alive_requests int
	handshake_timeout       time.Duration
mut:
	idle_conns &TlsIdleConnTracker = unsafe { nil }
pub mut:
	handler Handler
}

fn new_tls_handler_worker(wid int, ch chan &mbedtls.SSLConn, handler Handler, max_keep_alive_requests int, handshake_timeout time.Duration, idle_conns &TlsIdleConnTracker) thread {
	mut w := &TlsHandlerWorker{
		id:                      wid
		ch:                      ch
		handler:                 handler
		max_keep_alive_requests: max_keep_alive_requests
		handshake_timeout:       handshake_timeout
		idle_conns:              idle_conns
	}
	return spawn w.process_requests()
}

fn (mut w TlsHandlerWorker) process_requests() {
	for {
		mut conn := <-w.ch or { break }
		w.handle_conn(mut conn)
	}
}

fn (mut w TlsHandlerWorker) handle_conn(mut conn mbedtls.SSLConn) {
	// For H2 connections, serve_h2_conn_with_idle_tracker's serve() owns the
	// mark_idle/unmark_idle lifetime (it marks idle once and unmarks in its
	// own defer). Calling unmark_idle here a second time would race: after
	// serve()'s defer fires the OS can recycle conn.handle for a new
	// connection that has already called mark_idle, and this stale unmark
	// would silently evict it, preventing close_idle from ever shutting it
	// down and leaking the reader goroutine.
	mut is_h2 := false
	defer {
		if !is_h2 {
			w.idle_conns.unmark_idle(conn.handle)
		}
		// close_idle may have already closed conn.handle on Windows (a forced
		// closesocket to wake this worker from a blocked read). conn.shutdown()
		// both frees the mbedtls TLS resources (SSL ctx, config, certs, RNG,
		// ALPN alloc) AND closes the socket; we must always do the former but
		// must not close the socket a second time — that would race process-wide
		// SOCKET reuse and could close an unrelated socket. So relinquish socket
		// ownership and still call shutdown for the TLS cleanup. On non-Windows
		// was_force_closed is always false and the worker remains the sole closer.
		if w.idle_conns.was_force_closed(conn.handle) {
			conn.owns_socket = false
		}
		conn.shutdown() or {}
	}
	// Run the TLS handshake here on the worker thread (the accept thread only does
	// the raw TCP accept), so handshakes proceed in parallel up to worker_num and a
	// client that stalls mid-handshake can't wedge the accept loop or delay stop().
	// mark_idle first so close_idle can interrupt a stalled handshake by force-
	// closing the fd, and so a conn handed off while the server is shutting down is
	// closed (mark_idle returns false) rather than handshaked. On success, unmark
	// before the serve path does its own idle marking, keeping each handle tracked
	// once. On any early return the defer above performs the single shutdown.
	if !w.idle_conns.mark_idle(conn.handle) {
		return
	}
	conn.complete_handshake(w.handshake_timeout) or {
		$if debug {
			eprintln('TLS handshake failed: ${err}')
		}
		return
	}
	w.idle_conns.unmark_idle(conn.handle)
	// If the TLS handshake negotiated HTTP/2 via ALPN, switch to the HTTP/2
	// driver; otherwise fall through to the existing HTTP/1.1 path unchanged.
	if conn.negotiated_alpn() == 'h2' {
		is_h2 = true
		serve_h2_conn_with_idle_tracker(mut conn, mut w.handler, w.idle_conns, conn.handle) or {
			$if debug {
				eprintln('h2 server error: ${err}')
			}
		}
		return
	}
	mut reader := io.new_buffered_reader(reader: conn)
	defer {
		unsafe {
			reader.free()
		}
	}

	mut request_count := 0
	for {
		if !w.idle_conns.mark_idle(conn.handle) {
			return
		}
		mut req := parse_request(mut reader) or {
			if err !is io.Eof {
				$if debug {
					eprintln('error parsing TLS request: ${err}')
				}
			}
			return
		}
		w.idle_conns.unmark_idle(conn.handle)
		request_count++
		// `conn.ip` is the peer's IPv4 address as populated by mbedtls'
		// accept(); blank for IPv6, which is acceptable for keep-alive logic.
		if conn.ip != '' {
			req.header.add_custom('Remote-Addr', conn.ip) or {}
		}

		mut resp := w.handler.handle(req)
		normalize_server_response(mut resp, req)

		if !resp.header.contains(.content_length) {
			resp.header.set(.content_length, '${resp.body.len}')
		}

		max_reached := w.max_keep_alive_requests > 0 && request_count >= w.max_keep_alive_requests
		req_conn := (req.header.get(.connection) or { '' }).to_lower()
		resp_conn := (resp.header.get(.connection) or { '' }).to_lower()
		keep_alive := if max_reached {
			false
		} else if resp_conn == 'close' {
			false
		} else if resp_conn == 'keep-alive' {
			true
		} else if req_conn == 'close' {
			false
		} else if req_conn == 'keep-alive' {
			true
		} else {
			req.version == .v1_1
		}
		if max_reached || !resp.header.contains(.connection) {
			if keep_alive {
				resp.header.set(.connection, 'keep-alive')
			} else {
				resp.header.set(.connection, 'close')
			}
		}

		conn.write(resp.bytes()) or {
			$if debug {
				eprintln('error sending TLS response: ${err}')
			}
			return
		}

		if !keep_alive {
			return
		}
	}
}