/*-
 * SPDX-License-Identifier: BSD-2-Clause
 *
 * Copyright (c) 2024 Beckhoff Automation GmbH & Co. KG
 *
 * Redistribution and use in source and binary forms, with or without
 * modification, are permitted provided that the following conditions
 * are met:
 * 1. Redistributions of source code must retain the above copyright
 *    notice, this list of conditions and the following disclaimer.
 * 2. Redistributions in binary form must reproduce the above copyright
 *    notice, this list of conditions and the following disclaimer in the
 *    documentation and/or other materials provided with the distribution.
 *
 * THIS SOFTWARE IS PROVIDED BY THE AUTHOR AND CONTRIBUTORS ``AS IS'' AND
 * ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
 * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
 * ARE DISCLAIMED.  IN NO EVENT SHALL THE AUTHOR OR CONTRIBUTORS BE LIABLE
 * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
 * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS
 * OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION)
 * HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT
 * LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY
 * OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF
 * SUCH DAMAGE.
 */

#include <sys/stat.h>
#include <dirent.h>
#include <errno.h>
#include <fcntl.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <unistd.h>
#include <pwd.h>

#define	PAM_SM_SESSION

#include <security/pam_appl.h>
#include <security/pam_modules.h>
#include <security/pam_mod_misc.h>

#define	BASE_RUNTIME_DIR_PREFIX	"/var/run/xdg"
#define	RUNTIME_DIR_PREFIX	runtime_dir_prefix != NULL ? runtime_dir_prefix : BASE_RUNTIME_DIR_PREFIX

#define	RUNTIME_DIR_PREFIX_MODE	0711
#define	RUNTIME_DIR_MODE	0700	/* XDG spec */

#define	XDG_MAX_SESSION		100 /* Arbitrary limit because we need one */

static int
_pam_xdg_open(pam_handle_t *pamh, int flags __unused,
    int argc __unused, const char *argv[] __unused)
{
	struct passwd *passwd;
	const char *user;
	const char *runtime_dir_prefix;
	struct stat sb;
	char *runtime_dir = NULL;
	char *xdg_session_file;
	int rv, rt_dir_prefix, rt_dir, session_file, i;

	session_file = -1;
	rt_dir_prefix = -1;
	runtime_dir_prefix = openpam_get_option(pamh, "runtime_dir_prefix");

	/* Get user info */
	rv = pam_get_item(pamh, PAM_USER, (const void **)&user);
	if (rv != PAM_SUCCESS || user == NULL) {
		PAM_VERBOSE_ERROR("Can't get user information");
		goto out;
	}
	if ((passwd = getpwnam(user)) == NULL) {
		PAM_VERBOSE_ERROR("Can't get user information");
		rv = PAM_SESSION_ERR;
		goto out;
	}

	/* Open or create the base xdg directory */
	rt_dir_prefix = open(RUNTIME_DIR_PREFIX, O_DIRECTORY | O_NOFOLLOW);
	if (rt_dir_prefix < 0) {
		rt_dir_prefix = mkdir(RUNTIME_DIR_PREFIX, RUNTIME_DIR_PREFIX_MODE);
		if (rt_dir_prefix != 0) {
			PAM_VERBOSE_ERROR("Can't mkdir %s", RUNTIME_DIR_PREFIX);
			rv = PAM_SESSION_ERR;
			goto out;
		}
		rt_dir_prefix = open(RUNTIME_DIR_PREFIX, O_DIRECTORY | O_NOFOLLOW);
	}

	/* Open or create the user xdg directory */
	rt_dir = openat(rt_dir_prefix, user, O_DIRECTORY | O_NOFOLLOW);
	if (rt_dir < 0) {
		rt_dir = mkdirat(rt_dir_prefix, user, RUNTIME_DIR_MODE);
		if (rt_dir != 0) {
			PAM_VERBOSE_ERROR("mkdir: %s/%s (%d)", RUNTIME_DIR_PREFIX, user, rt_dir);
			rv = PAM_SESSION_ERR;
			goto out;
		}
		rv = fchownat(rt_dir_prefix, user, passwd->pw_uid, passwd->pw_gid, 0);
		if (rv != 0) {
			PAM_VERBOSE_ERROR("fchownat: %s/%s (%d)", RUNTIME_DIR_PREFIX, user, rv);
			rv = unlinkat(rt_dir_prefix, user, AT_REMOVEDIR);
			if (rv == -1)
				PAM_VERBOSE_ERROR("unlinkat: %s/%s (%d)", RUNTIME_DIR_PREFIX, user, errno);
			rv = PAM_SESSION_ERR;
			goto out;
		}
	} else {
		close(rt_dir);
		/* Check that the already create dir is correctly owned */
		rv = fstatat(rt_dir_prefix, user, &sb, 0);
		if (rv == -1) {
			PAM_VERBOSE_ERROR("fstatat %s/%s failed (%d)", RUNTIME_DIR_PREFIX, user, errno);
			rv = PAM_SESSION_ERR;
			goto out;
		}
		if (sb.st_uid != passwd->pw_uid ||
		  sb.st_gid != passwd->pw_gid) {
			PAM_VERBOSE_ERROR("%s/%s isn't owned by %d:%d\n", RUNTIME_DIR_PREFIX, user, passwd->pw_uid, passwd->pw_gid);
			rv = PAM_SESSION_ERR;
			goto out;
		}
		/* Test directory mode */
		if ((sb.st_mode & 0x1FF) != RUNTIME_DIR_MODE) {
			PAM_VERBOSE_ERROR("%s/%s have wrong mode\n", RUNTIME_DIR_PREFIX, user);
			rv = PAM_SESSION_ERR;
			goto out;
		}
	}

	/* Setup the environment variable */
	rv = asprintf(&runtime_dir, "XDG_RUNTIME_DIR=%s/%s", RUNTIME_DIR_PREFIX, user);
	if (rv < 0) {
		PAM_VERBOSE_ERROR("asprintf failed %d\n", rv);
		rv = PAM_SESSION_ERR;
		goto out;
	}
	rv = pam_putenv(pamh, runtime_dir);
	if (rv != PAM_SUCCESS) {
		PAM_VERBOSE_ERROR("pam_putenv: failed (%d)", rv);
		rv = PAM_SESSION_ERR;
		goto out;
	}

	/* Setup the session count file */
	for (i = 0; i < XDG_MAX_SESSION; i++) {
		rv = asprintf(&xdg_session_file, "%s/xdg_session.%d", user, i);
		if (rv < 0) {
			PAM_VERBOSE_ERROR("asprintf failed %d\n", rv);
			rv = PAM_SESSION_ERR;
			goto out;
		}
		rv = 0;
		session_file = openat(rt_dir_prefix, xdg_session_file, O_CREAT | O_EXCL, RUNTIME_DIR_MODE);
		free(xdg_session_file);
		if (session_file >= 0)
			break;
	}
	if (session_file < 0) {
		PAM_VERBOSE_ERROR("Too many sessions");
		rv = PAM_SESSION_ERR;
		goto out;
	}

out:
	if (session_file >= 0)
		close(session_file);
	if (rt_dir_prefix >= 0)
		close(rt_dir_prefix);

	if (runtime_dir)
		free(runtime_dir);
	return (rv);
}

static int
remove_dir(int fd)
{
	DIR *dirp;
	struct dirent *dp;

	dirp = fdopendir(fd);
	if (dirp == NULL)
		return (-1);

	while ((dp = readdir(dirp)) != NULL) {
		if (dp->d_type == DT_DIR) {
			int dirfd;

			if (strcmp(dp->d_name, ".") == 0 ||
			    strcmp(dp->d_name, "..") == 0)
				continue;
			dirfd = openat(fd, dp->d_name, 0);
			remove_dir(dirfd);
			close(dirfd);
			unlinkat(fd, dp->d_name, AT_REMOVEDIR);
			continue;
		}
		unlinkat(fd, dp->d_name, 0);
	}
	closedir(dirp);

	return (0);
}

static int
_pam_xdg_close(pam_handle_t *pamh __unused, int flags __unused,
    int argc __unused, const char *argv[] __unused)
{
	struct passwd *passwd;
	const char *user;
	const char *runtime_dir_prefix;
	struct stat sb;
	char *xdg_session_file;
	int rv, rt_dir_prefix, rt_dir, session_file, i;

	rt_dir = -1;
	rt_dir_prefix = -1;
	runtime_dir_prefix = openpam_get_option(pamh, "runtime_dir_prefix");

	/* Get user info */
	rv = pam_get_item(pamh, PAM_USER, (const void **)&user);
	if (rv != PAM_SUCCESS || user == NULL) {
		PAM_VERBOSE_ERROR("Can't get user information");
		goto out;
	}
	if ((passwd = getpwnam(user)) == NULL) {
		PAM_VERBOSE_ERROR("Can't get user information");
		rv = PAM_SESSION_ERR;
		goto out;
	}

	/* Open the xdg base directory */
	rt_dir_prefix = open(RUNTIME_DIR_PREFIX, O_DIRECTORY | O_NOFOLLOW);
	if (rt_dir_prefix < 0) {
		PAM_VERBOSE_ERROR("open: %s failed (%d)\n", runtime_dir_prefix, rt_dir_prefix);
		rv = PAM_SESSION_ERR;
		goto out;
	}
	/* Check that the already created dir is correctly owned */
	rv = fstatat(rt_dir_prefix, user, &sb, 0);
	if (rv == -1) {
		PAM_VERBOSE_ERROR("fstatat %s/%s failed (%d)", RUNTIME_DIR_PREFIX, user, errno);
		rv = PAM_SESSION_ERR;
		goto out;
	}
	if (sb.st_uid != passwd->pw_uid ||
	    sb.st_gid != passwd->pw_gid) {
		PAM_VERBOSE_ERROR("%s/%s isn't owned by %d:%d\n", RUNTIME_DIR_PREFIX, user, passwd->pw_uid, passwd->pw_gid);
		rv = PAM_SESSION_ERR;
		goto out;
	}
	/* Test directory mode */
	if ((sb.st_mode & 0x1FF) != RUNTIME_DIR_MODE) {
		PAM_VERBOSE_ERROR("%s/%s have wrong mode\n", RUNTIME_DIR_PREFIX, user);
		rv = PAM_SESSION_ERR;
		goto out;
	}

	/* Open the user xdg directory */
	rt_dir = openat(rt_dir_prefix, user, O_DIRECTORY | O_NOFOLLOW);
	if (rt_dir < 0) {
		PAM_VERBOSE_ERROR("openat: %s/%s failed (%d)\n", RUNTIME_DIR_PREFIX, user, rt_dir_prefix);
		rv = PAM_SESSION_ERR;
		goto out;
	}

	/* Get the last session file created */
	for (i = XDG_MAX_SESSION; i >= 0; i--) {
		rv = asprintf(&xdg_session_file, "%s/xdg_session.%d", user, i);
		if (rv < 0) {
			PAM_VERBOSE_ERROR("asprintf failed %d\n", rv);
			rv = PAM_SESSION_ERR;
			goto out;
		}
		rv = 0;
		session_file = openat(rt_dir_prefix, xdg_session_file, 0);
		if (session_file >= 0) {
			unlinkat(rt_dir_prefix, xdg_session_file, 0);
			free(xdg_session_file);
			break;
		}
		free(xdg_session_file);
	}
	if (session_file < 0) {
		PAM_VERBOSE_ERROR("Can't find session number\n");
		rv = PAM_SESSION_ERR;
		goto out;
	}
	close(session_file);

	/* Final cleanup if last user session */
	if (i == 0) {
		remove_dir(rt_dir);
		if (unlinkat(rt_dir_prefix, user, AT_REMOVEDIR) != 0) {
			PAM_VERBOSE_ERROR("Can't cleanup %s/%s\n", runtime_dir_prefix, user);
			rv = PAM_SESSION_ERR;
			goto out;
		}
	}

	rv = PAM_SUCCESS;
out:
	if (rt_dir >= 0)
		close(rt_dir);
	if (rt_dir_prefix >= 0)
		close(rt_dir_prefix);
	return (rv);
}

PAM_EXTERN int
pam_sm_open_session(pam_handle_t *pamh, int flags,
    int argc, const char *argv[])
{

	return (_pam_xdg_open(pamh, flags, argc, argv));
}

PAM_EXTERN int
pam_sm_close_session(pam_handle_t *pamh, int flags,
    int argc, const char *argv[])
{

	return (_pam_xdg_close(pamh, flags, argc, argv));
}

PAM_MODULE_ENTRY("pam_xdg");