io: allow slurping files with a given size
[libjh.git] / io.c
1 // Copyright (2013) Jann Horn <jann@thejh.net>
2 // This code is licensed under the AGPLv3.
3
4 HEADER #include <sys/types.h>
5 #include <unistd.h>
6 #include <errno.h>
7 #include <stdlib.h>
8 #include <stdbool.h>
9 #include <sys/stat.h>
10 #include <fcntl.h>
11 #include <string.h>
12 #include <assert.h>
13 #include <sys/mman.h>
14
15 // Wrapper for `read` that retries on partial reads.
16 // If last_res is non-NULL, it will be filled with the
17 // result of the last read() call (but all positive numbers
18 // become 1). In other words:
19 //  - -1: error
20 //  - 0: stream ended
21 //  - 1: no problems occured
22 PUBLIC_FN ssize_t read_nointr(int fd, void *buf, size_t count, int *last_res) {
23   errno = 0;
24   size_t done = 0;
25   while (done < count) {
26     ssize_t part_res = read(fd, buf+done, count-done);
27     if (part_res <= 0) {
28       if (last_res) *last_res = part_res;
29       if (done) return done;
30       return part_res;
31     }
32     done += part_res;
33   }
34   if (last_res) *last_res = 1;
35   return done;
36 }
37
38 // Wrapper for `write` that retries on partial writes.
39 // If last_res is non-NULL, it will be filled with the
40 // result of the last write() call (but all positive numbers
41 // become 1). In other words:
42 //  - -1: error
43 //  - 0: stream ended
44 //  - 1: no problems occured
45 PUBLIC_FN ssize_t write_nointr(int fd, void *buf, size_t count, int *last_res) {
46   errno = 0;
47   size_t done = 0;
48   while (done < count) {
49     ssize_t part_res = write(fd, buf+done, count-done);
50     if (part_res <= 0) {
51       if (last_res) *last_res = part_res;
52       if (done) return done;
53       return part_res;
54     }
55     done += part_res;
56   }
57   if (last_res) *last_res = 1;
58   return done;
59 }
60
61 // Read all data from the given file descriptor. Tries fstat()+read()
62 // first, but if fstat() doesn't work, falls back to multiple read()s.
63 // Specify JH_NO_STAT to prevent the fstat() call.
64 // The return value is a malloc'd buffer.
65 // The buffer will be null-terminated, so you can read text files with
66 // this and can specify len_out as NULL.
67 PUBLIC_CONST JH_SLURP_NO_STAT 1
68 PUBLIC_CONST JH_SLURP_REALLOC 2 /* realloc result block if it saves RAM */
69 PUBLIC_CONST JH_SLURP_8BYTE_PAD 4 /* pad buffer with eight nullbytes, not one */
70 PUBLIC_CONST JH_SLURP_FIXSIZE 8 /* len_out is actually input and specifies
71 HEADER                             the expected length. if there is more data,
72 HEADER                             that might be ignored silently. */
73 PUBLIC_FN void *slurp_fd(int fd, size_t *len_out, int flags) {
74   int errno_;
75   
76   // Let's just guess that the file is 1023 bytes. Will become 1024 with
77   // the nullbyte.
78   size_t size_guess = 1023;
79   bool trusted_guess = false; /* can we rely on the guess? */
80   
81   if (flags&JH_SLURP_FIXSIZE) {
82     size_guess = *len_out;
83     trusted_guess = true;
84     // if we want a fixed size, of course we don't want to stat
85     flags |= JH_SLURP_NO_STAT;
86   }
87   
88   int padlen = (flags&JH_SLURP_8BYTE_PAD)?8:1;
89   
90   // If we can determine the exact size, we don't have to guess. So try
91   // to determine the exact size.
92   if (!(flags&JH_SLURP_NO_STAT)) {
93     struct stat st;
94     if (fstat(fd, &st) == 0) {
95       if (st.st_size > 0) {
96         size_guess = st.st_size;
97         trusted_guess = true;
98       }
99     }
100   }
101   
102   char *buf = NULL;
103   int done = 0;
104   
105   while (1) {
106     buf = realloc(buf, size_guess+padlen);
107     if (buf == NULL) return NULL;
108     int last_res;
109     ssize_t read_res = read_nointr(fd, buf+done, size_guess-done, &last_res);
110     if (last_res == -1) { errno_=errno; free(buf); errno=errno_; return NULL; }
111     done += read_res;
112     if (last_res == 0 || trusted_guess) {
113       // out
114       if (padlen == 1) {
115         buf[done] = '\0';
116       } else {
117         *(uint64_t*)(buf+done) = 0;
118       }
119       if (done != size_guess && (flags&JH_SLURP_REALLOC)) {
120         // Well, it'd be weird if shrinking could fail... but meh, I can't find
121         // an explicit statement about this being disallowed, so try staying on
122         // the safe side.
123         char *buf_ = realloc(buf, done+padlen);
124         if (buf_) buf = buf_;
125       }
126       if (done != size_guess && (flags&JH_SLURP_FIXSIZE)) {
127         free(buf);
128         // not the most correct error message ever, but whatever.
129         errno = EFBIG;
130         return NULL;
131       }
132       if (len_out) *len_out = done;
133       return buf;
134     }
135     size_guess<<=1; // try two times the buffer size
136   }
137 }
138
139 PUBLIC_FN void *slurp_file(char *path, size_t *len_out, int flags) {
140   int fd = open(path, O_RDONLY|O_CLOEXEC);
141   if (fd == -1) return NULL;
142   char *res = slurp_fd(fd, len_out, flags);
143   int errno_ = errno;
144   close(fd);
145   errno = errno_;
146   return res;
147 }
148
149 // Write data into a file. len can be -1; that means that
150 // the real length is strlen(buf).
151 PUBLIC_FN int write_file(char *path, void *buf, ssize_t len, int open_flags) {
152   if (len == -1) len = strlen(buf);
153   
154   int fd = open(path, open_flags|O_CLOEXEC|O_WRONLY, 0777);
155   if (fd == -1) return 1;
156   ssize_t write_res = write_nointr(fd, buf, len, NULL);
157   int write_errno = errno;
158   int close_res = close(fd);
159   if (write_res != len) {
160     errno = write_errno;
161     return 1;
162   }
163   if (close_res) return 1;
164   return 0;
165 }
166
167 PUBLIC_FN void *mmap_file(void *addr, size_t length, int prot, int flags, char *path, off_t offset) {
168   assert((flags&MAP_ANONYMOUS) == 0);
169   assert(prot&(PROT_EXEC|PROT_READ|PROT_WRITE));
170   int open_flags = O_CLOEXEC;
171   if ((prot & (PROT_READ|PROT_EXEC)) && !(prot&PROT_WRITE)) open_flags |= O_RDONLY;
172   else if ((prot & (PROT_READ|PROT_EXEC))) open_flags |= O_RDWR;
173   else open_flags |= O_WRONLY;
174   int fd = open(path, open_flags);
175   if (fd == -1) return NULL;
176   void *r = mmap(addr, length, prot, flags, fd, offset);
177   int errno_ = errno;
178   close(fd);
179   errno = errno_;
180   return (r==MAP_FAILED) ? NULL : r;
181 }