Teach UploadPack "ref-prefix" in "ls-refs"

Add support for the "ref-prefix" parameter in the "ls-refs" command in
the fetch-pack/upload-pack protocol v2.

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

2 files changed, 62 insertions, 3 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 88ac49a51b..a50f170ffa 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
@@ -461,4 +461,43 @@ public class UploadPackTest {
 		assertThat(pckIn.readString(), is(tag.toObjectId().getName() + " refs/tags/tag"));
 		assertTrue(pckIn.readString() == PacketLineIn.END);
 	}
+
+	@Test
+	public void testV2LsRefsRefPrefix() throws Exception {
+		RevCommit tip = remote.commit().message("message").create();
+		remote.update("master", tip);
+		remote.update("other", tip);
+		remote.update("yetAnother", tip);
+
+		ByteArrayInputStream recvStream = uploadPackV2(
+			"command=ls-refs\n",
+			PacketLineIn.DELIM,
+			"ref-prefix refs/heads/maste",
+			"ref-prefix refs/heads/other",
+			PacketLineIn.END);
+		PacketLineIn pckIn = new PacketLineIn(recvStream);
+
+		assertThat(pckIn.readString(), is(tip.toObjectId().getName() + " refs/heads/master"));
+		assertThat(pckIn.readString(), is(tip.toObjectId().getName() + " refs/heads/other"));
+		assertTrue(pckIn.readString() == PacketLineIn.END);
+	}
+
+	@Test
+	public void testV2LsRefsRefPrefixNoSlash() throws Exception {
+		RevCommit tip = remote.commit().message("message").create();
+		remote.update("master", tip);
+		remote.update("other", tip);
+
+		ByteArrayInputStream recvStream = uploadPackV2(
+			"command=ls-refs\n",
+			PacketLineIn.DELIM,
+			"ref-prefix refs/heads/maste",
+			"ref-prefix r",
+			PacketLineIn.END);
+		PacketLineIn pckIn = new PacketLineIn(recvStream);
+
+		assertThat(pckIn.readString(), is(tip.toObjectId().getName() + " refs/heads/master"));
+		assertThat(pckIn.readString(), is(tip.toObjectId().getName() + " refs/heads/other"));
+		assertTrue(pckIn.readString() == PacketLineIn.END);
+	}
 }
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 07023286e9..355b1b9df2 100644
--- a/org.eclipse.jgit/src/org/eclipse/jgit/transport/UploadPack.java
+++ b/org.eclipse.jgit/src/org/eclipse/jgit/transport/UploadPack.java
@@ -69,6 +69,7 @@ import java.text.MessageFormat;
 import java.util.ArrayList;
 import java.util.Collection;
 import java.util.Collections;
+import java.util.HashMap;
 import java.util.HashSet;
 import java.util.List;
 import java.util.Map;
@@ -868,8 +869,9 @@ public class UploadPack {
 
 	private void lsRefsV2() throws IOException {
 		PacketLineOutRefAdvertiser adv = new PacketLineOutRefAdvertiser(pckOut);
-		Map<String, Ref> refs = getAdvertisedOrDefaultRefs();
 		String line;
+		ArrayList<String> refPrefixes = new ArrayList<>();
+		boolean needToFindSymrefs = false;
 
 		adv.setUseProtocolV2(true);
 
@@ -882,7 +884,9 @@ public class UploadPack {
 				if (line.equals("peel")) {
 					adv.setDerefTags(true);
 				} else if (line.equals("symrefs")) {
-					findSymrefs(adv, refs);
+					needToFindSymrefs = true;
+				} else if (line.startsWith("ref-prefix ")) {
+					refPrefixes.add(line.substring("ref-prefix ".length()));
 				} else {
 					throw new PackProtocolException("unexpected " + line);
 				}
@@ -892,7 +896,23 @@ public class UploadPack {
 		}
 		rawOut.stopBuffering();
 
-		adv.send(refs);
+		Map<String, Ref> refsToSend;
+		if (refPrefixes.isEmpty()) {
+			refsToSend = getAdvertisedOrDefaultRefs();
+		} else {
+			refsToSend = new HashMap<>();
+			for (String refPrefix : refPrefixes) {
+				for (Ref ref : db.getRefDatabase().getRefsByPrefix(refPrefix)) {
+					refsToSend.put(ref.getName(), ref);
+				}
+			}
+		}
+
+		if (needToFindSymrefs) {
+			findSymrefs(adv, refsToSend);
+		}
+
+		adv.send(refsToSend);
 		adv.end();
 	}