// SPDX-License-Identifier: GPL-2.0-only
/*
 * Copyright (C) 2024 ARM Limited.
 */

#define _GNU_SOURCE

#include <stdio.h>
#include <stdlib.h>
#include <stdbool.h>
#include <errno.h>
#include <fcntl.h>
#include <signal.h>
#include <string.h>
#include <unistd.h>

#include <sys/socket.h>

#include <linux/kernel.h>
#include <linux/if_alg.h>

#define DATA_SIZE (16 * 4096)

static int base, sock;

static int digest_len;
static char *ref;
static char *digest;
static char *alg_name;

static struct iovec data_iov;
static int zerocopy[2];
static int sigs;
static int iter;

static void handle_exit_signal(int sig, siginfo_t *info, void *context)
{
	printf("Terminated by signal %d, iterations=%d, signals=%d\n",
	       sig, iter, sigs);
	exit(0);
}

static void handle_kick_signal(int sig, siginfo_t *info, void *context)
{
	sigs++;
}

static char *drivers[] = {
	"crct10dif-arm64-ce",
	/* "crct10dif-arm64-neon", - Same priority as generic */
	"sha1-ce",
	"sha224-arm64",
	"sha224-arm64-neon",
	"sha224-ce",
	"sha256-arm64",
	"sha256-arm64-neon",
	"sha256-ce",
	"sha384-ce",
	"sha512-ce",
	"sha3-224-ce",
	"sha3-256-ce",
	"sha3-384-ce",
	"sha3-512-ce",
	"sm3-ce",
	"sm3-neon",
};

static bool create_socket(void)
{
	FILE *proc;
	struct sockaddr_alg addr;
	char buf[1024];
	char *c, *driver_name;
	bool is_shash, match;
	int ret, i;

	ret = socket(AF_ALG, SOCK_SEQPACKET, 0);
	if (ret < 0) {
		if (errno == EAFNOSUPPORT) {
			printf("AF_ALG not supported\n");
			return false;
		}

		printf("Failed to create AF_ALG socket: %s (%d)\n",
		       strerror(errno), errno);
		return false;
	}
	base = ret;

	memset(&addr, 0, sizeof(addr));
	addr.salg_family = AF_ALG;
	strncpy((char *)addr.salg_type, "hash", sizeof(addr.salg_type));

	proc = fopen("/proc/crypto", "r");
	if (!proc) {
		printf("Unable to open /proc/crypto\n");
		return false;
	}

	driver_name = NULL;
	is_shash = false;
	match = false;

	/* Look through /proc/crypto for a driver with kernel mode FP usage */
	while (!match) {
		c = fgets(buf, sizeof(buf), proc);
		if (!c) {
			if (feof(proc)) {
				printf("Nothing found in /proc/crypto\n");
				return false;
			}
			continue;
		}

		/* Algorithm descriptions are separated by a blank line */
		if (*c == '\n') {
			if (is_shash && driver_name) {
				for (i = 0; i < ARRAY_SIZE(drivers); i++) {
					if (strcmp(drivers[i],
						   driver_name) == 0) {
						match = true;
					}
				}
			}

			if (!match) {
				digest_len = 0;

				free(driver_name);
				driver_name = NULL;

				free(alg_name);
				alg_name = NULL;

				is_shash = false;
			}
			continue;
		}

		/* Remove trailing newline */
		c = strchr(buf, '\n');
		if (c)
			*c = '\0';

		/* Find the field/value separator and start of the value */
		c = strchr(buf, ':');
		if (!c)
			continue;
		c += 2;

		if (strncmp(buf, "digestsize", strlen("digestsize")) == 0)
			sscanf(c, "%d", &digest_len);

		if (strncmp(buf, "name", strlen("name")) == 0)
			alg_name = strdup(c);

		if (strncmp(buf, "driver", strlen("driver")) == 0)
			driver_name = strdup(c);

		if (strncmp(buf, "type", strlen("type")) == 0)
			if (strncmp(c, "shash", strlen("shash")) == 0)
				is_shash = true;
	}

	strncpy((char *)addr.salg_name, alg_name,
		sizeof(addr.salg_name) - 1);

	ret = bind(base, (struct sockaddr *)&addr, sizeof(addr));
	if (ret < 0) {
		printf("Failed to bind %s: %s (%d)\n",
		       addr.salg_name, strerror(errno), errno);
		return false;
	}

	ret = accept(base, NULL, 0);
	if (ret < 0) {
		printf("Failed to accept %s: %s (%d)\n",
		       addr.salg_name, strerror(errno), errno);
		return false;
	}

	sock = ret;

	ret = pipe(zerocopy);
	if (ret != 0) {
		printf("Failed to create zerocopy pipe: %s (%d)\n",
		       strerror(errno), errno);
		return false;
	}

	ref = malloc(digest_len);
	if (!ref) {
		printf("Failed to allocated %d byte reference\n", digest_len);
		return false;
	}

	digest = malloc(digest_len);
	if (!digest) {
		printf("Failed to allocated %d byte digest\n", digest_len);
		return false;
	}

	return true;
}

static bool compute_digest(void *buf)
{
	struct iovec iov;
	int ret, wrote;

	iov = data_iov;
	while (iov.iov_len) {
		ret = vmsplice(zerocopy[1], &iov, 1, SPLICE_F_GIFT);
		if (ret < 0) {
			printf("Failed to send buffer: %s (%d)\n",
			       strerror(errno), errno);
			return false;
		}

		wrote = ret;
		ret = splice(zerocopy[0], NULL, sock, NULL, wrote, 0);
		if (ret < 0) {
			printf("Failed to splice buffer: %s (%d)\n",
			       strerror(errno), errno);
		} else if (ret != wrote) {
			printf("Short splice: %d < %d\n", ret, wrote);
		}

		iov.iov_len -= wrote;
		iov.iov_base += wrote;
	}

reread:
	ret = recv(sock, buf, digest_len, 0);
	if (ret == 0) {
		printf("No digest returned\n");
		return false;
	}
	if (ret != digest_len) {
		if (errno == -EAGAIN)
			goto reread;
		printf("Failed to get digest: %s (%d)\n",
		       strerror(errno), errno);
		return false;
	}

	return true;
}

int main(void)
{
	char *data;
	struct sigaction sa;
	int ret;

	/* Ensure we have unbuffered output */
	setvbuf(stdout, NULL, _IOLBF, 0);

	/* The parent will communicate with us via signals */
	memset(&sa, 0, sizeof(sa));
	sa.sa_sigaction = handle_exit_signal;
	sa.sa_flags = SA_RESTART | SA_SIGINFO;
	sigemptyset(&sa.sa_mask);
	ret = sigaction(SIGTERM, &sa, NULL);
	if (ret < 0)
		printf("Failed to install SIGTERM handler: %s (%d)\n",
		       strerror(errno), errno);

	sa.sa_sigaction = handle_kick_signal;
	ret = sigaction(SIGUSR2, &sa, NULL);
	if (ret < 0)
		printf("Failed to install SIGUSR2 handler: %s (%d)\n",
		       strerror(errno), errno);

	data = malloc(DATA_SIZE);
	if (!data) {
		printf("Failed to allocate data buffer\n");
		return EXIT_FAILURE;
	}
	memset(data, 0, DATA_SIZE);

	data_iov.iov_base = data;
	data_iov.iov_len = DATA_SIZE;

	/*
	 * If we can't create a socket assume it's a lack of system
	 * support and fall back to a basic FPSIMD test for the
	 * benefit of fp-stress.
	 */
	if (!create_socket()) {
		execl("./fpsimd-test", "./fpsimd-test", NULL);
		printf("Failed to fall back to fspimd-test: %d (%s)\n",
			errno, strerror(errno));
		return EXIT_FAILURE;
	}

	/*
	 * Compute a reference digest we hope is repeatable, we do
	 * this at runtime partly to make it easier to play with
	 * parameters.
	 */
	if (!compute_digest(ref)) {
		printf("Failed to compute reference digest\n");
		return EXIT_FAILURE;
	}

	printf("AF_ALG using %s\n", alg_name);

	while (true) {
		if (!compute_digest(digest)) {
			printf("Failed to compute digest, iter=%d\n", iter);
			return EXIT_FAILURE;
		}

		if (memcmp(ref, digest, digest_len) != 0) {
			printf("Digest mismatch, iter=%d\n", iter);
			return EXIT_FAILURE;
		}

		iter++;
	}

	return EXIT_FAILURE;
}