httpzo / lib / range.ml

(* range.ml - Range request parsing and Content-Range response writing per RFC 7233 *)

open Base

(* Byte range specification *)
type byte_range =
  { kind : int      (* 0=Range, 1=Suffix, 2=Open *)
  ; start : int64
  ; end_ : int64
  }

(* Kind constants - internal *)
let kind_range = 0
let kind_suffix = 1
let kind_open = 2

(* Maximum ranges to parse *)
let max_ranges = 16

let empty = { kind = 0; start = 0L; end_ = 0L }

(* Query functions *)
let is_range r = r.kind = kind_range
let is_suffix r = r.kind = kind_suffix
let is_open r = r.kind = kind_open

type parse_status =
  | Valid
  | Invalid

(* Resolved byte range *)
type resolved =
  { start : int64
  ; end_ : int64
  ; length : int64
  }

let empty_resolved = { start = 0L; end_ = 0L; length = 0L }

type eval_result =
  | Full_content
  | Single_range
  | Multiple_ranges
  | Not_satisfiable

(* Skip whitespace *)
let skip_ws buf ~pos ~len =
  let p = ref pos in
  while !p < len && (
    let c = Base_bigstring.unsafe_get buf !p in
    Char.equal c ' ' || Char.equal c '\t'
  ) do
    Int.incr p
  done;
  !p
;;

(* Parse a non-negative int64, returns (value, end_pos, valid) *)
let parse_int64 buf ~pos ~len =
  let start = pos in
  let p = ref pos in
  let acc = ref 0L in
  let valid = ref true in
  while !valid && !p < len do
    let c = Base_bigstring.unsafe_get buf !p in
    if Char.is_digit c then (
      let digit = Int64.of_int (Char.to_int c - 48) in
      acc := Int64.(!acc * 10L + digit);
      Int.incr p
    ) else
      valid := false
  done;
  if !p > start then (!acc, !p, true) else (0L, pos, false)
;;

(* Parse a single range-spec *)
let parse_range_spec buf ~pos ~len =
  let pos = skip_ws buf ~pos ~len in
  if pos >= len then (false, empty, pos)
  else
    let c = Base_bigstring.unsafe_get buf pos in
    if Char.equal c '-' then
      (* Suffix range: -500 *)
      let (suffix, end_pos, valid) = parse_int64 buf ~pos:(pos + 1) ~len in
      if (not valid) || Int64.(suffix = 0L) then
        (false, empty, end_pos)
      else
        (true, { kind = kind_suffix; start = suffix; end_ = 0L }, end_pos)
    else
      (* Start-end or start- *)
      let (start, after_start, valid) = parse_int64 buf ~pos ~len in
      if not valid then (false, empty, after_start)
      else if after_start >= len then (false, empty, after_start)
      else if not (Char.equal (Base_bigstring.unsafe_get buf after_start) '-') then
        (false, empty, after_start)
      else
        let after_dash = after_start + 1 in
        if after_dash >= len || (
          let c = Base_bigstring.unsafe_get buf after_dash in
          Char.equal c ',' || Char.equal c ' ' || Char.equal c '\t'
        ) then
          (* Open range: start- *)
          (true, { kind = kind_open; start; end_ = 0L }, after_dash)
        else
          (* Closed range: start-end *)
          let (end_val, end_pos, end_valid) = parse_int64 buf ~pos:after_dash ~len in
          if (not end_valid) || Int64.(end_val < start) then
            (false, empty, end_pos)
          else
            (true, { kind = kind_range; start; end_ = end_val }, end_pos)
;;

(* Parse Range header into array - internal implementation working on buffer region *)
let parse_region buf ~off ~len (ranges : byte_range array) =
  let end_pos = off + len in
  (* Look for "=" to split unit and ranges *)
  let eq_pos = ref off in
  while !eq_pos < end_pos && not (Char.equal (Base_bigstring.unsafe_get buf !eq_pos) '=') do
    Int.incr eq_pos
  done;
  if !eq_pos >= end_pos then (Invalid, 0)
  else
    (* Check for "bytes" unit *)
    let unit_len = !eq_pos - off in
    if unit_len <> 5 then (Invalid, 0)
    else
      let is_bytes =
        let c0 = Base_bigstring.unsafe_get buf off in
        let c1 = Base_bigstring.unsafe_get buf (off + 1) in
        let c2 = Base_bigstring.unsafe_get buf (off + 2) in
        let c3 = Base_bigstring.unsafe_get buf (off + 3) in
        let c4 = Base_bigstring.unsafe_get buf (off + 4) in
        (Char.equal c0 'b' || Char.equal c0 'B') &&
        (Char.equal c1 'y' || Char.equal c1 'Y') &&
        (Char.equal c2 't' || Char.equal c2 'T') &&
        (Char.equal c3 'e' || Char.equal c3 'E') &&
        (Char.equal c4 's' || Char.equal c4 'S')
      in
      if not is_bytes then (Invalid, 0)
      else
        (* Parse comma-separated range specs *)
        let pos = ref (!eq_pos + 1) in
        let count = ref 0 in
        let valid = ref true in
        while !valid && !pos < end_pos && !count < max_ranges do
          pos := skip_ws buf ~pos:!pos ~len:end_pos;
          if !pos >= end_pos then
            valid := false
          else
            let (ok, range, after_range) = parse_range_spec buf ~pos:!pos ~len:end_pos in
            if ok then (
              Array.unsafe_set ranges !count range;
              Int.incr count
            );
            pos := skip_ws buf ~pos:after_range ~len:end_pos;
            if !pos < end_pos then
              if Char.equal (Base_bigstring.unsafe_get buf !pos) ',' then
                Int.incr pos
              else
                valid := false
        done;
        if !count > 0 then (Valid, !count) else (Invalid, 0)
;;

(* Parse Range header into array - from buffer and span *)
let parse buf (sp : Span.t) (ranges : byte_range array) =
  parse_region buf ~off:(Span.off sp) ~len:(Span.len sp) ranges
;;

(* Parse Range header from string - creates local buffer *)
let parse_string (s : string) (ranges : byte_range array) =
  let len = String.length s in
  let buf = Base_bigstring.create len in
  for i = 0 to len - 1 do
    Base_bigstring.unsafe_set buf i (String.unsafe_get s i)
  done;
  (* Bind result to prevent tail call - local buf must stay in scope *)
  let (status, count) = parse_region buf ~off:0 ~len ranges in
  (status, count)
;;

(* Resolve a single range *)
let resolve_range range ~resource_length =
  let res_len = resource_length in
  if Int64.(res_len <= 0L) then (false, empty_resolved)
  else
    let kind = range.kind in
    let start_val = range.start in
    let end_val = range.end_ in
    if kind = kind_range then
      (* Range: start-end *)
      if Int64.(start_val >= res_len) then (false, empty_resolved)
      else
        let end_clamped = Int64.min end_val Int64.(res_len - 1L) in
        let length = Int64.(end_clamped - start_val + 1L) in
        (true, { start = start_val; end_ = end_clamped; length })
    else if kind = kind_suffix then
      (* Suffix: -N (last N bytes) *)
      let suffix = start_val in  (* stored in start field *)
      if Int64.(suffix <= 0L) then (false, empty_resolved)
      else
        let start = Int64.max 0L Int64.(res_len - suffix) in
        let end_ = Int64.(res_len - 1L) in
        let length = Int64.(end_ - start + 1L) in
        (true, { start; end_; length })
    else
      (* Open: start- *)
      if Int64.(start_val >= res_len) then (false, empty_resolved)
      else
        let end_ = Int64.(res_len - 1L) in
        let length = Int64.(end_ - start_val + 1L) in
        (true, { start = start_val; end_; length })
;;

(* Evaluate ranges *)
let evaluate (ranges : byte_range array) ~count ~resource_length (out : resolved array) =
  if count = 0 then (Full_content, 0)
  else
    let resolved_count = ref 0 in
    for i = 0 to count - 1 do
      let (valid, r) = resolve_range (Array.unsafe_get ranges i) ~resource_length in
      if valid then (
        Array.unsafe_set out !resolved_count r;
        Int.incr resolved_count
      )
    done;
    if !resolved_count = 0 then (Not_satisfiable, 0)
    else if !resolved_count = 1 then (Single_range, 1)
    else (Multiple_ranges, !resolved_count)
;;

(* Response writing helpers *)

let write_accept_ranges dst ~off =
  let off = Buf_write.string dst ~off "Accept-Ranges: bytes" in
  Buf_write.crlf dst ~off
;;

let write_accept_ranges_none dst ~off =
  let off = Buf_write.string dst ~off "Accept-Ranges: none" in
  Buf_write.crlf dst ~off
;;

(* Content-Range: bytes start-end/total *)
let write_content_range dst ~off ~start ~end_ ~total =
  let off = Buf_write.string dst ~off "Content-Range: bytes " in
  let off = Buf_write.int64 dst ~off start in
  let off = Buf_write.char dst ~off '-' in
  let off = Buf_write.int64 dst ~off end_ in
  let off = Buf_write.char dst ~off '/' in
  let off = Buf_write.int64 dst ~off total in
  Buf_write.crlf dst ~off
;;

let write_content_range_resolved dst ~off r ~total =
  write_content_range dst ~off ~start:r.start ~end_:r.end_ ~total
;;

(* Content-Range: bytes * /total (for 416 responses) *)
let write_content_range_unsatisfiable dst ~off ~total =
  let off = Buf_write.string dst ~off "Content-Range: bytes */" in
  let off = Buf_write.int64 dst ~off total in
  Buf_write.crlf dst ~off
;;

(* Multipart helpers *)

let write_multipart_boundary dst ~off ~boundary =
  let off = Buf_write.string dst ~off "--" in
  let off = Buf_write.string dst ~off boundary in
  Buf_write.crlf dst ~off
;;

let write_multipart_final dst ~off ~boundary =
  let off = Buf_write.string dst ~off "--" in
  let off = Buf_write.string dst ~off boundary in
  let off = Buf_write.string dst ~off "--" in
  Buf_write.crlf dst ~off
;;

(* Generate a random boundary string *)
let generate_boundary () =
  let chars = "0123456789abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ" in
  let len = 24 in
  let buf = Bytes.create len in
  for i = 0 to len - 1 do
    let idx = Random.int (String.length chars) in
    Bytes.set buf i (String.get chars idx)
  done;
  Bytes.to_string buf
;;