Merge pull request #29 from frenetic-lang/checksum

Compute TCP checksum

Author: Marco Canini

lib/Packet.ml

@@ -86,6 +86,18 @@ type nwTos = int8
 
 type tpPort = int16
 
+let mk_pseudo_header (src : nwAddr) (dst : nwAddr) (proto : int) (len : int) =
+  (* XXX(seliopou): pseudo_header's allocated on every call. Given the usage
+   * pattern of this library, though, would it be safe to allocate once and
+   * reuse? *)
+  let pseudo_header = Cstruct.create 12 in
+  Cstruct.BE.set_uint32 pseudo_header 0  src;
+  Cstruct.BE.set_uint32 pseudo_header 4  dst;
+  Cstruct.set_uint8     pseudo_header 8  0;
+  Cstruct.set_uint8     pseudo_header 9  proto;
+  Cstruct.BE.set_uint16 pseudo_header 10 len;
+  pseudo_header
+
 module Tcp = struct
 
   module Flags = struct
@@ -202,6 +214,16 @@ module Tcp = struct
     let bits = Cstruct.shift bits sizeof_tcp in 
     Cstruct.blit pkt.payload 0 bits 0 (Cstruct.len pkt.payload)
 
+
+  let checksum (bits : Cstruct.t) (src : nwAddr) (dst : nwAddr) (pkt : t) =
+    let length = len pkt in
+    let pseudo_header = mk_pseudo_header src dst 0x6 length in
+    set_tcp_chksum bits 0;
+    let chksum = Checksum.ones_complement_list
+      (if (length mod 2) = 0
+        then [pseudo_header; Cstruct.sub bits 0 length]
+        else [pseudo_header; Cstruct.sub bits 0 length; Cstruct.of_string "\x00"]) in
+    set_tcp_chksum bits chksum
 end
 
 module Udp = struct
@@ -883,7 +905,8 @@ module Ip = struct
     let bits = Cstruct.shift bits header_len in
     match pkt.tp with
       | Tcp tcp -> 
-        Tcp.marshal bits tcp
+        Tcp.marshal bits tcp;
+        Tcp.checksum bits pkt.src pkt.dst tcp
       | Udp udp ->
         Udp.marshal bits udp
       | Icmp icmp -> 
@@ -1118,15 +1141,15 @@ let setNwTos pkt nwTos =
 
 let tp_setTpSrc tp src = match tp with 
   | Ip.Tcp tcp ->
-    Ip.Tcp { tcp with Tcp.src = src } (* JNF: checksum? *)
+    Ip.Tcp { tcp with Tcp.src = src }
   | Ip.Udp udp ->
     Ip.Udp { udp with Udp.src = src }
   | tp -> 
     tp
 
 let tp_setTpDst tp dst = match tp with 
   | Ip.Tcp tcp ->
-    Ip.Tcp { tcp with Tcp.dst = dst } (* JNF: checksum? *)
+    Ip.Tcp { tcp with Tcp.dst = dst }
   | Ip.Udp udp ->
     Ip.Udp { udp with Udp.dst = dst }
   | tp -> 

test/Test.ml

@@ -17,14 +17,14 @@ module RoundTrip = struct
   let unparsable_eq (l1, b1) (l2, b2) =
     l1 = l2 && compare (Cstruct.to_string b1) (Cstruct.to_string b2) = 0
 
-  let udp_eq e1 e2 =
+  let udp_eq ?(chksum=false) e1 e2 =
     let open Udp in
     e1.src = e2.src &&
     e1.dst = e2.dst &&
-    e1.chksum = e2.chksum &&
+    ((not chksum) || e1.chksum = e2.chksum) &&
     compare (Cstruct.to_string e1.payload) (Cstruct.to_string e2.payload) = 0
 
-  let tcp_eq e1 e2 =
+  let tcp_eq ?(chksum=false) e1 e2 =
     let open Tcp in
     e1.src = e2.src &&
     e1.dst = e2.dst &&
@@ -33,11 +33,11 @@ module RoundTrip = struct
     e1.offset = e2.offset &&
     e1.flags = e2.flags &&
     e1.window = e2.window &&
-    e1.chksum = e2.chksum &&
     e1.urgent = e2.urgent &&
+    ((not chksum) || e1.chksum = e2.chksum) &&
     compare (Cstruct.to_string e1.payload) (Cstruct.to_string e2.payload) = 0
 
-  let ip_eq e1 e2 =
+  let ip_eq ?(chksum=false) e1 e2 =
     let open Ip in
     e1.tos = e2.tos &&
     e1.ident = e2.ident &&
@@ -50,13 +50,13 @@ module RoundTrip = struct
       | Unparsable u1, Unparsable u2 ->
         unparsable_eq u1 u2
       | Udp p1, Udp p2 ->
-        udp_eq p1 p2
+        udp_eq ~chksum p1 p2
       | Tcp p1, Tcp p2 ->
-        tcp_eq p1 p2
+        tcp_eq ~chksum p1 p2
       | _, _ ->
         e1 = e2
 
-  let dl_eq e1 e2 =
+  let dl_eq ?(chksum=false) e1 e2 =
     e1.dlSrc = e2.dlSrc &&
     e1.dlDst = e2.dlDst &&
     e1.dlVlan = e2.dlVlan &&
@@ -65,12 +65,16 @@ module RoundTrip = struct
       | Unparsable u1, Unparsable u2 ->
         unparsable_eq u1 u2
       | Ip nw1, Ip nw2 ->
-        ip_eq nw1 nw2
+        ip_eq ~chksum nw1 nw2
       | _, _ ->
         e1 = e2
 
-  let prop_roundtrip parse marshal e =
-    dl_eq (parse (marshal e)) e
+  let prop_roundtrip ?(chksum=false) parse marshal e =
+    dl_eq ~chksum (parse (marshal e)) e
+
+  let prop_roundtrip2 parse marshal e =
+    let e' = parse (marshal e) in
+    prop_roundtrip ~chksum:true parse marshal e'
 
   TEST "Roundtrip property for unparsable Ethernet frames" =
     (packet_quickCheck (Arb.arbitrary_packet Arb.arbitrary_dl_unparsable)
@@ -92,13 +96,13 @@ module RoundTrip = struct
     let udp = Gen.map_gen (fun x -> Ip.Udp(x))
           (Arb.arbitrary_udp (Arb.arbitrary_payload 65507)) in
     (packet_quickCheck (mk_ip udp)
-      (prop_roundtrip parse marshal))
+      (prop_roundtrip2 parse marshal))
 
   TEST "Roundtrip property for TCP packets" =
     let tcp = Gen.map_gen (fun x -> Ip.Tcp(x))
           (Arb.arbitrary_tcp (Arb.arbitrary_payload (65507 - 128))) in
     (packet_quickCheck (mk_ip tcp)
-      (prop_roundtrip parse marshal))
+      (prop_roundtrip2 parse marshal))
 
 end