/*
 * suspend.c
 *
 * A simple user space suspend handler for swsusp.
 *
 * Copyright (C) 2005 Rafael J. Wysocki <rjw@sisk.pl>
 *
 * This file is released under the GPLv2.
 *
 */

#include <sys/types.h>
#include <sys/stat.h>
#include <sys/ioctl.h>
#include <sys/mman.h>
#include <syscall.h>
#include <fcntl.h>
#include <unistd.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <errno.h>

#include "swsusp.h"

static int dev;
static char buffer[PAGE_SIZE];

static inline unsigned int check_free_swap(void)
{
	int error;
	unsigned int free_swap;

	error = ioctl(dev, SNAPSHOT_IOCAVAIL_SWAP, &free_swap);
	if (!error)
		return free_swap;
	else
		printf("Error: errno = %d\n", errno);
	return 0;
}

static inline unsigned long get_swap_page(void)
{
	int error;
	unsigned long offset;

	error = ioctl(dev, SNAPSHOT_IOCGET_SWAP_PAGE, &offset);
	if (!error)
		return offset;
	return 0;
}

static inline int free_swap_page(unsigned long offset)
{
	return ioctl(dev, SNAPSHOT_IOCFREE_SWAP_PAGE, offset);
}

/**
 *	write_page - Write one page to a fresh swap location.
 *	@fd:	File handle of the resume partition
 *	@buf:	Pointer to the area we're writing.
 *	@offp:	Place to store the swap offset we used.
 */

static int write_page(int fd, void *buf, unsigned long *offp)
{
	unsigned long swap_offset;
	off_t offset;
	int res = -ENOSPC;
	ssize_t cnt = 0;

	swap_offset = get_swap_page();
	if (swap_offset) {
		offset = swap_offset * PAGE_SIZE;
		if (lseek(fd, offset, SEEK_SET) == offset) 
			cnt = write(fd, buf, PAGE_SIZE);
		if (cnt == PAGE_SIZE) {
			*offp = swap_offset;
			res = 0;
		} else {
			free_swap_page(swap_offset);
			if (cnt < 0)
				res = cnt;
		}
	}
	return res;
}

static inline void free_swap_map(struct swap_map_page *swap_map)
{
	struct swap_map_page *swp;

	while (swap_map) {
		swp = swap_map->next;
		free(swap_map);
		swap_map = swp;
	}
}

static struct swap_map_page *alloc_swap_map(unsigned int nr_pages)
{
	struct swap_map_page *swap_map, *swp;
	unsigned n = 0;

	if (!nr_pages)
		return NULL;

	printf("alloc_swap_map(): nr_pages = %d\n", nr_pages);
	swap_map = malloc(PAGE_SIZE);
	memset(swap_map, 0, PAGE_SIZE);
	swp = swap_map;
	for (n = MAP_PAGE_SIZE; n < nr_pages; n += MAP_PAGE_SIZE) {
		swp->next = malloc(PAGE_SIZE);
		memset(swp->next, 0, PAGE_SIZE);
		swp = swp->next;
		if (!swp) {
			free_swap_map(swap_map);
			return NULL;
		}
	}
	return swap_map;
}

/**
 *	reverse_swap_map - reverse the order of pages in the swap map
 *	@swap_map
 */

static inline struct swap_map_page *reverse_swap_map(struct swap_map_page *swap_map)
{
	struct swap_map_page *prev, *next;

	prev = NULL;
	while (swap_map) {
		next = swap_map->next;
		swap_map->next = prev;
		prev = swap_map;
		swap_map = next;
	}
	return prev;
}

/**
 *	free_swap_map_entries - free the swap entries allocated to store
 *	the swap map @swap_map (this is only called in case of an error)
 */
static inline void free_swap_map_entries(struct swap_map_page *swap_map)
{
	while (swap_map) {
		if (swap_map->next_swap)
			free_swap_page(swap_map->next_swap);
		swap_map = swap_map->next;
	}
}

/**
 *	save_swap_map - save the swap map used for tracing the data pages
 *	stored in the swap
 */

static int save_swap_map(int fd, struct swap_map_page *swap_map, unsigned long *start)
{
	unsigned long offset = 0;
	int error;

	while (swap_map) {
		swap_map->next_swap = offset;
		if ((error = write_page(fd, swap_map, &offset)))
			return error;
		swap_map = swap_map->next;
	}
	*start = offset;
	return 0;
}

/**
 *	free_image_entries - free the swap entries allocated to store
 *	the image data pages (this is only called in case of an error)
 */

static inline void free_image_entries(struct swap_map_page *swp)
{
	unsigned k;

	while (swp) {
		for (k = 0; k < MAP_PAGE_SIZE; k++)
			if (swp->entries[k])
				free_swap_page(swp->entries[k]);
		swp = swp->next;
	}
}

static inline void init_swap_map_handle(struct swap_map_handle *handle,
                                        int fd, struct swap_map_page *map)
{
	handle->cur = map;
	handle->k = 0;
	handle->fd = fd;
}

static inline int swap_map_write_page(struct swap_map_handle *handle, void *buf)
{
	int error;

	error = write_page(handle->fd, buf, handle->cur->entries + handle->k);
	if (error)
		return error;
	if (++handle->k >= MAP_PAGE_SIZE) {
		handle->cur = handle->cur->next;
		handle->k = 0;
	}
	return 0;
}

/**
 *	save_image - save the suspend image data
 */

static int save_image(struct swap_map_handle *handle,
                      unsigned int nr_pages)
{
	unsigned int m;
	int ret;
	int error = 0;

	printf("suspend: Saving image data pages (%u pages) ...     ", nr_pages);
	m = nr_pages / 100;
	if (!m)
		m = 1;
	nr_pages = 0;
	do {
		ret = read(dev, buffer, PAGE_SIZE);
		if (ret > 0) {
			error = swap_map_write_page(handle, buffer);
			if (error)
				break;
			if (!(nr_pages % m))
				printf("\b\b\b\b%3d%%", nr_pages / m);
			nr_pages++;
		}
	} while (ret > 0);
	if (!error)
		printf(" done\n");
	return error;
}

/**
 *	enough_swap - Make sure we have enough swap to save the image.
 *
 *	Returns TRUE or FALSE after checking the total amount of swap
 *	space avaiable from the resume partition.
 */

static int enough_swap(unsigned int nr_pages)
{
	unsigned int free_swap = check_free_swap();

	printf("suspend: Free swap pages: %u\n", free_swap);
	return free_swap > nr_pages;
}

static int mark_swapfiles(int fd, unsigned long start)
{
	int error = 0;

	lseek(fd, 0, SEEK_SET);
	if (read(fd, &swsusp_header, PAGE_SIZE) < PAGE_SIZE)
		return -EIO;
	if (!memcmp("SWAP-SPACE", swsusp_header.sig, 10) ||
	    !memcmp("SWAPSPACE2", swsusp_header.sig, 10)) {
		memcpy(swsusp_header.orig_sig, swsusp_header.sig, 10);
		memcpy(swsusp_header.sig, SWSUSP_SIG, 10);
		swsusp_header.image = start;
		lseek(fd, 0, SEEK_SET);
		if (write(fd, &swsusp_header, PAGE_SIZE) < PAGE_SIZE)
			error = -EIO;
	} else {
		printf("suspend: Device is not a swap space.\n");
		error = -ENODEV;
	}
	return error;
}

/**
 *	write_image - Write entire image and metadata.
 */

int write_image(char *resume_dev_name)
{
	struct swap_map_page *swap_map;
	struct swap_map_handle handle;
	struct swsusp_info *header;
	unsigned long start;
	int fd;
	int error;

	fd = open(resume_dev_name, O_RDWR | O_SYNC);
	if (fd < 0) {
		printf("suspend: Could not open resume device\n");
		return error;
	}
	error = read(dev, buffer, PAGE_SIZE);
	if (error < PAGE_SIZE)
		return error < 0 ? error : -EFAULT;
	header = (struct swsusp_info *)buffer;
	if (!enough_swap(header->pages)) {
		printf("suspend: Not enough free swap\n");
		return -ENOSPC;
	}
	swap_map = alloc_swap_map(header->pages);
	if (!swap_map)
		return -ENOMEM;
	init_swap_map_handle(&handle, fd, swap_map);

	error = swap_map_write_page(&handle, header);
	if (!error)
		error = save_image(&handle, header->pages - 1);
	if (error)
		goto Free_image_entries;

	swap_map = reverse_swap_map(swap_map);
	error = save_swap_map(fd, swap_map, &start);
	if (error)
		goto Free_map_entries;

	printf( "S" );
	error = mark_swapfiles(fd, start);
	printf( "|\n" );
	if (error)
		goto Free_map_entries;

Free_swap_map:
	free_swap_map(swap_map);
	fsync(fd);
	close(fd);
	return error;

Free_map_entries:
	free_swap_map_entries(swap_map);
Free_image_entries:
	free_image_entries(swap_map);
	goto Free_swap_map;
}

int main(int argc, char *argv[])
{
	char *snapshot_device_name, *resume_device_name;
	int image_size;
	int in_suspend;

	resume_device_name = argc <= 2 ? RESUME_DEVICE : argv[2];
	snapshot_device_name = argc <= 1 ? SNAPSHOT_DEVICE : argv[1];

	sync();
	setvbuf(stdout, NULL, _IONBF, 0);
	setvbuf(stderr, NULL, _IONBF, 0);

	if (mlockall(MCL_CURRENT | MCL_FUTURE)) {
		fprintf(stderr, "Could not lock myself\n");
		return 1;
	}

	image_size = DEFAULT_IMAGE_SIZE;
	while (image_size >= 0) {
		dev = open(snapshot_device_name, O_RDONLY);
		if (dev < 0)
			return -ENOENT;
		if (set_image_size(dev, image_size)) {
			fprintf(stderr, "Could not set the image size\n");
			return errno;
		}
		if (freeze(dev)) {
			fprintf(stderr, "Could not freeze processes\n");
			return errno;
		}
		if (!atomic_snapshot(dev, &in_suspend)) {
			if (!in_suspend) {
				unfreeze(dev);
				close(dev);
				return 0;
			}
			if (!write_image(resume_device_name))
				power_off();
			else
				image_size -= DEFAULT_IMAGE_SIZE;
		}
		unfreeze(dev);
		close(dev);
	}

	return 0;
}