Skip to content

Commit

Permalink
feat: allow max HTTP response size to be configured in the manifest (#16
Browse files Browse the repository at this point in the history
)
  • Loading branch information
zshipko authored Feb 8, 2024
1 parent 7713c8b commit b290563
Show file tree
Hide file tree
Showing 3 changed files with 46 additions and 10 deletions.
17 changes: 13 additions & 4 deletions bin/main.ml
Original file line number Diff line number Diff line change
Expand Up @@ -29,11 +29,13 @@ let print_timing ~time name f =
out

let main file func_name input loop timeout_ms allowed_paths allowed_hosts config
memory_max log_level log_file wasi stdin time =
memory_max http_max log_level log_file wasi stdin time =
let input = if stdin then read_stdin () else input in
let allowed_paths = split_allowed_paths allowed_paths in
let config = split_config config in
let memory = Manifest.{ max_pages = memory_max } in
let memory =
Manifest.{ max_pages = memory_max; max_http_response_bytes = http_max }
in
let manifest =
print_timing ~time "loaded manifest" @@ fun () ->
try
Expand Down Expand Up @@ -92,6 +94,13 @@ let memory_max =
let doc = "Max number of memory pages." in
Arg.(value & opt (some int) None & info [ "memory-max" ] ~docv:"PAGES" ~doc)

let http_max =
let doc = "Max number of bytes allowed in an HTTP response." in
Arg.(
value
& opt (some int) None
& info [ "http-response-max" ] ~docv:"BYTES" ~doc)

let timeout =
let doc = "Plugin timeout in milliseconds." in
Arg.(value & opt int 30000 & info [ "timeout"; "t" ] ~docv:"MILLIS" ~doc)
Expand Down Expand Up @@ -147,8 +156,8 @@ let time =
let main_t =
Term.(
const main $ file $ func_name $ input $ loop $ timeout $ allowed_paths
$ allowed_hosts $ config $ memory_max $ log_level $ log_file $ wasi $ stdin
$ time)
$ allowed_hosts $ config $ memory_max $ http_max $ log_level $ log_file
$ wasi $ stdin $ time)

let cmd = Cmd.v (Cmd.info "extism-call") main_t
let () = exit (Cmd.eval cmd)
27 changes: 24 additions & 3 deletions manifest/extism_manifest.ml
Original file line number Diff line number Diff line change
Expand Up @@ -5,7 +5,10 @@ type base64 = string
let yojson_of_base64 x = `String (Base64.encode_exn x)
let base64_of_yojson j = Yojson.Safe.Util.to_string j

type memory_options = { max_pages : int option [@yojson.option] }
type memory_options = {
max_pages : int option; [@yojson.option]
max_http_response_bytes : int option; [@yojson.option]
}
[@@deriving yojson]

type dict = (string * string) list
Expand Down Expand Up @@ -98,11 +101,29 @@ let of_file filename =
t_of_yojson j

let with_config config t = { t with config = Some config }
let with_memory_max max t = { t with memory = Some { max_pages = Some max } }

let with_memory_max max t =
match t.memory with
| None ->
{
t with
memory = Some { max_http_response_bytes = None; max_pages = Some max };
}
| Some m -> { t with memory = Some { m with max_pages = Some max } }

let with_http_response_max_bytes max t =
match t.memory with
| None ->
{
t with
memory = Some { max_http_response_bytes = Some max; max_pages = None };
}
| Some m ->
{ t with memory = Some { m with max_http_response_bytes = Some max } }

let%test "rountrip" =
let config = [ ("a", Some "b"); ("b", Some "c") ] in
let memory = { max_pages = Some 5 } in
let memory = { max_pages = Some 5; max_http_response_bytes = Some 9999 } in
let t =
create ~config ~memory ~allowed_hosts:[ "example.com" ]
~allowed_paths:[ ("a", "b") ]
Expand Down
12 changes: 9 additions & 3 deletions manifest/extism_manifest.mli
Original file line number Diff line number Diff line change
Expand Up @@ -25,8 +25,11 @@

type memory_options = {
max_pages : int option;
(** [max_pages] can be used to limit the total number of pages used by the
runtime *)
(** [max_pages] can be used to limit the total number of pages used by the
runtime *)
max_http_response_bytes: int option;
(** [max_http_response_bytes] can be used to limit the size of the response returned by
[extism_http_request] *)
}
[@@deriving yojson]
(** Memory options *)
Expand Down Expand Up @@ -128,4 +131,7 @@ val with_config : config -> t -> t
(** Returns a new {!t} with the [config] field updated *)

val with_memory_max : int -> t -> t
(** Returns a new {!t} with [memory.max_pages] updates *)
(** Returns a new {!t} with [memory.max_pages] updated *)

val with_http_response_max_bytes : int -> t -> t
(** Returns a new {!t} with [memory.max_http_response_bytes] updated *)

0 comments on commit b290563

Please sign in to comment.