| 1 | /* rio.c is a simple stream-oriented I/O abstraction that provides an interface |
| 2 | * to write code that can consume/produce data using different concrete input |
| 3 | * and output devices. For instance the same rdb.c code using the rio abstraction |
| 4 | * can be used to read and write the RDB format using in-memory buffers or files. |
| 5 | * |
| 6 | * A rio object provides the following methods: |
| 7 | * read: read from stream. |
| 8 | * write: write to stream. |
| 9 | * tell: get the current offset. |
| 10 | * |
| 11 | * It is also possible to set a 'checksum' method that is used by rio.c in order |
| 12 | * to compute a checksum of the data written or read, or to query the rio object |
| 13 | * for the current checksum. */ |
| 14 | |
| 15 | #include "fmacros.h" |
| 16 | #include <string.h> |
| 17 | #include <stdio.h> |
| 18 | #include "rio.h" |
| 19 | #include "util.h" |
| 20 | |
| 21 | uint64_t crc64(uint64_t crc, const unsigned char *s, uint64_t l); |
| 22 | |
| 23 | /* Returns 1 or 0 for success/failure. */ |
| 24 | static size_t rioBufferWrite(rio *r, const void *buf, size_t len) { |
| 25 | r->io.buffer.ptr = sdscatlen(r->io.buffer.ptr,(char*)buf,len); |
| 26 | r->io.buffer.pos += len; |
| 27 | return 1; |
| 28 | } |
| 29 | |
| 30 | /* Returns 1 or 0 for success/failure. */ |
| 31 | static size_t rioBufferRead(rio *r, void *buf, size_t len) { |
| 32 | if (sdslen(r->io.buffer.ptr)-r->io.buffer.pos < len) |
| 33 | return 0; /* not enough buffer to return len bytes. */ |
| 34 | memcpy(buf,r->io.buffer.ptr+r->io.buffer.pos,len); |
| 35 | r->io.buffer.pos += len; |
| 36 | return 1; |
| 37 | } |
| 38 | |
| 39 | /* Returns read/write position in buffer. */ |
| 40 | static off_t rioBufferTell(rio *r) { |
| 41 | return r->io.buffer.pos; |
| 42 | } |
| 43 | |
| 44 | /* Returns 1 or 0 for success/failure. */ |
| 45 | static size_t rioFileWrite(rio *r, const void *buf, size_t len) { |
| 46 | return fwrite(buf,len,1,r->io.file.fp); |
| 47 | } |
| 48 | |
| 49 | /* Returns 1 or 0 for success/failure. */ |
| 50 | static size_t rioFileRead(rio *r, void *buf, size_t len) { |
| 51 | return fread(buf,len,1,r->io.file.fp); |
| 52 | } |
| 53 | |
| 54 | /* Returns read/write position in file. */ |
| 55 | static off_t rioFileTell(rio *r) { |
| 56 | return ftello(r->io.file.fp); |
| 57 | } |
| 58 | |
| 59 | static const rio rioBufferIO = { |
| 60 | rioBufferRead, |
| 61 | rioBufferWrite, |
| 62 | rioBufferTell, |
| 63 | NULL, /* update_checksum */ |
| 64 | 0, /* current checksum */ |
| 65 | { { NULL, 0 } } /* union for io-specific vars */ |
| 66 | }; |
| 67 | |
| 68 | static const rio rioFileIO = { |
| 69 | rioFileRead, |
| 70 | rioFileWrite, |
| 71 | rioFileTell, |
| 72 | NULL, /* update_checksum */ |
| 73 | 0, /* current checksum */ |
| 74 | { { NULL, 0 } } /* union for io-specific vars */ |
| 75 | }; |
| 76 | |
| 77 | void rioInitWithFile(rio *r, FILE *fp) { |
| 78 | *r = rioFileIO; |
| 79 | r->io.file.fp = fp; |
| 80 | } |
| 81 | |
| 82 | void rioInitWithBuffer(rio *r, sds s) { |
| 83 | *r = rioBufferIO; |
| 84 | r->io.buffer.ptr = s; |
| 85 | r->io.buffer.pos = 0; |
| 86 | } |
| 87 | |
| 88 | /* This function can be installed both in memory and file streams when checksum |
| 89 | * computation is needed. */ |
| 90 | void rioGenericUpdateChecksum(rio *r, const void *buf, size_t len) { |
| 91 | r->cksum = crc64(r->cksum,buf,len); |
| 92 | } |
| 93 | |
| 94 | /* ------------------------------ Higher level interface --------------------------- |
| 95 | * The following higher level functions use lower level rio.c functions to help |
| 96 | * generating the Redis protocol for the Append Only File. */ |
| 97 | |
| 98 | /* Write multi bulk count in the format: "*<count>\r\n". */ |
| 99 | size_t rioWriteBulkCount(rio *r, char prefix, int count) { |
| 100 | char cbuf[128]; |
| 101 | int clen; |
| 102 | |
| 103 | cbuf[0] = prefix; |
| 104 | clen = 1+ll2string(cbuf+1,sizeof(cbuf)-1,count); |
| 105 | cbuf[clen++] = '\r'; |
| 106 | cbuf[clen++] = '\n'; |
| 107 | if (rioWrite(r,cbuf,clen) == 0) return 0; |
| 108 | return clen; |
| 109 | } |
| 110 | |
| 111 | /* Write binary-safe string in the format: "$<count>\r\n<payload>\r\n". */ |
| 112 | size_t rioWriteBulkString(rio *r, const char *buf, size_t len) { |
| 113 | size_t nwritten; |
| 114 | |
| 115 | if ((nwritten = rioWriteBulkCount(r,'$',len)) == 0) return 0; |
| 116 | if (len > 0 && rioWrite(r,buf,len) == 0) return 0; |
| 117 | if (rioWrite(r,"\r\n",2) == 0) return 0; |
| 118 | return nwritten+len+2; |
| 119 | } |
| 120 | |
| 121 | /* Write a long long value in format: "$<count>\r\n<payload>\r\n". */ |
| 122 | size_t rioWriteBulkLongLong(rio *r, long long l) { |
| 123 | char lbuf[32]; |
| 124 | unsigned int llen; |
| 125 | |
| 126 | llen = ll2string(lbuf,sizeof(lbuf),l); |
| 127 | return rioWriteBulkString(r,lbuf,llen); |
| 128 | } |
| 129 | |
| 130 | /* Write a double value in the format: "$<count>\r\n<payload>\r\n" */ |
| 131 | size_t rioWriteBulkDouble(rio *r, double d) { |
| 132 | char dbuf[128]; |
| 133 | unsigned int dlen; |
| 134 | |
| 135 | dlen = snprintf(dbuf,sizeof(dbuf),"%.17g",d); |
| 136 | return rioWriteBulkString(r,dbuf,dlen); |
| 137 | } |