From 7ac2b65cb7e1e6163fd34e5a76c657a4d1940ccc Mon Sep 17 00:00:00 2001
From: Christophe Raffalli <christophe@raffalli.eu>
Date: Sat, 8 Jul 2023 03:58:22 -1000
Subject: [PATCH] fixed ktls test + bug in stubs

---
 src/ssl_stubs.c      |  2 +-
 tests/ssl_io_ktls.ml | 27 ++++++++++++++++++---------
 tests/util_ktls.ml   | 22 +++++++++++++++++-----
 3 files changed, 36 insertions(+), 15 deletions(-)

diff --git a/src/ssl_stubs.c b/src/ssl_stubs.c
index 36ccb61..df5a0e5 100644
--- a/src/ssl_stubs.c
+++ b/src/ssl_stubs.c
@@ -548,7 +548,7 @@ CAMLprim value caml_ssl_ktls_send_available(value out_fd) {
 
 CAMLprim value caml_ssl_ktls_recv_available(value out_fd) {
   CAMLparam1(out_fd);
-  int r = BIO_get_ktls_recv(SSL_get_wbio(SSL_val(out_fd)));
+  int r = BIO_get_ktls_recv(SSL_get_rbio(SSL_val(out_fd)));
   CAMLreturn(Val_int(r));
 }
 
diff --git a/tests/ssl_io_ktls.ml b/tests/ssl_io_ktls.ml
index d95abc8..8c8ce2d 100644
--- a/tests/ssl_io_ktls.ml
+++ b/tests/ssl_io_ktls.ml
@@ -8,8 +8,6 @@ let test_verify () =
 
   let context = Ssl.create_context ~ktls:true TLSv1_2 Client_context in
   let ssl = Ssl.open_connection_with_context context addr in
-  assert(Ssl.ktls_send_available ssl);
-  assert(Ssl.ktls_recv_available ssl);
   let verify_result =
     try
       Ssl.verify ssl;
@@ -36,14 +34,11 @@ let test_set_host () =
   let domain = Unix.domain_of_sockaddr addr in
   let sock = Unix.socket domain Unix.SOCK_STREAM 0 in
   let ssl = Ssl.embed_socket sock context in
-  Ssl.set_host ssl "localhost";
   Unix.connect sock addr;
   Ssl.connect ssl;
   let verify_result =
     try
       Ssl.verify ssl;
-      assert(Ssl.ktls_send_available ssl);
-      assert(Ssl.ktls_recv_available ssl);
       ""
     with
     | e -> Printexc.to_string e
@@ -65,13 +60,26 @@ let test_read_write () =
 
   let context = Ssl.create_context ~ktls:true TLSv1_2 Client_context in
   let ssl = Ssl.open_connection_with_context context addr in
-  assert(Ssl.ktls_send_available ssl);
-  assert(Ssl.ktls_recv_available ssl);
   let send_msg = "send" in
   let write_buf = Bytes.create (String.length send_msg) in
-  Unix.single_write (Ssl.file_descr_of_socket ssl) write_buf 0 4 |> ignore;
+  Util.write ssl write_buf 0 4 |> ignore;
   let read_buf = Bytes.create 8 in
-  Unix.read (Ssl.file_descr_of_socket ssl) read_buf 0 8 |> ignore;
+  Util.read ssl read_buf 0 8 |> ignore;
+  Ssl.shutdown_connection ssl;
+  check string "received message" "received" (Bytes.to_string read_buf)
+
+(* test to very that unix read/write are compatible with Ssl's*)
+let test_read_write2 () =
+  let addr = Unix.ADDR_INET (Unix.inet_addr_of_string "127.0.0.1", 11345) in
+  Util.server_thread addr (Some (fun _ -> "received")) |> ignore;
+
+  let context = Ssl.create_context ~ktls:true TLSv1_2 Client_context in
+  let ssl = Ssl.open_connection_with_context context addr in
+  let send_msg = "send" in
+  let write_buf = Bytes.create (String.length send_msg) in
+  Ssl.write ssl write_buf 0 4 |> ignore;
+  let read_buf = Bytes.create 8 in
+  Ssl.read ssl read_buf 0 8 |> ignore;
   Ssl.shutdown_connection ssl;
   check string "received message" "received" (Bytes.to_string read_buf)
 
@@ -82,5 +90,6 @@ let () =
       , [ test_case "Verify" `Quick test_verify
         ; test_case "Set host" `Quick test_set_host
         ; test_case "Read write" `Quick test_read_write
+        ; test_case "Read write2" `Quick test_read_write2
         ] )
     ]
diff --git a/tests/util_ktls.ml b/tests/util_ktls.ml
index cd47887..2bed3a3 100644
--- a/tests/util_ktls.ml
+++ b/tests/util_ktls.ml
@@ -13,18 +13,32 @@ type server_args =
   ; parser : (string -> string) option
   }
 
+let write ssl =
+  if Ssl.ktls_send_available ssl
+  then Unix.single_write (Ssl.file_descr_of_socket ssl)
+  else Ssl.write ssl
+
+let write_substring ssl =
+  if Ssl.ktls_send_available ssl
+  then Unix.write_substring (Ssl.file_descr_of_socket ssl)
+  else Ssl.write_substring ssl
+
+let read ssl =
+  if Ssl.ktls_recv_available ssl
+  then Unix.read (Ssl.file_descr_of_socket ssl)
+  else Ssl.read ssl
+
 let server_rw_loop ssl parser_func =
-  let fd = Ssl.file_descr_of_socket ssl in
   let rw_loop = ref true in
   while !rw_loop do
     try
       let read_buf = Bytes.create 256 in
-      let read_bytes = Unix.read fd read_buf 0 256 in
+      let read_bytes = read ssl read_buf 0 256 in
       if read_bytes > 0
       then (
         let input = Bytes.to_string read_buf in
         let response = parser_func input in
-        Unix.write_substring fd response 0 (String.length response) |> ignore;
+        write_substring ssl response 0 (String.length response) |> ignore;
         Ssl.close_notify ssl |> ignore;
         rw_loop := false)
     with
@@ -63,8 +77,6 @@ let server_listen args =
     let listen = Unix.accept socket in
     let ssl = embed_socket (fst listen) context in
     accept ssl;
-    assert(ktls_send_available ssl);
-    assert(ktls_recv_available ssl);
     (* Exit right away unless we need to rw *)
     (match args.parser with
     | Some parser_func -> server_rw_loop ssl parser_func