calculate TCP checksum

ip_packet.nim

@@ -1,6 +1,7 @@
 from nativesockets import ntohs, ntohl, htons, htonl
 from net import IpAddress, IpAddressFamily, Port, `$`
 from posix import InAddr, inet_ntoa
+import unittest
 
 type
   Ether_header {.importc: "struct ether_header", pure, final,
@@ -97,33 +98,84 @@ proc parseEthernetPacket*(input: string): IpPacket =
   else:
     result = IpPacket(protocol: other)
 
+proc tcpChecksum(buffer: string): uint16 =
+  let ip = cast[ptr Ip](buffer.cstring)
+  var tcp = cast[ptr uint16](cast[ByteAddress](ip) + ip.ip_hl.int * 4)
+  var tcpLen = ntohs(ip.ip_len) - ip.ip_hl.uint16 * 4
+  var checksum = 0.uint32
+  # Add source IP address
+  checksum = checksum + (ntohl(ip.ip_src.s_addr) shr 16 and 0xFFFF).uint16
+  checksum = checksum + (ntohl(ip.ip_src.s_addr) and 0xFFFF).uint16
+  # Add dest IP address
+  checksum = checksum + (ntohl(ip.ip_dst.s_addr) shr 16 and 0xFFFF).uint16
+  checksum = checksum + (ntohl(ip.ip_dst.s_addr) and 0xFFFF).uint16
+  # Add protocol and reserved
+  checksum = checksum + ip.ip_p.uint16
+  # Add length of the IP payload
+  checksum = checksum + tcpLen
+  # Add IP payload
+  while tcpLen > 1:
+    checksum = checksum + ntohs(tcp[])
+    tcp = cast[ptr uint16](cast[ByteAddress](tcp) + sizeof(uint16))
+    tcpLen = tcpLen - sizeof(uint16).uint16
+  # Add leftover single byte
+  if tcpLen != 0:
+    checksum = checksum + cast[ptr uint8](tcp)[]
+  # Fold 16-bit segments of the 32-bit checksum and add carry bit. Result is the
+  # one's complement.
+  checksum = (checksum and 0xFFFF).uint16 + (checksum shr 16).uint16
+  checksum = checksum + checksum shr 16
+  result = not checksum.uint16
+
 proc serialize*(packet: IpPacket): string =
   case packet.protocol
   of tcp:
+    result = newString(sizeof(Ip) + sizeof(Tcphdr))
+    zeroMem(result.cstring, result.len)
     let srcIp = InAddr(s_addr: cast[uint32](packet.tcpIpSrc.address_v4))
     let dstIp = InAddr(s_addr: cast[uint32](packet.tcpIpDst.address_v4))
-    var ipHeader = Ip(ip_hl: 5,
-                      ip_v: 4,
-                      ip_tos: 0,
-                      ip_len: sizeof(Ip).uint16 + sizeof(TcpHdr).uint16,
-                      ip_id: htons(54321), # FIXME: random number
-                      ip_off: 0,
-                      ip_ttl: 64,
-                      ip_p: 6.cuchar,
-                      ip_sum: htons(0), # done by kernel
-                      ip_src: srcIp,
-                      ip_dst: dstIp)
-    var tcpHeader = Tcphdr(th_sport: htons(packet.tcpPortSrc.uint16),
-                           th_dport: htons(packet.tcpPortDst.uint16),
-                           th_seq: htonl(packet.tcpSeqNumber),
-                           th_ack: 0,
-                           th_off: 5,
-                           th_flags: cast[uint8](packet.tcpFlags),
-                           th_win: 1452 * 10,
-                           th_sum: 0,
-                           th_urp: 0)
-    result = newString(sizeof(Ip) + sizeof(Tcphdr))
-    copyMem(addr result[0], addr ipHeader, sizeof(Ip))
-    copyMem(addr result[sizeof(Ip)], addr tcpHeader, sizeof(Tcphdr))
+
+    let ipHeader = cast[ptr Ip](addr result[0])
+    ipHeader.ip_hl = 5
+    ipHeader.ip_v = 4
+    ipHeader.ip_tos = 0
+    ipHeader.ip_len = htons(sizeof(Ip).uint16 + sizeof(TcpHdr).uint16)
+    ipHeader.ip_id = htons(54321) # FIXME = random number
+    ipHeader.ip_off = 0
+    ipHeader.ip_ttl = 64
+    ipHeader.ip_p = 6.cuchar
+    ipHeader.ip_src = srcIp
+    ipHeader.ip_dst = dstIp
+
+    let tcpHeader = cast[ptr Tcphdr](addr result[sizeof(Ip)])
+    tcpHeader.th_sport = htons(packet.tcpPortSrc.uint16)
+    tcpHeader.th_dport = htons(packet.tcpPortDst.uint16)
+    tcpHeader.th_seq = htonl(packet.tcpSeqNumber)
+    tcpHeader.th_ack = 0
+    tcpHeader.th_off = 5
+    tcpHeader.th_flags = cast[uint8](packet.tcpFlags)
+    tcpHeader.th_win = htons(1452 * 10)
+    tcpHeader.th_urp = 0
+    tcpHeader.th_sum = htons(tcpChecksum(result))
+
   else:
     raise newException(ValueError, "protocol not supported")
+
+suite "ip_packet tests":
+  setup:
+    var
+      ipHeader = Ip(ip_hl: 5, ip_v: 4, ip_tos: 0, ip_len: htons(40),
+                    ip_id: htons(54321), ip_off: 0, ip_ttl: 64, ip_p: 6.cuchar,
+                    ip_src: InAddr(s_addr: htonl(0x5abb2bd9.uint32)),
+                    ip_dst: InAddr(s_addr: htonl(0x0a000069.uint32)))
+      tcpHeader = Tcphdr(th_sport: htons(4321), th_dport: htons(1234),
+                         th_seq: htonl(345364), th_ack: 0, th_off: 5,
+                         th_flags: cast[uint8]({SYN}), th_win: htons(1452 * 10),
+                         th_urp: 0)
+
+  test "tcpChecksum":
+    var buffer = newString(sizeof(Ip) + sizeof(Tcphdr))
+    zeroMem(addr buffer[0], buffer.len)
+    copyMem(addr buffer[0], addr ipHeader, sizeof(Ip))
+    copyMem(addr buffer[sizeof(Ip)], addr tcpHeader, sizeof(Tcphdr))
+    check(tcpChecksum(buffer) == 35681)