From b2724409540e0744070f7d7bef12d501079fb386 Mon Sep 17 00:00:00 2001 From: Jose Marcial Vieira Bisneto Date: Fri, 16 Nov 2018 11:33:52 -0300 Subject: [PATCH] Handle protocol message with a length while reading. --- net/connectionstream.c | 19 ++++++++++++++++++- 1 file changed, 18 insertions(+), 1 deletion(-) diff --git a/net/connectionstream.c b/net/connectionstream.c index a036e60..ccca82b 100644 --- a/net/connectionstream.c +++ b/net/connectionstream.c @@ -73,7 +73,24 @@ int libp2p_net_connection_read(void* stream_context, struct StreamMessage** msg, uint8_t* result_buffer = NULL; int current_size = 0; while (1) { - int retVal = socket_read(ctx->socket_descriptor, (char*)&buffer[0], 4096, 0, timeout_secs); + int retVal; + if (current_size) { + retVal = socket_read(ctx->socket_descriptor, (char*)&buffer[0], sizeof(buffer), 0, timeout_secs); + } else { + // initial reading, see if it is a protocol request with length. + retVal = socket_read(ctx->socket_descriptor, (char*)&buffer[0], 2, 0, timeout_secs); + if (retVal == 2) { + if (buffer[0] > 0 && buffer[1] == '/') { + // if the second byte is a slash, it may be a protocol message with a length in the first byte. + retVal = socket_read(ctx->socket_descriptor, (char*)&buffer[2], buffer[0]-1, 0, timeout_secs); + } else { + // if it isn't, read as much as possible. + retVal = socket_read(ctx->socket_descriptor, (char*)&buffer[2], sizeof(buffer)-2, 0, timeout_secs); + } + if (retVal > 0) + retVal += 2; // compensate the first read. + } + } ctx->last_comm_epoch = time(NULL); libp2p_logger_debug("connectionstream", "Retrieved %d bytes from socket %d.\n", retVal, ctx->socket_descriptor); if (retVal < 1) { // get out of the loop