CA-289145: Close socket if error occurs when connecting

Signed-off-by: Yang Qian <[email protected]>

Author: krizex

lib/open_uri.ml

@@ -15,36 +15,43 @@
  *
  *)
 
+open Xapi_stdext_pervasives.Pervasiveext
+
 let handle_socket f s =
   try
-    let result = f s in
-    Unix.close s;
-    result
+    f s
   with e ->
     Backtrace.is_important e;
-    Unix.close s;
     raise e
 
 let open_tcp f host port =
-      let host_entry = Unix.gethostbyname host in
-      let sockaddr = Unix.ADDR_INET(host_entry.Unix.h_addr_list.(0), port) in
-      let s = Unix.socket Unix.PF_INET Unix.SOCK_STREAM 0 in
+  let host_entry = Unix.gethostbyname host in
+  let sockaddr = Unix.ADDR_INET(host_entry.Unix.h_addr_list.(0), port) in
+  let s = Unix.socket Unix.PF_INET Unix.SOCK_STREAM 0 in
+  finally (fun () ->
       Unix.connect s sockaddr;
       handle_socket f s
+    ) (fun () ->
+      Unix.close s
+    )
 
 let with_open_uri uri f =
   match Uri.scheme uri with
   | Some "http" ->
     begin match Uri.host uri, Uri.port uri with
-    | Some host, Some port -> open_tcp f host port
-    | Some host, None      -> open_tcp f host 80
-    | _, _ -> failwith (Printf.sprintf "Failed to parse host and port from URI: %s" (Uri.to_string uri))
+      | Some host, Some port -> open_tcp f host port
+      | Some host, None      -> open_tcp f host 80
+      | _, _ -> failwith (Printf.sprintf "Failed to parse host and port from URI: %s" (Uri.to_string uri))
     end
   | Some "file" ->
     let filename = Uri.path_and_query uri in
     let sockaddr = Unix.ADDR_UNIX filename in
     let s = Unix.socket Unix.PF_UNIX Unix.SOCK_STREAM 0 in
-    Unix.connect s sockaddr;
-    handle_socket f s
+    finally (fun () ->
+        Unix.connect s sockaddr;
+        handle_socket f s
+      ) (fun () ->
+        Unix.close s
+      )
   | Some x -> failwith (Printf.sprintf "Unsupported URI scheme: %s" x)
   | None -> failwith (Printf.sprintf "Failed to parse URI: %s" (Uri.to_string uri))

lib/posix_channel.ml

@@ -190,8 +190,14 @@ let receive protocols =
   | V4V_proxy(_, _) -> assert false (* weight is 0 above *)
   | TCP_proxy(ip, port) ->
     let s = Unix.socket Unix.PF_INET Unix.SOCK_STREAM 0 in
-    Unix.connect s (Unix.ADDR_INET(Unix.inet_addr_of_string ip, port));
-    s
+    begin
+      try
+        Unix.connect s (Unix.ADDR_INET(Unix.inet_addr_of_string ip, port));
+        s
+      with e ->
+        Unix.close s;
+        raise e
+    end
   | Unix_sendmsg(_, path, token) ->
     let s = Unix.socket Unix.PF_UNIX Unix.SOCK_STREAM 0 in
     finally