From 49c03b167e060862df569b12fc6222b4b0221bfe Mon Sep 17 00:00:00 2001
From: Swrup <swrup@protonmail.com>
Date: Mon, 4 Apr 2022 10:31:08 +0200
Subject: [PATCH] check image mime type

---
 src/babillard.ml            | 13 -------------
 src/babillard_page.eml.html |  2 +-
 src/db.ml                   | 32 ++++++++++++++++++++++++++++----
 src/dune                    |  3 +++
 src/user.ml                 |  9 ++++-----
 5 files changed, 36 insertions(+), 23 deletions(-)

diff --git a/src/babillard.ml b/src/babillard.ml
index b38df76..6ccb0a0 100644
--- a/src/babillard.ml
+++ b/src/babillard.ml
@@ -273,19 +273,6 @@ let () =
       (Array.map (fun query -> Db.exec query ()) tables)
   then Dream.error (fun log -> log "can't create babillard's tables")
 
-let clean_image image =
-  let name, alt, content = image in
-  let name =
-    match name with
-    | Some name -> Dream.html_escape name
-    | None ->
-      (* make up random name if no name was given *)
-      Uuidm.to_string (Uuidm.v4_gen random_state ())
-  in
-  if not (is_valid_image content) then Error "invalid image"
-  else if String.length alt > 1000 then Error "Image description too long"
-  else Ok (name, alt, content)
-
 let make_thumbnail content =
   let open Bos in
   (* jpp *)
diff --git a/src/babillard_page.eml.html b/src/babillard_page.eml.html
index 2589689..b5cee28 100644
--- a/src/babillard_page.eml.html
+++ b/src/babillard_page.eml.html
@@ -43,7 +43,7 @@ let f request =
           <input name="tags" type="text" class="form-control" id="tags" aria-labelledby="tags-label"></input>
           <br />
           <label for="file" id="file-label" class="form-label">Picture:</label>
-          <input id="file" name="file" aria-describedby="file-label" type="file" accept="image/*">
+          <input id="file" name="file" aria-describedby="file-label" type="file" accept="image/png,image/jpeg,image/webp">
           <br />
           <label for="alt" id="alt-label" class="form-label">Image description:</label>
           <input name="alt" type="text" class="form-control" id="alt" aria-labelledby="alt-label"></input>
diff --git a/src/db.ml b/src/db.ml
index e43af41..f263afd 100644
--- a/src/db.ml
+++ b/src/db.ml
@@ -23,10 +23,6 @@ let () =
   if Result.is_error (Db.exec set_foreign_keys_on ()) then
     Dream.error (fun log -> log "can't set foreign_keys on")
 
-(* TODO do image validation: length and MIME types with conan*)
-(* TODO do the same for text input: check length, forbidden chars and have a forbidden words filter*)
-let is_valid_image _content = true
-
 let () =
   let query =
     Caqti_request.exec Caqti_type.unit
@@ -38,3 +34,31 @@ let () =
   | Error _e ->
     Format.eprintf "db error@\n";
     exit 1
+
+let mime_database = Conan.Process.database ~tree:Conan_light.tree
+
+let mime contents =
+  match Conan_string.run ~database:mime_database contents with
+  | Ok m -> Conan.Metadata.mime m
+  | Error _ -> None
+
+let clean_image image =
+  let name, alt, content = image in
+  let name =
+    match name with
+    | Some name -> Dream.html_escape name
+    | None ->
+      (* make up random name if no name was given *)
+      Uuidm.to_string (Uuidm.v4_gen random_state ())
+  in
+  if String.length name > 1000 then Error "Image name too long"
+  else if String.length alt > 1000 then Error "Image description too long"
+  else if String.length content > 4200000 then Error "Image size too big"
+  else
+    match mime content with
+    | None -> Error "invalid image type"
+    | Some mime -> (
+      match mime with
+      | "image/jpeg" | "image/png" | "image/webp" -> Ok (name, alt, content)
+      | _unsupported_mime_type ->
+        Error (Format.sprintf "unsupported image type: %s" mime) )
diff --git a/src/dune b/src/dune
index 27ae89a..600b3fb 100644
--- a/src/dune
+++ b/src/dune
@@ -26,6 +26,9 @@
   caqti
   caqti.blocking
   caqti-driver-sqlite3
+  conan
+  conan.string
+  conan-database.light
   directories
   dream
   emile
diff --git a/src/user.ml b/src/user.ml
index b29381a..83fabbc 100644
--- a/src/user.ml
+++ b/src/user.ml
@@ -255,11 +255,10 @@ let get_avatar user_id =
 let upload_avatar files user_id =
   match files with
   | [] -> Error "No file provided"
-  | [ (_, content) ] ->
-    if not (is_valid_image content) then Error "Invalid image"
-    else
-      let^ () = Db.exec Q.upload_avatar (content, user_id) in
-      Ok ()
+  | [ (name_opt, content) ] ->
+    let* _name, _alt, content = clean_image (name_opt, "avatar", content) in
+    let^ () = Db.exec Q.upload_avatar (content, user_id) in
+    Ok ()
   | _files -> Error "More than one file provided"
 
 let is_admin user_id =