ucw/fb-socket.c

   1 /*
   2  *      UCW Library -- Fast Buffered I/O on Sockets with Timeouts
   3  *
   4  *      (c) 2008 Martin Mares <mj@ucw.cz>
   5  *      (c) 2017 Pavel Charvat <pchar@ucw.cz>
   6  *
   7  *      This software may be freely distributed and used according to the terms
   8  *      of the GNU Lesser General Public License.
   9  */
  10
  11 #include <ucw/lib.h>
  12 #include <ucw/fastbuf.h>
  13 #include <ucw/fb-socket.h>
  14
  15 #include <stdio.h>
  16 #include <unistd.h>
  17 #include <poll.h>
  18 #include <errno.h>
  19
  20 struct fb_sock {
  21   struct fastbuf fb;
  22   struct fbsock_params par;
  23   byte buf[0];
  24 };
  25
  26 #define FB_SOCK(f) ((struct fb_sock *)(f))
  27
  28 static int
  29 fbs_refill(struct fastbuf *f)
  30 {
  31   struct fbsock_params *p = &FB_SOCK(f)->par;
  32   struct pollfd pf = {
  33       .fd = p->fd,
  34       .events = POLLIN
  35   };
  36
  37   for (;;)
  38     {
  39       int e = poll(&pf, 1, p->timeout_ms);
  40       if (e < 0)
  41         {
  42           p->err(p->data, FBSOCK_READ, "read error");
  43           return 0;
  44         }
  45       if (!e)
  46         {
  47           p->err(p->data, FBSOCK_READ | FBSOCK_TIMEOUT, "read timeout");
  48           return 0;
  49         }
  50
  51       f->bptr = f->buffer;
  52       int l = read(p->fd, f->buffer, f->bufend-f->buffer);
  53       if (l < 0)
  54         {
  55           if (errno == EINTR || errno == EAGAIN)
  56             continue;
  57           p->err(p->data, FBSOCK_READ, "read error");
  58           return 0;
  59         }
  60       f->bstop = f->buffer + l;
  61       f->pos += l;
  62       return l;
  63     }
  64 }
  65
  66 static void
  67 fbs_spout(struct fastbuf *f)
  68 {
  69   struct fbsock_params *p = &FB_SOCK(f)->par;
  70   struct pollfd pf = {
  71       .fd = p->fd,
  72       .events = POLLOUT,
  73   };
  74
  75   int l = f->bptr - f->bstop;
  76   f->bptr = f->bstop;
  77   f->pos += l;
  78   char *buf = f->bstop;
  79
  80   while (l)
  81     {
  82       int e = poll(&pf, 1, p->timeout_ms);
  83       if (e < 0)
  84         {
  85           p->err(p->data, FBSOCK_WRITE, "write error");
  86           return;
  87         }
  88       if (!e)
  89         {
  90           p->err(p->data, FBSOCK_WRITE | FBSOCK_TIMEOUT, "write timeout");
  91           return;
  92         }
  93
  94       e = write(p->fd, buf, l);
  95       if (e < 0)
  96         {
  97           if (errno == EINTR || errno == EAGAIN)
  98             continue;
  99           p->err(p->data, FBSOCK_WRITE, "write error");
 100           return;
 101         }
 102       buf += e;
 103       l -= e;
 104     }
 105 }
 106
 107 static void
 108 fbs_close(struct fastbuf *f)
 109 {
 110   if (!FB_SOCK(f)->par.fd_is_shared)
 111     close(FB_SOCK(f)->par.fd);
 112   xfree(f);
 113 }
 114
 115 struct fastbuf *
 116 fbsock_create(struct fbsock_params *p)
 117 {
 118   struct fb_sock *F = xmalloc(sizeof(*F) + p->bufsize);
 119   struct fastbuf *f = &F->fb;
 120
 121   bzero(F, sizeof(*F));
 122   F->par = *p;
 123   f->buffer = F->buf;
 124   f->bptr = f->bstop = f->buffer;
 125   f->bufend = f->buffer + p->bufsize;
 126   f->name = "<socket>";
 127   f->refill = fbs_refill;
 128   f->spout = fbs_spout;
 129   f->close = fbs_close;
 130   f->can_overwrite_buffer = 1;
 131   return f;
 132 }
 133
 134 #ifdef TEST
 135
 136 #include <stdlib.h>
 137
 138 static void test_err(void *x UNUSED, uint flags, char *msg UNUSED)
 139 {
 140   if (flags & FBSOCK_READ)
 141     printf("READ");
 142   else if (flags & FBSOCK_WRITE)
 143     printf("WRITE");
 144   if (flags & FBSOCK_TIMEOUT)
 145     printf(" TIMEOUT\n");
 146   else
 147     printf(" ERROR\n");
 148   exit(0);
 149 }
 150
 151 int main(void)
 152 {
 153   int fd[2];
 154   if (pipe(fd) < 0)
 155     ASSERT(0);
 156
 157   struct fbsock_params p = {
 158       .fd = fd[0],
 159       .bufsize = 16,
 160       .timeout_ms = 100,
 161       .err = test_err
 162   };
 163   struct fastbuf *f = fbsock_create(&p);
 164
 165   bputsn(f, "Oook!");           // This fits in PIPE_BUF
 166   bflush(f);
 167
 168   char buf[256];
 169   if (!bgets(f, buf, sizeof(buf)))
 170     die("bgets failed");
 171   if (strcmp(buf, "Oook!"))
 172     die("Misread input");
 173
 174   bgets(f, buf, sizeof(buf));
 175   puts("WRONG");
 176   exit(0);
 177 }
 178
 179 #endif