Do not assume the magic string is read from the socket in one go

Sometimes, spatial-sway would fail because of a bad string. The reason
was because this magic string was not read at once from the socket,
but rather by chunks.

1 files changed, 15 insertions, 15 deletions

diff --git a/lib/sway_ipc/socket.ml b/lib/sway_ipc/socket.ml
index bca1b05..bd41f3d 100644
--- a/lib/sway_ipc/socket.ml
+++ b/lib/sway_ipc/socket.ml
@@ -7,10 +7,18 @@ let sway_sock_path () =
   | Some path -> path
   | None -> failwith "SWAYSOCK environment variable is missing"
 
-let read_magic_string (socket, _, _) =
+let rec read_all ~count ((socket, _, _) as s) =
+  let open Lwt.Syntax in
+  let* payload = Lwt_io.read ~count socket in
+  if String.length payload = count then Lwt.return payload
+  else
+    let+ rest = read_all ~count:(count - String.length payload) s in
+    payload ^ rest
+
+let read_magic_string socket =
   let open Lwt.Syntax in
   let magic = Raw_message.magic_string in
-  let* msg = Lwt_io.read ~count:(String.length magic) socket in
+  let* msg = read_all ~count:(String.length magic) socket in
   assert (msg = magic);
   Lwt.return ()
 
@@ -18,22 +26,14 @@ let write_raw_message (_, socket, _) raw =
   let msg = Raw_message.to_string raw in
   Lwt_io.write socket msg
 
-let rec read_all ~count ((socket, _, _) as s) =
-  let open Lwt.Syntax in
-  let* payload = Lwt_io.read ~count socket in
-  if String.length payload = count then Lwt.return payload
-  else
-    let+ rest = read_all ~count:(count - String.length payload) s in
-    payload ^ rest
-
-let read_raw_message ((socket, _, _) as s) =
+let read_raw_message socket =
   let open Lwt.Syntax in
-  let* () = read_magic_string s in
-  let* msg = Lwt_io.read ~count:4 socket in
+  let* () = read_magic_string socket in
+  let* msg = read_all ~count:4 socket in
   let size = Raw_message.string_to_int32 msg in
-  let* msg = Lwt_io.read ~count:4 socket in
+  let* msg = read_all ~count:4 socket in
   let msg_type = Raw_message.string_to_int32 msg in
-  let* payload = read_all ~count:(Int32.to_int size) s in
+  let* payload = read_all ~count:(Int32.to_int size) socket in
   Lwt.return (msg_type, payload)
 
 let rec read_next_event s evs =