From 73060a463cfabc0512a44083e89f6f4de2f772ce Mon Sep 17 00:00:00 2001
From: Brad Rydzewski <brad.rydzewski@gmail.com>
Date: Tue, 23 Sep 2014 20:47:40 -0700
Subject: [PATCH] added limit to streaming output to avoid crashes

---
 shared/build/build.go  |  2 +-
 shared/build/writer.go | 15 +++++++++++++++
 2 files changed, 16 insertions(+), 1 deletion(-)

diff --git a/shared/build/build.go b/shared/build/build.go
index 2b74c8a66..9c314a180 100644
--- a/shared/build/build.go
+++ b/shared/build/build.go
@@ -395,7 +395,7 @@ func (b *Builder) run() error {
 
 	// attach to the container
 	go func() {
-		b.dockerClient.Containers.Attach(run.ID, &writer{b.Stdout})
+		b.dockerClient.Containers.Attach(run.ID, &writer{b.Stdout, 0})
 	}()
 
 	// start the container
diff --git a/shared/build/writer.go b/shared/build/writer.go
index 2f5fe9e1c..fce230af4 100644
--- a/shared/build/writer.go
+++ b/shared/build/writer.go
@@ -12,12 +12,17 @@ var (
 	// the prefix used to determine if this is
 	// data that should be stripped from the output
 	prefix = []byte("#DRONE:")
+
+	// default limit to use when streaming build output.
+	DefaultLimit = 2000000
 )
 
 // custom writer to intercept the build
 // output
 type writer struct {
 	io.Writer
+
+	length int
 }
 
 // Write appends the contents of p to the buffer. It will
@@ -25,6 +30,16 @@ type writer struct {
 // output, and will alter the output accordingly.
 func (w *writer) Write(p []byte) (n int, err error) {
 
+	// ensure we haven't exceeded the limit
+	if w.length > DefaultLimit {
+		w.Writer.Write([]byte("Truncating build output ..."))
+		return len(p), nil
+	}
+
+	// track the number of bytes written to the
+	// buffer so that we can limit it.
+	w.length += len(p)
+
 	lines := strings.Split(string(p), "\n")
 	for i, line := range lines {