wimlib.net Git - wimlib/blob - 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__ */