svn commit: r256968 - head/usr.sbin/pkg

[Baptiste Daroussin]

Author: bapt
Date: Wed Oct 23 14:06:07 2013
New Revision: 256968
URL: http://svnweb.freebsd.org/changeset/base/256968

Log:
  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
  
  Obtained from:	pkg git repo
  MFC after:	3 days

Modified:
  head/usr.sbin/pkg/dns_utils.c
  head/usr.sbin/pkg/dns_utils.h
  head/usr.sbin/pkg/pkg.c

Modified: head/usr.sbin/pkg/dns_utils.c
==============================================================================
--- head/usr.sbin/pkg/dns_utils.c	Wed Oct 23 14:04:09 2013	(r256967)
+++ head/usr.sbin/pkg/dns_utils.c	Wed Oct 23 14:06:07 2013	(r256968)
@@ -1,5 +1,5 @@
 /*-
- * Copyright (c) 2012 Baptiste Daroussin <bapt at FreeBSD.org>
+ * Copyright (c) 2012-2013 Baptiste Daroussin <bapt at 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)
+{
+	unsigned int r, l;
+	struct dns_srvinfo *da, *db;
+
+       	da = *(struct dns_srvinfo **)a;
+	db = *(struct dns_srvinfo **)b;
+
+	l = da->priority;
+	r = db->priority;
+
+	return ((l > r) - (l < r));
+}
+
+static int
+srv_final_cmp(const void *a, const void *b)
+{
+	unsigned int r, l, wr, wl;
+	int res;
+	struct dns_srvinfo *da, *db;
+
+       	da = *(struct dns_srvinfo **)a;
+	db = *(struct dns_srvinfo **)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];
 

Modified: head/usr.sbin/pkg/dns_utils.h
==============================================================================
--- head/usr.sbin/pkg/dns_utils.h	Wed Oct 23 14:04:09 2013	(r256967)
+++ head/usr.sbin/pkg/dns_utils.h	Wed Oct 23 14:06:07 2013	(r256968)
@@ -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;
 };

Modified: head/usr.sbin/pkg/pkg.c
==============================================================================
--- head/usr.sbin/pkg/pkg.c	Wed Oct 23 14:04:09 2013	(r256967)
+++ head/usr.sbin/pkg/pkg.c	Wed Oct 23 14:06:07 2013	(r256968)
@@ -191,8 +191,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) {