1 files changed, 53 insertions, 18 deletions

diff --git a/ecp/src/ecp/vconn/vconn.c b/ecp/src/ecp/vconn/vconn.c
index c0e269f..780faab 100644
--- a/ecp/src/ecp/vconn/vconn.c
+++ b/ecp/src/ecp/vconn/vconn.c
@@ -72,16 +72,24 @@ static ssize_t handle_next(ECPConnection *conn, unsigned char *msg, size_t msg_s
     return ECP_SIZE_ECDH_PUB;
 }
 
-static ssize_t handle_relay(ECPConnection *conn, unsigned char *msg, size_t msg_size, ECP2Buffer *bufs) {
+static ssize_t handle_relay(ECPConnection *conn, unsigned char *_msg, size_t msg_size, ECP2Buffer *bufs) {
     ECPSocket *sock = conn->sock;
     ECPConnection *conn_next;
     ECPBuffer payload;
     unsigned char idx, _idx;
-    size_t _msg_size = msg_size;
+    unsigned char *msg = _msg;
+    uint16_t pkt_size;
+    size_t _pkt_size;
     ssize_t rv;
     int _rv;
 
-    if (msg_size < ECP_MIN_PKT) return ECP_ERR_SIZE;
+    if (msg_size < sizeof(uint16_t)) return ECP_ERR_SIZE;
+
+    pkt_size  = (uint16_t)msg[0] << 8;
+    pkt_size |= (uint16_t)msg[1];
+    msg += sizeof(uint16_t);
+
+    if (msg_size < sizeof(uint16_t) + pkt_size) return ECP_ERR_SIZE;
 
     idx = msg[ECP_SIZE_PROTO];
     if (idx == ECP_KEYID_INV) return ECP_ERR_KEYID;
@@ -154,8 +162,10 @@ static ssize_t handle_relay(ECPConnection *conn, unsigned char *msg, size_t msg_
             if (_idx == ECP_KEYID_PERMA) {
                 /* this is init reply */
                 msg[ECP_SIZE_PROTO] = (ECP_KEYID_PERMA << 4) | ECP_KEYID_NOKEY;
-                memmove(msg+ECP_SIZE_PROTO+1, msg+ECP_SIZE_PROTO+1+ECP_SIZE_ECDH_PUB, msg_size-(ECP_SIZE_PROTO+1+ECP_SIZE_ECDH_PUB));
-                _msg_size -= ECP_SIZE_ECDH_PUB;
+                memmove(msg+ECP_SIZE_PROTO+1, msg+ECP_SIZE_PROTO+1+ECP_SIZE_ECDH_PUB, pkt_size-(ECP_SIZE_PROTO+1+ECP_SIZE_ECDH_PUB));
+                pkt_size -= ECP_SIZE_ECDH_PUB;
+                _msg[0] = pkt_size >> 8;
+                _msg[1] = pkt_size;
             }
 
             break;
@@ -165,15 +175,22 @@ static ssize_t handle_relay(ECPConnection *conn, unsigned char *msg, size_t msg_
             return ECP_ERR_CTYPE;
     }
 
-    payload.buffer = msg - ECP_SIZE_MTYPE;
+    payload.buffer = ecp_msg_get_pld(ECP_MTYPE_RELAY, _msg);
     payload.size = bufs->payload->size - (payload.buffer - bufs->payload->buffer);
 
     ecp_pld_set_type(payload.buffer, payload.size, ECP_MTYPE_EXEC);
-    rv = ecp_pld_send(conn_next, bufs->packet, &payload, ECP_SIZE_MTYPE + _msg_size, ECP_SEND_FLAG_REPLY);
 
-    ecp_conn_refcount_dec(conn_next);
+    if (conn_next->parent) {
+        _pkt_size = 0;
+    } else {
+        _pkt_size = ECP_SIZE_VCONN_PKT;
+    }
 
+    rv = ecp_pld_send_wsize(conn_next, bufs->packet, &payload, ECP_SIZE_MTYPE + sizeof(uint16_t) + pkt_size, ECP_SEND_FLAG_REPLY, _pkt_size);
+    ecp_conn_refcount_dec(conn_next);
     if (rv < 0) return rv;
+
+    /* payload buffer is reused, so we must return msg_size */
     return msg_size;
 }
 
@@ -181,21 +198,31 @@ static ssize_t handle_relay(ECPConnection *conn, unsigned char *msg, size_t msg_
 
 static ssize_t handle_exec(ECPConnection *conn, unsigned char *msg, size_t msg_size, ECP2Buffer *bufs) {
     ECPBuffer *packet;
+    uint16_t pkt_size;
     ssize_t rv;
 
-    if (bufs == NULL) return ECP_ERR;
+    if (msg_size < sizeof(uint16_t)) return ECP_ERR_SIZE;
+
+    pkt_size  = (uint16_t)msg[0] << 8;
+    pkt_size |= (uint16_t)msg[1];
+    msg += sizeof(uint16_t);
+
+    if (msg_size < sizeof(uint16_t) + pkt_size) return ECP_ERR_SIZE;
 
     packet = bufs->packet;
-    if (packet->size < msg_size) return ECP_ERR_SIZE;
+    if (packet->size < pkt_size) return ECP_ERR_SIZE;
 
-    memcpy(packet->buffer, msg, msg_size);
-    rv = ecp_pkt_handle(conn->sock, conn, NULL, bufs, msg_size);
-    return rv;
+    memcpy(packet->buffer, msg, pkt_size);
+    rv = ecp_pkt_handle(conn->sock, conn, NULL, bufs, pkt_size);
+    if (rv < 0) return rv;
+
+    /* payload buffer can be reused, so we must return msg_size */
+    return msg_size;
 }
 
 ssize_t ecp_vconn_pack_parent(ECPConnection *conn, ECPBuffer *payload, ECPBuffer *packet, size_t pkt_size, ecp_tr_addr_t *addr) {
     unsigned char *msg;
-    size_t hdr_size;
+    size_t hdr_size, _pkt_size;
     ssize_t rv;
     int _rv;
 
@@ -206,10 +233,19 @@ ssize_t ecp_vconn_pack_parent(ECPConnection *conn, ECPBuffer *payload, ECPBuffer
     if (msg == NULL) return ECP_ERR;
 
     hdr_size = msg - payload->buffer;
-    if (payload->size < pkt_size+hdr_size) return ECP_ERR_SIZE;
+    if (payload->size < hdr_size + sizeof(uint16_t) + pkt_size) return ECP_ERR_SIZE;
+
+    msg[0] = pkt_size >> 8;
+    msg[1] = pkt_size;
+    memcpy(msg + sizeof(uint16_t), packet->buffer, pkt_size);
+
+    if (conn->parent) {
+        _pkt_size = 0;
+    } else {
+        _pkt_size = ECP_SIZE_VCONN_PKT;
+    }
 
-    memcpy(msg, packet->buffer, pkt_size);
-    rv = ecp_pack_conn(conn, packet, ECP_KEYID_INV, ECP_KEYID_INV, 0, NULL, NULL, payload, pkt_size+hdr_size, addr);
+    rv = ecp_pack_conn(conn, packet, _pkt_size, ECP_KEYID_INV, ECP_KEYID_INV, 0, NULL, NULL, payload, hdr_size + sizeof(uint16_t) + pkt_size, addr);
     return rv;
 }
 
@@ -272,7 +308,6 @@ int ecp_vconn_handle_open(ECPConnection *conn, ECP2Buffer *bufs) {
     /* we should release incoming packet before sending next open packet */
     ecp_tr_release(bufs->packet, 1);
     rv = _ecp_conn_open(vconn->next, conn, NULL, 1);
-    if (rv) rv = ECP_ERR_NEXT;
 
     return rv;
 }