// SPDX-License-Identifier: GPL-2.0-only
/*
 * Copyright (C) 2015-2019 ARM Limited.
 * Original author: Dave Martin <Dave.Martin@arm.com>
 */
#define _GNU_SOURCE
#include <assert.h>
#include <errno.h>
#include <limits.h>
#include <stddef.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <getopt.h>
#include <unistd.h>
#include <sys/auxv.h>
#include <sys/prctl.h>
#include <asm/hwcap.h>
#include <asm/sigcontext.h>

static int inherit = 0;
static int no_inherit = 0;
static int force = 0;
static unsigned long vl;
static int set_ctl = PR_SVE_SET_VL;
static int get_ctl = PR_SVE_GET_VL;

static const struct option options[] = {
 { "force", no_argument, NULL, 'f' },
 { "inherit", no_argument, NULL, 'i' },
 { "max", no_argument, NULL, 'M' },
 { "no-inherit", no_argument, &no_inherit, 1 },
 { "sme", no_argument, NULL, 's' },
 { "help", no_argument, NULL, '?' },
 {}
};

static char const *program_name;

static int parse_options(int argc, char **argv)
{
 int c;
 char *rest;

 program_name = strrchr(argv[0], '/');
 if (program_name)
  ++program_name;
 else
  program_name = argv[0];

 while ((c = getopt_long(argc, argv, "Mfhi", options, NULL)) != -1)
  switch (c) {
  case 'M': vl = SVE_VL_MAX; break;
  case 'f': force = 1; break;
  case 'i': inherit = 1; break;
  case 's': set_ctl = PR_SME_SET_VL;
    get_ctl = PR_SME_GET_VL;
    break;
  case 0:  break;
  default: goto error;
  }

 if (inherit && no_inherit)
  goto error;

 if (!vl) {
  /* vector length */
  if (optind >= argc)
   goto error;

  errno = 0;
  vl = strtoul(argv[optind], &rest, 0);
  if (*rest) {
   vl = ULONG_MAX;
   errno = EINVAL;
  }
  if (vl == ULONG_MAX && errno) {
   fprintf(stderr, "%s: %s: %s\n",
    program_name, argv[optind], strerror(errno));
   goto error;
  }

  ++optind;
 }

 /* command */
 if (optind >= argc)
  goto error;

 return 0;

error:
 fprintf(stderr,
  "Usage: %s [-f | --force] "
  "[-i | --inherit | --no-inherit] "
  "{-M | --max | <vector length>} "
  "<command> [<arguments> ...]\n",
  program_name);
 return -1;
}

int main(int argc, char **argv)
{
 int ret = 126; /* same as sh(1) command-not-executable error */
 long flags;
 char *path;
 int t, e;

 if (parse_options(argc, argv))
  return 2; /* same as sh(1) builtin incorrect-usage */

 if (vl & ~(vl & PR_SVE_VL_LEN_MASK)) {
  fprintf(stderr, "%s: Invalid vector length %lu\n",
   program_name, vl);
  return 2; /* same as sh(1) builtin incorrect-usage */
 }

 if (!(getauxval(AT_HWCAP) & HWCAP_SVE)) {
  fprintf(stderr, "%s: Scalable Vector Extension not present\n",
   program_name);

  if (!force)
   goto error;

  fputs("Going ahead anyway (--force):  "
        "This is a debug option.  Don't rely on it.\n",
        stderr);
 }

 flags = PR_SVE_SET_VL_ONEXEC;
 if (inherit)
  flags |= PR_SVE_VL_INHERIT;

 t = prctl(set_ctl, vl | flags);
 if (t < 0) {
  fprintf(stderr, "%s: PR_SVE_SET_VL: %s\n",
   program_name, strerror(errno));
  goto error;
 }

 t = prctl(get_ctl);
 if (t == -1) {
  fprintf(stderr, "%s: PR_SVE_GET_VL: %s\n",
   program_name, strerror(errno));
  goto error;
 }
 flags = PR_SVE_VL_LEN_MASK;
 flags = t & ~flags;

 assert(optind < argc);
 path = argv[optind];

 execvp(path, &argv[optind]);
 e = errno;
 if (errno == ENOENT)
  ret = 127; /* same as sh(1) not-found error */
 fprintf(stderr, "%s: %s: %s\n", program_name, path, strerror(e));

error:
 return ret;  /* same as sh(1) not-executable error */
}