More safety checks
[blerg.git] / database / database.c
1 /* Blerg is (C) 2011 The Dominion of Awesome, and is distributed under a
2  * BSD-style license.  Please see the COPYING file for details.
3  */
4 #include <stdint.h>
5 #include <stdlib.h>
6 #include <string.h>
7 #include <stdio.h>
8 #include <time.h>
9 #include <unistd.h>
10 #include <sys/stat.h>
11 #include <sys/types.h>
12 #include <sys/mman.h>
13 #include <sys/file.h>
14 #include <fcntl.h>
15 #include "database.h"
16 #include "subscription.h"
17 #include "util.h"
18 #include "config.h"
19
20 #define CHECK_VALID_BLERG(r)                               \
21         if (blerg == NULL) {                               \
22                 fprintf(stderr, "Invalid struct blerg\n"); \
23                 return r;                                  \
24         }
25
26 uint64_t blerg_get_record_count(struct blerg *blerg) {
27         uint64_t count;
28         flock(blerg->meta_fd, LOCK_SH);
29         count = blerg->meta->sequence;
30         flock(blerg->meta_fd, LOCK_UN);
31         return count;
32 }
33
34 /* Returns last usable record */
35 uint64_t blerg_increment_record_count(struct blerg *blerg) {
36         uint64_t count;
37         flock(blerg->meta_fd, LOCK_EX);
38         count = blerg->meta->sequence++;
39         flock(blerg->meta_fd, LOCK_UN);
40         return count;
41 }
42
43 void blerg_segment_close(struct blerg *blerg) {
44         if (blerg->data != NULL)
45                 munmap((void *)blerg->data, blerg->data_size);
46         if (blerg->data_fd != -1)
47                 close(blerg->data_fd);
48         if (blerg->index != NULL)
49                 munmap((void *)blerg->index, RECORDS_PER_SEGMENT * sizeof(struct record));
50         if (blerg->index_fd != -1)
51                 close(blerg->index_fd);
52 }
53
54 int blerg_remap_data(struct blerg *blerg) {
55         struct stat st;
56
57         if (blerg->data != NULL)
58                 munmap(blerg->data, blerg->data_size);
59         fstat(blerg->data_fd, &st);
60         if (st.st_size == 0) {
61                 /* Can't map an empty data file. */
62                 return 1;
63         }
64         blerg->data = (char *) mmap(NULL, st.st_size, PROT_READ, MAP_SHARED, blerg->data_fd, 0);
65         if (blerg->data == MAP_FAILED) {
66                 perror("Could not remap data");
67                 return 0;
68         }
69         blerg->data_size = st.st_size;
70         return 1;
71 }
72
73 int blerg_segment_switch(struct blerg *blerg, int new_segment) {
74         char filename[512];
75         uint64_t max_sequence_no = blerg_get_record_count(blerg);
76         struct stat st;
77
78         if (blerg->index != NULL && blerg->data != NULL && new_segment == blerg->current_segment) {
79                 return 1;
80         }
81
82         if (new_segment > max_sequence_no / RECORDS_PER_SEGMENT) {
83                 fprintf(stderr, "Cannot switch to sequence beyond last record\n");
84                 return 0;
85         }
86         if (new_segment < 0) {
87                 fprintf(stderr, "Cannot switch to negative segment\n");
88                 return 0;
89         }
90
91         blerg_segment_close(blerg);
92
93         /* Load and map the index */
94         snprintf(filename, 512, "%s/index%d", blerg->base_path, new_segment);
95         blerg->index_fd = open(filename, O_RDWR | O_CREAT, 0600);
96         if (blerg->index_fd == -1) {
97                 perror("Could not open index");
98                 goto open_failed_index_open;
99         }
100         flock(blerg->index_fd, LOCK_EX);
101         fstat(blerg->index_fd, &st);
102         if (st.st_size == 0) {
103                 /* ftruncate() means never having to say you're sorry.  Sorry
104                    in this case meaning "allocating disk space for a 1MB file
105                    full or zeroes". */
106                 ftruncate(blerg->index_fd, RECORDS_PER_SEGMENT * sizeof(struct record));
107         }
108         flock(blerg->index_fd, LOCK_UN);
109
110         blerg->index = (struct record *) mmap(NULL, RECORDS_PER_SEGMENT * sizeof(struct record), PROT_READ | PROT_WRITE, MAP_SHARED, blerg->index_fd, 0);
111         if (blerg->index == MAP_FAILED) {
112                 perror("Could not mmap index");
113                 goto open_failed_index_mmap;
114         }
115
116         /* Load data file */
117         sprintf(filename, "%s/data%d", blerg->base_path, new_segment);
118         blerg->data_fd = open(filename, O_RDWR | O_APPEND | O_CREAT, 0600);
119         if (blerg->data_fd == -1) {
120                 perror("Could not open data");
121                 goto open_failed_data_open;
122         }
123
124         if (!blerg_remap_data(blerg)) {
125                 goto open_failed_data_mmap;
126         }
127
128         blerg->current_segment = new_segment;
129
130         return 1;
131
132 open_failed_data_mmap:
133         close(blerg->data_fd);
134 open_failed_data_open:
135         munmap((void *)blerg->index, RECORDS_PER_SEGMENT * sizeof(struct record));
136 open_failed_index_mmap:
137         close(blerg->index_fd);
138 open_failed_index_open:
139         return 0;
140 }
141
142 int blerg_exists(const char *name) {
143         int namelen = strlen(name);
144         char filename[512];
145
146         if (!valid_name(name)) {
147                 fprintf(stderr, "Invalid name\n");
148                 return 0;
149         }
150
151         snprintf(filename, 512, "%s/%s", DATA_PATH, name);
152         if (access(filename, F_OK) == -1)
153                 return 0;
154         else
155                 return 1;
156 }
157
158 struct blerg *blerg_open(const char *name) {
159         int namelen = strlen(name);
160         char filename[512];
161         struct stat st;
162         uint64_t sequence;
163
164         if (!valid_name(name)) {
165                 fprintf(stderr, "Invalid name\n");
166                 return NULL;
167         }
168         struct blerg *blerg = malloc(sizeof(struct blerg));
169         if (!blerg) {
170                 perror("Cannot allocate memory for blerg");
171                 goto open_failed_blerg_malloc;
172         }
173         blerg->name = malloc(namelen + 1);
174         memcpy(blerg->name, name, namelen + 1);
175         blerg->meta_fd = blerg->index_fd = blerg->data_fd = -1;
176         blerg->meta = NULL;
177         blerg->index = NULL;
178         blerg->data = NULL;
179
180         /* Make the directory if it doesn't exist */
181         blerg->base_path = malloc(512);
182         snprintf(blerg->base_path, 512, "%s/%s", DATA_PATH, name);
183         if (access(blerg->base_path, F_OK) == -1)
184                 mkdir(blerg->base_path, 0755);
185
186         /* Open and map metadata */
187         snprintf(filename, 512, "%s/meta", blerg->base_path);
188         blerg->meta_fd = open(filename, O_RDWR | O_CREAT, 0600);
189         if (blerg->meta_fd == -1) {
190                 perror("Could not open metadata");
191                 goto open_failed_meta_open;
192         }
193         fstat(blerg->meta_fd, &st);
194         if (st.st_size < sizeof(struct meta)) {
195                 /* Extend the file if sizeof(struct meta) is larger than the
196                    file. This allows seamless upgrades as long as struct meta
197                    only adds members. */
198                 posix_fallocate(blerg->meta_fd, 0, sizeof(struct meta));
199         }
200         blerg->meta = (struct meta *) mmap(NULL, sizeof(struct meta), PROT_READ | PROT_WRITE, MAP_SHARED, blerg->meta_fd, 0);
201         if (blerg->meta == MAP_FAILED) {
202                 perror("Could not map metadata");
203                 goto open_failed_meta_mmap;
204         }
205
206         /* Open and map index and data for the current segment */
207         blerg->current_segment = blerg_get_record_count(blerg) / RECORDS_PER_SEGMENT;
208         if (!blerg_segment_switch(blerg, blerg->current_segment)) {
209                 fprintf(stderr, "Could not switch segment\n");
210                 goto open_failed_segment_switch;
211         }
212
213         return blerg;
214
215 open_failed_segment_switch:
216         munmap((void *)blerg->meta, sizeof(struct meta));
217 open_failed_meta_mmap:
218         close(blerg->meta_fd);
219 open_failed_meta_open:
220         free(blerg->name);
221         free(blerg);
222 open_failed_blerg_malloc:
223         return NULL;
224 }
225
226 int blerg_close(struct blerg *blerg) {
227         CHECK_VALID_BLERG(0)
228         blerg_segment_close(blerg);
229         munmap((void *)blerg->meta, sizeof(struct meta));
230         close(blerg->meta_fd);
231         free(blerg->base_path);
232         free(blerg->name);
233         free(blerg);
234         return 1;
235 }
236
237 int blerg_store(struct blerg *blerg, const char *data, int len) {
238         struct stat st;
239         int n;
240
241         CHECK_VALID_BLERG(-1)
242
243         if (len > MAX_RECORD_SIZE || len <= 0) {
244                 fprintf(stderr, "len out of bounds\n");
245                 return -1;
246         }
247
248         flock(blerg->index_fd, LOCK_EX);
249         flock(blerg->data_fd, LOCK_EX);
250
251         uint64_t record = blerg_get_record_count(blerg);
252         if (record == -1) {  /* Intentional signed-unsigned coercion */
253                 fprintf(stderr, "Could not find free record\n");
254                 return -1;
255         }
256         int segment = record / RECORDS_PER_SEGMENT;
257         if (segment != blerg->current_segment)
258                 blerg_segment_switch(blerg, segment);
259         int seg_rec = record % RECORDS_PER_SEGMENT;
260
261         /* Get the position for the new data */
262         fstat(blerg->data_fd, &st);
263         int curpos = st.st_size;
264
265         /* Write data to the data log */
266         n = write(blerg->data_fd, data, len);
267         if (n < len) {
268                 perror("Could not write data");
269                 /* Truncate anything we may have written */
270                 ftruncate(blerg->data_fd, curpos);
271                 return -1;
272         }
273
274         /* Update the index */
275         blerg->index[seg_rec].flags = 0x0001;
276         blerg->index[seg_rec].offset = curpos;
277         blerg->index[seg_rec].length = len;
278         blerg->index[seg_rec].timestamp = time(NULL);
279
280         /* And finally increment the record count */
281         blerg_increment_record_count(blerg);
282
283         flock(blerg->data_fd, LOCK_UN);
284         flock(blerg->index_fd, LOCK_UN);
285
286         /* Now do those dirty microblogging deeds */
287         tag_scan(blerg->name, data, len, record);
288         subscription_notify(blerg->name, record);
289
290         return record;
291 }
292
293 int blerg_fetch(struct blerg *blerg, int record, char **data, int *length) {
294         CHECK_VALID_BLERG(0)
295         if (record < 0 || record >= blerg_get_record_count(blerg)) {
296                 fprintf(stderr, "Invalid record\n");
297                 return 0;
298         }
299         if (data == NULL || length == NULL) {
300                 fprintf(stderr, "data or length is null\n");
301                 return 0;
302         }
303
304         int segment = record / RECORDS_PER_SEGMENT;
305         if (segment != blerg->current_segment)
306                 blerg_segment_switch(blerg, segment);
307         int seg_rec = record % RECORDS_PER_SEGMENT;
308
309         if ((blerg->index[seg_rec].flags & 0x1) == 0) {
310                 fprintf(stderr, "Invalid record\n");
311                 return 0;
312         }
313
314         int rec_offset = blerg->index[seg_rec].offset;
315         int rec_length = blerg->index[seg_rec].length;
316         if (rec_offset >= blerg->data_size) {
317                 /* We're accessing an out-of-bounds record in our mmap.  Remap
318                    and recheck. */
319                 if (!blerg_remap_data(blerg)) {
320                         return 0;
321                 }
322                 if (rec_offset >= blerg->data_size) {
323                         fprintf(stderr, "Record offset outside of data!?");
324                         return 0;
325                 }
326         }
327
328         *data = malloc(rec_length);
329         if (*data == NULL) {
330                 perror("Could not allocate string in fetch");
331                 return 0;
332         }
333
334         memcpy(*data, blerg->data + rec_offset, rec_length);
335
336         *length = rec_length;
337
338         return 1;
339 }
340
341 time_t blerg_get_timestamp(struct blerg *blerg, int record) {
342         CHECK_VALID_BLERG(0)
343         if (record < 0 || record >= blerg_get_record_count(blerg)) {
344                 fprintf(stderr, "Invalid record\n");
345                 return 0;
346         }
347
348         int segment = record / RECORDS_PER_SEGMENT;
349         if (segment != blerg->current_segment)
350                 blerg_segment_switch(blerg, segment);
351         int seg_rec = record % RECORDS_PER_SEGMENT;
352
353         if ((blerg->index[seg_rec].flags & 0x1) == 0) {
354                 fprintf(stderr, "Invalid record\n");
355                 return 0;
356         }
357
358         return blerg->index[seg_rec].timestamp;
359 }
360
361 int blerg_set_subscription_mark(struct blerg *blerg) {
362         CHECK_VALID_BLERG(0)
363         blerg->meta->subscription_mark = subscription_count_items(blerg->name);
364         return 1;
365 }
366
367 uint64_t blerg_get_subscription_mark(struct blerg *blerg) {
368         CHECK_VALID_BLERG(0)
369         return blerg->meta->subscription_mark;
370 }