1 files changed, 46 insertions, 12 deletions

diff --git a/src/rio.h b/src/rio.h
index c996c54f6..9576335e8 100644
--- a/src/rio.h
+++ b/src/rio.h
@@ -1,6 +1,6 @@
 /*
  * Copyright (c) 2009-2012, Pieter Noordhuis <pcnoordhuis at gmail dot com>
- * Copyright (c) 2009-2012, Salvatore Sanfilippo <antirez at gmail dot com>
+ * Copyright (c) 2009-2019, Salvatore Sanfilippo <antirez at gmail dot com>
  * All rights reserved.
  *
  * Redistribution and use in source and binary forms, with or without
@@ -35,6 +35,10 @@
 #include <stdio.h>
 #include <stdint.h>
 #include "sds.h"
+#include "connection.h"
+
+#define RIO_FLAG_READ_ERROR (1<<0)
+#define RIO_FLAG_WRITE_ERROR (1<<1)
 
 struct _rio {
     /* Backend functions.
@@ -51,8 +55,8 @@ struct _rio {
      * computation. */
     void (*update_cksum)(struct _rio *, const void *buf, size_t len);
 
-    /* The current checksum */
-    uint64_t cksum;
+    /* The current checksum and flags (see RIO_FLAG_*) */
+    uint64_t cksum, flags;
 
     /* number of bytes read or written */
     size_t processed_bytes;
@@ -73,14 +77,20 @@ struct _rio {
             off_t buffered; /* Bytes written since last fsync. */
             off_t autosync; /* fsync after 'autosync' bytes written. */
         } file;
-        /* Multiple FDs target (used to write to N sockets). */
+        /* Connection object (used to read from socket) */
+        struct {
+            connection *conn;   /* Connection */
+            off_t pos;    /* pos in buf that was returned */
+            sds buf;      /* buffered data */
+            size_t read_limit;  /* don't allow to buffer/read more than that */
+            size_t read_so_far; /* amount of data read from the rio (not buffered) */
+        } conn;
+        /* FD target (used to write to pipe). */
         struct {
-            int *fds;       /* File descriptors. */
-            int *state;     /* Error state of each fd. 0 (if ok) or errno. */
-            int numfds;
+            int fd;       /* File descriptor. */
             off_t pos;
             sds buf;
-        } fdset;
+        } fd;
     } io;
 };
 
@@ -91,11 +101,14 @@ typedef struct _rio rio;
  * if needed. */
 
 static inline size_t rioWrite(rio *r, const void *buf, size_t len) {
+    if (r->flags & RIO_FLAG_WRITE_ERROR) return 0;
     while (len) {
         size_t bytes_to_write = (r->max_processing_chunk && r->max_processing_chunk < len) ? r->max_processing_chunk : len;
         if (r->update_cksum) r->update_cksum(r,buf,bytes_to_write);
-        if (r->write(r,buf,bytes_to_write) == 0)
+        if (r->write(r,buf,bytes_to_write) == 0) {
+            r->flags |= RIO_FLAG_WRITE_ERROR;
             return 0;
+        }
         buf = (char*)buf + bytes_to_write;
         len -= bytes_to_write;
         r->processed_bytes += bytes_to_write;
@@ -104,10 +117,13 @@ static inline size_t rioWrite(rio *r, const void *buf, size_t len) {
 }
 
 static inline size_t rioRead(rio *r, void *buf, size_t len) {
+    if (r->flags & RIO_FLAG_READ_ERROR) return 0;
     while (len) {
         size_t bytes_to_read = (r->max_processing_chunk && r->max_processing_chunk < len) ? r->max_processing_chunk : len;
-        if (r->read(r,buf,bytes_to_read) == 0)
+        if (r->read(r,buf,bytes_to_read) == 0) {
+            r->flags |= RIO_FLAG_READ_ERROR;
             return 0;
+        }
         if (r->update_cksum) r->update_cksum(r,buf,bytes_to_read);
         buf = (char*)buf + bytes_to_read;
         len -= bytes_to_read;
@@ -124,11 +140,29 @@ static inline int rioFlush(rio *r) {
     return r->flush(r);
 }
 
+/* This function allows to know if there was a read error in any past
+ * operation, since the rio stream was created or since the last call
+ * to rioClearError(). */
+static inline int rioGetReadError(rio *r) {
+    return (r->flags & RIO_FLAG_READ_ERROR) != 0;
+}
+
+/* Like rioGetReadError() but for write errors. */
+static inline int rioGetWriteError(rio *r) {
+    return (r->flags & RIO_FLAG_WRITE_ERROR) != 0;
+}
+
+static inline void rioClearErrors(rio *r) {
+    r->flags &= ~(RIO_FLAG_READ_ERROR|RIO_FLAG_WRITE_ERROR);
+}
+
 void rioInitWithFile(rio *r, FILE *fp);
 void rioInitWithBuffer(rio *r, sds s);
-void rioInitWithFdset(rio *r, int *fds, int numfds);
+void rioInitWithConn(rio *r, connection *conn, size_t read_limit);
+void rioInitWithFd(rio *r, int fd);
 
-void rioFreeFdset(rio *r);
+void rioFreeFd(rio *r);
+void rioFreeConn(rio *r, sds* out_remainingBufferedData);
 
 size_t rioWriteBulkCount(rio *r, char prefix, long count);
 size_t rioWriteBulkString(rio *r, const char *buf, size_t len);