Fix an overflow bug in String.concat (#833)

Fix an overflow bug in String.concat and make the function faster.

Changes

@@ -61,6 +61,9 @@ Next version (tbd):
 - GPR#810: check for integer overflow in Array.concat
   (Jeremy Yallop)
 
+- GPR#805, GPR#815, GPR#833: check for integer overflow in String.concat
+  (Jeremy Yallop,
+   review by Damien Doligez, Alain Frisch, Daniel Bünzli, Fabrice Le Fessant)
 
 OCaml 4.04.0:
 -------------

stdlib/bytes.ml

@@ -97,23 +97,28 @@ let iter f a =
 let iteri f a =
   for i = 0 to length a - 1 do f i (unsafe_get a i) done
 
-let concat sep l =
-  match l with
-    [] -> empty
+let ensure_ge x y = if x >= y then x else invalid_arg "Bytes.concat"
+
+let rec sum_lengths acc seplen = function
+  | [] -> acc
+  | hd :: [] -> length hd + acc
+  | hd :: tl -> sum_lengths (ensure_ge (length hd + seplen + acc) acc) seplen tl
+
+let rec unsafe_blits dst pos sep seplen = function
+    [] -> dst
+  | hd :: [] ->
+    unsafe_blit hd 0 dst pos (length hd); dst
   | hd :: tl ->
-      let num = ref 0 and len = ref 0 in
-      List.iter (fun s -> incr num; len := !len + length s) l;
-      let r = create (!len + length sep * (!num - 1)) in
-      unsafe_blit hd 0 r 0 (length hd);
-      let pos = ref(length hd) in
-      List.iter
-        (fun s ->
-          unsafe_blit sep 0 r !pos (length sep);
-          pos := !pos + length sep;
-          unsafe_blit s 0 r !pos (length s);
-          pos := !pos + length s)
-        tl;
-      r
+    unsafe_blit hd 0 dst pos (length hd);
+    unsafe_blit sep 0 dst (pos + length hd) seplen;
+    unsafe_blits dst (pos + length hd + seplen) sep seplen tl
+
+let concat sep = function
+    [] -> empty
+  | l -> let seplen = length sep in
+          unsafe_blits 
+            (create (sum_lengths 0 seplen l))
+            0 sep seplen l
 
 let cat s1 s2 =
   let l1 = length s1 in

stdlib/string.ml

@@ -44,23 +44,28 @@ let fill =
 let blit =
   B.blit_string
 
-let concat sep l =
-  match l with
-  | [] -> ""
+let ensure_ge x y = if x >= y then x else invalid_arg "String.concat"
+
+let rec sum_lengths acc seplen = function
+  | [] -> acc
+  | hd :: [] -> length hd + acc
+  | hd :: tl -> sum_lengths (ensure_ge (length hd + seplen + acc) acc) seplen tl
+
+let rec unsafe_blits dst pos sep seplen = function
+    [] -> dst
+  | hd :: [] ->
+    unsafe_blit hd 0 dst pos (length hd); dst
   | hd :: tl ->
-      let num = ref 0 and len = ref 0 in
-      List.iter (fun s -> incr num; len := !len + length s) l;
-      let r = B.create (!len + length sep * (!num - 1)) in
-      unsafe_blit hd 0 r 0 (length hd);
-      let pos = ref(length hd) in
-      List.iter
-        (fun s ->
-          unsafe_blit sep 0 r !pos (length sep);
-          pos := !pos + length sep;
-          unsafe_blit s 0 r !pos (length s);
-          pos := !pos + length s)
-        tl;
-      Bytes.unsafe_to_string r
+    unsafe_blit hd 0 dst pos (length hd);
+    unsafe_blit sep 0 dst (pos + length hd) seplen;
+    unsafe_blits dst (pos + length hd + seplen) sep seplen tl
+
+let concat sep = function
+    [] -> ""
+  | l -> let seplen = length sep in bts @@
+          unsafe_blits 
+            (B.create (sum_lengths 0 seplen l))
+            0 sep seplen l
 
 let iter f s =
   B.iter f (bos s)

testsuite/tests/lib-string/test_string.ml

@@ -36,3 +36,17 @@ let () =
     check_split ' ' (String.sub s 0 i)
   done
 ;;
+
+(* GPR#805/815/833 *)
+
+let ()  =
+  if Sys.word_size = 32 then begin
+    let big = String.make Sys.max_string_length 'x' in
+    let push x l = l := x :: !l in
+    let (+=) a b = a := !a + b in
+    let sz, l = ref 0, ref [] in
+    while !sz >= 0 do push big l; sz += Sys.max_string_length done;
+    while !sz <= 0 do push big l; sz += Sys.max_string_length done;
+    try ignore (String.concat "" !l); assert false
+    with Invalid_argument _ -> ()
+  end