4 files changed, 222 insertions, 9 deletions

diff --git a/usr.sbin/pkg/Makefile b/usr.sbin/pkg/Makefile
index 360e301..c2ca0a3 100644
--- a/usr.sbin/pkg/Makefile
+++ b/usr.sbin/pkg/Makefile
@@ -1,7 +1,7 @@
 # $FreeBSD$
 
 PROG=	pkg
-SRC=	pkg.c
+SRCS=	pkg.c dns_utils.c
 
 NO_MAN=	yes
 DPADD=	${LIBARCHIVE} ${LIBELF} ${LIBFETCH}
diff --git a/usr.sbin/pkg/dns_utils.c b/usr.sbin/pkg/dns_utils.c
new file mode 100644
index 0000000..e88bf98
--- /dev/null
+++ b/usr.sbin/pkg/dns_utils.c
@@ -0,0 +1,136 @@
+/*-
+ * Copyright (c) 2012 Baptiste Daroussin <bapt@FreeBSD.org>
+ * All rights reserved.
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions
+ * are met:
+ * 1. Redistributions of source code must retain the above copyright
+ *    notice, this list of conditions and the following disclaimer,
+ *    without modification, immediately at the beginning of the file.
+ * 2. Redistributions in binary form must reproduce the above copyright
+ *    notice, this list of conditions and the following disclaimer in the
+ *    documentation and/or other materials provided with the distribution.
+ *
+ * THIS SOFTWARE IS PROVIDED BY THE AUTHOR ``AS IS'' AND ANY EXPRESS OR
+ * IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES
+ * OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE DISCLAIMED.
+ * IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR ANY DIRECT, INDIRECT,
+ * INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT
+ * NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
+ * DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY
+ * THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
+ * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF
+ * THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+ */
+
+#include <sys/cdefs.h>
+__FBSDID("$FreeBSD$");
+
+#include <stdlib.h>
+#include <string.h>
+#include <netinet/in.h>
+#include <resolv.h>
+
+#include "dns_utils.h"
+
+typedef union {
+	HEADER hdr;
+	unsigned char buf[1024];
+} dns_query;
+
+struct dns_srvinfo *
+dns_getsrvinfo(const char *zone)
+{
+	struct dns_srvinfo **res, *first;
+	unsigned char *end, *p;
+	char host[MAXHOSTNAMELEN];
+	dns_query q;
+	int len, qdcount, ancount, n, i;
+	unsigned int type, class, ttl, priority, weight, port;
+
+	if ((len = res_query(zone, C_IN, T_SRV, q.buf, sizeof(q.buf))) == -1 ||
+	    len < (int)sizeof(HEADER))
+		return (NULL);
+
+	qdcount = ntohs(q.hdr.qdcount);
+	ancount = ntohs(q.hdr.ancount);
+
+	end = q.buf + len;
+	p = q.buf + sizeof(HEADER);
+
+	while(qdcount > 0 && p < end) {
+		qdcount--;
+		if((len = dn_expand(q.buf, end, p, host, MAXHOSTNAMELEN)) < 0)
+			return (NULL);
+		p += len + NS_QFIXEDSZ;
+	}
+
+	res = malloc(sizeof(struct dns_srvinfo) * ancount);
+	if (res == NULL)
+		return (NULL);
+	memset(res, 0, sizeof(struct dns_srvinfo) * ancount);
+
+	n = 0;
+	while (ancount > 0 && p < end) {
+		ancount--;
+		len = dn_expand(q.buf, end, p, host, MAXHOSTNAMELEN);
+		if (len < 0) {
+			for (i = 0; i < n; i++)
+				free(res[i]);
+			free(res);
+			return NULL;
+		}
+
+		p += len;
+
+		NS_GET16(type, p);
+		NS_GET16(class, p);
+		NS_GET32(ttl, p);
+		NS_GET16(len, p);
+
+		if (type != T_SRV) {
+			p += len;
+			continue;
+		}
+
+		NS_GET16(priority, p);
+		NS_GET16(weight, p);
+		NS_GET16(port, p);
+
+		len = dn_expand(q.buf, end, p, host, MAXHOSTNAMELEN);
+		if (len < 0) {
+			for (i = 0; i < n; i++)
+				free(res[i]);
+			free(res);
+			return (NULL);
+		}
+
+		res[n] = malloc(sizeof(struct dns_srvinfo));
+		if (res[n] == NULL) {
+			for (i = 0; i < n; i++)
+				free(res[i]);
+			free(res);
+			return (NULL);
+		}
+		res[n]->type = type;
+		res[n]->class = class;
+		res[n]->ttl = ttl;
+		res[n]->priority = priority;
+		res[n]->weight = weight;
+		res[n]->port = port;
+		res[n]->next = NULL;
+		strlcpy(res[n]->host, host, MAXHOSTNAMELEN);
+
+		p += len;
+		n++;
+	}
+
+	for (i = 0; i < n - 1; i++)
+		res[i]->next = res[i + 1];
+
+	first = res[0];
+	free(res);
+
+	return (first);
+}
diff --git a/usr.sbin/pkg/dns_utils.h b/usr.sbin/pkg/dns_utils.h
new file mode 100644
index 0000000..0f3367b
--- /dev/null
+++ b/usr.sbin/pkg/dns_utils.h
@@ -0,0 +1,45 @@
+/*-
+ * Copyright (c) 2012 Baptiste Daroussin <bapt@FreeBSD.org>
+ * All rights reserved.
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions
+ * are met:
+ * 1. Redistributions of source code must retain the above copyright
+ *    notice, this list of conditions and the following disclaimer.
+ * 2. Redistributions in binary form must reproduce the above copyright
+ *    notice, this list of conditions and the following disclaimer in the
+ *    documentation and/or other materials provided with the distribution.
+ * 
+ * THIS SOFTWARE IS PROVIDED BY THE AUTHOR AND CONTRIBUTORS ``AS IS'' AND
+ * ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+ * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
+ * ARE DISCLAIMED.  IN NO EVENT SHALL THE AUTHOR OR CONTRIBUTORS BE LIABLE
+ * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
+ * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS
+ * OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION)
+ * HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT
+ * LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY
+ * OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF
+ * SUCH DAMAGE.
+ *
+ * $FreeBSD$
+ */
+
+#ifndef DNS_UTILS_H
+#define DNS_UTILS_H
+struct dns_srvinfo {
+	unsigned int type;
+	unsigned int class;
+	unsigned int ttl;
+	unsigned int priority;
+	unsigned int weight;
+	unsigned int port;
+	char host[MAXHOSTNAMELEN];
+	struct dns_srvinfo *next;
+};
+
+struct dns_srvinfo *
+	dns_getsrvinfo(const char *zone);
+
+#endif
diff --git a/usr.sbin/pkg/pkg.c b/usr.sbin/pkg/pkg.c
index 8dc7bcf..061008c 100644
--- a/usr.sbin/pkg/pkg.c
+++ b/usr.sbin/pkg/pkg.c
@@ -48,9 +48,10 @@ __FBSDID("$FreeBSD$");
 #include <unistd.h>
 
 #include "elf_tables.h"
+#include "dns_utils.h"
 
 #define _LOCALBASE "/usr/local"
-#define _PKGS_URL "http://pkgbeta.FreeBSD.org"
+#define _PKGS_URL "http://pkg.FreeBSD.org"
 
 static const char *
 elf_corres_to_string(struct _elf_corres *m, int e)
@@ -281,16 +282,20 @@ install_pkg_static(char *path, char *pkgpath)
 static int
 bootstrap_pkg(void)
 {
+	struct url *u;
 	FILE *remote;
 	FILE *config;
 	char *site;
+	struct dns_srvinfo *mirrors, *current;
+	/* To store _https._tcp. + hostname + \0 */
+	char zone[MAXHOSTNAMELEN + 13];
 	char url[MAXPATHLEN];
 	char conf[MAXPATHLEN];
 	char abi[BUFSIZ];
 	char tmppkg[MAXPATHLEN];
 	char buf[10240];
 	char pkgstatic[MAXPATHLEN];
-	int fd, retry, ret;
+	int fd, retry, ret, max_retry;
 	struct url_stat st;
 	off_t done, r;
 	time_t now;
@@ -298,9 +303,11 @@ bootstrap_pkg(void)
 
 	done = 0;
 	last = 0;
+	max_retry = 3;
 	ret = -1;
 	remote = NULL;
 	config = NULL;
+	current = mirrors = NULL;
 
 	printf("Bootstrapping pkg please wait\n");
 
@@ -324,12 +331,37 @@ bootstrap_pkg(void)
 		return (-1);
 	}
 
-	retry = 3;
-	do {
-		remote = fetchXGetURL(url, &st, "");
-		if (remote == NULL)
-			sleep(1);
-	} while (remote == NULL && retry-- > 0);
+	retry = max_retry;
+
+	u = fetchParseURL(url);
+	while (remote == NULL) {
+		if (retry == max_retry) {
+			if (strcmp(u->scheme, "file") != 0) {
+				snprintf(zone, sizeof(zone),
+				    "_%s._tcp.%s", u->scheme, u->host);
+				printf("%s\n", zone);
+				mirrors = dns_getsrvinfo(zone);
+				current = mirrors;
+			}
+		}
+
+		if (mirrors != NULL)
+			strlcpy(u->host, current->host, sizeof(u->host));
+
+		remote = fetchXGet(u, &st, "");
+		if (remote == NULL) {
+			--retry;
+			if (retry <= 0)
+				goto fetchfail;
+			if (mirrors == NULL) {
+				sleep(1);
+			} else {
+				current = current->next;
+				if (current == NULL)
+					current = mirrors;
+			}
+		}
+	}
 
 	if (remote == NULL)
 		goto fetchfail;