pkg(7): rework the arguments handling

Rework the arguments handling around using getopt_long:
* add long option support
* add -4 and -6 support to enforce ipv4 or ipv6

While here fix a regression which occured between FreeBSD 12.1 and
FreeBSD 12.2 where pkg bootstrap -y stopped working

PR:		252270
MFC after:	2 weeks
Submitted by:	evilham <contact@evilham.com>
Differential Revision:	https://reviews.freebsd.org/D27860

2 files changed, 102 insertions, 40 deletions

diff --git a/usr.sbin/pkg/pkg.7 b/usr.sbin/pkg/pkg.7
index f7116f5da828..c672306ac08d 100644
--- a/usr.sbin/pkg/pkg.7
+++ b/usr.sbin/pkg/pkg.7
@@ -24,7 +24,7 @@
 .\"
 .\" $FreeBSD$
 .\"
-.Dd November 6, 2020
+.Dd December 31, 2020
 .Dt PKG 7
 .Os
 .Sh NAME
@@ -36,12 +36,15 @@
 .Nm
 add
 .Op Fl f
+.Op Fl y
 .Ao Pa pkg.txz Ac
 .Nm
 .Fl N
 .Nm
+.Op Fl 4 | Fl 6
 bootstrap
 .Op Fl f
+.Op Fl y
 .Sh DESCRIPTION
 .Nm
 is the package management tool.
@@ -65,7 +68,7 @@ is not installed yet, it will be fetched, have its signature verified,
 installed, and then have the original command forwarded to it.
 If already installed, the command requested will be forwarded to the real
 .Xr pkg 8 .
-.It Nm Li add Oo Fl f Oc Ao Pa pkg.txz Ac
+.It Nm Li add Oo Fl f Oc Oo Fl y Oc Ao Pa pkg.txz Ac
 Install
 .Xr pkg 8
 from a local package instead of fetching from remote.
@@ -76,21 +79,37 @@ If the
 flag is specified, then
 .Xr pkg 8
 will be installed regardless if it is already installed.
+If the
+.Fl y
+flag is specified, no confirmation will be asked when bootstrapping
+.Xr pkg 8 .
 .It Nm Fl N
 Do not bootstrap, just determine if
 .Xr pkg 8
 is actually installed or not.
 Returns 0 and the number of packages installed
 if it is, otherwise 1.
-.It Nm Li bootstrap Op Fl f
+.It Nm Oo Fl 4 | Fl 6 Oc Li bootstrap Oo Fl f Oc Oo Fl y Oc
 Attempt to bootstrap and do not forward anything to
 .Xr pkg 8
 after it is installed.
+With
+.Fl 4
+and
+.Fl 6 ,
+.Nm
+will force IPv4 or IPv6 respectively to fetch
+.Xr pkg 8
+and its signatures as needed.
 If the
 .Fl f
 flag is specified, then
 .Xr pkg 8
 will be fetched and installed regardless if it is already installed.
+If the
+.Fl y
+flag is specified, no confirmation will be asked when bootstrapping
+.Xr pkg 8 .
 .El
 .Sh CONFIGURATION
 Configuration varies in whether it is in a repository configuration file
diff --git a/usr.sbin/pkg/pkg.c b/usr.sbin/pkg/pkg.c
index a35efd8f91a8..788fdb39ebb9 100644
--- a/usr.sbin/pkg/pkg.c
+++ b/usr.sbin/pkg/pkg.c
@@ -43,6 +43,7 @@ __FBSDID("$FreeBSD$");
 #include <errno.h>
 #include <fcntl.h>
 #include <fetch.h>
+#include <getopt.h>
 #include <libutil.h>
 #include <paths.h>
 #include <stdbool.h>
@@ -173,7 +174,7 @@ install_pkg_static(const char *path, const char *pkgpath, bool force)
 }
 
 static int
-fetch_to_fd(const char *url, char *path)
+fetch_to_fd(const char *url, char *path, const char *fetchOpts)
 {
 	struct url *u;
 	struct dns_srvinfo *mirrors, *current;
@@ -225,7 +226,7 @@ fetch_to_fd(const char *url, char *path)
 			u->port = current->port;
 		}
 
-		remote = fetchXGet(u, &st, "");
+		remote = fetchXGet(u, &st, fetchOpts);
 		if (remote == NULL) {
 			--retry;
 			if (retry <= 0)
@@ -828,7 +829,7 @@ cleanup:
 }
 
 static int
-bootstrap_pkg(bool force)
+bootstrap_pkg(bool force, const char *fetchOpts)
 {
 	int fd_pkg, fd_sig;
 	int ret;
@@ -865,7 +866,7 @@ bootstrap_pkg(bool force)
 	snprintf(tmppkg, MAXPATHLEN, "%s/pkg.txz.XXXXXX",
 	    getenv("TMPDIR") ? getenv("TMPDIR") : _PATH_TMP);
 
-	if ((fd_pkg = fetch_to_fd(url, tmppkg)) == -1)
+	if ((fd_pkg = fetch_to_fd(url, tmppkg, fetchOpts)) == -1)
 		goto fetchfail;
 
 	if (signature_type != NULL &&
@@ -877,7 +878,7 @@ bootstrap_pkg(bool force)
 			snprintf(url, MAXPATHLEN, "%s/Latest/pkg.txz.sig",
 			    packagesite);
 
-			if ((fd_sig = fetch_to_fd(url, tmpsig)) == -1) {
+			if ((fd_sig = fetch_to_fd(url, tmpsig, fetchOpts)) == -1) {
 				fprintf(stderr, "Signature for pkg not "
 				    "available.\n");
 				goto fetchfail;
@@ -893,7 +894,7 @@ bootstrap_pkg(bool force)
 			snprintf(url, MAXPATHLEN, "%s/Latest/pkg.txz.pubkeysig",
 			    packagesite);
 
-			if ((fd_sig = fetch_to_fd(url, tmpsig)) == -1) {
+			if ((fd_sig = fetch_to_fd(url, tmpsig, fetchOpts)) == -1) {
 				fprintf(stderr, "Signature for pkg not "
 				    "available.\n");
 				goto fetchfail;
@@ -948,6 +949,14 @@ static const char non_interactive_message[] =
 "Please set ASSUME_ALWAYS_YES=yes environment variable to be able to bootstrap "
 "in non-interactive (stdin not being a tty)\n";
 
+static const char args_bootstrap_message[] =
+"Too many arguments\n"
+"Usage: pkg [-4|-6] bootstrap [-f] [-y]\n";
+
+static const char args_add_message[] =
+"Too many arguments\n"
+"Usage: pkg add [-f] [-y] {pkg.txz}\n";
+
 static int
 pkg_query_yes_no(void)
 {
@@ -1037,27 +1046,74 @@ main(int argc, char *argv[])
 {
 	char pkgpath[MAXPATHLEN];
 	const char *pkgarg;
-	int i;
-	bool bootstrap_only, force, yes;
+	bool activation_test, add_pkg, bootstrap_only, force, yes;
+	signed char ch;
+	const char *fetchOpts;
+	char *command;
 
+	activation_test = false;
+	add_pkg = false;
 	bootstrap_only = false;
+	command = NULL;
+	fetchOpts = "";
 	force = false;
 	pkgarg = NULL;
 	yes = false;
 
+	struct option longopts[] = {
+		{ "force",		no_argument,		NULL,	'f' },
+		{ "only-ipv4",		no_argument,		NULL,	'4' },
+		{ "only-ipv6",		no_argument,		NULL,	'6' },
+		{ "yes",		no_argument,		NULL,	'y' },
+		{ NULL,			0,			NULL,	0   },
+	};
+
 	snprintf(pkgpath, MAXPATHLEN, "%s/sbin/pkg", getlocalbase());
 
-	if (argc > 1 && strcmp(argv[1], "bootstrap") == 0) {
-		bootstrap_only = true;
-		if (argc > 3) {
-			fprintf(stderr, "Too many arguments\nUsage: pkg bootstrap [-f]\n");
-			exit(EXIT_FAILURE);
-		}
-		if (argc == 3 && strcmp(argv[2], "-f") == 0) {
+	while ((ch = getopt_long(argc, argv, "-:fyN46", longopts, NULL)) != -1) {
+		switch (ch) {
+		case 'f':
 			force = true;
-		} else if (argc == 3) {
-			fprintf(stderr, "Invalid argument specified\nUsage: pkg bootstrap [-f]\n");
-			exit(EXIT_FAILURE);
+			break;
+		case 'N':
+			activation_test = true;
+			break;
+		case 'y':
+			yes = true;
+			break;
+		case '4':
+			fetchOpts = "4";
+			break;
+		case '6':
+			fetchOpts = "6";
+			break;
+		case 1:
+			// Non-option arguments, first one is the command
+			if (command == NULL) {
+				command = argv[optind-1];
+				if (strcmp(command, "add") == 0) {
+					add_pkg = true;
+				}
+				else if (strcmp(command, "bootstrap") == 0) {
+					bootstrap_only = true;
+				}
+			}
+			// bootstrap doesn't accept other arguments
+			else if (bootstrap_only) {
+				fprintf(stderr, args_bootstrap_message);
+				exit(EXIT_FAILURE);
+			}
+			// For add, we accept exactly one further argument
+			else if (add_pkg && pkgarg != NULL) {
+				fprintf(stderr, args_add_message);
+				exit(EXIT_FAILURE);
+			}
+			else if (add_pkg) {
+				pkgarg = argv[optind-1];
+			}
+			break;
+		default:
+			break;
 		}
 	}
 
@@ -1065,19 +1121,14 @@ main(int argc, char *argv[])
 		/* 
 		 * To allow 'pkg -N' to be used as a reliable test for whether
 		 * a system is configured to use pkg, don't bootstrap pkg
-		 * when that argument is given as argv[1].
+		 * when that that option is passed.
 		 */
-		if (argv[1] != NULL && strcmp(argv[1], "-N") == 0)
+		if (activation_test)
 			errx(EXIT_FAILURE, "pkg is not installed");
 
 		config_init();
 
-		if (argc > 1 && strcmp(argv[1], "add") == 0) {
-			if (argc > 2 && strcmp(argv[2], "-f") == 0) {
-				force = true;
-				pkgarg = argv[3];
-			} else
-				pkgarg = argv[2];
+		if (add_pkg) {
 			if (pkgarg == NULL) {
 				fprintf(stderr, "Path to pkg.txz required\n");
 				exit(EXIT_FAILURE);
@@ -1095,16 +1146,8 @@ main(int argc, char *argv[])
 		 * not tty. Check the environment to see if user has answer
 		 * tucked in there already.
 		 */
-		config_bool(ASSUME_ALWAYS_YES, &yes);
-		if (!yes) {
-			for (i = 1; i < argc; i++) {
-				if (strcmp(argv[i], "-y") == 0 ||
-				    strcmp(argv[i], "--yes") == 0) {
-					yes = true;
-					break;
-				}
-			}
-		}
+		if (!yes)
+			config_bool(ASSUME_ALWAYS_YES, &yes);
 		if (!yes) {
 			if (!isatty(fileno(stdin))) {
 				fprintf(stderr, non_interactive_message);
@@ -1115,7 +1158,7 @@ main(int argc, char *argv[])
 			if (pkg_query_yes_no() == 0)
 				exit(EXIT_FAILURE);
 		}
-		if (bootstrap_pkg(force) != 0)
+		if (bootstrap_pkg(force, fetchOpts) != 0)
 			exit(EXIT_FAILURE);
 		config_finish();