Add :max-size option to byte-array-store

ring-core/src/ring/middleware/multipart_params/byte_array.clj

@@ -1,17 +1,68 @@
 (ns ring.middleware.multipart-params.byte-array
   "A multipart storage engine for storing uploads as in-memory byte arrays."
-  (:import [java.io InputStream]
+  (:import [java.io InputStream OutputStream]
+           [java.lang IndexOutOfBoundsException System]
            [org.apache.commons.io IOUtils]))
 
+(def ^:private ^:const byte-array-type (type (byte-array 0)))
+
+(defn- bounded-buffer
+  "Creates a ByteArrayOutputStream-like buffer which enters a failure state
+  if it is supplied with more than max-size bytes of data. It returns a
+  vector containing the buffer, which subclasses OutputStream, and a function
+  which collects the data in the buffer into a single byte array. If too much
+  data is provided to the buffer, then the collect function will return
+  :ring.middleware.multipart-params.byte-array/too-large."
+  [max-size]
+  ;; State consists of total length and a list of [arr i j n] vectors, where
+  ;; arr is an array, [i, i + n) is a range in the array, and j is the
+  ;; corresponding position in the stream.
+  (let [state (atom [0 []])]
+    [(proxy [OutputStream] []
+       (write [v & more]
+         (let [[arr i n] (if (= (type v) byte-array-type)
+                           (if (empty? more)
+                             [v 0 (count v)] ; write(byte[])
+                             (vec (conj more v))) ; write(byte[], int, int)
+                           [(byte-array [v]) 0 1])] ; write(int)
+           (if (> n (- (count arr) i))
+             (throw (IndexOutOfBoundsException.)))
+           (swap! state (fn [[j arrs]]
+                          (let [next (+ j n)]
+                            (if (> next max-size)
+                              nil
+                              [next (conj arrs [arr i j n])])))))))
+     (fn []
+       (let [s @state]
+         (if (nil? s)
+           ::too-large
+           (let [[len arrs] s
+                 out (byte-array len)
+                 idx (atom 0)]
+             (run! (fn [[arr i j n]]
+                     (System/arraycopy arr i out j n))
+                   arrs)
+             out))))]))
+
 (defn byte-array-store
   "Returns a function that stores multipart file parameters as an array of
-  bytes. The multipart parameters will be stored as maps with the following
-  keys:
+  bytes. Accepts the following options:
+
+  :max-size - maximum size of files to accept, in bytes
+
+  The multipart parameters will be stored as maps with the following keys:
 
   :filename     - the name of the uploaded file
   :content-type - the content type of the uploaded file
   :bytes        - an array of bytes containing the uploaded content"
-  []
-  (fn [item]
-    (-> (select-keys item [:filename :content-type])
-        (assoc :bytes (IOUtils/toByteArray ^InputStream (:stream item))))))
+  {:arglists '([] [options])}
+  ([] (byte-array-store {}))
+  ([options]
+   (fn [item]
+     (-> (select-keys item [:filename :content-type])
+         (assoc :bytes (if (contains? options :max-size)
+                         (let [[s f] (bounded-buffer (:max-size options))]
+                           (.transferTo ^InputStream (:stream item) s)
+                           (f))
+                         (IOUtils/toByteArray
+                          ^InputStream (:stream item))))))))

ring-core/test/ring/middleware/multipart_params/test/byte_array.clj

@@ -12,3 +12,24 @@
     (is (= (:filename result) "foo.txt"))
     (is (= (:content-type result) "text/plain"))
     (is (= (String. (:bytes result)) "foo"))))
+
+(deftest test-byte-array-store-max-size-success
+  (let [store  (byte-array-store {:max-size 100})
+        result (store
+                {:filename "foo.txt"
+                 :content-type "text/plain"
+                 :stream (string-input-stream "foo")})]
+    (is (= (:filename result) "foo.txt"))
+    (is (= (:content-type result) "text/plain"))
+    (is (= (String. (:bytes result)) "foo"))))
+
+(deftest test-byte-array-store-max-size-failure
+  (let [store  (byte-array-store {:max-size 1})
+        result (store
+                {:filename "foo.txt"
+                 :content-type "text/plain"
+                 :stream (string-input-stream "foo")})]
+    (is (= (:filename result) "foo.txt"))
+    (is (= (:content-type result) "text/plain"))
+    (is (= (:bytes result)
+           :ring.middleware.multipart-params.byte-array/too-large))))