From 7b59670d6032500c8633fa9ddc5a72130696cf44 Mon Sep 17 00:00:00 2001
From: Simon Cruanes <simon.cruanes.2007@m4x.org>
Date: Sun, 17 Nov 2019 16:52:42 -0600
Subject: [PATCH] feat(bin): accept `PUT` method for uploading up to max size

---
 src/bin/http_of_dir.ml | 64 +++++++++++++++++++++++++++++++++++-------
 1 file changed, 54 insertions(+), 10 deletions(-)

diff --git a/src/bin/http_of_dir.ml b/src/bin/http_of_dir.ml
index 93350e87..92a1e75e 100644
--- a/src/bin/http_of_dir.ml
+++ b/src/bin/http_of_dir.ml
@@ -2,6 +2,20 @@
 module S = Tiny_httpd
 module Pf = Printf
 
+type config = {
+  mutable addr: string;
+  mutable port: int;
+  mutable upload: bool;
+  mutable max_upload_size: int;
+}
+
+let default_config () : config = {
+  addr="127.0.0.1";
+  port=8080;
+  upload=true;
+  max_upload_size = 10 * 1024 * 1024;
+}
+
 let contains_dot_dot s =
   try
     String.iteri
@@ -38,10 +52,37 @@ let same_path a b =
   Filename.dirname a = Filename.dirname b &&
   Filename.basename a = Filename.basename b
 
-let serve ~addr ~port (dir:string) : _ result = 
-  let server = S.create ~addr ~port () in
+let serve ~config (dir:string) : _ result = 
+  let server = S.create ~addr:config.addr ~port:config.port () in
+  if config.upload then (
+    S.add_path_handler server ~meth:`PUT "/%s"
+      ~accept:(fun req ->
+          match S.Request.get_header_int req "Content-Length" with
+          | Some n when n > config.max_upload_size ->
+            Error (403, "max upload size is " ^ string_of_int config.max_upload_size)
+          | Some _ when contains_dot_dot req.S.Request.path ->
+            Error (403, "invalid path (contains '..')")
+          | Some _ -> Ok ()
+          | None ->
+            Error (403, "must know size before hand: max upload size is " ^
+                        string_of_int config.max_upload_size)
+        )
+      (fun path req ->
+         let fpath = Filename.concat dir path in
+         let oc =
+           try open_out fpath
+           with e ->
+             S.Response.fail_raise ~code:403 "cannot upload to %S: %s"
+               path (Printexc.to_string e)
+         in
+         output_string oc req.S.Request.body;
+         flush oc;
+         close_out oc;
+         S.Response.make (Ok "upload successful")
+      )
+  );
   S.add_path_handler server ~meth:`GET "/%s"
-    (fun _req path () ->
+    (fun path _req ->
        let f = Filename.concat dir path in
        if contains_dot_dot f then (
          S.Response.fail ~code:403 "Path is forbidden";
@@ -65,17 +106,20 @@ let serve ~addr ~port (dir:string) : _ result =
 
 
 let main () =
-  let addr_ = ref "127.0.0.1" in
-  let port_ = ref 8080 in
+  let config = default_config () in
   let dir_ = ref "." in
   Arg.parse (Arg.align [
-      "--addr", Set_string addr_, " address to listen on";
-      "-a", Set_string addr_, " alias to --listen";
-      "--port", Set_int port_, " port to listen on";
-      "-p", Set_int port_, " alias to --port";
+      "--addr", String (fun s -> config.addr <- s), " address to listen on";
+      "-a", String (fun s -> config.addr <- s), " alias to --listen";
+      "--port", Int (fun x -> config.port <- x), " port to listen on";
+      "-p", Int (fun x -> config.port <- x), " alias to --port";
       "--dir", Set_string dir_, " directory to serve (default: \".\")";
+      "--no-upload", Unit (fun () -> config.upload <- false), " disable file uploading";
+      "--max-upload", Int (fun i -> config.max_upload_size <- 1024 * 1024 * i),
+      "maximum size of files that can be uploaded, in MB";
+      "--debug", Unit (fun () -> S._enable_debug true), " debug mode";
     ]) (fun _ -> raise (Arg.Bad "no positional arguments")) "http_of_dir [options]";
-  match serve ~addr:!addr_ ~port:!port_ !dir_ with
+  match serve ~config !dir_ with
   | Ok () -> ()
   | Error e ->
     raise e