Sanitize username inputs in the database layer
[blerg.git] / database / database.c
1 #include <stdint.h>
2 #include <stdlib.h>
3 #include <string.h>
4 #include <stdio.h>
5 #include <time.h>
6 #include <unistd.h>
7 #include <sys/stat.h>
8 #include <sys/types.h>
9 #include <sys/mman.h>
10 #include <sys/file.h>
11 #include <fcntl.h>
12 #include "database.h"
13 #include "util.h"
14 #include "config.h"
15
16 uint64_t blerg_get_record_count(struct blerg *blerg) {
17         uint64_t count;
18         flock(blerg->meta_fd, LOCK_SH);
19         count = blerg->meta->sequence;
20         flock(blerg->meta_fd, LOCK_UN);
21         return count;
22 }
23
24 /* Returns last usable record */
25 uint64_t blerg_increment_record_count(struct blerg *blerg) {
26         uint64_t count;
27         flock(blerg->meta_fd, LOCK_EX);
28         count = blerg->meta->sequence++;
29         flock(blerg->meta_fd, LOCK_UN);
30         return count;
31 }
32
33 void blerg_segment_close(struct blerg *blerg) {
34         if (blerg->data != NULL)
35                 munmap((void *)blerg->data, blerg->data_size);
36         if (blerg->data_fd != -1)
37                 close(blerg->data_fd);
38         if (blerg->index != NULL)
39                 munmap((void *)blerg->index, RECORDS_PER_SEGMENT * sizeof(struct record));
40         if (blerg->index_fd != -1)
41                 close(blerg->index_fd);
42 }
43
44 int blerg_segment_switch(struct blerg *blerg, int new_segment) {
45         char filename[512];
46         uint64_t max_sequence = blerg_get_record_count(blerg);
47         struct stat st;
48
49         if (new_segment > max_sequence / RECORDS_PER_SEGMENT) {
50                 fprintf(stderr, "Cannot switch to sequence beyond last record\n");
51                 return 0;
52         }
53
54         blerg_segment_close(blerg);
55
56         /* Load and map the index */
57         snprintf(filename, 512, "%s/index%d", blerg->base_path, new_segment);
58         blerg->index_fd = open(filename, O_RDWR | O_CREAT, 0600);
59         if (blerg->index_fd == -1) {
60                 perror("Could not open index");
61                 goto open_failed_index_open;
62         }
63         flock(blerg->index_fd, LOCK_EX);
64         fstat(blerg->index_fd, &st);
65         if (st.st_size == 0) {
66                 int i;
67                 struct record r;
68                 memset((void *)&r, 0, sizeof(struct record));
69                 for (i = 0; i < RECORDS_PER_SEGMENT; i++) {
70                         write(blerg->index_fd, &r, sizeof(struct record));
71                 }
72         }
73         flock(blerg->index_fd, LOCK_UN);
74
75         blerg->index = (struct record *) mmap(NULL, RECORDS_PER_SEGMENT * sizeof(struct record), PROT_READ | PROT_WRITE, MAP_SHARED, blerg->index_fd, 0);
76         if (blerg->index == MAP_FAILED) {
77                 perror("Could not mmap index");
78                 goto open_failed_index_mmap;
79         }
80
81         /* Load data file */
82         sprintf(filename, "%s/data%d", blerg->base_path, new_segment);
83         blerg->data_fd = open(filename, O_RDWR | O_APPEND | O_CREAT, 0600);
84         fstat(blerg->data_fd, &st);
85         blerg->data_size = st.st_size;
86         if (blerg->data_fd == -1) {
87                 perror("Could not open data");
88                 goto open_failed_data_open;
89         }
90
91         if (blerg->data_size > 0) {
92                 blerg->data = (char *) mmap(NULL, blerg->data_size, PROT_READ, MAP_SHARED, blerg->data_fd, 0);
93                 if (blerg->data == MAP_FAILED) {
94                         perror("Could not mmap data");
95                         goto open_failed_data_mmap;
96                 }
97         }
98
99         return 1;
100
101 open_failed_data_mmap:
102         close(blerg->data_fd);
103 open_failed_data_open:
104         munmap((void *)blerg->index, sizeof(RECORDS_PER_SEGMENT * sizeof(struct record)));
105 open_failed_index_mmap:
106         close(blerg->index_fd);
107 open_failed_index_open:
108         return 0;
109 }
110
111 int blerg_exists(const char *name) {
112         int namelen = strlen(name);
113         char filename[512];
114
115         if (!valid_name(name)) {
116                 fprintf(stderr, "Invalid name\n");
117                 return 0;
118         }
119
120         snprintf(filename, 512, "%s/%s", DATA_PATH, name);
121         if (access(filename, F_OK) == -1)
122                 return 0;
123         else
124                 return 1;
125 }
126
127 struct blerg *blerg_open(const char *name) {
128         int namelen = strlen(name);
129         char filename[512];
130         struct stat st;
131         uint64_t sequence;
132
133         if (!valid_name(name)) {
134                 fprintf(stderr, "Invalid name\n");
135                 return NULL;
136         }
137         struct blerg *blerg = malloc(sizeof(struct blerg));
138         if (!blerg) {
139                 perror("Cannot allocate memory for blerg");
140                 goto open_failed_blerg_malloc;
141         }
142         blerg->name = malloc(namelen + 1);
143         memcpy(blerg->name, name, namelen + 1);
144         blerg->meta_fd = blerg->index_fd = blerg->data_fd = -1;
145         blerg->meta = NULL;
146         blerg->index = NULL;
147         blerg->data = NULL;
148
149         /* Make the directory if it doesn't exist */
150         blerg->base_path = malloc(512);
151         snprintf(blerg->base_path, 512, "%s/%s", DATA_PATH, name);
152         if (access(blerg->base_path, F_OK) == -1)
153                 mkdir(blerg->base_path, 0755);
154
155         /* Open and map metadata */
156         snprintf(filename, 512, "%s/meta", blerg->base_path);
157         blerg->meta_fd = open(filename, O_RDWR | O_CREAT, 0600);
158         if (blerg->meta_fd == -1) {
159                 perror("Could not open metadata");
160                 goto open_failed_meta_open;
161         }
162         fstat(blerg->meta_fd, &st);
163         if (st.st_size == 0) {
164                 char *buf = (char *) malloc(sizeof(struct meta));
165                 memset(buf, 0, sizeof(struct meta));
166                 write(blerg->meta_fd, buf, sizeof(struct meta));
167                 free(buf);
168         }
169         blerg->meta = (struct meta *) mmap(NULL, sizeof(struct meta), PROT_READ | PROT_WRITE, MAP_SHARED, blerg->meta_fd, 0);
170         if (blerg->meta == MAP_FAILED) {
171                 perror("Could not map metadata");
172                 goto open_failed_meta_mmap;
173         }
174
175         /* Open and map index and data for the current segment */
176         blerg->current_segment = blerg_get_record_count(blerg) / RECORDS_PER_SEGMENT;
177         if (!blerg_segment_switch(blerg, blerg->current_segment)) {
178                 fprintf(stderr, "Could not switch segment\n");
179                 goto open_failed_segment_switch;
180         }
181
182         return blerg;
183
184 open_failed_segment_switch:
185         munmap((void *)blerg->meta, sizeof(struct meta));
186 open_failed_meta_mmap:
187         close(blerg->meta_fd);
188 open_failed_meta_open:
189         free(blerg->name);
190         free(blerg);
191 open_failed_blerg_malloc:
192         return NULL;
193 }
194
195 int blerg_close(struct blerg *blerg) {
196         blerg_segment_close(blerg);
197         munmap((void *)blerg->meta, sizeof(struct meta));
198         close(blerg->meta_fd);
199         free(blerg->base_path);
200         free(blerg->name);
201         free(blerg);
202         return 1;
203 }
204
205 int blerg_store(struct blerg *blerg, const char *data, int len) {
206         if (len > MAX_RECORD_SIZE) {
207                 fprintf(stderr, "len > 64K\n");
208                 return -1;
209         }
210
211         flock(blerg->index_fd, LOCK_EX);
212         flock(blerg->data_fd, LOCK_EX);
213
214         uint64_t record = blerg_increment_record_count(blerg);
215         if (record == -1) {
216                 fprintf(stderr, "Could not find free record\n");
217                 return -1;
218         }
219         int segment = record / RECORDS_PER_SEGMENT;
220         if (segment != blerg->current_segment)
221                 blerg_segment_switch(blerg, segment);
222         int seg_rec = record % RECORDS_PER_SEGMENT;
223
224         /* Get the position for the new data */
225         FILE *datafile = fdopen(dup(blerg->data_fd), "a");
226         fseek(datafile, 0, SEEK_END);
227         int curpos = ftell(datafile);
228         fclose(datafile);
229
230         int bytes = 0;
231         do {
232                 int n = write(blerg->data_fd, data + bytes, len);
233                 if (n == -1) {
234                         perror("Could not write data");
235                         /* Truncate anything we may have written */
236                         ftruncate(blerg->data_fd, curpos);
237                         return -1;
238                 }
239                 bytes += n;
240         } while (bytes < len);
241         blerg->index[seg_rec].flags = 0x0001;
242         blerg->index[seg_rec].offset = curpos;
243         blerg->index[seg_rec].length = len;
244         blerg->index[seg_rec].timestamp = time(NULL);
245
246         tag_scan(blerg->name, data, len, record);
247
248         flock(blerg->data_fd, LOCK_UN);
249         flock(blerg->index_fd, LOCK_UN);
250
251         return record;
252 }
253
254 int blerg_fetch(struct blerg *blerg, int record, char **data, int *length) {
255         if (record < 0) {
256                 fprintf(stderr, "Invalid record\n");
257                 return 0;
258         }
259
260         int segment = record / RECORDS_PER_SEGMENT;
261         if (segment != blerg->current_segment)
262                 blerg_segment_switch(blerg, segment);
263         int seg_rec = record % RECORDS_PER_SEGMENT;
264
265         if ((blerg->index[seg_rec].flags & 0x1) == 0) {
266                 fprintf(stderr, "Invalid record\n");
267                 return 0;
268         }
269
270         int rec_offset = blerg->index[seg_rec].offset;
271         int rec_length = blerg->index[seg_rec].length;
272         if (rec_offset >= blerg->data_size) {
273                 /* We're accessing an out-of-bounds record in our mmap.
274                    Recheck size and remap. */
275                 struct stat st;
276                 fstat(blerg->data_fd, &st);
277                 blerg->data_size = st.st_size;
278                 if (rec_offset > blerg->data_size) {
279                         fprintf(stderr, "Record offset outside of data!?");
280                         return 0;
281                 }
282
283                 munmap(blerg->data, blerg->data_size);
284                 blerg->data = (char *) mmap(NULL, blerg->data_size, PROT_READ, MAP_SHARED, blerg->data_fd, 0);
285                 if (blerg->data == MAP_FAILED) {
286                         perror("Could not remap data");
287                         return 0;
288                 }
289         }
290
291         *data = malloc(rec_length);
292         if (*data == NULL) {
293                 perror("Could not allocate string in fetch");
294                 return 0;
295         }
296
297         memcpy(*data, blerg->data + rec_offset, rec_length);
298
299         *length = rec_length;
300
301         return 1;
302 }
303
304 time_t blerg_get_timestamp(struct blerg *blerg, int record) {
305         if (record < 0) {
306                 fprintf(stderr, "Invalid record\n");
307                 return 0;
308         }
309
310         int segment = record / RECORDS_PER_SEGMENT;
311         if (segment != blerg->current_segment)
312                 blerg_segment_switch(blerg, segment);
313         int seg_rec = record % RECORDS_PER_SEGMENT;
314
315         if ((blerg->index[seg_rec].flags & 0x1) == 0) {
316                 fprintf(stderr, "Invalid record\n");
317                 return 0;
318         }
319
320         return blerg->index[seg_rec].timestamp;
321 }