Fix incorrect HeaderV2_IPV4 declaration, and various tweaks

Author: ba0f3

proxyproto.nimble

@@ -1,6 +1,6 @@
 # Package
 
-version       = "0.1.2"
+version       = "0.1.3"
 author        = "Huy Doan"
 description   = "PROXY Protocol enabler for aged programs"
 license       = "MIT"

src/proxyproto.nim

@@ -6,12 +6,11 @@ from dynlib import LibHandle, symAddr
 
 const
   v1sig = "PROXY".cstring
-  v2sig = "\c\n\c\n\x00\c\nQUIT\n"
-
-proc c_memcmp(a, b: pointer, size: csize): cint {.importc: "memcmp", header: "<string.h>", noSideEffect.}
+  # v2       \c         \n    \c    \n    \0    \c    \n    Q     U     I     T     \n
+  v2sig = [0x0D.byte, 0x0A, 0x0D, 0x0A, 0x00, 0x0D, 0x0A, 0x51, 0x55, 0x49, 0x54, 0x0A]
 
 type
-  HeaderV2_IPV4 {.union.} = object
+  HeaderV2_IPV4 = object
     src_addr: uint32
     dst_addr: uint32
     src_port: uint16
@@ -51,8 +50,9 @@ type
 
 proc pp_handshake*(fd: SocketHandle, sa: ptr SockAddr, sl: ptr Socklen): int  =
   var
-    size: int
     hdr: Header
+    src: Sockaddr_in
+    src6: Sockaddr_in6
 
   while true:
     result = recv(fd, addr hdr, sizeof(hdr), 0)
@@ -61,21 +61,16 @@ proc pp_handshake*(fd: SocketHandle, sa: ptr SockAddr, sl: ptr Socklen): int  =
   if result == -1:
     return if (errno == EAGAIN): 0 else: -1
 
-  if result >= 16 and c_memcmp(addr hdr.v2, v2sig.cstring, 12) == 0 and (hdr.v2.ver_cmd and 0xF0) == 0x20:
-    size = 16 + ntohs(hdr.v2.length).int
-    if result < size:
-      return -1
+  if result >= 16 and hdr.v2.sig == v2sig and (hdr.v2.ver_cmd and 0xF0) == 0x20 and result >= 16 + ntohs(hdr.v2.length).int:
     case hdr.v2.ver_cmd and 0xF
     of 0x1: # PROXY command
       case hdr.v2.fam
       of 0x11: # TCPv4
-        var src: Sockaddr_in
         src.sin_family = AF_INET.TSa_Family
         src.sin_addr.s_addr = hdr.v2.address.ip4.src_addr
         src.sin_port = hdr.v2.address.ip4.src_port
         copyMem(sa, addr src, sl[])
       of 0x21: # TCPv6
-        var src6: Sockaddr_in6
         src6.sin6_family = AF_INET6.TSa_Family
         copyMem(addr src6.sin6_addr, addr hdr.v2.address.ip6.src_addr, 16)
         src6.sin6_port = hdr.v2.address.ip6.src_port
@@ -86,7 +81,7 @@ proc pp_handshake*(fd: SocketHandle, sa: ptr SockAddr, sl: ptr Socklen): int  =
       return -1
     else:
       return -1
-  elif result >= 8 and c_memcmp(addr hdr, v1sig, 5) == 0:
+  elif result >= 8 and cmpMem(addr hdr, v1sig, 5) == 0:
     let backslash_r = hdr.v1.find('\r')
     if backslash_r == -1 or hdr.v1[backslash_r + 1] != '\n':
       return -1
@@ -95,65 +90,59 @@ proc pp_handshake*(fd: SocketHandle, sa: ptr SockAddr, sl: ptr Socklen): int  =
       tmp: array[40, char]
       step = SRC_IP
       idx = 0
-    zeroMem(addr tmp, sizeof(tmp))
+      c: char
     if hdr.v1[9] == '4':
-      var src: Sockaddr_in
       src.sin_family = AF_INET.TSa_Family
       for i in 11..backslash_r:
-        var c = hdr.v1[i]
+        c = hdr.v1[i]
         if c == ' ' or c == '\0':
           case step
           of SRC_IP:
             step = DST_IP
+            zeroMem(addr tmp, sizeof(tmp))
             src.sin_addr.s_addr = cast[uint32](parseIPAddress($cast[cstring](addr tmp)).address_v4)
-          of DST_IP:
-            # temporary ignore destination ip address
+          of DST_IP: # temporary ignore destination ip address
             step = SRC_PORT
           of SRC_PORT:
             step = DST_PORT
+            zeroMem(addr tmp, sizeof(tmp))
             src.sin_port = htons(parseInt($cast[cstring](addr tmp)).uint16)
             copyMem(sa, addr src, sl[])
             return backslash_r
-          of DST_PORT:
-            # temporary ignore destination port
+          of DST_PORT: # temporary ignore destination port
             break
-          zeroMem(addr tmp, sizeof(tmp))
           idx = 0
         else:
           tmp[idx] = c
           inc(idx)
     elif hdr.v1[0] == '6':
-      var src6: Sockaddr_in6
       src6.sin6_family = AF_INET6.TSa_Family
       for i in 11..backslash_r:
-        var c = hdr.v1[i]
+        c = hdr.v1[i]
         if c == ' ' or c == '\0':
           case step
           of SRC_IP:
             step = DST_IP
             var ip = parseIPAddress($cast[cstring](addr tmp))
+            zeroMem(addr tmp, sizeof(tmp))
             copyMem(addr src6.sin6_addr, addr ip.address_v6, 16)
-          of DST_IP:
-            # temporary ignore destination ip address
+          of DST_IP: # temporary ignore destination ip address
             step = SRC_PORT
           of SRC_PORT:
             step = DST_PORT
             src6.sin6_port = htons(parseInt($cast[cstring](addr tmp)).uint16)
+            zeroMem(addr tmp, sizeof(tmp))
             copyMem(sa, addr src6, sl[])
             return backslash_r
-          of DST_PORT:
-            # temporary ignore destination port
+          of DST_PORT: # temporary ignore destination port
             break
-          zeroMem(addr tmp, sizeof(tmp))
           idx = 0
         else:
           tmp[idx] = c
           inc(idx)
-    else:
-      # not supported protocol
+    else: # not supported protocol
       return -1
-  else:
-    ##  Wrong protocol
+  else: #  Wrong protocol
     return -1
 
 when isMainModule: