Improve random number generation
[wimlib] / src / util.c
1 /*
2  * util.c - utility functions
3  */
4
5 /*
6  * Copyright (C) 2012-2016 Eric Biggers
7  *
8  * This file is free software; you can redistribute it and/or modify it under
9  * the terms of the GNU Lesser General Public License as published by the Free
10  * Software Foundation; either version 3 of the License, or (at your option) any
11  * later version.
12  *
13  * This file is distributed in the hope that it will be useful, but WITHOUT
14  * ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or FITNESS
15  * FOR A PARTICULAR PURPOSE. See the GNU Lesser General Public License for more
16  * details.
17  *
18  * You should have received a copy of the GNU Lesser General Public License
19  * along with this file; if not, see http://www.gnu.org/licenses/.
20  */
21
22 #ifdef HAVE_CONFIG_H
23 #  include "config.h"
24 #endif
25
26 #include <errno.h>
27 #include <fcntl.h>
28 #include <limits.h>
29 #include <stdio.h>
30 #include <stdlib.h>
31 #include <string.h>
32 #ifdef HAVE_SYS_SYSCTL_H
33 #  include <sys/types.h>
34 #  include <sys/sysctl.h>
35 #endif
36 #ifdef HAVE_SYS_SYSCALL_H
37 #  include <sys/syscall.h>
38 #endif
39 #include <unistd.h>
40
41 #include "wimlib.h"
42 #include "wimlib/assert.h"
43 #include "wimlib/error.h"
44 #include "wimlib/timestamp.h"
45 #include "wimlib/util.h"
46 #include "wimlib/xml.h"
47
48 /*******************
49  * Memory allocation
50  *******************/
51
52 static void *(*wimlib_malloc_func) (size_t)         = malloc;
53 static void  (*wimlib_free_func)   (void *)         = free;
54 static void *(*wimlib_realloc_func)(void *, size_t) = realloc;
55
56 void *
57 wimlib_malloc(size_t size)
58 {
59         void *ptr;
60
61 retry:
62         ptr = (*wimlib_malloc_func)(size);
63         if (unlikely(!ptr)) {
64                 if (size == 0) {
65                         size = 1;
66                         goto retry;
67                 }
68         }
69         return ptr;
70 }
71
72 void
73 wimlib_free_memory(void *ptr)
74 {
75         (*wimlib_free_func)(ptr);
76 }
77
78 void *
79 wimlib_realloc(void *ptr, size_t size)
80 {
81         if (size == 0)
82                 size = 1;
83         return (*wimlib_realloc_func)(ptr, size);
84 }
85
86 void *
87 wimlib_calloc(size_t nmemb, size_t size)
88 {
89         size_t total_size = nmemb * size;
90         void *p;
91
92         if (size != 0 && nmemb > SIZE_MAX / size) {
93                 errno = ENOMEM;
94                 return NULL;
95         }
96
97         p = MALLOC(total_size);
98         if (p)
99                 p = memset(p, 0, total_size);
100         return p;
101 }
102
103 char *
104 wimlib_strdup(const char *str)
105 {
106         return memdup(str, strlen(str) + 1);
107 }
108
109 #ifdef __WIN32__
110 wchar_t *
111 wimlib_wcsdup(const wchar_t *str)
112 {
113         return memdup(str, (wcslen(str) + 1) * sizeof(wchar_t));
114 }
115 #endif
116
117 void *
118 wimlib_aligned_malloc(size_t size, size_t alignment)
119 {
120         wimlib_assert(is_power_of_2(alignment));
121
122         void *ptr = MALLOC(sizeof(void *) + alignment - 1 + size);
123         if (ptr) {
124                 void *orig_ptr = ptr;
125                 ptr = (void *)ALIGN((uintptr_t)ptr + sizeof(void *), alignment);
126                 ((void **)ptr)[-1] = orig_ptr;
127         }
128         return ptr;
129 }
130
131 void
132 wimlib_aligned_free(void *ptr)
133 {
134         if (ptr)
135                 FREE(((void **)ptr)[-1]);
136 }
137
138 void *
139 memdup(const void *mem, size_t size)
140 {
141         void *ptr = MALLOC(size);
142         if (ptr)
143                 ptr = memcpy(ptr, mem, size);
144         return ptr;
145 }
146
147 /* API function documented in wimlib.h  */
148 WIMLIBAPI int
149 wimlib_set_memory_allocator(void *(*malloc_func)(size_t),
150                             void (*free_func)(void *),
151                             void *(*realloc_func)(void *, size_t))
152 {
153         wimlib_malloc_func  = malloc_func  ? malloc_func  : malloc;
154         wimlib_free_func    = free_func    ? free_func    : free;
155         wimlib_realloc_func = realloc_func ? realloc_func : realloc;
156
157         xml_set_memory_allocator(wimlib_malloc_func, wimlib_free_func,
158                                  wimlib_realloc_func);
159         return 0;
160 }
161
162 /*******************
163  * String utilities
164  *******************/
165
166 #ifndef HAVE_MEMPCPY
167 void *mempcpy(void *dst, const void *src, size_t n)
168 {
169         return memcpy(dst, src, n) + n;
170 }
171 #endif
172
173 /**************************
174  * Random number generation
175  **************************/
176
177 #ifndef __WIN32__
178 /*
179  * Generate @n cryptographically secure random bytes (thread-safe)
180  *
181  * This is the UNIX version.  It uses the Linux getrandom() system call if
182  * available; otherwise, it falls back to reading from /dev/urandom.
183  */
184 void
185 get_random_bytes(void *p, size_t n)
186 {
187         if (n == 0)
188                 return;
189 #ifdef HAVE_NR_GETRANDOM
190         static bool getrandom_unavailable;
191
192         if (getrandom_unavailable)
193                 goto try_dev_urandom;
194         do {
195                 int res = syscall(__NR_getrandom, p, n, 0);
196                 if (unlikely(res < 0)) {
197                         if (errno == ENOSYS) {
198                                 getrandom_unavailable = true;
199                                 goto try_dev_urandom;
200                         }
201                         if (errno == EINTR)
202                                 continue;
203                         ERROR_WITH_ERRNO("getrandom() failed");
204                         wimlib_assert(0);
205                         res = 0;
206                 }
207                 p += res;
208                 n -= res;
209         } while (n != 0);
210         return;
211
212 try_dev_urandom:
213         ;
214 #endif /* HAVE_NR_GETRANDOM */
215         int fd = open("/dev/urandom", O_RDONLY);
216         if (fd < 0) {
217                 ERROR_WITH_ERRNO("Unable to open /dev/urandom");
218                 wimlib_assert(0);
219         }
220         do {
221                 int res = read(fd, p, min(n, INT_MAX));
222                 if (unlikely(res < 0)) {
223                         if (errno == EINTR)
224                                 continue;
225                         ERROR_WITH_ERRNO("Error reading from /dev/urandom");
226                         wimlib_assert(0);
227                         res = 0;
228                 }
229                 p += res;
230                 n -= res;
231         } while (n != 0);
232         close(fd);
233 }
234 #endif /* !__WIN32__ */
235
236 /*
237  * Generate @n cryptographically secure random alphanumeric characters
238  * (thread-safe)
239  *
240  * This is implemented on top of get_random_bytes().  For efficiency the calls
241  * to get_random_bytes() are batched.
242  */
243 void
244 get_random_alnum_chars(tchar *p, size_t n)
245 {
246         u32 r[64];
247         int r_idx = 0;
248         int r_end = 0;
249
250         for (; n != 0; p++, n--) {
251                 tchar x;
252
253                 if (r_idx >= r_end) {
254                         r_idx = 0;
255                         r_end = min(n, ARRAY_LEN(r));
256                         get_random_bytes(r, r_end * sizeof(r[0]));
257                 }
258
259                 STATIC_ASSERT(sizeof(r[0]) == sizeof(u32));
260                 while (unlikely(r[r_idx] >= UINT32_MAX - (UINT32_MAX % 62)))
261                         get_random_bytes(&r[r_idx], sizeof(r[0]));
262
263                 x = r[r_idx++] % 62;
264
265                 if (x < 26)
266                         *p = 'a' + x;
267                 else if (x < 52)
268                         *p = 'A' + x - 26;
269                 else
270                         *p = '0' + x - 52;
271         }
272 }
273
274 /************************
275  * System information
276  ************************/
277
278 #ifndef __WIN32__
279 unsigned
280 get_available_cpus(void)
281 {
282         long n = sysconf(_SC_NPROCESSORS_ONLN);
283         if (n < 1 || n >= UINT_MAX) {
284                 WARNING("Failed to determine number of processors; assuming 1.");
285                 return 1;
286         }
287         return n;
288 }
289 #endif /* !__WIN32__ */
290
291 #ifndef __WIN32__
292 u64
293 get_available_memory(void)
294 {
295 #if defined(_SC_PAGESIZE) && defined(_SC_PHYS_PAGES)
296         long page_size = sysconf(_SC_PAGESIZE);
297         long num_pages = sysconf(_SC_PHYS_PAGES);
298         if (page_size <= 0 || num_pages <= 0)
299                 goto default_size;
300         return ((u64)page_size * (u64)num_pages);
301 #else
302         int mib[2] = {CTL_HW, HW_MEMSIZE};
303         u64 memsize;
304         size_t len = sizeof(memsize);
305         if (sysctl(mib, ARRAY_LEN(mib), &memsize, &len, NULL, 0) < 0 || len != 8)
306                 goto default_size;
307         return memsize;
308 #endif
309
310 default_size:
311         WARNING("Failed to determine available memory; assuming 1 GiB");
312         return (u64)1 << 30;
313 }
314 #endif /* !__WIN32__ */