thodg/got

diff --git a/got/got.1 b/got/got.1
index 21ccb8a..869f13d 100644
--- a/got/got.1
+++ b/got/got.1
@@ -136,7 +136,7 @@ follows the globbing rules documented in
 .It Cm im
 Short alias for
 .Cm import .
-.It Cm clone Oo Fl a Oc Oo Fl q Oc Oo Fl v Oc Ar repository-URL Op Ar directory
+.It Cm clone Oo Fl a Oc Oo Fl b Ar branch Oc Oo Fl q Oc Oo Fl v Oc Ar repository-URL Op Ar directory
 Clone a Git repository at the specified
 .Ar repository-URL
 into the specified
@@ -203,8 +203,25 @@ are as follows:
 .Bl -tag -width Ds
 .It Fl a
 Fetch all branches from the remote repository.
-If this option is not specified, a branch resolved via the repository's HEAD
-reference will be fetched.
+If this option is not specified, a branch resolved via the remote
+repository's HEAD reference will be fetched.
+Cannot be used together with the
+.Fl b
+option.
+.It Fl b Ar branch
+Fetch the specified
+.Ar branch
+from the remote repository.
+This option may be specified multiple times to build a list of branches
+to fetch.
+If the branch corresponding to the remote repository's HEAD reference is not
+in this list, the cloned repository's HEAD reference will be set to the first
+branch which was fetched.
+If this option is not specified, a branch resolved via the remote
+repository's HEAD reference will be fetched.
+Cannot be used together with the
+.Fl a
+option.
 .It Fl m
 Create the cloned repository as a mirror of the original repository.
 This is useful if the cloned repository will not be used to store
@@ -250,11 +267,11 @@ The maximum is 3.
 .It Cm cl
 Short alias for
 .Cm clone .
-.It Cm fetch Oo Fl a Oc Oo Fl r Ar repository-path Oc Oo Fl q Oc Oo Fl v Oc Op Ar remote-repository-name
+.It Cm fetch Oo Fl a Oc Oo Fl b Ar branch Oc Oo Fl r Ar repository-path Oc Oo Fl q Oc Oo Fl v Oc Op Ar remote-repository
 Fetch new changes from a remote repository.
 If no
-.Ar remote-repository-name
-is specified the name
+.Ar remote-repository
+is specified,
 .Dq origin
 will be used.
 The remote repository's URL is obtained from the corresponding entry in the
@@ -295,8 +312,22 @@ are as follows:
 .Bl -tag -width Ds
 .It Fl a
 Fetch all branches from the remote repository.
-If this option is not specified, a branch resolved via the repository's HEAD
-reference will be fetched.
+If this option is not specified, a branch resolved via the remote
+repository's HEAD reference will be fetched.
+Cannot be used together with the
+.Fl b
+option.
+.It Fl b Ar branch
+Fetch the specified
+.Ar branch
+from the remote repository.
+This option may be specified multiple times to build a list of branches
+to fetch.
+If this option is not specified, a branch resolved via the remote
+repository's HEAD reference will be fetched.
+Cannot be used together with the
+.Fl a
+option.
 .It Fl r Ar repository-path
 Use the repository at the specified path.
 If not specified, assume the repository is located at or above the current
diff --git a/got/got.c b/got/got.c
index c73a14b..4ba89fc 100644
--- a/got/got.c
+++ b/got/got.c
@@ -811,8 +811,8 @@ done:
 __dead static void
 usage_clone(void)
 {
-	fprintf(stderr, "usage: %s clone [-a] [-m] [-q] [-v] repository-url "
-	    "[directory]\n", getprogname());
+	fprintf(stderr, "usage: %s clone [-a] [-b branch] [-m] [-q] [-v] "
+	    "repository-url [directory]\n", getprogname());
 	exit(1);
 }
 
@@ -885,6 +885,21 @@ fetch_progress(void *arg, const char *message, off_t packfile_size,
 }
 
 static const struct got_error *
+create_head_ref(struct got_reference *target_ref, struct got_repository *repo)
+{
+	const struct got_error *err;
+	struct got_reference *head_symref;
+
+	err = got_ref_alloc_symref(&head_symref, GOT_REF_HEAD, target_ref);
+	if (err)
+		return err;
+
+	err = got_ref_write(head_symref, repo);
+	got_ref_close(head_symref);
+	return err;
+}
+
+static const struct got_error *
 cmd_clone(int argc, char *argv[])
 {
 	const struct got_error *error = NULL;
@@ -893,7 +908,7 @@ cmd_clone(int argc, char *argv[])
 	char *default_destdir = NULL, *id_str = NULL;
 	const char *repo_path;
 	struct got_repository *repo = NULL;
-	struct got_pathlist_head refs, symrefs;
+	struct got_pathlist_head refs, symrefs, wanted_branches;
 	struct got_pathlist_entry *pe;
 	struct got_object_id *pack_hash = NULL;
 	int ch, fetchfd = -1;
@@ -908,12 +923,19 @@ cmd_clone(int argc, char *argv[])
 
 	TAILQ_INIT(&refs);
 	TAILQ_INIT(&symrefs);
+	TAILQ_INIT(&wanted_branches);
 
-	while ((ch = getopt(argc, argv, "amvq")) != -1) {
+	while ((ch = getopt(argc, argv, "ab:mvq")) != -1) {
 		switch (ch) {
 		case 'a':
 			fetch_all_branches = 1;
 			break;
+		case 'b':
+			error = got_pathlist_append(&wanted_branches,
+			    optarg, NULL);
+			if (error)
+				return error;
+			break;
 		case 'm':
 			mirror_references = 1;
 			break;
@@ -934,6 +956,9 @@ cmd_clone(int argc, char *argv[])
 	argc -= optind;
 	argv += optind;
 
+	if (fetch_all_branches && !TAILQ_EMPTY(&wanted_branches))
+		errx(1, "-a and -b options are mutually exclusive\n");
+
 	uri = argv[0];
 
 	if (argc == 1)
@@ -1023,7 +1048,8 @@ cmd_clone(int argc, char *argv[])
 	fpa.verbosity = verbosity;
 	error = got_fetch_pack(&pack_hash, &refs, &symrefs,
 	    GOT_FETCH_DEFAULT_REMOTE_NAME, mirror_references,
-	    fetch_all_branches, fetchfd, repo, fetch_progress, &fpa);
+	    fetch_all_branches, &wanted_branches, fetchfd,
+	    repo, fetch_progress, &fpa);
 	if (error)
 		goto done;
 
@@ -1088,19 +1114,41 @@ cmd_clone(int argc, char *argv[])
 			goto done;
 		}
 
-		error = got_ref_alloc_symref(&head_symref,
-		    GOT_REF_HEAD, target_ref);
+		if (verbosity >= 0)
+			printf("Setting %s to %s\n", refname, target);
+		error = create_head_ref(target_ref, repo);
 		got_ref_close(target_ref);
 		if (error)
 			goto done;
+	}
+	if (pe == NULL) {
+		/*
+		 * We failed to set the HEAD reference. If we asked for
+		 * a set of wanted branches use the first of one of those
+		 * which could be fetched instead.
+		 */
+		 TAILQ_FOREACH(pe, &wanted_branches, entry) {
+			const char *target = pe->path;
+			struct got_reference *target_ref;
 
-		if (verbosity >= 0)
-			printf("Setting %s to %s\n", GOT_REF_HEAD,
-			    got_ref_get_symref_target(head_symref));
+			error = got_ref_open(&target_ref, repo, target, 0);
+			if (error) {
+				if (error->code == GOT_ERR_NOT_REF) {
+					error = NULL;
+					continue;
+				}
+				goto done;
+			}
 
-		error = got_ref_write(head_symref, repo);
-		if (error)
-			goto done;
+			if (verbosity >= 0)
+				printf("Setting %s to %s\n", GOT_REF_HEAD,
+				    got_ref_get_name(target_ref));
+			error = create_head_ref(target_ref, repo);
+			got_ref_close(target_ref);
+			if (error)
+				goto done;
+			break;
+		}
 	}
 
 	/* Create a config file git-fetch(1) can understand. */
@@ -1187,6 +1235,7 @@ done:
 		free(pe->data);
 	}
 	got_pathlist_free(&symrefs);
+	got_pathlist_free(&wanted_branches);
 	free(pack_hash);
 	free(proto);
 	free(host);
@@ -1268,8 +1317,8 @@ done:
 __dead static void
 usage_fetch(void)
 {
-	fprintf(stderr, "usage: %s fetch [-a] [-r repository-path] [-q] [-v] "
-	    "[remote-repository-name]\n", getprogname());
+	fprintf(stderr, "usage: %s fetch [-a] [-b branch] [-r repository-path] "
+	    "[-q] [-v] [remote-repository-name]\n", getprogname());
 	exit(1);
 }
 
@@ -1286,7 +1335,7 @@ cmd_fetch(int argc, char *argv[])
 	char *id_str = NULL;
 	struct got_repository *repo = NULL;
 	struct got_worktree *worktree = NULL;
-	struct got_pathlist_head refs, symrefs;
+	struct got_pathlist_head refs, symrefs, wanted_branches;
 	struct got_pathlist_entry *pe;
 	struct got_object_id *pack_hash = NULL;
 	int i, ch, fetchfd = -1;
@@ -1295,12 +1344,19 @@ cmd_fetch(int argc, char *argv[])
 
 	TAILQ_INIT(&refs);
 	TAILQ_INIT(&symrefs);
+	TAILQ_INIT(&wanted_branches);
 
-	while ((ch = getopt(argc, argv, "ar:vq")) != -1) {
+	while ((ch = getopt(argc, argv, "ab:r:vq")) != -1) {
 		switch (ch) {
 		case 'a':
 			fetch_all_branches = 1;
 			break;
+		case 'b':
+			error = got_pathlist_append(&wanted_branches,
+			    optarg, NULL);
+			if (error)
+				return error;
+			break;
 		case 'r':
 			repo_path = realpath(optarg, NULL);
 			if (repo_path == NULL)
@@ -1325,6 +1381,9 @@ cmd_fetch(int argc, char *argv[])
 	argc -= optind;
 	argv += optind;
 
+	if (fetch_all_branches && !TAILQ_EMPTY(&wanted_branches))
+		errx(1, "-a and -b options are mutually exclusive\n");
+
 	if (argc == 0)
 		remote_name = GOT_FETCH_DEFAULT_REMOTE_NAME;
 	else if (argc == 1)
@@ -1427,8 +1486,8 @@ cmd_fetch(int argc, char *argv[])
 	fpa.last_p_resolved = -1;
 	fpa.verbosity = verbosity;
 	error = got_fetch_pack(&pack_hash, &refs, &symrefs, remote->name,
-	    remote->mirror_references, fetch_all_branches, fetchfd, repo,
-	    fetch_progress, &fpa);
+	    remote->mirror_references, fetch_all_branches, &wanted_branches,
+	    fetchfd, repo, fetch_progress, &fpa);
 	if (error)
 		goto done;
 
@@ -1526,6 +1585,7 @@ done:
 		free(pe->data);
 	}
 	got_pathlist_free(&symrefs);
+	got_pathlist_free(&wanted_branches);
 	free(id_str);
 	free(cwd);
 	free(repo_path);
diff --git a/include/got_fetch.h b/include/got_fetch.h
index 14f29c4..a7b9590 100644
--- a/include/got_fetch.h
+++ b/include/got_fetch.h
@@ -61,4 +61,5 @@ typedef const struct got_error *(*got_fetch_progress_cb)(void *,
  */
 const struct got_error *got_fetch_pack(struct got_object_id **,
 	struct got_pathlist_head *, struct got_pathlist_head *, const char *,
-	int, int, int, struct got_repository *, got_fetch_progress_cb, void *);
+	int, int, struct got_pathlist_head *, int, struct got_repository *,
+	got_fetch_progress_cb, void *);
diff --git a/lib/fetch.c b/lib/fetch.c
index e36d5b8..d9e55b8 100644
--- a/lib/fetch.c
+++ b/lib/fetch.c
@@ -387,7 +387,8 @@ check_pack_hash(int fd, size_t sz, uint8_t *hcomp)
 const struct got_error*
 got_fetch_pack(struct got_object_id **pack_hash, struct got_pathlist_head *refs,
     struct got_pathlist_head *symrefs, const char *remote_name,
-    int mirror_references, int fetch_all_branches, int fetchfd,
+    int mirror_references, int fetch_all_branches,
+    struct got_pathlist_head *wanted_branches, int fetchfd,
     struct got_repository *repo,
     got_fetch_progress_cb progress_cb, void *progress_arg)
 {
@@ -554,7 +555,7 @@ got_fetch_pack(struct got_object_id **pack_hash, struct got_pathlist_head *refs,
 		goto done;
 	}
 	err = got_privsep_send_fetch_req(&fetchibuf, nfetchfd, &have_refs,
-	    fetch_all_branches);
+	    fetch_all_branches, wanted_branches);
 	if (err != NULL)
 		goto done;
 	nfetchfd = -1;
diff --git a/lib/got_lib_privsep.h b/lib/got_lib_privsep.h
index 3df98dc..ca92bcd 100644
--- a/lib/got_lib_privsep.h
+++ b/lib/got_lib_privsep.h
@@ -110,6 +110,8 @@ enum got_imsg_type {
 
 	/* Messages related to networking. */
 	GOT_IMSG_FETCH_REQUEST,
+	GOT_IMSG_FETCH_HAVE_REF,
+	GOT_IMSG_FETCH_WANTED_BRANCH,
 	GOT_IMSG_FETCH_OUTFD,
 	GOT_IMSG_FETCH_SYMREFS,
 	GOT_IMSG_FETCH_REF,
@@ -239,17 +241,26 @@ struct got_imsg_tag_object {
 	 */
 } __attribute__((__packed__));
 
-/* Structures for GOT_IMSG_FETCH_REQUEST data. */
+/* Structure for GOT_IMSG_FETCH_HAVE_REF data. */
 struct got_imsg_fetch_have_ref {
 	uint8_t id[SHA1_DIGEST_LENGTH];
 	size_t name_len;
 	/* Followed by name_len data bytes. */
 } __attribute__((__packed__));
 
+/* Structure for GOT_IMSG_FETCH_WANTED_BRANCH data. */
+struct got_imsg_fetch_wanted_branch {
+	size_t name_len;
+	/* Followed by name_len data bytes. */
+} __attribute__((__packed__));
+
+/* Structure for GOT_IMSG_FETCH_REQUEST data. */
 struct got_imsg_fetch_request {
 	int fetch_all_branches;
 	size_t n_have_refs;
-	/* Followed by n_have_refs times of got_imsg_fetch_have_ref data. */
+	size_t n_wanted_branches;
+	/* Followed by n_have_refs GOT_IMSG_FETCH_HAVE_REF messages. */
+	/* Followed by n_wanted_branches times GOT_IMSG_FETCH_WANTED_BRANCH. */
 } __attribute__((__packed__));
 
 /* Structures for GOT_IMSG_FETCH_SYMREFS data. */
@@ -345,7 +356,7 @@ struct got_imsg_remote {
 	int mirror_references;
 
 	/* Followed by name_len + url_len data bytes. */
-};
+} __attribute__((__packed__));
 
 /*
  * Structure for GOT_IMSG_GITCONFIG_REMOTES data.
@@ -390,7 +401,7 @@ const struct got_error *got_privsep_send_index_pack_done(struct imsgbuf *);
 const struct got_error *got_privsep_recv_index_progress(int *, int *, int *,
     int *, int *, struct imsgbuf *ibuf);
 const struct got_error *got_privsep_send_fetch_req(struct imsgbuf *, int,
-    struct got_pathlist_head *, int);
+    struct got_pathlist_head *, int, struct got_pathlist_head *);
 const struct got_error *got_privsep_send_fetch_outfd(struct imsgbuf *, int);
 const struct got_error *got_privsep_send_fetch_symrefs(struct imsgbuf *,
     struct got_pathlist_head *);
diff --git a/lib/privsep.c b/lib/privsep.c
index c03ab3d..805dd3e 100644
--- a/lib/privsep.c
+++ b/lib/privsep.c
@@ -412,73 +412,107 @@ got_privsep_send_obj(struct imsgbuf *ibuf, struct got_object *obj)
 
 const struct got_error *
 got_privsep_send_fetch_req(struct imsgbuf *ibuf, int fd,
-   struct got_pathlist_head *have_refs, int fetch_all_branches)
+   struct got_pathlist_head *have_refs, int fetch_all_branches,
+   struct got_pathlist_head *wanted_branches)
 {
 	const struct got_error *err = NULL;
 	struct ibuf *wbuf;
-	size_t len, n_have_refs = 0;
+	size_t len;
 	struct got_pathlist_entry *pe;
-
+	struct got_imsg_fetch_request fetchreq;
+
+	memset(&fetchreq, 0, sizeof(fetchreq));
+	fetchreq.fetch_all_branches = fetch_all_branches;
+	TAILQ_FOREACH(pe, have_refs, entry)
+		fetchreq.n_have_refs++;
+	TAILQ_FOREACH(pe, wanted_branches, entry)
+		fetchreq.n_wanted_branches++;
 	len = sizeof(struct got_imsg_fetch_request);
-	TAILQ_FOREACH(pe, have_refs, entry) {
-		len += sizeof(struct got_imsg_fetch_have_ref) + pe->path_len;
-		n_have_refs++;
-	}
 	if (len >= MAX_IMSGSIZE - IMSG_HEADER_SIZE) {
 		close(fd);
 		return got_error(GOT_ERR_NO_SPACE);
 	}
 
-	wbuf = imsg_create(ibuf, GOT_IMSG_FETCH_REQUEST, 0, 0, len);
-	if (wbuf == NULL) {
-		close(fd);
-		return got_error_from_errno("imsg_create FETCH_REQUEST");
-	}
+	if (imsg_compose(ibuf, GOT_IMSG_FETCH_REQUEST, 0, 0, fd,
+	    &fetchreq, sizeof(fetchreq)) == -1)
+		return got_error_from_errno(
+		    "imsg_compose FETCH_SERVER_PROGRESS");
 
-	/* Keep in sync with struct got_imsg_fetch_request definition! */
-	if (imsg_add(wbuf, &fetch_all_branches, sizeof(fetch_all_branches))
-	    == -1) {
-		err = got_error_from_errno("imsg_add FETCH_REQUEST");
-		ibuf_free(wbuf);
-		close(fd);
-		return err;
-	}
-	if (imsg_add(wbuf, &n_have_refs, sizeof(n_have_refs)) == -1) {
-		err = got_error_from_errno("imsg_add FETCH_REQUEST");
-		ibuf_free(wbuf);
+	err = flush_imsg(ibuf);
+	if (err) {
 		close(fd);
 		return err;
 	}
+	fd = -1;
 
 	TAILQ_FOREACH(pe, have_refs, entry) {
 		const char *name = pe->path;
 		size_t name_len = pe->path_len;
 		struct got_object_id *id = pe->data;
 
+		len = sizeof(struct got_imsg_fetch_have_ref) + name_len;
+		wbuf = imsg_create(ibuf, GOT_IMSG_FETCH_HAVE_REF, 0, 0, len);
+		if (wbuf == NULL)
+			return got_error_from_errno("imsg_create FETCH_HAVE_REF");
+
 		/* Keep in sync with struct got_imsg_fetch_have_ref! */
 		if (imsg_add(wbuf, id->sha1, sizeof(id->sha1)) == -1) {
-			err = got_error_from_errno("imsg_add FETCH_REQUEST");
+			err = got_error_from_errno("imsg_add FETCH_HAVE_REF");
 			ibuf_free(wbuf);
-			close(fd);
 			return err;
 		}
 		if (imsg_add(wbuf, &name_len, sizeof(name_len)) == -1) {
-			err = got_error_from_errno("imsg_add FETCH_REQUEST");
+			err = got_error_from_errno("imsg_add FETCH_HAVE_REF");
 			ibuf_free(wbuf);
-			close(fd);
 			return err;
 		}
 		if (imsg_add(wbuf, name, name_len) == -1) {
-			err = got_error_from_errno("imsg_add FETCH_REQUEST");
+			err = got_error_from_errno("imsg_add FETCH_HAVE_REF");
 			ibuf_free(wbuf);
-			close(fd);
 			return err;
 		}
+
+		wbuf->fd = -1;
+		imsg_close(ibuf, wbuf);
+		err = flush_imsg(ibuf);
+		if (err)
+			return err;
 	}
 
-	wbuf->fd = fd;
-	imsg_close(ibuf, wbuf);
-	return flush_imsg(ibuf);
+	TAILQ_FOREACH(pe, wanted_branches, entry) {
+		const char *name = pe->path;
+		size_t name_len = pe->path_len;
+
+		len = sizeof(struct got_imsg_fetch_wanted_branch) + name_len;
+		wbuf = imsg_create(ibuf, GOT_IMSG_FETCH_WANTED_BRANCH, 0, 0,
+		    len);
+		if (wbuf == NULL)
+			return got_error_from_errno(
+			     "imsg_create FETCH_WANTED_BRANCH");
+
+		/* Keep in sync with struct got_imsg_fetch_wanted_branch! */
+		if (imsg_add(wbuf, &name_len, sizeof(name_len)) == -1) {
+			err = got_error_from_errno(
+			    "imsg_add FETCH_WANTED_BRANCH");
+			ibuf_free(wbuf);
+			return err;
+		}
+		if (imsg_add(wbuf, name, name_len) == -1) {
+			err = got_error_from_errno(
+			     "imsg_add FETCH_WANTED_BRANCH");
+			ibuf_free(wbuf);
+			return err;
+		}
+
+		wbuf->fd = -1;
+		imsg_close(ibuf, wbuf);
+		err = flush_imsg(ibuf);
+		if (err)
+			return err;
+	}
+
+	return NULL;
+
 }
 
 const struct got_error *
diff --git a/libexec/got-fetch-pack/got-fetch-pack.c b/libexec/got-fetch-pack/got-fetch-pack.c
index 35e6b77..cb84af6 100644
--- a/libexec/got-fetch-pack/got-fetch-pack.c
+++ b/libexec/got-fetch-pack/got-fetch-pack.c
@@ -483,7 +483,7 @@ fetch_error(const char *buf, size_t len)
 static const struct got_error *
 fetch_pack(int fd, int packfd, struct got_object_id *packid,
     struct got_pathlist_head *have_refs, int fetch_all_branches,
-    struct imsgbuf *ibuf)
+    struct got_pathlist_head *wanted_branches, struct imsgbuf *ibuf)
 {
 	const struct got_error *err = NULL;
 	char buf[GOT_FETCH_PKTMAX];
@@ -561,10 +561,21 @@ fetch_pack(int fd, int packfd, struct got_object_id *packid,
 			    getprogname(), refname);
 
 		if (strncmp(refname, "refs/heads/", 11) == 0) {
-			if (default_branch != NULL &&
-			    !match_branch(refname, default_branch))
-				continue;
-			found_branch = 1;
+			if (fetch_all_branches) {
+				found_branch = 1;
+			} else if (!TAILQ_EMPTY(wanted_branches)) {
+				TAILQ_FOREACH(pe, wanted_branches, entry) {
+					if (match_branch(refname, pe->path))
+						break;
+				}
+				if (pe == NULL)
+					continue;
+				found_branch = 1;
+			} else if (default_branch != NULL) {
+				if (!match_branch(refname, default_branch))
+					continue;
+				found_branch = 1;
+			}
 		} else if (strncmp(refname, "refs/tags/", 10) != 0) {
 			if (chattygot) {
 				fprintf(stderr, "%s: ignoring '%s' which is "
@@ -837,11 +848,12 @@ main(int argc, char **argv)
 	struct imsgbuf ibuf;
 	struct imsg imsg;
 	struct got_pathlist_head have_refs;
+	struct got_pathlist_head wanted_branches;
 	struct got_pathlist_entry *pe;
-	struct got_imsg_fetch_request *fetch_req = NULL;
+	struct got_imsg_fetch_request fetch_req;
 	struct got_imsg_fetch_have_ref href;
-	size_t datalen, remain;
-	size_t offset;
+	struct got_imsg_fetch_wanted_branch wbranch;
+	size_t datalen;
 #if 0
 	static int attached;
 	while (!attached)
@@ -849,6 +861,7 @@ main(int argc, char **argv)
 #endif
 
 	TAILQ_INIT(&have_refs);
+	TAILQ_INIT(&wanted_branches);
 
 	if (getenv("GOT_FETCH_DEBUG") != NULL) {
 		fprintf(stderr, "%s being chatty!\n", getprogname());
@@ -876,43 +889,46 @@ main(int argc, char **argv)
 		goto done;
 	}
 	datalen = imsg.hdr.len - IMSG_HEADER_SIZE;
-	if (datalen < sizeof(struct got_imsg_fetch_request)) {
-		err = got_error(GOT_ERR_PRIVSEP_LEN);
-		goto done;
-	}
-	fetch_req = (struct got_imsg_fetch_request *)imsg.data;
-	if (datalen < sizeof(*fetch_req) +
-	    sizeof(struct got_imsg_fetch_have_ref) *
-	    fetch_req->n_have_refs) {
+	if (datalen < sizeof(fetch_req)) {
 		err = got_error(GOT_ERR_PRIVSEP_LEN);
 		goto done;
 	}
-	offset = sizeof(*fetch_req);
-	remain = datalen;
-	for (i = 0; i < fetch_req->n_have_refs; i++) {
+	memcpy(&fetch_req, imsg.data, sizeof(fetch_req));
+	fetchfd = imsg.fd;
+	imsg_free(&imsg);
+
+	for (i = 0; i < fetch_req.n_have_refs; i++) {
 		struct got_object_id *id;
 		char *refname;
 
-		if (remain < sizeof(href) || offset > datalen) {
+		if ((err = got_privsep_recv_imsg(&imsg, &ibuf, 0)) != 0) {
+			if (err->code == GOT_ERR_PRIVSEP_PIPE)
+				err = NULL;
+			goto done;
+		}
+		if (imsg.hdr.type == GOT_IMSG_STOP)
+			goto done;
+		if (imsg.hdr.type != GOT_IMSG_FETCH_HAVE_REF) {
+			err = got_error(GOT_ERR_PRIVSEP_MSG);
+			goto done;
+		}
+		datalen = imsg.hdr.len - IMSG_HEADER_SIZE;
+		if (datalen < sizeof(href)) {
 			err = got_error(GOT_ERR_PRIVSEP_LEN);
 			goto done;
 		}
-		memcpy(&href, imsg.data + offset, sizeof(href));
-		remain -= sizeof(href);
-		if (remain < href.name_len) {
+		memcpy(&href, imsg.data, sizeof(href));
+		if (datalen - sizeof(href) < href.name_len) {
 			err = got_error(GOT_ERR_PRIVSEP_LEN);
 			goto done;
 		}
-		remain -= href.name_len;
 		refname = malloc(href.name_len + 1);
 		if (refname == NULL) {
 			err = got_error_from_errno("malloc");
 			goto done;
 		}
-		offset += sizeof(href);
-		memcpy(refname, imsg.data + offset, href.name_len);
+		memcpy(refname, imsg.data + sizeof(href), href.name_len);
 		refname[href.name_len] = '\0';
-		offset += href.name_len;
 
 		id = malloc(sizeof(*id));
 		if (id == NULL) {
@@ -927,8 +943,50 @@ main(int argc, char **argv)
 			free(id);
 			goto done;
 		}
+
+		imsg_free(&imsg);
+	}
+
+	for (i = 0; i < fetch_req.n_wanted_branches; i++) {
+		char *refname;
+
+		if ((err = got_privsep_recv_imsg(&imsg, &ibuf, 0)) != 0) {
+			if (err->code == GOT_ERR_PRIVSEP_PIPE)
+				err = NULL;
+			goto done;
+		}
+		if (imsg.hdr.type == GOT_IMSG_STOP)
+			goto done;
+		if (imsg.hdr.type != GOT_IMSG_FETCH_WANTED_BRANCH) {
+			err = got_error(GOT_ERR_PRIVSEP_MSG);
+			goto done;
+		}
+		datalen = imsg.hdr.len - IMSG_HEADER_SIZE;
+		if (datalen < sizeof(wbranch)) {
+			err = got_error(GOT_ERR_PRIVSEP_LEN);
+			goto done;
+		}
+		memcpy(&wbranch, imsg.data, sizeof(wbranch));
+		if (datalen - sizeof(wbranch) < wbranch.name_len) {
+			err = got_error(GOT_ERR_PRIVSEP_LEN);
+			goto done;
+		}
+		refname = malloc(wbranch.name_len + 1);
+		if (refname == NULL) {
+			err = got_error_from_errno("malloc");
+			goto done;
+		}
+		memcpy(refname, imsg.data + sizeof(wbranch), wbranch.name_len);
+		refname[wbranch.name_len] = '\0';
+
+		err = got_pathlist_append(&wanted_branches, refname, NULL);
+		if (err) {
+			free(refname);
+			goto done;
+		}
+
+		imsg_free(&imsg);
 	}
-	fetchfd = imsg.fd;
 
 	if ((err = got_privsep_recv_imsg(&imsg, &ibuf, 0)) != 0) {
 		if (err->code == GOT_ERR_PRIVSEP_PIPE)
@@ -948,13 +1006,16 @@ main(int argc, char **argv)
 	packfd = imsg.fd;
 
 	err = fetch_pack(fetchfd, packfd, &packid, &have_refs,
-	    fetch_req->fetch_all_branches, &ibuf);
+	    fetch_req.fetch_all_branches, &wanted_branches, &ibuf);
 done:
 	TAILQ_FOREACH(pe, &have_refs, entry) {
 		free((char *)pe->path);
 		free(pe->data);
 	}
 	got_pathlist_free(&have_refs);
+	TAILQ_FOREACH(pe, &wanted_branches, entry)
+		free((char *)pe->path);
+	got_pathlist_free(&wanted_branches);
 	if (packfd != -1 && close(packfd) == -1 && err == NULL)
 		err = got_error_from_errno("close");
 	if (err != NULL)