helper: Support bundle2 responses to bundle2 pushes

Even when the client doesn't indicate supporting bundle2 replies,
mercurial servers send errors through bundle2 replies. So we're pretty
much forced to support bundle2 replies, and now send a replycaps part
unconditionally, so that we always get bundle2 replies to bundle2
pushes through native wire protocol.

cinnabar/hg/__init__.py

@@ -662,10 +662,9 @@ def revs():
         if repo.local():
             repo.local().ui.setconfig('server', 'validate', True)
         b2caps = bundle2caps(repo) if unbundle20 else {}
-        if b2caps and (repo.url().startswith(('http://', 'https://')) or
-                       not isinstance(repo, HelperRepo)):
-            b2caps['replycaps'] = True
         logging.getLogger('bundle2').debug('%r', b2caps)
+        if b2caps:
+            b2caps['replycaps'] = encodecaps({'error': ['abort']})
         cg = create_bundle(store, push_commits, b2caps)
         if not isinstance(repo, HelperRepo):
             cg = util.chunkbuffer(cg)

cinnabar/hg/bundle.py

@@ -443,8 +443,9 @@ def create_bundle(store, commits, bundle2caps={}):
         from mercurial.util import chunkbuffer
         yield 'HG20'
         yield '\0' * 4  # bundle parameters length: no params
-        if bundle2caps.get('replycaps'):
-            for chunk in bundlepart('REPLYCAPS'):
+        replycaps = bundle2caps.get('replycaps')
+        if replycaps:
+            for chunk in bundlepart('REPLYCAPS', data=chunkbuffer(replycaps)):
                 yield chunk
         for chunk in bundlepart('CHANGEGROUP',
                                 advisoryparams=(('version', version),),

hg-bundle.c

@@ -2,28 +2,53 @@
 #include "hg-bundle.h"
 #include <stdint.h>
 
-static size_t copy_data(uint32_t len, FILE *in, FILE *out)
+struct bundle_writer {
+	union {
+		FILE *file;
+		struct strbuf *buf;
+	} out;
+	int type;
+};
+
+#define WRITER_FILE 1
+#define WRITER_STRBUF 2
+
+static size_t write_data(const unsigned char *buf, size_t size,
+			 struct bundle_writer *out)
+{
+	switch (out->type) {
+	case WRITER_FILE:
+		return fwrite(buf, 1, size, out->out.file);
+	case WRITER_STRBUF:
+		strbuf_add(out->out.buf, buf, size);
+		return size;
+	default:
+		return 0;
+	}
+}
+
+static size_t copy_data(uint32_t len, FILE *in, struct bundle_writer *out)
 {
 	unsigned char buf[4096];
 	size_t ret = len;
 	while (len) {
 		uint32_t sz = len > sizeof(buf) ? sizeof(buf) : len;
 		fread(buf, 1, sz, in);
-		fwrite(buf, 1, sz, out);
+		write_data(buf, sz, out);
 		len -= sz;
 	}
 	return ret;
 }
 
-static size_t copy_chunk(int adjust, FILE *in, FILE *out)
+static size_t copy_chunk(int adjust, FILE *in, struct bundle_writer *out)
 {
 	unsigned char buf[4];
 	const unsigned char *p = buf;
 	uint32_t len;
 	size_t ret = 0;
 	//TODO: Check for errors, etc.
 	fread(buf, 1, 4, in);
-	fwrite(buf, 1, 4, out);
+	write_data(buf, 4, out);
 	len = get_be32(p);
 	if (len <= adjust)
 		//TODO: len != 0 is actually invalid
@@ -33,12 +58,12 @@ static size_t copy_chunk(int adjust, FILE *in, FILE *out)
 	return ret;
 }
 
-static size_t copy_changegroup_chunk(FILE *in, FILE *out)
+static size_t copy_changegroup_chunk(FILE *in, struct bundle_writer *out)
 {
 	return copy_chunk(4, in, out);
 }
 
-static void copy_changegroup(FILE *in, FILE *out)
+static void copy_changegroup(FILE *in, struct bundle_writer *out)
 {
 	/* changesets */
 	while (copy_changegroup_chunk(in, out)) {}
@@ -50,18 +75,18 @@ static void copy_changegroup(FILE *in, FILE *out)
 	}
 }
 
-static size_t copy_bundle2_chunk(FILE *in, FILE *out)
+static size_t copy_bundle2_chunk(FILE *in, struct bundle_writer *out)
 {
 	return copy_chunk(0, in, out);
 }
 
-void copy_bundle(FILE *in, FILE *out)
+static void copy_bundle_internal(FILE *in, struct bundle_writer *out)
 {
 	unsigned char buf[4];
 	const unsigned char *p = buf;
 	//TODO: Check for errors, etc.
 	fread(buf, 1, 4, in);
-	fwrite(buf, 1, 4, out);
+	write_data(buf, 4, out);
 	if (memcmp(buf, "HG20", 4)) {
 		copy_data(get_be32(p) - 4, in, out);
 		copy_changegroup(in, out);
@@ -74,3 +99,19 @@ void copy_bundle(FILE *in, FILE *out)
 		while (copy_bundle2_chunk(in, out)) {}
 	}
 }
+
+void copy_bundle(FILE *in, FILE *out)
+{
+	struct bundle_writer writer;
+	writer.type = WRITER_FILE;
+	writer.out.file = out;
+	copy_bundle_internal(in, &writer);
+}
+
+void copy_bundle_to_strbuf(FILE *in, struct strbuf *out)
+{
+	struct bundle_writer writer;
+	writer.type = WRITER_STRBUF;
+	writer.out.buf = out;
+	copy_bundle_internal(in, &writer);
+}

hg-bundle.h

@@ -1,8 +1,10 @@
 #ifndef BUNDLE_H
 #define BUNDLE_H
 
+#include "strbuf.h"
 #include <stdio.h>
 
 extern void copy_bundle(FILE *in, FILE *out);
+extern void copy_bundle_to_strbuf(FILE *in, struct strbuf *out);
 
 #endif

hg-connect-stdio.c

@@ -120,6 +120,7 @@ static void stdio_push_command(struct hg_connection *conn,
 			       struct strbuf *response, FILE *in, off_t len,
 			       const char *command, ...)
 {
+	int is_bundle2 = 0;
 	char buf[4096];
 	struct strbuf header = STRBUF_INIT;
 	va_list ap;
@@ -137,6 +138,13 @@ static void stdio_push_command(struct hg_connection *conn,
 	xwrite(conn->stdio.proc.in, header.buf, header.len);
 	strbuf_release(&header);
 
+        if (len > 4) {
+		char header[4] = { 0, };
+		fread(header, 4, 1, in);
+		fseek(in, 0L, SEEK_SET);
+		is_bundle2 = memcmp(header, "HG20", 4) == 0;
+	}
+
 	while (len) {
 		size_t read = sizeof(buf) > len ? len : sizeof(buf);
 		read = fread(buf, 1, read, in);
@@ -145,11 +153,15 @@ static void stdio_push_command(struct hg_connection *conn,
 	}
 
 	xwrite(conn->stdio.proc.in, "0\n", 2);
-	/* There are two responses, one for output, one for actual response. */
-	//TODO: actually handle output here
-	stdio_read_response(conn, &header);
-	strbuf_release(&header);
-	stdio_read_response(conn, response);
+	if (is_bundle2) {
+		copy_bundle_to_strbuf(conn->stdio.out, response);
+	} else {
+		/* There are two responses, one for output, one for actual response. */
+		//TODO: actually handle output here
+		stdio_read_response(conn, &header);
+		strbuf_release(&header);
+		stdio_read_response(conn, response);
+	}
 }
 
 static int stdio_finish(struct hg_connection *conn)