Fix frame serialization logic

inhabitedtype · Jan 5, 2019 · bba1257 · bba1257
1 parent 8382d0e
commit bba1257
Show file tree

Hide file tree

Showing 2 changed files with 29 additions and 8 deletions.
diff --git a/lib/websocket.ml b/lib/websocket.ml
@@ -291,12 +291,12 @@ module Frame = struct
     end
   ;;
 
-  let serialize_control faraday ~opcode =
-    serialize_headers faraday ~is_fin:true ~opcode ~payload_length:0
+  let serialize_control ?mask faraday ~opcode =
+    let opcode = (opcode :> Opcode.t) in
+    serialize_headers faraday ?mask ~is_fin:true ~opcode ~payload_length:0
 
   let schedule_serialize ?mask faraday ~is_fin ~opcode ~payload ~off ~len =
-    let payload_length = Bigstring.length payload in
-    serialize_headers faraday ?mask ~is_fin ~opcode ~payload_length;
+    serialize_headers faraday ?mask ~is_fin ~opcode ~payload_length:len;
     begin match mask with
     | None -> ()
     | Some mask -> apply_mask mask payload ~off ~len
@@ -305,8 +305,16 @@ module Frame = struct
   ;;
 
   let serialize_bytes ?mask faraday ~is_fin ~opcode ~payload ~off ~len =
-    let payload_length = Bytes.length payload in
-    serialize_headers faraday ?mask ~is_fin ~opcode ~payload_length;
+    serialize_headers faraday ?mask ~is_fin ~opcode ~payload_length:len;
+    begin match mask with
+    | None -> ()
+    | Some mask -> apply_mask_bytes mask payload ~off ~len
+    end;
+    Faraday.write_bytes faraday payload ~off ~len;
+  ;;
+
+  let schedule_serialize_bytes ?mask faraday ~is_fin ~opcode ~payload ~off ~len =
+    serialize_headers faraday ?mask ~is_fin ~opcode ~payload_length:len;
     begin match mask with
     | None -> ()
     | Some mask -> apply_mask_bytes mask payload ~off ~len

diff --git a/lib/websocket.mli b/lib/websocket.mli
@@ -83,7 +83,11 @@ module Frame : sig
 
   val parse : t Angstrom.t
 
-  val serialize_control : Faraday.t -> opcode:Opcode.standard_control -> unit
+  val serialize_control
+    : ?mask:int32
+    -> Faraday.t
+    -> opcode:Opcode.standard_control
+    -> unit
 
   val schedule_serialize 
     :  ?mask:int32
@@ -93,10 +97,19 @@ module Frame : sig
     -> payload:Bigstring.t
     -> off:int
     -> len:int
-    -> Faraday.t
     -> unit
 
   val schedule_serialize_bytes
+    :  ?mask:int32
+    -> Faraday.t
+    -> is_fin:bool
+    -> opcode:Opcode.t
+    -> payload:Bytes.t
+    -> off:int
+    -> len:int
+    -> unit
+
+  val serialize_bytes
     :  ?mask:int32
     -> Faraday.t
     -> is_fin:bool