1 files changed, 43 insertions, 30 deletions

diff --git a/ecp/src/ecp/core.c b/ecp/src/ecp/core.c
index 2741df2..33debdb 100644
--- a/ecp/src/ecp/core.c
+++ b/ecp/src/ecp/core.c
@@ -2245,6 +2245,7 @@ ssize_t ecp_msg_handle(ECPConnection *conn, ecp_seq_t seq, unsigned char mtype,
 ssize_t ecp_pld_handle_one(ECPConnection *conn, ecp_seq_t seq, unsigned char *payload, size_t pld_size, ECP2Buffer *bufs) {
     unsigned char mtype;
     unsigned char *msg;
+    unsigned char frag_tot;
     size_t hdr_size, msg_size;
     ssize_t rv;
     int _rv;
@@ -2252,13 +2253,20 @@ ssize_t ecp_pld_handle_one(ECPConnection *conn, ecp_seq_t seq, unsigned char *pa
     _rv = ecp_pld_get_type(payload, pld_size, &mtype);
     if (_rv) return _rv;
 
+    frag_tot = 0;
+    if (mtype & ECP_MTYPE_FLAG_FRAG) {
+        _rv = ecp_pld_get_frag(payload, pld_size, NULL, &frag_tot, NULL);
+        if (_rv) return _rv;
+    }
+
+    ecp_timer_pop(conn, mtype, frag_tot);
+
     msg = ecp_pld_get_msg(payload, pld_size);
     if (msg == NULL) return ECP_ERR;
     hdr_size = msg - payload;
     msg_size = pld_size - hdr_size;
 
     rv = ecp_msg_handle(conn, seq, mtype, msg, msg_size, bufs);
-    ecp_timer_pop(conn, mtype);
     if (rv < 0) {
         ecp_err_handle(conn, mtype, rv);
         return rv;
@@ -2417,35 +2425,40 @@ ssize_t ecp_unpack(ECPSocket *sock, ECPConnection *parent, ecp_tr_addr_t *addr,
 
     ecp_buf2nonce(&nonce_pkt, nonce_buf);
     if (is_open) {
+        ecp_nonce_t nonce_offset;
+
         if (ECP_NONCE_LTE(nonce_pkt, nonce_conn)) {
-            ecp_nonce_t nonce_offset = nonce_conn - nonce_pkt;
+            ecp_ack_t _nonce_map;
 
-            if (nonce_offset < ECP_SIZE_ACKB) {
-                ecp_ack_t nonce_mask = ((ecp_ack_t)1 << nonce_offset);
+            nonce_offset = nonce_conn - nonce_pkt;
+            if (nonce_offset >= ECP_SIZE_ACKB) {
+                rv = ECP_ERR_SEQ;
+                goto unpack_fin;
+            }
 
-                if (nonce_mask & nonce_map) _rv = ECP_ERR_SEQ;
-                if (!_rv) nonce_in = nonce_conn;
-            } else {
-                _rv = ECP_ERR_SEQ;
+            _nonce_map = ((ecp_ack_t)1 << nonce_offset);
+            if (nonce_map & _nonce_map) {
+                rv = ECP_ERR_SEQ;
+                goto unpack_fin;
             }
+
+            nonce_map |= _nonce_map;
+            nonce_in = nonce_conn;
         } else {
-            ecp_nonce_t nonce_offset = nonce_pkt - nonce_conn;
+            nonce_offset = nonce_pkt - nonce_conn;
 
-            if (nonce_offset < ECP_MAX_SEQ_FWD) {
-                if (nonce_offset < ECP_SIZE_ACKB) {
-                    nonce_map = nonce_map << nonce_offset;
-                } else {
-                    nonce_map = 0;
-                }
-                nonce_map |= 1;
-                nonce_in = nonce_pkt;
+            if (nonce_offset >= ECP_MAX_SEQ_FWD) {
+                rv = ECP_ERR_SEQ;
+                goto unpack_fin;
+            }
+
+            if (nonce_offset < ECP_SIZE_ACKB) {
+                nonce_map = nonce_map << nonce_offset;
             } else {
-                _rv = ECP_ERR_SEQ;
+                nonce_map = 0;
             }
-        }
-        if (_rv) {
-            rv = _rv;
-            goto unpack_fin;
+            nonce_map |= 1;
+            nonce_in = nonce_pkt;
         }
     }
 
@@ -2876,23 +2889,23 @@ int ecp_pld_set_type(unsigned char *pld, size_t pld_size, unsigned char mtype) {
     return ECP_OK;
 }
 
-int ecp_pld_get_frag(unsigned char *pld, size_t pld_size, unsigned char *frag_cnt, unsigned char *frag_tot, uint16_t *frag_size) {
+int ecp_pld_get_frag(unsigned char *pld, size_t pld_size, unsigned char *frag_cnt, unsigned char *frag_tot, uint16_t *frag_sz) {
     if (pld_size < ECP_SIZE_MTYPE) return ECP_ERR_SIZE;
 
     if (!(pld[0] & ECP_MTYPE_FLAG_FRAG)) return ECP_ERR;
     if (pld_size < (ECP_SIZE_MTYPE + ECP_SIZE_MT_FLAG(pld[0]))) return ECP_ERR_SIZE;
     if (pld[2] == 0) return ECP_ERR;
 
-    *frag_cnt = pld[1];
-    *frag_tot = pld[2];
-    *frag_size = \
+    if (frag_cnt) *frag_cnt = pld[1];
+    if (frag_tot) *frag_tot = pld[2];
+    if (frag_sz) *frag_sz = \
         ((uint16_t)pld[3] << 8) | \
         ((uint16_t)pld[4]);
 
     return ECP_OK;
 }
 
-int ecp_pld_set_frag(unsigned char *pld, size_t pld_size, unsigned char frag_cnt, unsigned char frag_tot, uint16_t frag_size) {
+int ecp_pld_set_frag(unsigned char *pld, size_t pld_size, unsigned char frag_cnt, unsigned char frag_tot, uint16_t frag_sz) {
     if (pld_size < ECP_SIZE_MTYPE) return ECP_ERR_SIZE;
 
     if (!(pld[0] & ECP_MTYPE_FLAG_FRAG)) return ECP_ERR;
@@ -2900,8 +2913,8 @@ int ecp_pld_set_frag(unsigned char *pld, size_t pld_size, unsigned char frag_cnt
 
     pld[1] = frag_cnt;
     pld[2] = frag_tot;
-    pld[3] = frag_size >> 8;
-    pld[4] = frag_size;
+    pld[3] = frag_sz >> 8;
+    pld[4] = frag_sz;
 
     return ECP_OK;
 }
@@ -2915,7 +2928,7 @@ int ecp_pld_get_pts(unsigned char *pld, size_t pld_size, ecp_pts_t *pts) {
     if (pld_size < (ECP_SIZE_MTYPE + ECP_SIZE_MT_FLAG(pld[0]))) return ECP_ERR_SIZE;
 
     offset = ECP_SIZE_MTYPE + ECP_SIZE_MT_FRAG(pld[0]);
-    *pts = \
+    if (pts) *pts = \
         ((ecp_pts_t)pld[offset]     << 24) | \
         ((ecp_pts_t)pld[offset + 1] << 16) | \
         ((ecp_pts_t)pld[offset + 2] << 8)  | \