Introduce a simple memory allocator for managing the memory pool used for

allocating buffers etc.

5 files changed, 173 insertions, 131 deletions

diff --git a/Makefile.am b/Makefile.am
index 7ac907d..61f6deb 100644
--- a/Makefile.am
+++ b/Makefile.am
@@ -99,6 +99,7 @@ libsuspend_common_a_SOURCES=\
 	splashy_funcs.h splashy_funcs.c \
 	fbsplash_funcs.h fbsplash_funcs.c \
 	bootsplash.h bootsplash.c \
+	memalloc.h memalloc.c \
 	whitelist.h whitelist.c \
 	s2ram.h s2ram.c
 
diff --git a/encrypt.h b/encrypt.h
index ab1b37d..d3feb02 100644
--- a/encrypt.h
+++ b/encrypt.h
@@ -15,6 +15,7 @@
 
 /* Maximum length of a passphrase, in characters */
 #define PASS_SIZE	128
+#define PASSBUF_SIZE	(2 * PASS_SIZE)
 /* Symmetric cipher used for image encryption, the size of its key and its
  * block, in bytes
  */
diff --git a/resume.c b/resume.c
index 6c1f4be..4adf67e 100644
--- a/resume.c
+++ b/resume.c
@@ -29,6 +29,7 @@
 #endif
 
 #include "swsusp.h"
+#include "memalloc.h"
 #include "config_parser.h"
 #include "md5.h"
 #include "splash.h"
@@ -49,6 +50,7 @@ static unsigned int decompress_buf_size;
 #endif
 #ifdef CONFIG_ENCRYPT
 static char do_decrypt;
+static char password[PASSBUF_SIZE];
 #else
 #define do_decrypt 0
 #endif
@@ -159,10 +161,6 @@ static inline int atomic_restore(int dev)
 	return ioctl(dev, SNAPSHOT_ATOMIC_RESTORE, 0);
 }
 
-static unsigned int page_size;
-static unsigned int buffer_size;
-static void *mem_pool;
-
 /**
  *	read_page - Read data from a swap location
  *	@fd:		File handle of the resume partition.
@@ -264,55 +262,70 @@ static int load_extents_page(struct swap_map_handle *handle)
 }
 
 /**
+ *	free_swap_reader - free memory allocated for loading the image
+ *	@handle:	Structure containing pointers to memory buffers to free.
+ */
+static void free_swap_reader(struct swap_map_handle *handle)
+{
+#ifdef CONFIG_COMPRESS
+	if (do_decompress) {
+		freemem(handle->lzo_work_buffer);
+		freemem(handle->read_buffer);
+	}
+#endif
+#ifdef CONFIG_ENCRYPT
+	if (do_decrypt)
+		 freemem(handle->decrypt_buffer);
+#endif
+	 freemem(handle->buffer);
+	 freemem(handle->extents);
+}
+
+/**
  *	init_swap_reader - initialize the structure used for loading the image
  *	@handle:	Structure to initialize.
  *	@fd:		File descriptor associated with the swap.
  *	@start:		Swap location (offset) of the first image page.
- *	@buf:		Memory pool to use for buffers.
  *	@image_size:	Total size of the image data.
  *
  *	Initialize buffers and related fields of @handle and load the first
  *	array of extents.
  */
 static int init_swap_reader(struct swap_map_handle *handle, int fd,
-				loff_t start, void *buf, loff_t image_size)
+				loff_t start, loff_t image_size)
 {
 	int error;
 
-	if (!start || !buf)
+	if (!start)
 		return -EINVAL;
 
 	handle->fd = fd;
 	handle->total_size = image_size;
 
-	handle->extents = buf;
-	buf += page_size;
+	handle->extents = getmem(page_size);
 
-	handle->buffer = buf;
-	handle->read_buffer = buf;
-	buf += buffer_size;
+	handle->buffer = getmem(buffer_size);
+	handle->read_buffer = handle->buffer;
 
 #ifdef CONFIG_ENCRYPT
-	if (do_decrypt) {
-		handle->decrypt_buffer = buf;
-		buf += page_size;
-	}
+	if (do_decrypt)
+		handle->decrypt_buffer = getmem(page_size);
 #endif
 
 #ifdef CONFIG_COMPRESS
 	if (do_decompress) {
-		handle->read_buffer = buf;
-		buf += decompress_buf_size;
-		handle->lzo_work_buffer = buf;
-		/* This buffer must hold at least LZO1X_1_MEM_COMPRESS bytes */
+		handle->read_buffer = getmem(decompress_buf_size);
+		handle->lzo_work_buffer = getmem(LZO1X_1_MEM_COMPRESS);
 	}
 #endif
 
 	/* Read the table of extents */
 	handle->next_extents = start;
 	error = load_extents_page(handle);
-	if (error)
+	if (error) {
+		free_swap_reader(handle);
 		return error;
+	}
 
 	if (verify_checksum)
 		md5_init_ctx(&handle->ctx);
@@ -506,14 +519,13 @@ static char *print_checksum(char * buf, unsigned char *checksum)
 
 #ifdef CONFIG_ENCRYPT
 static int decrypt_key(struct image_header_info *header, unsigned char *key,
-			unsigned char *ivec, void *buffer)
+			unsigned char *ivec)
 {
 	gcry_ac_handle_t rsa_hd;
 	gcry_ac_data_t rsa_data_set, key_set;
 	gcry_ac_key_t rsa_priv;
 	gcry_mpi_t mpi;
 	unsigned char *buf, *out, *key_buf, *ivec_buf;
-	char *pass_buf;
 	struct md5_ctx ctx;
 	struct RSA_data *rsa;
 	gcry_cipher_hd_t sym_hd;
@@ -530,16 +542,15 @@ static int decrypt_key(struct image_header_info *header, unsigned char *key,
 	if (ret)
 		goto Free_rsa;
 
-	pass_buf = buffer;
-	key_buf = (unsigned char *)pass_buf + PASS_SIZE;
-	ivec_buf = key_buf + PK_KEY_SIZE;
-	out = ivec_buf + PK_CIPHER_BLOCK;
+	key_buf = getmem(PK_KEY_SIZE);
+	ivec_buf = getmem(PK_CIPHER_BLOCK);
+	out = getmem(KEY_TEST_SIZE);
 	do {
-		splash.read_password(pass_buf, 0);
+		splash.read_password(password, 0);
 		memset(ivec_buf, 0, PK_CIPHER_BLOCK);
-		strncpy((char *)ivec_buf, pass_buf, PK_CIPHER_BLOCK);
+		strncpy(ivec_buf, password, PK_CIPHER_BLOCK);
 		md5_init_ctx(&ctx);
-		md5_process_bytes(pass_buf, strlen(pass_buf), &ctx);
+		md5_process_bytes(password, strlen(password), &ctx);
 		md5_finish_ctx(&ctx, key_buf);
 		ret = gcry_cipher_setkey(sym_hd, key_buf, PK_KEY_SIZE);
 		if (!ret)
@@ -548,8 +559,8 @@ static int decrypt_key(struct image_header_info *header, unsigned char *key,
 
 		if (!ret)
 			ret = gcry_cipher_encrypt(sym_hd,
-					out, KEY_TEST_SIZE,
-					KEY_TEST_DATA, KEY_TEST_SIZE);
+						out, KEY_TEST_SIZE,
+						KEY_TEST_DATA, KEY_TEST_SIZE);
 
 		if (ret)
 			break;
@@ -565,7 +576,7 @@ static int decrypt_key(struct image_header_info *header, unsigned char *key,
 		ret = gcry_cipher_open(&sym_hd, PK_CIPHER,
 				GCRY_CIPHER_MODE_CFB, GCRY_CIPHER_SECURE);
 	if (ret)
-		goto Free_rsa;
+		goto Free_buffers;
 
 	ret = gcry_cipher_setkey(sym_hd, key_buf, PK_KEY_SIZE);
 	if (!ret)
@@ -641,6 +652,11 @@ Destroy_data_set:
 Close_cypher:
 	gcry_cipher_close(sym_hd);
 
+Free_buffers:
+	freemem(out);
+	freemem(ivec_buf);
+	freemem(key_buf);
+
 Free_rsa:
 	gcry_ac_close(rsa_hd);
 
@@ -649,7 +665,7 @@ Free_rsa:
 #endif
 
 static int open_resume_dev(char *resume_dev_name, 
-			    struct swsusp_header *swsusp_header)
+                           struct swsusp_header *swsusp_header)
 {
 	ssize_t size = sizeof(struct swsusp_header);
 	unsigned int shift = (resume_offset + 1) * page_size - size;
@@ -713,14 +729,14 @@ static void pause_resume(int pause)
 static int read_image(int dev, int fd, struct swsusp_header *swsusp_header)
 {
 	static struct swap_map_handle handle;
-	static unsigned char orig_checksum[16], checksum[16];
+	static unsigned char orig_checksum[16], checksum[16], csum_buf[48];
+	struct image_header_info *header;
 	int ret, error = 0;
-	struct image_header_info *header = mem_pool;
-	char *buffer = (char *)mem_pool + page_size;
 	ssize_t size = sizeof(struct swsusp_header);
 	unsigned int shift = (resume_offset + 1) * page_size - size;
 	char c;
 
+	header = getmem(page_size);
 	error = read_page(fd, header, swsusp_header->image);
 	if (error)
 		goto Reboot_question;
@@ -730,8 +746,8 @@ static int read_image(int dev, int fd, struct swsusp_header *swsusp_header)
 
 	if (header->flags & IMAGE_CHECKSUM) {
 		memcpy(orig_checksum, header->checksum, 16);
-		print_checksum(buffer, orig_checksum);
-		printf("%s: MD5 checksum %s\n", my_name, buffer);
+		print_checksum(csum_buf, orig_checksum);
+		printf("%s: MD5 checksum %s\n", my_name, csum_buf);
 		verify_checksum = 1;
 	}
 	splash.progress(10);
@@ -759,12 +775,12 @@ static int read_image(int dev, int fd, struct swsusp_header *swsusp_header)
 
 		printf("%s: Encrypted image\n", my_name);
 		if (header->flags & IMAGE_USE_RSA) {
-			error = decrypt_key(header, key, ivec, buffer);
+			error = decrypt_key(header, key, ivec);
 		} else {
 			int j;
 
-			splash.read_password(buffer, 0);
-			encrypt_init(key, ivec, buffer);
+			splash.read_password(password, 0);
+			encrypt_init(key, ivec, password);
 			for (j = 0; j < CIPHER_BLOCK; j++)
 				ivec[j] ^= header->salt[j];
 		}
@@ -795,7 +811,7 @@ static int read_image(int dev, int fd, struct swsusp_header *swsusp_header)
 	if (error)
 		goto Reboot_question;
 
-	error = init_swap_reader(&handle, fd, header->map_start, buffer,
+	error = init_swap_reader(&handle, fd, header->map_start,
 					header->image_data_size);
 	if (!error) {
 		struct timeval begin, end;
@@ -809,13 +825,14 @@ static int read_image(int dev, int fd, struct swsusp_header *swsusp_header)
 				fprintf(stderr,
 					"%s: MD5 checksum does not match\n",
 					my_name);
-				print_checksum(buffer, checksum);
+				print_checksum(csum_buf, checksum);
 				fprintf(stderr,
 					"%s: Computed MD5 checksum %s\n",
-					my_name, buffer);
+					my_name, csum_buf);
 				error = -EINVAL;
 			}
 		}
+		free_swap_reader(&handle);
 		if (error)
 			goto Reboot_question;
 		gettimeofday(&end, NULL);
@@ -884,10 +901,14 @@ Reboot_question:
 		gcry_cipher_close(handle.cipher_handle);
 #endif
 
+	freemem(header);
+
 	if (error) {
-		sprintf(buffer, "%s: Error %d loading the image\nPress "
+		char message[SPLASH_GENERIC_MESSAGE_SIZE];
+
+		sprintf(message, "%s: Error %d loading the image\nPress "
 			ENTER_KEY_NAME " to continue\n", my_name, error);
-		splash.dialog(buffer);
+		splash.dialog(message);
 	} else if (header->resume_pause != 0) {
 		pause_resume(header->resume_pause);
 	} else {
@@ -1016,8 +1037,8 @@ int main(int argc, char *argv[])
 	else
 		splash_param = SPL_RESUME;
 
-	page_size = getpagesize();
-	buffer_size = BUFFER_PAGES * page_size;
+	get_page_and_buffer_sizes();
+
 	mem_size = 2 * page_size + buffer_size;
 #ifdef CONFIG_ENCRYPT
 	printf("%s: libgcrypt version: %s\n", my_name,
@@ -1030,16 +1051,16 @@ int main(int argc, char *argv[])
 	 * The formula below follows from the worst-case expansion calculation
 	 * for LZO1 (size / 16 + 67) and the fact that the size of the
 	 * compressed data must be stored in the buffer (sizeof(size_t)).
-	 * Additionally, we want the buffer size to be a multiple of page_size.
 	 */
 	decompress_buf_size = buffer_size +
-		((page_size + (buffer_size >> 4) + 66 + sizeof(size_t))
-			& ~(page_size - 1));
-	mem_size += decompress_buf_size + LZO1X_1_MEM_COMPRESS;
+			round_up_page_size((buffer_size >> 4) + 67 +
+						sizeof(size_t));
+	mem_size += decompress_buf_size +
+			round_up_page_size(LZO1X_1_MEM_COMPRESS);
 #endif
-	mem_pool = malloc(mem_size);
-	if (!mem_pool) {
-		error = errno;
+
+	error = init_memalloc(page_size, mem_size);
+	if (error) {
 		fprintf(stderr, "%s: Could not allocate memory\n", my_name);
 		return error;
 	}
@@ -1134,7 +1155,7 @@ Free:
 
 	close_printk();
 
-	free(mem_pool);
+	free_memalloc();
 
 	return error;
 }
diff --git a/suspend.c b/suspend.c
index c93d219..1cfbfe4 100644
--- a/suspend.c
+++ b/suspend.c
@@ -37,6 +37,7 @@
 #endif
 
 #include "swsusp.h"
+#include "memalloc.h"
 #include "config_parser.h"
 #include "md5.h"
 #include "splash.h"
@@ -70,7 +71,7 @@ static unsigned int compress_buf_size;
 static char do_encrypt;
 static char use_RSA;
 static char key_name[MAX_STR_LEN] = SUSPEND_KEY_FILE;
-static char password[PASS_SIZE];
+static char password[PASSBUF_SIZE];
 static unsigned long encrypt_buf_size;
 #else
 #define do_encrypt 0
@@ -195,10 +196,6 @@ static struct config_par parameters[] = {
 	}
 };
 
-static unsigned int page_size;
-/* This MUST be an multipe of page_size */
-static unsigned int buffer_size;
-static void *mem_pool;
 #ifdef CONFIG_ENCRYPT
 struct key_data *key_data;
 gcry_cipher_hd_t cipher_handle;
@@ -538,45 +535,55 @@ struct swap_map_handle {
 };
 
 /**
+ *	free_swap_writer - free memory allocated for saving the image
+ *	@handle:	Structure containing pointers to memory buffers to free.
+ */
+static void free_swap_writer(struct swap_map_handle *handle)
+{
+#ifdef CONFIG_COMPRESS
+	if (do_compress) {
+		freemem(handle->lzo_work_buffer);
+		freemem(handle->write_buffer);
+	}
+#endif
+#ifdef CONFIG_ENCRYPT
+	if (do_encrypt)
+		freemem(handle->encrypt_buffer);
+#endif
+	freemem(handle->buffer);
+	freemem(handle->extents);
+}
+
+/**
  *	init_swap_writer - initialize the structure used for saving the image
  *	@handle:	Structure to initialize.
  *	@dev:		Special device file to read image pages from.
  *	@fd:		File descriptor associated with the swap.
- *	@buf:		Memory pool to use for buffers.
  *
  *	It doesn't preallocate swap, so preallocate_swap() has to be called on
  *	@handle after this.
  */
-static int
-init_swap_writer(struct swap_map_handle *handle, int dev, int fd, void *buf)
+static int init_swap_writer(struct swap_map_handle *handle, int dev, int fd)
 {
 	loff_t offset;
 
-	if (!buf)
-		return -EINVAL;
-
-	handle->extents = buf;
-	buf += page_size;
+	handle->extents = getmem(page_size);
 
-	handle->buffer = buf;
-	handle->page_buffer = buf;
-	handle->write_buffer = buf;
-	buf += buffer_size;
+	handle->buffer = getmem(buffer_size);
+	handle->page_buffer = handle->buffer;
+	handle->write_buffer = handle->buffer;
 
 #ifdef CONFIG_ENCRYPT
 	if (do_encrypt) {
-		handle->encrypt_buffer = buf;
-		handle->encrypt_ptr = buf;
-		buf += encrypt_buf_size;
+		handle->encrypt_buffer = getmem(encrypt_buf_size);
+		handle->encrypt_ptr = handle->encrypt_buffer;
 	}
 #endif
 
 #ifdef CONFIG_COMPRESS
 	if (do_compress) {
-		handle->write_buffer = buf;
-		buf += compress_buf_size;
-		handle->lzo_work_buffer = buf;
-		/* This buffer must hold at least LZO1X_1_MEM_COMPRESS bytes */
+		handle->write_buffer = getmem(compress_buf_size);
+		handle->lzo_work_buffer = getmem(LZO1X_1_MEM_COMPRESS);
 	}
 #endif
 
@@ -587,8 +594,10 @@ init_swap_writer(struct swap_map_handle *handle, int dev, int fd, void *buf)
 	memset(handle->extents, 0, page_size);
 	handle->nr_extents = 0;
 	offset = get_swap_page(dev);
-	if (!offset)
+	if (!offset) {
+		free_swap_writer(handle);
 		return -ENOSPC;
+	}
 	handle->extents_spc = offset;
 
 	if (compute_checksum)
@@ -914,7 +923,6 @@ static int mark_swap(int fd, loff_t start)
 /**
  *	write_image - Write entire image and metadata.
  */
-
 int write_image(int snapshot_fd, int resume_fd)
 {
 	static struct swap_map_handle handle;
@@ -932,8 +940,7 @@ int write_image(int snapshot_fd, int resume_fd)
 	if (!start)
 		return -ENOSPC;
 
-	error = init_swap_writer(&handle, snapshot_fd, resume_fd,
-						mem_pool + page_size);
+	error = init_swap_writer(&handle, snapshot_fd, resume_fd);
 	if (error)
 		return error;
 
@@ -949,15 +956,23 @@ int write_image(int snapshot_fd, int resume_fd)
 		struct swsusp_info *image_header;
 		ssize_t ret;
 
-		image_header = (struct swsusp_info *)handle.page_buffer;
+		/*
+		 * Do it in such a way that save_image() will believe it has
+		 * already read the header page.
+		 */
+		image_header = handle.page_buffer;
 		ret = read(snapshot_fd, image_header, page_size);
-		if (ret < page_size)
-			return ret < 0 ? ret : -EFAULT;
+		if (ret < page_size) {
+			error = ret < 0 ? ret : -EFAULT;
+			goto Exit;
+		}
 		handle.page_buffer += page_size;
 		image_size = image_header->size;
 		nr_pages = image_header->pages;
-		if (!nr_pages)
-			return -ENODATA;
+		if (!nr_pages) {
+			error = -ENODATA;
+			goto Exit;
+		}
 		/* We have already read one page */
 		nr_pages--;
 	}
@@ -966,16 +981,18 @@ int write_image(int snapshot_fd, int resume_fd)
 	handle.swap_needed = image_size;
 	if (!enough_swap(&handle)) {
 		fprintf(stderr, "%s: Not enough free swap\n", my_name);
-		return -ENOSPC;
+		error = -ENOSPC;
+		goto Exit;
 	}
 	if (!preallocate_swap(&handle)) {
 		fprintf(stderr, "%s: Failed to allocate swap\n", my_name);
-		return -ENOSPC;
+		error = -ENOSPC;
+		goto Exit;
 	}
 	/* Shift handle.cur_offset for the first call to next_swap_page() */
 	handle.cur_offset -= page_size;
 
-	header  = mem_pool;
+	header  = getmem(page_size);
 	memset(header, 0, page_size);
 	header->pages = nr_pages;
 	header->flags = 0;
@@ -1027,7 +1044,7 @@ No_RSA:
 	if (error) {
 		fprintf(stderr,"%s: libgcrypt error: %s\n", my_name,
 			gcry_strerror(error));
-		goto Exit;
+		goto Free_header;
 	}
 
 Save_image:
@@ -1075,7 +1092,11 @@ Save_image:
 	if (!error)
 		printf( "|" );
 
+ Free_header:
+	freemem(header);
  Exit:
+	free_swap_writer(&handle);
+
 	printf("\n");
 	return error;
 }
@@ -1114,12 +1135,13 @@ static int reset_signature(int fd)
 	}
 	fsync(fd);
 	if (error) {
-		fprintf(stderr,
-			"reset_signature: Error %d resetting the image.\n"
-			"There should be valid image on disk. Powerdown and do normal resume.\n"
-			"Continuing with this booted system will lead to data corruption.\n", 
-			error);
-		while(1);
+		fprintf(stderr, "%s: Error %d resetting the image.\n"
+			"There should be valid image on disk. "
+			"Powerdown and carry out normal resume.\n"
+			"Continuing with this booted system "
+			"will lead to data corruption.\n", myname, error);
+		while(1)
+			sleep(10);
 	}
 	return error;
 }
@@ -1246,7 +1268,6 @@ Unfreeze:
  *	console_fd - get file descriptor for given file name and verify
  *	if that's a console descriptor (based on the code of openvt)
  */
-
 static inline int console_fd(const char *fname)
 {
 	int fd;
@@ -1255,7 +1276,8 @@ static inline int console_fd(const char *fname)
 	fd = open(fname, O_RDONLY);
 	if (fd < 0 && errno == EACCES)
 		fd = open(fname, O_WRONLY);
-	if (fd >= 0 && (ioctl(fd, KDGKBTYPE, &arg) || (arg != KB_101 && arg != KB_84))) {
+	if (fd >= 0 && (ioctl(fd, KDGKBTYPE, &arg)
+	    || (arg != KB_101 && arg != KB_84))) {
 		close(fd);
 		return -ENOTTY;
 	}
@@ -1273,11 +1295,10 @@ static int set_kmsg_redirect;
  *	the standard streams to it.  The number of the currently active
  *	virtual terminal is saved via @orig_vc
  */
-
 static int prepare_console(int *orig_vc, int *new_vc)
 {
 	int fd, error, vt = -1;
-	char *vt_name = mem_pool;
+	char vt_name[GENERIC_NAME_SIZE];
 	struct vt_stat vtstat;
 	char clear_vt, tiocl[2];
 
@@ -1353,9 +1374,8 @@ Close_fd:
 
 /**
  *	restore_console - switch to the virtual console that was active before
- *	suspend
+ *	                  suspend
  */
-
 static void restore_console(int fd, int orig_vc)
 {
 	int error;
@@ -1530,7 +1550,7 @@ static int lock_vt(void)
 	struct sigaction sa;
 	struct vt_mode vtm;
 	struct vt_stat vtstat;
-	char *vt_name = mem_pool;
+	char vt_name[GENERIC_NAME_SIZE];
 	int fd, error;
 
 	fd = console_fd("/dev/console");
@@ -1715,8 +1735,7 @@ int main(int argc, char *argv[])
 	do {
 		ret = open("/dev/null", O_RDWR);
 		if (ret < 0) {
-			ret = errno;
-			suspend_error("Could not open /dev/null.");
+			perror(argv[0]);
 			return ret;
 		}
 	} while (ret < 3);
@@ -1760,8 +1779,7 @@ int main(int argc, char *argv[])
 	if (resume_pause > RESUME_PAUSE_MAX)
 		resume_pause = RESUME_PAUSE_MAX;
 
-	page_size = getpagesize();
-	buffer_size = BUFFER_PAGES * page_size;
+	get_page_and_buffer_sizes();
 
 	mem_size = 2 * page_size + buffer_size;
 #ifdef CONFIG_COMPRESS
@@ -1770,13 +1788,13 @@ int main(int argc, char *argv[])
 		 * The formula below follows from the worst-case expansion
 		 * calculation for LZO1 (size / 16 + 67) and the fact that the
 		 * size of the compressed data must be stored in the buffer
-		 * (sizeof(size_t)).  Additionally, we want the buffer size to
-		 * be a multiple of page_size.
+		 * (sizeof(size_t)).
 		 */
 		compress_buf_size = buffer_size +
-			((page_size + (buffer_size >> 4) + 66 + sizeof(size_t))
-				& ~(page_size - 1));
-		mem_size += compress_buf_size + LZO1X_1_MEM_COMPRESS;
+			round_up_page_size((buffer_size >> 4) + 67 +
+						sizeof(size_t));
+		mem_size += compress_buf_size +
+				round_up_page_size(LZO1X_1_MEM_COMPRESS);
 	}
 #endif
 #ifdef CONFIG_ENCRYPT
@@ -1792,20 +1810,21 @@ int main(int argc, char *argv[])
 			do_encrypt = 0;
 		} else {
 			encrypt_buf_size = ENCRYPT_BUF_PAGES * page_size;
-			mem_size += encrypt_buf_size;
+			mem_size += encrypt_buf_size +
+				round_up_page_size(sizeof(struct key_data));
 		}
 	}
 #endif
-	mem_pool = malloc(mem_size);
-	if (!mem_pool) {
-		ret = errno;
-		suspend_error("Could not allocate memory.");
+
+	ret = init_memalloc(page_size, mem_size);
+	if (ret) {
+		fprintf(stderr, "%s: Could not allocate memory\n", my_name);
 		return ret;
 	}
+
 #ifdef CONFIG_ENCRYPT
 	if (do_encrypt) {
-		mem_size -= buffer_size;
-		key_data = (struct key_data *)((char *)mem_pool + mem_size);
+		key_data = getmem(sizeof(struct key_data));
 		generate_key();
 	}
 #endif
@@ -1908,10 +1927,8 @@ int main(int argc, char *argv[])
 		s2ram = !s2ram_hacks();
 #endif
 #ifdef CONFIG_ENCRYPT
-        if (do_encrypt && ! use_RSA) {
-                splash.read_password((char *)mem_pool,1);
-                strncpy(password,(char *)mem_pool,PASS_SIZE);
-        }
+        if (do_encrypt && ! use_RSA)
+                splash.read_password(password, 1);
 #endif
 
 	open_printk();
@@ -1959,10 +1976,12 @@ Umount:
 		umount(chroot_path);
 	}
 #ifdef CONFIG_ENCRYPT
-	if (do_encrypt)
+	if (do_encrypt) {
 		gcry_cipher_close(cipher_handle);
+		freemem(key_data);
+	}
 #endif
-	free(mem_pool);
+	free_memalloc();
 
 	return ret;
 }
diff --git a/swsusp.h b/swsusp.h
index 52e7cac..bcf1d56 100644
--- a/swsusp.h
+++ b/swsusp.h
@@ -159,8 +159,6 @@ struct buf_block {
 	char data[1];
 } __attribute__((packed));
 
-#define BUFFER_PAGES	32
-
 #define SNAPSHOT_DEVICE	"/dev/snapshot"
 #define RESUME_DEVICE ""
 
@@ -179,3 +177,5 @@ struct buf_block {
 #define ENTER_KEY_NAME	"ENTER"
 
 #define RESUME_PAUSE_MAX	60
+
+#define GENERIC_NAME_SIZE	256