// SPDX-License-Identifier: GPL-2.0-only /* * Copyright (c) 2025, Google LLC. * Pasha Tatashin <pasha.tatashin@soleen.com> */ #define _GNU_SOURCE #include <stdio.h>#include <stdlib.h>#include <string.h>#include <getopt.h>#include <fcntl.h>#include <unistd.h>#include <sys/ioctl.h>#include <sys/syscall.h>#include <sys/mman.h>#include <sys/types.h>#include <sys/resource.h>#include <sys/stat.h>#include <errno.h>#include <stdarg.h>#include <linux/unistd.h> #include <libliveupdate.h> int luo_open_device(void){	return open(LUO_DEVICE, O_RDWR);} int luo_ensure_nofile_limit(long min_limit){	struct rlimit hl; 	/* Allow to extra files to be used by test itself */	min_limit += 32; 	if (getrlimit(RLIMIT_NOFILE, &hl) < 0)		return -errno; 	if (hl.rlim_cur >= min_limit)		return 0; 	hl.rlim_cur = min_limit;	if (hl.rlim_cur > hl.rlim_max)		hl.rlim_max = hl.rlim_cur; 	if (setrlimit(RLIMIT_NOFILE, &hl) < 0)		return -errno; 	return 0;} int luo_create_session(int luo_fd, const char *name){	struct liveupdate_ioctl_create_session arg = { .size = sizeof(arg) }; 	snprintf((char *)arg.name, LIVEUPDATE_SESSION_NAME_LENGTH, "%.*s",		 LIVEUPDATE_SESSION_NAME_LENGTH - 1, name); 	if (ioctl(luo_fd, LIVEUPDATE_IOCTL_CREATE_SESSION, &arg))		return -errno; 	return arg.fd;} int luo_retrieve_session(int luo_fd, const char *name){	struct liveupdate_ioctl_retrieve_session arg = { .size = sizeof(arg) }; 	snprintf((char *)arg.name, LIVEUPDATE_SESSION_NAME_LENGTH, "%.*s",		 LIVEUPDATE_SESSION_NAME_LENGTH - 1, name); 	if (ioctl(luo_fd, LIVEUPDATE_IOCTL_RETRIEVE_SESSION, &arg))		return -errno; 	return arg.fd;} int luo_session_preserve_fd(int session_fd, int fd, __u64 token){	struct liveupdate_session_preserve_fd arg = {		.size = sizeof(arg),		.fd = fd,		.token = token,	}; 	if (ioctl(session_fd, LIVEUPDATE_SESSION_PRESERVE_FD, &arg))		return -errno; 	return 0;} int luo_session_retrieve_fd(int session_fd, __u64 token){	struct liveupdate_session_retrieve_fd arg = {		.size = sizeof(arg),		.token = token,	}; 	if (ioctl(session_fd, LIVEUPDATE_SESSION_RETRIEVE_FD, &arg))		return -errno; 	return arg.fd;} /* Helper function to get a session name via ioctl. */int luo_get_session_name(int session_fd, char *name, size_t name_len){	struct liveupdate_session_get_name args = {}; 	args.size = sizeof(args); 	if (ioctl(session_fd, LIVEUPDATE_SESSION_GET_NAME, &args))		return -errno; 	strncpy(name, (char *)args.name, name_len - 1);	name[name_len - 1] = '\0'; 	return 0;} int create_and_preserve_memfd(int session_fd, int token, const char *data){	long page_size = getpagesize();	void *map = MAP_FAILED;	int mfd = -1, ret = -1; 	mfd = memfd_create("test_mfd", 0);	if (mfd < 0)		return -errno; 	if (ftruncate(mfd, page_size) != 0)		goto out; 	map = mmap(NULL, page_size, PROT_WRITE, MAP_SHARED, mfd, 0);	if (map == MAP_FAILED)		goto out; 	snprintf(map, page_size, "%s", data);	munmap(map, page_size); 	ret = luo_session_preserve_fd(session_fd, mfd, token);	if (ret)		goto out; 	ret = 0;out:	if (ret != 0 && errno != 0)		ret = -errno;	if (mfd >= 0)		close(mfd);	return ret;} int restore_and_verify_memfd(int session_fd, int token,			     const char *expected_data){	long page_size = getpagesize();	void *map = MAP_FAILED;	int mfd = -1, ret = -1; 	mfd = luo_session_retrieve_fd(session_fd, token);	if (mfd < 0)		return mfd; 	map = mmap(NULL, page_size, PROT_READ, MAP_SHARED, mfd, 0);	if (map == MAP_FAILED)		goto out; 	if (expected_data && strcmp(expected_data, map) != 0) {		ksft_print_msg("Data mismatch! Expected '%s', Got '%s'\n",			       expected_data, (char *)map);		ret = -EINVAL;		goto out_munmap;	} 	ret = mfd;out_munmap:	munmap(map, page_size);out:	if (ret < 0 && errno != 0)		ret = -errno;	if (ret < 0 && mfd >= 0)		close(mfd);	return ret;} int luo_session_finish(int session_fd){	struct liveupdate_session_finish arg = { .size = sizeof(arg) }; 	if (ioctl(session_fd, LIVEUPDATE_SESSION_FINISH, &arg) < 0)		return -errno; 	return 0;} void create_state_file(int luo_fd, const char *session_name, int token,		       int next_stage){	char buf[32];	int state_session_fd; 	state_session_fd = luo_create_session(luo_fd, session_name);	if (state_session_fd < 0)		fail_exit("luo_create_session for state tracking"); 	snprintf(buf, sizeof(buf), "%d", next_stage);	if (create_and_preserve_memfd(state_session_fd, token, buf) < 0)		fail_exit("create_and_preserve_memfd for state tracking"); 	/*	 * DO NOT close session FD, otherwise it is going to be unpreserved	 */} void restore_and_read_stage(int state_session_fd, int token, int *stage){	char buf[32] = {0};	int mfd; 	mfd = restore_and_verify_memfd(state_session_fd, token, NULL);	if (mfd < 0)		fail_exit("failed to restore state memfd"); 	if (read(mfd, buf, sizeof(buf) - 1) < 0)		fail_exit("failed to read state mfd"); 	*stage = atoi(buf); 	close(mfd);} void daemonize_and_wait(void){	pid_t pid; 	ksft_print_msg("[STAGE 1] Forking persistent child to hold sessions...\n"); 	pid = fork();	if (pid < 0)		fail_exit("fork failed"); 	if (pid > 0) {		ksft_print_msg("[STAGE 1] Child PID: %d. Resources are pinned.\n", pid);		ksft_print_msg("[STAGE 1] You may now perform kexec reboot.\n");		exit(EXIT_SUCCESS);	} 	/* Detach from terminal so closing the window doesn't kill us */	if (setsid() < 0)		fail_exit("setsid failed"); 	close(STDIN_FILENO);	close(STDOUT_FILENO);	close(STDERR_FILENO); 	/* Change dir to root to avoid locking filesystems */	if (chdir("/") < 0)		exit(EXIT_FAILURE); 	while (1)		sleep(60);} static int parse_stage_args(int argc, char *argv[]){	int stage = 1;	int opt; 	optind = 1;	while ((opt = getopt(argc, argv, "s:")) != -1) {		switch (opt) {		case 's':			stage = atoi(optarg);			if (stage != 1 && stage != 2)				fail_exit("Invalid stage argument");			break;		default:			fail_exit("Unknown argument");		}	} 	return stage;} int luo_test(int argc, char *argv[],	     const char *state_session_name,	     luo_test_stage1_fn stage1,	     luo_test_stage2_fn stage2){	int target_stage = parse_stage_args(argc, argv);	int luo_fd = luo_open_device();	int state_session_fd;	int detected_stage; 	if (luo_fd < 0) {		ksft_exit_skip("Failed to open %s. Is the luo module loaded?\n",			       LUO_DEVICE);	} 	state_session_fd = luo_retrieve_session(luo_fd, state_session_name);	if (state_session_fd == -ENOENT)		detected_stage = 1;	else if (state_session_fd >= 0)		detected_stage = 2;	else		fail_exit("Failed to check for state session"); 	if (target_stage != detected_stage) {		ksft_exit_fail_msg("Stage mismatch Requested stage %d, but system is in stage %d.\n"				   "(State session %s: %s)\n",				   target_stage, detected_stage, state_session_name,				   (detected_stage == 2) ? "EXISTS" : "MISSING");	} 	if (target_stage == 1)		stage1(luo_fd);	else		stage2(luo_fd, state_session_fd); 	return 0;}