Support protocol v2 want-ref in UploadPack

UploadPack already allows the client to send wanted OIDs as "want"
lines. Extend UploadPack to also allow the client to send wanted ref
names as "want-ref" lines when the fetch is done using protocol v2.

The corresponding Git commit is 516e2b76bd ("upload-pack: implement
ref-in-want", 2018-06-28).

To support a two-stage rollout, two configuration variables are
provided: uploadpack.allowrefinwant (default "false") allows clients to
specify "want-ref" in their requests, and uploadpack.advertiserefinwant
(default "true") makes UploadPack advertise this capability. If
uploadpack.allowrefinwant is true but uploadpack.advertiserefinwant is
false, UploadPack will not advertise that it supports "want-ref", but it
will support it.

Change-Id: I3c24077949640d453af90d81a7f48ce4b8ac9833
Signed-off-by: Jonathan Tan <jonathantanmy@google.com>

4 files changed, 279 insertions, 0 deletions

diff --git a/org.eclipse.jgit.test/tst/org/eclipse/jgit/transport/UploadPackTest.java b/org.eclipse.jgit.test/tst/org/eclipse/jgit/transport/UploadPackTest.java
index 9f49073eae..7ef713bf45 100644
--- a/org.eclipse.jgit.test/tst/org/eclipse/jgit/transport/UploadPackTest.java
+++ b/org.eclipse.jgit.test/tst/org/eclipse/jgit/transport/UploadPackTest.java
@@ -1,11 +1,13 @@
 package org.eclipse.jgit.transport;
 
+import static org.hamcrest.Matchers.containsString;
 import static org.hamcrest.Matchers.hasItems;
 import static org.hamcrest.Matchers.is;
 import static org.hamcrest.Matchers.theInstance;
 import static org.junit.Assert.assertFalse;
 import static org.junit.Assert.assertThat;
 import static org.junit.Assert.assertTrue;
+import static org.junit.Assert.fail;
 
 import java.util.Arrays;
 import java.util.Collections;
@@ -505,6 +507,51 @@ public class UploadPackTest {
 	}
 
 	@Test
+	public void testV2CapabilitiesRefInWant() throws Exception {
+		server.getConfig().setBoolean("uploadpack", null, "allowrefinwant", true);
+		ByteArrayInputStream recvStream =
+			uploadPackV2Setup(null, null, PacketLineIn.END);
+		PacketLineIn pckIn = new PacketLineIn(recvStream);
+
+		assertThat(pckIn.readString(), is("version 2"));
+		assertThat(
+			Arrays.asList(pckIn.readString(), pckIn.readString()),
+			// TODO(jonathantanmy) This check overspecifies the
+			// order of the capabilities of "fetch".
+			hasItems("ls-refs", "fetch=ref-in-want shallow"));
+		assertTrue(pckIn.readString() == PacketLineIn.END);
+	}
+
+	@Test
+	public void testV2CapabilitiesRefInWantNotAdvertisedIfUnallowed() throws Exception {
+		server.getConfig().setBoolean("uploadpack", null, "allowrefinwant", false);
+		ByteArrayInputStream recvStream =
+			uploadPackV2Setup(null, null, PacketLineIn.END);
+		PacketLineIn pckIn = new PacketLineIn(recvStream);
+
+		assertThat(pckIn.readString(), is("version 2"));
+		assertThat(
+			Arrays.asList(pckIn.readString(), pckIn.readString()),
+			hasItems("ls-refs", "fetch=shallow"));
+		assertTrue(pckIn.readString() == PacketLineIn.END);
+	}
+
+	@Test
+	public void testV2CapabilitiesRefInWantNotAdvertisedIfAdvertisingForbidden() throws Exception {
+		server.getConfig().setBoolean("uploadpack", null, "allowrefinwant", true);
+		server.getConfig().setBoolean("uploadpack", null, "advertiserefinwant", false);
+		ByteArrayInputStream recvStream =
+			uploadPackV2Setup(null, null, PacketLineIn.END);
+		PacketLineIn pckIn = new PacketLineIn(recvStream);
+
+		assertThat(pckIn.readString(), is("version 2"));
+		assertThat(
+			Arrays.asList(pckIn.readString(), pckIn.readString()),
+			hasItems("ls-refs", "fetch=shallow"));
+		assertTrue(pckIn.readString() == PacketLineIn.END);
+	}
+
+	@Test
 	@SuppressWarnings("boxing")
 	public void testV2EmptyRequest() throws Exception {
 		ByteArrayInputStream recvStream = uploadPackV2(PacketLineIn.END);
@@ -1161,6 +1208,178 @@ public class UploadPackTest {
 			PacketLineIn.END);
 	}
 
+	@Test
+	public void testV2FetchWantRefIfNotAllowed() throws Exception {
+		RevCommit one = remote.commit().message("1").create();
+		remote.update("one", one);
+
+		try {
+			uploadPackV2(
+				"command=fetch\n",
+				PacketLineIn.DELIM,
+				"want-ref refs/heads/one\n",
+				"done\n",
+				PacketLineIn.END);
+		} catch (PackProtocolException e) {
+			assertThat(
+				e.getMessage(),
+				containsString("unexpected want-ref refs/heads/one"));
+			return;
+		}
+		fail("expected PackProtocolException");
+	}
+
+	@Test
+	public void testV2FetchWantRef() throws Exception {
+		RevCommit one = remote.commit().message("1").create();
+		RevCommit two = remote.commit().message("2").create();
+		RevCommit three = remote.commit().message("3").create();
+		remote.update("one", one);
+		remote.update("two", two);
+		remote.update("three", three);
+
+		server.getConfig().setBoolean("uploadpack", null, "allowrefinwant", true);
+
+		ByteArrayInputStream recvStream = uploadPackV2(
+			"command=fetch\n",
+			PacketLineIn.DELIM,
+			"want-ref refs/heads/one\n",
+			"want-ref refs/heads/two\n",
+			"done\n",
+			PacketLineIn.END);
+		PacketLineIn pckIn = new PacketLineIn(recvStream);
+		assertThat(pckIn.readString(), is("wanted-refs"));
+		assertThat(
+				Arrays.asList(pckIn.readString(), pckIn.readString()),
+				hasItems(
+					one.toObjectId().getName() + " refs/heads/one",
+					two.toObjectId().getName() + " refs/heads/two"));
+		assertThat(pckIn.readString(), theInstance(PacketLineIn.DELIM));
+		assertThat(pckIn.readString(), is("packfile"));
+		parsePack(recvStream);
+
+		assertTrue(client.hasObject(one.toObjectId()));
+		assertTrue(client.hasObject(two.toObjectId()));
+		assertFalse(client.hasObject(three.toObjectId()));
+	}
+
+	@Test
+	public void testV2FetchBadWantRef() throws Exception {
+		RevCommit one = remote.commit().message("1").create();
+		remote.update("one", one);
+
+		server.getConfig().setBoolean("uploadpack", null, "allowrefinwant", true);
+
+		try {
+			uploadPackV2(
+				"command=fetch\n",
+				PacketLineIn.DELIM,
+				"want-ref refs/heads/one\n",
+				"want-ref refs/heads/nonExistentRef\n",
+				"done\n",
+				PacketLineIn.END);
+		} catch (PackProtocolException e) {
+			assertThat(
+				e.getMessage(),
+				containsString("Invalid ref name: refs/heads/nonExistentRef"));
+			return;
+		}
+		fail("expected PackProtocolException");
+	}
+
+	@Test
+	public void testV2FetchMixedWantRef() throws Exception {
+		RevCommit one = remote.commit().message("1").create();
+		RevCommit two = remote.commit().message("2").create();
+		RevCommit three = remote.commit().message("3").create();
+		remote.update("one", one);
+		remote.update("two", two);
+		remote.update("three", three);
+
+		server.getConfig().setBoolean("uploadpack", null, "allowrefinwant", true);
+
+		ByteArrayInputStream recvStream = uploadPackV2(
+			"command=fetch\n",
+			PacketLineIn.DELIM,
+			"want-ref refs/heads/one\n",
+			"want " + two.toObjectId().getName() + "\n",
+			"done\n",
+			PacketLineIn.END);
+		PacketLineIn pckIn = new PacketLineIn(recvStream);
+		assertThat(pckIn.readString(), is("wanted-refs"));
+		assertThat(
+				pckIn.readString(),
+				is(one.toObjectId().getName() + " refs/heads/one"));
+		assertThat(pckIn.readString(), theInstance(PacketLineIn.DELIM));
+		assertThat(pckIn.readString(), is("packfile"));
+		parsePack(recvStream);
+
+		assertTrue(client.hasObject(one.toObjectId()));
+		assertTrue(client.hasObject(two.toObjectId()));
+		assertFalse(client.hasObject(three.toObjectId()));
+	}
+
+	@Test
+	public void testV2FetchWantRefWeAlreadyHave() throws Exception {
+		RevCommit one = remote.commit().message("1").create();
+		remote.update("one", one);
+
+		server.getConfig().setBoolean("uploadpack", null, "allowrefinwant", true);
+
+		ByteArrayInputStream recvStream = uploadPackV2(
+			"command=fetch\n",
+			PacketLineIn.DELIM,
+			"want-ref refs/heads/one\n",
+			"have " + one.toObjectId().getName(),
+			"done\n",
+			PacketLineIn.END);
+		PacketLineIn pckIn = new PacketLineIn(recvStream);
+
+		// The client still needs to know the hash of the object that
+		// refs/heads/one points to, even though it already has the
+		// object ...
+		assertThat(pckIn.readString(), is("wanted-refs"));
+		assertThat(
+				pckIn.readString(),
+				is(one.toObjectId().getName() + " refs/heads/one"));
+		assertThat(pckIn.readString(), theInstance(PacketLineIn.DELIM));
+
+		// ... but the client does not need the object itself.
+		assertThat(pckIn.readString(), is("packfile"));
+		parsePack(recvStream);
+		assertFalse(client.hasObject(one.toObjectId()));
+	}
+
+	@Test
+	public void testV2FetchWantRefAndDeepen() throws Exception {
+		RevCommit parent = remote.commit().message("parent").create();
+		RevCommit child = remote.commit().message("x").parent(parent).create();
+		remote.update("branch1", child);
+
+		server.getConfig().setBoolean("uploadpack", null, "allowrefinwant", true);
+
+		ByteArrayInputStream recvStream = uploadPackV2(
+			"command=fetch\n",
+			PacketLineIn.DELIM,
+			"want-ref refs/heads/branch1\n",
+			"deepen 1\n",
+			"done\n",
+			PacketLineIn.END);
+		PacketLineIn pckIn = new PacketLineIn(recvStream);
+
+		// wanted-refs appears first, then shallow-info.
+		assertThat(pckIn.readString(), is("wanted-refs"));
+		assertThat(pckIn.readString(), is(child.toObjectId().getName() + " refs/heads/branch1"));
+		assertThat(pckIn.readString(), theInstance(PacketLineIn.DELIM));
+		assertThat(pckIn.readString(), is("shallow-info"));
+		assertThat(pckIn.readString(), is("shallow " + child.toObjectId().getName()));
+		assertThat(pckIn.readString(), theInstance(PacketLineIn.DELIM));
+		assertThat(pckIn.readString(), is("packfile"));
+		parsePack(recvStream);
+		assertTrue(client.hasObject(child.toObjectId()));
+		assertFalse(client.hasObject(parent.toObjectId()));
+	}
+
 	private static class RejectAllRefFilter implements RefFilter {
 		@Override
 		public Map<String, Ref> filter(Map<String, Ref> refs) {
diff --git a/org.eclipse.jgit/src/org/eclipse/jgit/transport/GitProtocolConstants.java b/org.eclipse.jgit/src/org/eclipse/jgit/transport/GitProtocolConstants.java
index 10cd775304..86740d42c5 100644
--- a/org.eclipse.jgit/src/org/eclipse/jgit/transport/GitProtocolConstants.java
+++ b/org.eclipse.jgit/src/org/eclipse/jgit/transport/GitProtocolConstants.java
@@ -167,6 +167,13 @@ public class GitProtocolConstants {
 	public static final String OPTION_FILTER = "filter"; //$NON-NLS-1$
 
 	/**
+	 * The client specified a want-ref expression.
+	 *
+	 * @since 5.1
+	 */
+	public static final String OPTION_WANT_REF = "want-ref"; //$NON-NLS-1$
+
+	/**
 	 * The client supports atomic pushes. If this option is used, the server
 	 * will update all refs within one atomic transaction.
 	 *
@@ -231,6 +238,13 @@ public class GitProtocolConstants {
 	public static final String CAPABILITY_PUSH_OPTIONS = "push-options"; //$NON-NLS-1$
 
 	/**
+	 * The server supports the client specifying ref names.
+	 *
+	 * @since 5.1
+	 */
+	public static final String CAPABILITY_REF_IN_WANT = "ref-in-want"; //$NON-NLS-1$
+
+	/**
 	 * The server supports listing refs using protocol v2.
 	 *
 	 * @since 5.0
diff --git a/org.eclipse.jgit/src/org/eclipse/jgit/transport/TransferConfig.java b/org.eclipse.jgit/src/org/eclipse/jgit/transport/TransferConfig.java
index 4ae1ccb420..6b8d5c598e 100644
--- a/org.eclipse.jgit/src/org/eclipse/jgit/transport/TransferConfig.java
+++ b/org.eclipse.jgit/src/org/eclipse/jgit/transport/TransferConfig.java
@@ -126,6 +126,7 @@ public class TransferConfig {
 	private final boolean allowInvalidPersonIdent;
 	private final boolean safeForWindows;
 	private final boolean safeForMacOS;
+	private final boolean allowRefInWant;
 	private final boolean allowTipSha1InWant;
 	private final boolean allowReachableSha1InWant;
 	private final boolean allowFilter;
@@ -180,6 +181,7 @@ public class TransferConfig {
 			ignore.add(ObjectChecker.ErrorType.ZERO_PADDED_FILEMODE);
 		}
 
+		allowRefInWant = rc.getBoolean("uploadpack", "allowrefinwant", false);
 		allowTipSha1InWant = rc.getBoolean(
 				"uploadpack", "allowtipsha1inwant", false);
 		allowReachableSha1InWant = rc.getBoolean(
@@ -262,6 +264,14 @@ public class TransferConfig {
 	}
 
 	/**
+	 * @return true if clients are allowed to specify a "want-ref" line
+	 * @since 5.1
+	 */
+	public boolean isAllowRefInWant() {
+		return allowRefInWant;
+	}
+
+	/**
 	 * Get {@link org.eclipse.jgit.transport.RefFilter} respecting configured
 	 * hidden refs.
 	 *
diff --git a/org.eclipse.jgit/src/org/eclipse/jgit/transport/UploadPack.java b/org.eclipse.jgit/src/org/eclipse/jgit/transport/UploadPack.java
index aedc7a6db2..c093e8b1db 100644
--- a/org.eclipse.jgit/src/org/eclipse/jgit/transport/UploadPack.java
+++ b/org.eclipse.jgit/src/org/eclipse/jgit/transport/UploadPack.java
@@ -45,6 +45,7 @@ package org.eclipse.jgit.transport;
 
 import static org.eclipse.jgit.lib.Constants.R_TAGS;
 import static org.eclipse.jgit.lib.RefDatabase.ALL;
+import static org.eclipse.jgit.transport.GitProtocolConstants.CAPABILITY_REF_IN_WANT;
 import static org.eclipse.jgit.transport.GitProtocolConstants.COMMAND_FETCH;
 import static org.eclipse.jgit.transport.GitProtocolConstants.COMMAND_LS_REFS;
 import static org.eclipse.jgit.transport.GitProtocolConstants.OPTION_AGENT;
@@ -62,6 +63,7 @@ import static org.eclipse.jgit.transport.GitProtocolConstants.OPTION_SHALLOW;
 import static org.eclipse.jgit.transport.GitProtocolConstants.OPTION_SIDE_BAND;
 import static org.eclipse.jgit.transport.GitProtocolConstants.OPTION_SIDE_BAND_64K;
 import static org.eclipse.jgit.transport.GitProtocolConstants.OPTION_THIN_PACK;
+import static org.eclipse.jgit.transport.GitProtocolConstants.OPTION_WANT_REF;
 
 import java.io.ByteArrayOutputStream;
 import java.io.EOFException;
@@ -77,6 +79,7 @@ import java.util.HashSet;
 import java.util.List;
 import java.util.Map;
 import java.util.Set;
+import java.util.TreeMap;
 
 import org.eclipse.jgit.annotations.Nullable;
 import org.eclipse.jgit.errors.CorruptObjectException;
@@ -959,9 +962,27 @@ public class UploadPack {
 
 		boolean includeTag = false;
 		boolean filterReceived = false;
+		TreeMap<String, ObjectId> wantedRefs = new TreeMap<>();
 		while ((line = pckIn.readString()) != PacketLineIn.END) {
 			if (line.startsWith("want ")) { //$NON-NLS-1$
 				wantIds.add(ObjectId.fromString(line.substring(5)));
+			} else if (transferConfig.isAllowRefInWant() &&
+					line.startsWith(OPTION_WANT_REF + " ")) { //$NON-NLS-1$
+				String refName = line.substring(OPTION_WANT_REF.length() + 1);
+				Ref ref = db.getRefDatabase().exactRef(refName);
+				if (ref == null) {
+					throw new PackProtocolException(
+							MessageFormat.format(JGitText.get().invalidRefName,
+								refName));
+				}
+				ObjectId oid = ref.getObjectId();
+				if (oid == null) {
+					throw new PackProtocolException(
+							MessageFormat.format(JGitText.get().invalidRefName,
+								refName));
+				}
+				wantedRefs.put(refName, oid);
+				wantIds.add(oid);
 			} else if (line.startsWith("have ")) { //$NON-NLS-1$
 				peerHas.add(ObjectId.fromString(line.substring(5)));
 			} else if (line.equals("done")) { //$NON-NLS-1$
@@ -1062,6 +1083,18 @@ public class UploadPack {
 		}
 
 		if (doneReceived || okToGiveUp()) {
+			if (!wantedRefs.isEmpty()) {
+				if (sectionSent) {
+					pckOut.writeDelim();
+				}
+				pckOut.writeString("wanted-refs\n"); //$NON-NLS-1$
+				for (Map.Entry<String, ObjectId> entry : wantedRefs.entrySet()) {
+					pckOut.writeString(entry.getValue().getName() + ' ' +
+							entry.getKey() + '\n');
+				}
+				sectionSent = true;
+			}
+
 			if (shallowCommits != null) {
 				if (sectionSent)
 					pckOut.writeDelim();
@@ -1125,9 +1158,12 @@ public class UploadPack {
 		ArrayList<String> caps = new ArrayList<>();
 		caps.add("version 2"); //$NON-NLS-1$
 		caps.add(COMMAND_LS_REFS);
+		boolean advertiseRefInWant = transferConfig.isAllowRefInWant() &&
+			db.getConfig().getBoolean("uploadpack", null, "advertiserefinwant", true);
 		caps.add(
 				COMMAND_FETCH + '=' +
 				(transferConfig.isAllowFilter() ? OPTION_FILTER + ' ' : "") + //$NON-NLS-1$
+				(advertiseRefInWant ? CAPABILITY_REF_IN_WANT + ' ' : "") + //$NON-NLS-1$
 				OPTION_SHALLOW);
 		return caps;
 	}