wimlib.net Git - wimlib/blob - src/util.c

   1 /*
   2  * util.c - utility functions
   3  */
   4
   5 /*
   6  * Copyright 2012-2023 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 https://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
  47 /*******************
  48  * Memory allocation
  49  *******************/
  50
  51 static void *(*wimlib_malloc_func) (size_t)         = malloc;
  52 static void  (*wimlib_free_func)   (void *)         = free;
  53 static void *(*wimlib_realloc_func)(void *, size_t) = realloc;
  54
  55 void *
  56 wimlib_malloc(size_t size)
  57 {
  58         void *ptr;
  59
  60 retry:
  61         ptr = (*wimlib_malloc_func)(size);
  62         if (unlikely(!ptr)) {
  63                 if (size == 0) {
  64                         size = 1;
  65                         goto retry;
  66                 }
  67         }
  68         return ptr;
  69 }
  70
  71 void
  72 wimlib_free_memory(void *ptr)
  73 {
  74         (*wimlib_free_func)(ptr);
  75 }
  76
  77 void *
  78 wimlib_realloc(void *ptr, size_t size)
  79 {
  80         if (size == 0)
  81                 size = 1;
  82         return (*wimlib_realloc_func)(ptr, size);
  83 }
  84
  85 void *
  86 wimlib_calloc(size_t nmemb, size_t size)
  87 {
  88         size_t total_size = nmemb * size;
  89         void *p;
  90
  91         if (size != 0 && nmemb > SIZE_MAX / size) {
  92                 errno = ENOMEM;
  93                 return NULL;
  94         }
  95
  96         p = MALLOC(total_size);
  97         if (p)
  98                 p = memset(p, 0, total_size);
  99         return p;
 100 }
 101
 102 char *
 103 wimlib_strdup(const char *str)
 104 {
 105         return memdup(str, strlen(str) + 1);
 106 }
 107
 108 #ifdef _WIN32
 109 wchar_t *
 110 wimlib_wcsdup(const wchar_t *str)
 111 {
 112         return memdup(str, (wcslen(str) + 1) * sizeof(wchar_t));
 113 }
 114 #endif
 115
 116 void *
 117 wimlib_aligned_malloc(size_t size, size_t alignment)
 118 {
 119         wimlib_assert(is_power_of_2(alignment));
 120
 121         void *ptr = MALLOC(sizeof(void *) + alignment - 1 + size);
 122         if (ptr) {
 123                 void *orig_ptr = ptr;
 124                 ptr = (void *)ALIGN((uintptr_t)ptr + sizeof(void *), alignment);
 125                 ((void **)ptr)[-1] = orig_ptr;
 126         }
 127         return ptr;
 128 }
 129
 130 void
 131 wimlib_aligned_free(void *ptr)
 132 {
 133         if (ptr)
 134                 FREE(((void **)ptr)[-1]);
 135 }
 136
 137 void *
 138 memdup(const void *mem, size_t size)
 139 {
 140         void *ptr = MALLOC(size);
 141         if (ptr)
 142                 ptr = memcpy(ptr, mem, size);
 143         return ptr;
 144 }
 145
 146 /* API function documented in wimlib.h  */
 147 WIMLIBAPI int
 148 wimlib_set_memory_allocator(void *(*malloc_func)(size_t),
 149                             void (*free_func)(void *),
 150                             void *(*realloc_func)(void *, size_t))
 151 {
 152         wimlib_malloc_func  = malloc_func  ? malloc_func  : malloc;
 153         wimlib_free_func    = free_func    ? free_func    : free;
 154         wimlib_realloc_func = realloc_func ? realloc_func : realloc;
 155         return 0;
 156 }
 157
 158 /*******************
 159  * String utilities
 160  *******************/
 161
 162 #ifndef HAVE_MEMPCPY
 163 void *mempcpy(void *dst, const void *src, size_t n)
 164 {
 165         return memcpy(dst, src, n) + n;
 166 }
 167 #endif
 168
 169 /**************************
 170  * Random number generation
 171  **************************/
 172
 173 #ifndef _WIN32
 174 /*
 175  * Generate @n cryptographically secure random bytes (thread-safe)
 176  *
 177  * This is the UNIX version.  It uses the Linux getrandom() system call if
 178  * available; otherwise, it falls back to reading from /dev/urandom.
 179  */
 180 void
 181 get_random_bytes(void *p, size_t n)
 182 {
 183         if (n == 0)
 184                 return;
 185 #ifdef __NR_getrandom
 186         static bool getrandom_unavailable;
 187
 188         if (getrandom_unavailable)
 189                 goto try_dev_urandom;
 190         do {
 191                 int res = syscall(__NR_getrandom, p, n, 0);
 192                 if (unlikely(res < 0)) {
 193                         if (errno == ENOSYS) {
 194                                 getrandom_unavailable = true;
 195                                 goto try_dev_urandom;
 196                         }
 197                         if (errno == EINTR)
 198                                 continue;
 199                         ERROR_WITH_ERRNO("getrandom() failed");
 200                         wimlib_assert(0);
 201                         res = 0;
 202                 }
 203                 p += res;
 204                 n -= res;
 205         } while (n != 0);
 206         return;
 207
 208 try_dev_urandom:
 209         ;
 210 #endif /* __NR_getrandom */
 211         int fd = open("/dev/urandom", O_RDONLY);
 212         if (fd < 0) {
 213                 ERROR_WITH_ERRNO("Unable to open /dev/urandom");
 214                 wimlib_assert(0);
 215         }
 216         do {
 217                 int res = read(fd, p, min(n, INT_MAX));
 218                 if (unlikely(res < 0)) {
 219                         if (errno == EINTR)
 220                                 continue;
 221                         ERROR_WITH_ERRNO("Error reading from /dev/urandom");
 222                         wimlib_assert(0);
 223                         res = 0;
 224                 }
 225                 p += res;
 226                 n -= res;
 227         } while (n != 0);
 228         close(fd);
 229 }
 230 #endif /* !_WIN32 */
 231
 232 /*
 233  * Generate @n cryptographically secure random alphanumeric characters
 234  * (thread-safe)
 235  *
 236  * This is implemented on top of get_random_bytes().  For efficiency the calls
 237  * to get_random_bytes() are batched.
 238  */
 239 void
 240 get_random_alnum_chars(tchar *p, size_t n)
 241 {
 242         u32 r[64];
 243         int r_idx = 0;
 244         int r_end = 0;
 245
 246         for (; n != 0; p++, n--) {
 247                 tchar x;
 248
 249                 if (r_idx >= r_end) {
 250                         r_idx = 0;
 251                         r_end = min(n, ARRAY_LEN(r));
 252                         get_random_bytes(r, r_end * sizeof(r[0]));
 253                 }
 254
 255                 STATIC_ASSERT(sizeof(r[0]) == sizeof(u32));
 256                 while (unlikely(r[r_idx] >= UINT32_MAX - (UINT32_MAX % 62)))
 257                         get_random_bytes(&r[r_idx], sizeof(r[0]));
 258
 259                 x = r[r_idx++] % 62;
 260
 261                 if (x < 26)
 262                         *p = 'a' + x;
 263                 else if (x < 52)
 264                         *p = 'A' + x - 26;
 265                 else
 266                         *p = '0' + x - 52;
 267         }
 268 }
 269
 270 /************************
 271  * System information
 272  ************************/
 273
 274 #ifndef _WIN32
 275 unsigned
 276 get_available_cpus(void)
 277 {
 278         long n = sysconf(_SC_NPROCESSORS_ONLN);
 279         if (n < 1 || n >= UINT_MAX) {
 280                 WARNING("Failed to determine number of processors; assuming 1.");
 281                 return 1;
 282         }
 283         return n;
 284 }
 285 #endif /* !_WIN32 */
 286
 287 #ifndef _WIN32
 288 u64
 289 get_available_memory(void)
 290 {
 291 #if defined(_SC_PAGESIZE) && defined(_SC_PHYS_PAGES)
 292         long page_size = sysconf(_SC_PAGESIZE);
 293         long num_pages = sysconf(_SC_PHYS_PAGES);
 294         if (page_size <= 0 || num_pages <= 0)
 295                 goto default_size;
 296         return ((u64)page_size * (u64)num_pages);
 297 #else
 298         int mib[2] = {CTL_HW, HW_MEMSIZE};
 299         u64 memsize;
 300         size_t len = sizeof(memsize);
 301         if (sysctl(mib, ARRAY_LEN(mib), &memsize, &len, NULL, 0) < 0 || len != 8)
 302                 goto default_size;
 303         return memsize;
 304 #endif
 305
 306 default_size:
 307         WARNING("Failed to determine available memory; assuming 1 GiB");
 308         return (u64)1 << 30;
 309 }
 310 #endif /* !_WIN32 */