Refactor SideBandOutputStream to be buffered

Instead of relying on our callers to wrap us up inside of a
BufferedOutputStream and using the proper block sizing, do the
buffering directly inside of SideBandOutputStream.  This ensures
we don't get large write-throughs from BufferedOutputStream that
might overflow the configured packet size.

The constructor of SideBandOutputStream is also beefed up to check
its arguments and ensure they are within acceptable ranges for the
current side-band protocol.

Change-Id: Ic14567327d03c9e972f9734b8228178bc448867d
Signed-off-by: Shawn O. Pearce <spearce@spearce.org>

4 files changed, 92 insertions, 45 deletions

diff --git a/org.eclipse.jgit/src/org/eclipse/jgit/transport/PacketLineOut.java b/org.eclipse.jgit/src/org/eclipse/jgit/transport/PacketLineOut.java
index 81dd4f6a15..51506b20aa 100644
--- a/org.eclipse.jgit/src/org/eclipse/jgit/transport/PacketLineOut.java
+++ b/org.eclipse.jgit/src/org/eclipse/jgit/transport/PacketLineOut.java
@@ -1,5 +1,5 @@
 /*
- * Copyright (C) 2008-2009, Google Inc.
+ * Copyright (C) 2008-2010, Google Inc.
  * Copyright (C) 2008-2009, Robin Rosenberg <robin.rosenberg@dewire.com>
  * Copyright (C) 2008, Shawn O. Pearce <spearce@spearce.org>
  * and other copyright owners as documented in the project's IP log.
@@ -105,14 +105,6 @@ public class PacketLineOut {
 		out.write(packet);
 	}
 
-	void writeChannelPacket(final int channel, final byte[] buf, int off,
-			int len) throws IOException {
-		formatLength(len + 5);
-		lenbuffer[4] = (byte) channel;
-		out.write(lenbuffer, 0, 5);
-		out.write(buf, off, len);
-	}
-
 	/**
 	 * Write a packet end marker, sometimes referred to as a flush command.
 	 * <p>
@@ -149,6 +141,10 @@ public class PacketLineOut {
 			'7', '8', '9', 'a', 'b', 'c', 'd', 'e', 'f' };
 
 	private void formatLength(int w) {
+		formatLength(lenbuffer, w);
+	}
+
+	static void formatLength(byte[] lenbuffer, int w) {
 		int o = 3;
 		while (o >= 0 && w != 0) {
 			lenbuffer[o--] = hexchar[w & 0xf];
diff --git a/org.eclipse.jgit/src/org/eclipse/jgit/transport/SideBandOutputStream.java b/org.eclipse.jgit/src/org/eclipse/jgit/transport/SideBandOutputStream.java
index 5e50fd89b3..6e0a52627e 100644
--- a/org.eclipse.jgit/src/org/eclipse/jgit/transport/SideBandOutputStream.java
+++ b/org.eclipse.jgit/src/org/eclipse/jgit/transport/SideBandOutputStream.java
@@ -1,5 +1,5 @@
 /*
- * Copyright (C) 2008, Google Inc.
+ * Copyright (C) 2008-2010, Google Inc.
  * and other copyright owners as documented in the project's IP log.
  *
  * This program and the accompanying materials are made available
@@ -47,11 +47,10 @@ import java.io.IOException;
 import java.io.OutputStream;
 
 /**
- * Multiplexes data and progress messages
+ * Multiplexes data and progress messages.
  * <p>
- * To correctly use this class you must wrap it in a BufferedOutputStream with a
- * buffer size no larger than either {@link #SMALL_BUF} or {@link #MAX_BUF},
- * minus {@link #HDR_SIZE}.
+ * This stream is buffered at packet sizes, so the caller doesn't need to wrap
+ * it in yet another buffered stream.
  */
 class SideBandOutputStream extends OutputStream {
 	static final int CH_DATA = SideBandInputStream.CH_DATA;
@@ -66,34 +65,93 @@ class SideBandOutputStream extends OutputStream {
 
 	static final int HDR_SIZE = 5;
 
-	private final int channel;
+	private final OutputStream out;
 
-	private final PacketLineOut pckOut;
+	private final byte[] buffer;
 
-	private byte[] singleByteBuffer;
+	/**
+	 * Number of bytes in {@link #buffer} that are valid data.
+	 * <p>
+	 * Initialized to {@link #HDR_SIZE} if there is no application data in the
+	 * buffer, as the packet header always appears at the start of the buffer.
+	 */
+	private int cnt;
 
-	SideBandOutputStream(final int chan, final PacketLineOut out) {
-		channel = chan;
-		pckOut = out;
+	/**
+	 * Create a new stream to write side band packets.
+	 *
+	 * @param chan
+	 *            channel number to prefix all packets with, so the remote side
+	 *            can demultiplex the stream and get back the original data.
+	 *            Must be in the range [0, 255].
+	 * @param sz
+	 *            maximum size of a data packet within the stream. The remote
+	 *            side needs to agree to the packet size to prevent buffer
+	 *            overflows. Must be in the range [HDR_SIZE + 1, MAX_BUF).
+	 * @param os
+	 *            stream that the packets are written onto. This stream should
+	 *            be attached to a SideBandInputStream on the remote side.
+	 */
+	SideBandOutputStream(final int chan, final int sz, final OutputStream os) {
+		if (chan <= 0 || chan > 255)
+			throw new IllegalArgumentException("channel " + chan
+					+ " must be in range [0, 255]");
+		if (sz <= HDR_SIZE)
+			throw new IllegalArgumentException("packet size " + sz
+					+ " must be >= " + HDR_SIZE);
+		else if (MAX_BUF < sz)
+			throw new IllegalArgumentException("packet size " + sz
+					+ " must be <= " + MAX_BUF);
+
+		out = os;
+		buffer = new byte[sz];
+		buffer[4] = (byte) chan;
+		cnt = HDR_SIZE;
 	}
 
 	@Override
 	public void flush() throws IOException {
-		if (channel != CH_DATA)
-			pckOut.flush();
+		if (HDR_SIZE < cnt)
+			writeBuffer();
+		out.flush();
 	}
 
 	@Override
-	public void write(final byte[] b, final int off, final int len)
-			throws IOException {
-		pckOut.writeChannelPacket(channel, b, off, len);
+	public void write(final byte[] b, int off, int len) throws IOException {
+		while (0 < len) {
+			int capacity = buffer.length - cnt;
+			if (cnt == HDR_SIZE && capacity < len) {
+				// Our block to write is bigger than the packet size,
+				// stream it out as-is to avoid unnecessary copies.
+				PacketLineOut.formatLength(buffer, buffer.length);
+				out.write(buffer, 0, HDR_SIZE);
+				out.write(b, off, capacity);
+				off += capacity;
+				len -= capacity;
+
+			} else {
+				if (capacity == 0)
+					writeBuffer();
+
+				int n = Math.min(len, capacity);
+				System.arraycopy(b, off, buffer, cnt, n);
+				cnt += n;
+				off += n;
+				len -= n;
+			}
+		}
 	}
 
 	@Override
 	public void write(final int b) throws IOException {
-		if (singleByteBuffer == null)
-			singleByteBuffer = new byte[1];
-		singleByteBuffer[0] = (byte) b;
-		write(singleByteBuffer);
+		if (cnt == buffer.length)
+			writeBuffer();
+		buffer[cnt++] = (byte) b;
+	}
+
+	private void writeBuffer() throws IOException {
+		PacketLineOut.formatLength(buffer, cnt);
+		out.write(buffer, 0, cnt);
+		cnt = HDR_SIZE;
 	}
 }
diff --git a/org.eclipse.jgit/src/org/eclipse/jgit/transport/SideBandProgressMonitor.java b/org.eclipse.jgit/src/org/eclipse/jgit/transport/SideBandProgressMonitor.java
index 89d338c897..efce7b1da7 100644
--- a/org.eclipse.jgit/src/org/eclipse/jgit/transport/SideBandProgressMonitor.java
+++ b/org.eclipse.jgit/src/org/eclipse/jgit/transport/SideBandProgressMonitor.java
@@ -1,5 +1,5 @@
 /*
- * Copyright (C) 2008, Google Inc.
+ * Copyright (C) 2008-2010, Google Inc.
  * and other copyright owners as documented in the project's IP log.
  *
  * This program and the accompanying materials are made available
@@ -43,7 +43,7 @@
 
 package org.eclipse.jgit.transport;
 
-import java.io.BufferedOutputStream;
+import java.io.OutputStream;
 import java.io.OutputStreamWriter;
 import java.io.PrintWriter;
 
@@ -66,12 +66,8 @@ class SideBandProgressMonitor implements ProgressMonitor {
 
 	private int totalWork;
 
-	SideBandProgressMonitor(final PacketLineOut pckOut) {
-		final int bufsz = SideBandOutputStream.SMALL_BUF
-				- SideBandOutputStream.HDR_SIZE;
-		out = new PrintWriter(new OutputStreamWriter(new BufferedOutputStream(
-				new SideBandOutputStream(SideBandOutputStream.CH_PROGRESS,
-						pckOut), bufsz), Constants.CHARSET));
+	SideBandProgressMonitor(final OutputStream os) {
+		out = new PrintWriter(new OutputStreamWriter(os, Constants.CHARSET));
 	}
 
 	public void start(final int totalTasks) {
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 b76b22b77e..39c4243bad 100644
--- a/org.eclipse.jgit/src/org/eclipse/jgit/transport/UploadPack.java
+++ b/org.eclipse.jgit/src/org/eclipse/jgit/transport/UploadPack.java
@@ -43,9 +43,6 @@
 
 package org.eclipse.jgit.transport;
 
-import static org.eclipse.jgit.transport.BasePackFetchConnection.MultiAck;
-
-import java.io.BufferedOutputStream;
 import java.io.EOFException;
 import java.io.IOException;
 import java.io.InputStream;
@@ -70,6 +67,7 @@ import org.eclipse.jgit.revwalk.RevFlagSet;
 import org.eclipse.jgit.revwalk.RevObject;
 import org.eclipse.jgit.revwalk.RevTag;
 import org.eclipse.jgit.revwalk.RevWalk;
+import org.eclipse.jgit.transport.BasePackFetchConnection.MultiAck;
 import org.eclipse.jgit.transport.RefAdvertiser.PacketLineOutRefAdvertiser;
 import org.eclipse.jgit.util.io.InterruptTimer;
 import org.eclipse.jgit.util.io.TimeoutInputStream;
@@ -556,13 +554,12 @@ public class UploadPack {
 			int bufsz = SideBandOutputStream.SMALL_BUF;
 			if (options.contains(OPTION_SIDE_BAND_64K))
 				bufsz = SideBandOutputStream.MAX_BUF;
-			bufsz -= SideBandOutputStream.HDR_SIZE;
-
-			packOut = new BufferedOutputStream(new SideBandOutputStream(
-					SideBandOutputStream.CH_DATA, pckOut), bufsz);
 
+			packOut = new SideBandOutputStream(SideBandOutputStream.CH_DATA,
+					bufsz, rawOut);
 			if (progress)
-				pm = new SideBandProgressMonitor(pckOut);
+				pm = new SideBandProgressMonitor(new SideBandOutputStream(
+						SideBandOutputStream.CH_PROGRESS, bufsz, rawOut));
 		}
 
 		final PackWriter pw;