1 files changed, 29 insertions, 32 deletions

diff --git a/src/ppkt.c b/src/ppkt.c
index dc011a4..7c06f27 100644
--- a/src/ppkt.c
+++ b/src/ppkt.c
@@ -6,7 +6,7 @@
 #include <netinet/ip.h>
 #include <netinet/ip_icmp.h>
 
-enum ptype ppkt_type_to_enum(struct ppkt_header * pheader)
+enum ptype ppkt_type_to_enum(struct ppkt_header const * pheader)
 {
     enum ptype pt = (enum ptype)pheader->type;
 
@@ -22,7 +22,7 @@ enum ptype ppkt_type_to_enum(struct ppkt_header * pheader)
 
 int ppkt_process_icmp(struct psock * sock)
 {
-    if (sock->current.peer.ss_family == AF_INET) {
+    if (sock->current.peer_sockaddr.ss_family == AF_INET) {
         sock->current.pkt_buf.icmphdr.checksum = ntohs(sock->current.pkt_buf.icmphdr.checksum);
         sock->current.pkt_buf.icmphdr.un.echo.id = ntohs(sock->current.pkt_buf.icmphdr.un.echo.id);
         sock->current.pkt_buf.icmphdr.un.echo.sequence = ntohs(sock->current.pkt_buf.icmphdr.un.echo.sequence);
@@ -71,52 +71,49 @@ static size_t ppkt_data_size(union ppkt_body * pbody, enum ptype type)
 
 int ppkt_process_ppkt(struct psock * sock)
 {
-    size_t const min_pkt_siz = sizeof(struct iphdr) +
-                               sizeof(sock->current.pkt_buf.icmphdr);
-    size_t data_siz = sizeof(sock->current.pkt_buf.pheader);
+    size_t const min_pkt_siz = sizeof(struct iphdr) + sizeof(sock->current.pkt_buf.icmphdr);
 
-    if (sock->current.peer.ss_family == AF_INET) {
-        if (sock->current.bytes_read < min_pkt_siz + data_siz)
-        {
-            return -1;
-        }
-    } else {
+    if (sock->current.peer_sockaddr.ss_family != AF_INET) {
+        return -1;
+    }
+
+    if (sock->current.bytes_read < min_pkt_siz + sizeof(sock->current.pkt_buf.pheader)) {
         return -1;
     }
 
     sock->current.pkt_buf.pheader.total_size = ntohs(sock->current.pkt_buf.pheader.total_size);
-    if (sock->current.pkt_buf.pheader.total_size != sock->current.bytes_read - min_pkt_siz)
-    {
+    if (sock->current.pkt_buf.pheader.total_size != sock->current.bytes_read - min_pkt_siz) {
         return -1;
     }
 
-    switch (ppkt_type_to_enum(&sock->current.pkt_buf.pheader)) {
+    enum ptype packet_type = ppkt_type_to_enum(&sock->current.pkt_buf.pheader);
+    size_t packet_body_size = ppkt_size(packet_type);
+    if (packet_body_size == 0 ||
+        packet_body_size > sock->current.bytes_read - (min_pkt_siz + sizeof(sock->current.pkt_buf.pheader))) {
+        return -1;
+    }
+
+    switch (packet_type) {
         case PTYPE_INVALID:
             return -1;
         case PTYPE_AUTH_REQUEST:
-            sock->current.pkt_buf.pbody.auth_request.magic =
-                ntohl(sock->current.pkt_buf.pbody.auth_request.magic);
+            sock->current.pkt_buf.pbody.auth_request.magic = ntohl(sock->current.pkt_buf.pbody.auth_request.magic);
             sock->current.pkt_buf.pbody.auth_request.authdata_siz =
                 ntohs(sock->current.pkt_buf.pbody.auth_request.authdata_siz);
-
-            data_siz += sizeof(sock->current.pkt_buf.pbody.auth_request);
-            data_siz += sock->current.pkt_buf.pbody.auth_request.authdata_siz;
             break;
         case PTYPE_AUTH_RESPONSE:
             break;
     }
-    if (sock->current.bytes_read != min_pkt_siz + data_siz)
-    {
+
+    size_t packet_data_size = ppkt_data_size(&sock->current.pkt_buf.pbody, packet_type);
+    if (sock->current.bytes_read != min_pkt_siz + packet_body_size + packet_data_size) {
         return -1;
     }
 
     return 0;
 }
 
-static void ppkt_init_pkt(struct pdesc * desc,
-                          struct ppkt_buffer * pkt_buf,
-                          enum ptype type,
-                          size_t data_siz)
+static void ppkt_init_pkt(struct pdesc * desc, struct ppkt_buffer * pkt_buf, enum ptype type, size_t data_siz)
 {
     pkt_buf->icmphdr.un.echo.id = desc->identifier;
     pkt_buf->icmphdr.un.echo.sequence = ++desc->sequence;
@@ -134,22 +131,19 @@ static void ppkt_init_pkt(struct pdesc * desc,
     }
     pkt_buf->icmphdr.code = 0;
     pkt_buf->pheader.type = type;
-    pkt_buf->pheader.total_size = ppkt_size(type) +
-                                  ppkt_data_size(&pkt_buf->pbody, type);
+    pkt_buf->pheader.total_size = ppkt_size(type) + ppkt_data_size(&pkt_buf->pbody, type);
 }
 
 static void ppkt_finalize_pkt(struct ppkt_buffer * const pkt_buf)
 {
     pkt_buf->icmphdr.un.echo.id = htons(pkt_buf->icmphdr.un.echo.id);
     pkt_buf->icmphdr.un.echo.sequence = htons(pkt_buf->icmphdr.un.echo.sequence);
-    switch (pkt_buf->pheader.type)
-    {
+    switch (pkt_buf->pheader.type) {
         case PTYPE_INVALID:
             break;
         case PTYPE_AUTH_REQUEST:
             pkt_buf->pbody.auth_request.magic = htonl(PTUNNEL_MAGIC);
-            pkt_buf->pbody.auth_request.authdata_siz =
-                htons(pkt_buf->pbody.auth_request.authdata_siz);
+            pkt_buf->pbody.auth_request.authdata_siz = htons(pkt_buf->pbody.auth_request.authdata_siz);
             break;
         case PTYPE_AUTH_RESPONSE:
             break;
@@ -160,7 +154,10 @@ static void ppkt_finalize_pkt(struct ppkt_buffer * const pkt_buf)
     pkt_buf->icmphdr.checksum = icmp_checksum_iovec(pkt_buf->iovec, pkt_buf->iovec_used);
 }
 
-void ppkt_prepare_auth_request(struct pdesc * desc, struct ppkt_buffer * pkt_buf, uint8_t * authdata, size_t authdata_siz)
+void ppkt_prepare_auth_request(struct pdesc * desc,
+                               struct ppkt_buffer * pkt_buf,
+                               uint8_t * authdata,
+                               size_t authdata_siz)
 {
     ppkt_init_pkt(desc, pkt_buf, PTYPE_AUTH_REQUEST, authdata_siz);