1 files changed, 26 insertions, 3 deletions

diff --git a/src/dgram/dgram_link.c b/src/dgram/dgram_link.c
index c1f4845..d26b85e 100644
--- a/src/dgram/dgram_link.c
+++ b/src/dgram/dgram_link.c
@@ -35,6 +35,7 @@
 #include <net/if.h>
 #include <stdlib.h>
 #include <string.h>
+#include <poll.h>
 #include "if_ether.h"
 #include "dgram.h"
 #include "crc.h"
@@ -159,12 +160,13 @@ tx_len_t
 dgram_squeak(struct onet_link *link, mac_addr_t dst)
 {
     struct dgram_params params;
-    uint8_t pad[8];
+    uint8_t pad[SQUEAK_NPAD];
 
     if (link == NULL) {
         return -EINVAL;
     }
 
+    memset(pad, 0, sizeof(pad));
     params.dst = dst;
     params.buf = pad;
     params.len = sizeof(pad);
@@ -175,11 +177,14 @@ dgram_squeak(struct onet_link *link, mac_addr_t dst)
 rx_len_t
 dgram_recv(struct onet_link *link, void *buf, uint16_t len)
 {
+    const size_t MAX_MISS = 8;
     socklen_t addr_len;
     struct sockaddr_ll saddr;
     struct onet_dgram *o1p_hdr;
     struct ether_hdr *hdr;
-    size_t dgram_len, recv_len;
+    struct pollfd pfd;
+    size_t dgram_len, miss_count = 0;
+    ssize_t recv_len;
     uint32_t crc;
     uint16_t proto;
     mac_addr_t dest_mac, src_mac;
@@ -204,11 +209,25 @@ dgram_recv(struct onet_link *link, void *buf, uint16_t len)
     saddr.sll_ifindex = link->iface_idx;
     saddr.sll_halen = HW_ADDR_LEN;
 
+    /* Setup poll descriptor */
+    pfd.fd = link->sockfd;
+    pfd.events = POLLIN;
+
     /*
      * Wait until we get a packet for us with the right
      * protocol ID.
      */
     for (;;) {
+        if (miss_count > MAX_MISS) {
+            free(p);
+            return -ETIMEDOUT;
+        }
+
+        if (poll(&pfd, 1, RECV_POLL_MS) < 0) {
+            free(p);
+            return -ETIMEDOUT;
+        }
+
         recv_len = recvfrom(
             link->sockfd, p, dgram_len,
             0, (struct sockaddr *)&saddr,
@@ -221,19 +240,23 @@ dgram_recv(struct onet_link *link, void *buf, uint16_t len)
         src_mac = mac_swap(hdr->source);
 
         if (proto != PROTO_ID) {
+            ++miss_count;
             continue;
         }
 
         o1p_hdr = DGRAM_HDR(p);
         crc = crc32(o1p_hdr, sizeof(*o1p_hdr) - sizeof(crc));
         if (crc != o1p_hdr->crc32) {
+            ++miss_count;
             continue;
         }
 
+        link->last_recv = src_mac;
+
         /* Is this a squeak? */
         if (o1p_hdr->type == OTYPE_SQUEAK) {
             squeak_back(link, src_mac, dest_mac);
-            continue;
+            break;
         }
 
         /* If this is for everyone, take it */