Browse Source

Use one buffer for connsok line receives without relying on MSG_PEEK and abort appropriately when recv() returns 0

master
Con Kolivas 11 years ago
parent
commit
b57ffe23d6
  1. 57
      src/ckpool.c
  2. 2
      src/ckpool.h

57
src/ckpool.c

@ -124,24 +124,31 @@ bool ping_main(ckpool_t *ckp)
return true; return true;
} }
/* Peek in a socket, and then receive only one line at a time, allocing enough /* Read from a socket into cs->buf till we get an '\n', converting it to '\0'
* memory in *buf */ * and storing how much extra data we've received, to be moved to the beginning
* of the buffer for use on the next receive. */
int read_socket_line(connsock_t *cs, int timeout) int read_socket_line(connsock_t *cs, int timeout)
{ {
size_t buflen = 0, bufofs = 0;
int ret = -1, bufsiz;
char *eom = NULL; char *eom = NULL;
tv_t tv_timeout; tv_t tv_timeout;
size_t buflen;
int ret = -1;
fd_set rd; fd_set rd;
dealloc(cs->buf);
if (unlikely(cs->fd < 0)) if (unlikely(cs->fd < 0))
return ret; goto out;
if (unlikely(!cs->buf))
cs->buf = ckzalloc(PAGESIZE);
else if (cs->buflen) {
memmove(cs->buf, cs->buf + cs->bufofs, cs->buflen);
cs->buf[cs->buflen] = '\0';
cs->bufofs = cs->buflen;
cs->buflen = 0;
} else
cs->bufofs = 0;
bufsiz = PAGESIZE;
while (!eom) { while (!eom) {
char readbuf[PAGESIZE] = {}; char readbuf[PAGESIZE] = {};
int extralen;
FD_ZERO(&rd); FD_ZERO(&rd);
FD_SET(cs->fd, &rd); FD_SET(cs->fd, &rd);
@ -155,33 +162,27 @@ int read_socket_line(connsock_t *cs, int timeout)
LOGERR("Select failed in read_socket_line"); LOGERR("Select failed in read_socket_line");
goto out; goto out;
} }
ret = recv(cs->fd, readbuf, bufsiz - 2, MSG_PEEK); ret = recv(cs->fd, readbuf, PAGESIZE - 4, 0);
if (ret < 0) { if (ret < 1) {
LOGERR("Failed to recv in read_socket_line"); LOGERR("Failed to recv in read_socket_line");
goto out; goto out;
} }
if (!ret) buflen = cs->bufofs + ret + 1;
continue;
eom = strchr(readbuf, '\n');
if (eom)
extralen = eom - readbuf + 1;
else
extralen = ret;
buflen += extralen + 1;
align_len(&buflen); align_len(&buflen);
cs->buf = realloc(cs->buf, buflen); cs->buf = realloc(cs->buf, buflen);
if (unlikely(!cs->buf)) if (unlikely(!cs->buf))
quit(1, "Failed to alloc buf of %d bytes in read_socket_line", (int)buflen); quit(1, "Failed to alloc rem of %d bytes in read_socket_line", (int)buflen);
ret = recv(cs->fd, cs->buf + bufofs, extralen, 0); memcpy(cs->buf + cs->bufofs, readbuf, ret);
if (ret < 0) { eom = strchr(cs->buf, '\n');
LOGERR("Failed to recv %d bytes in read_socket_line", (int)buflen); cs->bufofs += ret;
goto out; cs->buf[cs->bufofs] = '\0';
} }
bufofs += ret; if (eom < cs->buf + cs->bufofs) {
} cs->buflen = cs->buf + cs->bufofs - eom - 1;
eom = cs->buf + bufofs; cs->bufofs = eom - cs->buf + 1;
} else
cs->buflen = cs->bufofs = 0;
eom[0] = '\0'; eom[0] = '\0';
ret = bufofs + 1;
out: out:
if (ret < 1) if (ret < 1)
dealloc(cs->buf); dealloc(cs->buf);

2
src/ckpool.h

@ -35,6 +35,8 @@ struct connsock {
char *port; char *port;
char *auth; char *auth;
char *buf; char *buf;
int bufofs;
int buflen;
}; };
typedef struct connsock connsock_t; typedef struct connsock connsock_t;

Loading…
Cancel
Save