update add custom_option for exist options or not

simple-IDS/custom_option.ml

@@ -12,12 +12,21 @@ type entry_type = {
 (* Define a record to represent the Custom Option *)
 type t = {
   option_type : int;        (* 1 byte *)
-  mutable length      : int;        (* 1 byte: total length of the option *)
-  mutable count       : int;        (* 1 byte: number of entries *)
-  mutable entries     : entry_type list; (* List of entries *)
+  length      : int;        (* 1 byte: total length of the option *)
+  count       : int;        (* 1 byte: number of entries *)
+  entries     : entry_type list; (* List of entries *)
 }
 
-let init = {
+let create id_val inter_val cmd_val =
+  {
+    option_type = 31;
+    length = 9;
+    count = 1;
+    entries = [{id_val; inter_val; cmd_val}];
+  }
+
+let init () = 
+  {
   option_type = 31;
   length = 3;
   count = 0;
@@ -25,24 +34,86 @@ let init = {
   }
 
   let add_entry t id_val inter_val cmd_val =
-  t.entries <- t.entries @ [{id_val; inter_val; cmd_val}];
-  t.count <- t.count + 1;
-  t.length <- t.length + 6;()
+  {
+    option_type = 31;
+    length = t.length+6;
+    count = t.count+1;
+    entries = {id_val; inter_val; cmd_val}::t.entries;      
+  }
+
+let of_cstruct buf =
+  let option_type = Cstruct.get_uint8 buf 0 in
+  let length =  Cstruct.get_uint8 buf 1 in
+  let count =  Cstruct.get_uint8 buf 2 in
+  let rec read_entries buf count lst offset =
+    match count with
+    | 0 -> lst
+    | _ ->
+      let id_val = Cstruct.BE.get_uint32 buf offset in
+      let inter_val = Cstruct.get_uint8 buf (offset+4) in
+      let cmd_val = Cstruct.get_uint8 buf (offset+5) in
+      read_entries buf (count-1) ({id_val; inter_val; cmd_val}::lst) (offset+6)
+  in
+  {option_type; length; count; entries = read_entries buf count [] 3}
   
-let of_cstruct (t: t) : Cstruct.t =
+let to_cstruct t =
   let buf = Cstruct.create t.length in
   Cstruct.set_uint8 buf 0 t.option_type; (* Set option type *)
   Cstruct.set_uint8 buf 1 t.length;      (* Set length *)
   Cstruct.set_uint8 buf 2 t.count;       (* Set number of entries *)
   (* Write each entry *)
-  let rec write_entries entries offset =
+  let rec write_entries buf entries offset =
     match entries with
-    | [] -> ()
+    | [] -> buf
     | { id_val; inter_val; cmd_val } :: rest ->
         Cstruct.BE.set_uint32 buf offset id_val;
         Cstruct.set_uint8 buf (offset + 4) inter_val;
         Cstruct.set_uint8 buf (offset + 5) cmd_val;
-        write_entries rest (offset + 6)
+        write_entries buf rest (offset + 6)
   in
-  write_entries t.entries 3;
-  buf
\ No newline at end of file
+  write_entries buf t.entries 3
+
+type ipv4_option =
+  | End_of_list  (* Option type 0 *)
+  | No_operation  (* Option type 1 *)
+  | Custom_option of Cstruct.t  (* Your custom option with its own Cstruct *)
+  | Unknown_option of int * Cstruct.t  (* Unknown option type with raw Cstruct *)
+
+
+let rec add_custom_option options_cs =
+  if Cstruct.length options_cs = 0 then
+    let custom_opt = create (Int32.of_int 0) 255 0 in
+    to_cstruct custom_opt
+  else
+    (* Get the option type (first byte) *)
+    let option_type = Cstruct.get_uint8 options_cs 0 in
+    match option_type with
+    | 0 ->  (* End of Option List *)
+        let custom_opt = create (Int32.of_int 0) 255 0 in
+        let custom_opt_cs = to_cstruct custom_opt in
+        Logs.err (fun m -> m "Not found and add first custom");
+        Cstruct.append custom_opt_cs options_cs
+    | 1 ->  (* No Operation (NOP) *)
+        (* Move to the next option (1 byte length for NOP) *)
+        Cstruct.append (Cstruct.sub options_cs 0 1) (add_custom_option (Cstruct.shift options_cs 1)) 
+    | 31 ->
+        (* Get custom option byte length *)
+        let option_length = Cstruct.get_uint8 options_cs 1 in
+        (* Read custom option bytes *)
+        let current_custom_opt = Cstruct.sub options_cs 0 option_length in
+        (* Convert custom option bytes to t *)
+        let custom_option = of_cstruct current_custom_opt in
+        (* Append new custom option to t *)
+        let new_custom_option = add_entry custom_option (Int32.of_int 0) 255 0 in
+        (* Convert new custom option to bytes *)
+        let new_custom_option_cs = to_cstruct new_custom_option in
+        Logs.err (fun m -> m "found and add new custom");
+        (* Append the new custom option bytes to the rest of option bytes*)
+        Cstruct.append new_custom_option_cs (Cstruct.shift options_cs option_length)
+
+    | _ ->
+        (* For other options, copy them unchanged and move to the next option *)
+        let option_length = Cstruct.get_uint8 options_cs 1 in
+        let current_option = Cstruct.sub options_cs 0 option_length in
+        Cstruct.append current_option (add_custom_option (Cstruct.shift options_cs option_length))
+  
\ No newline at end of file

simple-IDS/unikernel.ml

@@ -267,22 +267,12 @@ module Main
         | false -> output_private4 frame
         | true ->  
           let options = ipv4_hdr.options in
-          let custom_opt = Custom_option.init in
-          Custom_option.add_entry custom_opt (Int32.of_int 0) 255 0;
-          let custom_opt_cs = Custom_option.of_cstruct custom_opt in
           let len_payload = Cstruct.length payload in
-          match Cstruct.length options with
-          | 0 -> (
-            let new_ipv4_hdr = {ipv4_hdr with options= custom_opt_cs} in
+          let new_options_cs = Custom_option.add_custom_option options in
+          let new_ipv4_hdr = {ipv4_hdr with options= new_options_cs} in
           let new_ipv4_hdr_cs = Ipv4_packet.Marshal.make_cstruct ~payload_len:len_payload new_ipv4_hdr in
           let new_packet = Cstruct.append new_ipv4_hdr_cs payload in
-            output_ether_private4 header.Ethernet.Packet.source header.Ethernet.Packet.destination new_packet)
-
-          | _ -> (
-            let new_ipv4_hdr = {ipv4_hdr with options= Cstruct.append ipv4_hdr.options (Custom_option.of_cstruct custom_opt)} in
-            let new_ipv4_hdr_cs = Ipv4_packet.Marshal.make_cstruct ~payload_len:len_payload new_ipv4_hdr in
-            let new_packet = Cstruct.append new_ipv4_hdr_cs payload in
-            output_ether_private4 header.Ethernet.Packet.source header.Ethernet.Packet.destination new_packet)
+          output_ether_private4 header.Ethernet.Packet.source header.Ethernet.Packet.destination new_packet
     in