MFC: r256968, r256971, r256978

Improve SRV records support for the pkg(8) bootstrap:
- order srv records by priorities
- for all entries of the same priority, order randomly respect the weight
- select the port where to fetch from respect the port provided in the SRV
record

Allow to bootstrap by doing pkg add ./a/path/to/a/pkg_package.txz

Approved by:	re (glebius)

3 files changed, 108 insertions, 3 deletions

diff --git a/usr.sbin/pkg/dns_utils.c b/usr.sbin/pkg/dns_utils.c
index cbdb3d5..ce88048 100644
--- a/usr.sbin/pkg/dns_utils.c
+++ b/usr.sbin/pkg/dns_utils.c
@@ -1,5 +1,5 @@
 /*-
- * Copyright (c) 2012 Baptiste Daroussin <bapt@FreeBSD.org>
+ * Copyright (c) 2012-2013 Baptiste Daroussin <bapt@FreeBSD.org>
  * All rights reserved.
  *
  * Redistribution and use in source and binary forms, with or without
@@ -39,6 +39,77 @@ typedef union {
 	unsigned char buf[1024];
 } dns_query;
 
+static int
+srv_priority_cmp(const void *a, const void *b)
+{
+	const struct dns_srvinfo *da, *db;
+	unsigned int r, l;
+
+	da = *(struct dns_srvinfo * const *)a;
+	db = *(struct dns_srvinfo * const *)b;
+
+	l = da->priority;
+	r = db->priority;
+
+	return ((l > r) - (l < r));
+}
+
+static int
+srv_final_cmp(const void *a, const void *b)
+{
+	const struct dns_srvinfo *da, *db;
+	unsigned int r, l, wr, wl;
+	int res;
+
+	da = *(struct dns_srvinfo * const *)a;
+	db = *(struct dns_srvinfo * const *)b;
+
+	l = da->priority;
+	r = db->priority;
+
+	res = ((l > r) - (l < r));
+
+	if (res == 0) {
+		wl = da->finalweight;
+		wr = db->finalweight;
+		res = ((wr > wl) - (wr < wl));
+	}
+
+	return (res);
+}
+
+static void
+compute_weight(struct dns_srvinfo **d, int first, int last)
+{
+	int i, j, totalweight;
+	int *chosen;
+
+	chosen = malloc(sizeof(int) * (last - first + 1));
+	totalweight = 0;
+	
+	for (i = 0; i <= last; i++)
+		totalweight += d[i]->weight;
+
+	if (totalweight == 0)
+		return;
+
+	for (i = 0; i <= last; i++) {
+		for (;;) {
+			chosen[i] = random() % (d[i]->weight * 100 / totalweight);
+			for (j = 0; j < i; j++) {
+				if (chosen[i] == chosen[j])
+					break;
+			}
+			if (j == i) {
+				d[i]->finalweight = chosen[i];
+				break;
+			}
+		}
+	}
+
+	free(chosen);
+}
+
 struct dns_srvinfo *
 dns_getsrvinfo(const char *zone)
 {
@@ -46,7 +117,7 @@ dns_getsrvinfo(const char *zone)
 	unsigned char *end, *p;
 	char host[MAXHOSTNAMELEN];
 	dns_query q;
-	int len, qdcount, ancount, n, i;
+	int len, qdcount, ancount, n, i, f, l;
 	unsigned int type, class, ttl, priority, weight, port;
 
 	if ((len = res_query(zone, C_IN, T_SRV, q.buf, sizeof(q.buf))) == -1 ||
@@ -125,6 +196,21 @@ dns_getsrvinfo(const char *zone)
 		n++;
 	}
 
+	qsort(res, n, sizeof(res[0]), srv_priority_cmp);
+
+	priority = f = l = 0;
+	for (i = 0; i < n; i++) {
+		if (res[i]->priority != priority) {
+			if (f != l)
+				compute_weight(res, f, l);
+			f = i;
+			priority = res[i]->priority;
+		}
+		l = i;
+	}
+
+	qsort(res, n, sizeof(res[0]), srv_final_cmp);
+
 	for (i = 0; i < n - 1; i++)
 		res[i]->next = res[i + 1];
 
diff --git a/usr.sbin/pkg/dns_utils.h b/usr.sbin/pkg/dns_utils.h
index 0f3367b..b1418bb 100644
--- a/usr.sbin/pkg/dns_utils.h
+++ b/usr.sbin/pkg/dns_utils.h
@@ -35,6 +35,7 @@ struct dns_srvinfo {
 	unsigned int priority;
 	unsigned int weight;
 	unsigned int port;
+	unsigned int finalweight;
 	char host[MAXHOSTNAMELEN];
 	struct dns_srvinfo *next;
 };
diff --git a/usr.sbin/pkg/pkg.c b/usr.sbin/pkg/pkg.c
index 0aef3da..34ab26f 100644
--- a/usr.sbin/pkg/pkg.c
+++ b/usr.sbin/pkg/pkg.c
@@ -34,6 +34,7 @@ __FBSDID("$FreeBSD$");
 #include <archive_entry.h>
 #include <err.h>
 #include <errno.h>
+#include <fcntl.h>
 #include <fetch.h>
 #include <paths.h>
 #include <stdbool.h>
@@ -191,8 +192,10 @@ bootstrap_pkg(void)
 			}
 		}
 
-		if (mirrors != NULL)
+		if (mirrors != NULL) {
 			strlcpy(u->host, current->host, sizeof(u->host));
+			u->port = current->port;
+		}
 
 		remote = fetchXGet(u, &st, "");
 		if (remote == NULL) {
@@ -295,7 +298,9 @@ int
 main(__unused int argc, char *argv[])
 {
 	char pkgpath[MAXPATHLEN];
+	char pkgstatic[MAXPATHLEN];
 	bool yes = false;
+	int fd, ret;
 
 	snprintf(pkgpath, MAXPATHLEN, "%s/sbin/pkg",
 	    getenv("LOCALBASE") ? getenv("LOCALBASE") : _LOCALBASE);
@@ -309,6 +314,19 @@ main(__unused int argc, char *argv[])
 		if (argv[1] != NULL && strcmp(argv[1], "-N") == 0)
 			errx(EXIT_FAILURE, "pkg is not installed");
 
+		if (argc > 2 && strcmp(argv[1], "add") == 0 &&
+		    access(argv[2], R_OK) == 0) {
+			fd = open(argv[2], O_RDONLY);
+			if (fd == -1)
+				err(EXIT_FAILURE, "Unable to open %s", argv[2]);
+
+			if ((ret = extract_pkg_static(fd, pkgstatic, MAXPATHLEN)) == 0)
+				ret = install_pkg_static(pkgstatic, argv[2]);
+			close(fd);
+			if (ret != 0)
+				exit(EXIT_FAILURE);
+			exit(EXIT_SUCCESS);
+		}
 		/*
 		 * Do not ask for confirmation if either of stdin or stdout is
 		 * not tty. Check the environment to see if user has answer