diff --git a/bin/main.ml b/bin/main.ml index 62cab54..ff66aa3 100644 --- a/bin/main.ml +++ b/bin/main.ml @@ -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 @@ -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) @@ -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) diff --git a/manifest/extism_manifest.ml b/manifest/extism_manifest.ml index a217a1b..7a01b7b 100644 --- a/manifest/extism_manifest.ml +++ b/manifest/extism_manifest.ml @@ -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 @@ -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") ] diff --git a/manifest/extism_manifest.mli b/manifest/extism_manifest.mli index b766147..debae00 100644 --- a/manifest/extism_manifest.mli +++ b/manifest/extism_manifest.mli @@ -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 *) @@ -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 *)