Adding upstream version 1.2.0.

Signed-off-by: Daniel Baumann <daniel@debian.org>

limio/limit.go

@@ -0,0 +1,31 @@
+package limio
+
+import "io"
+
+type limitedWriter struct {
+	W io.Writer
+	N int64
+}
+
+func (l *limitedWriter) Write(p []byte) (n int, err error) {
+	if l.N <= 0 {
+		return 0, io.ErrShortWrite
+	}
+	if int64(len(p)) > l.N {
+		p = p[0:l.N]
+		err = io.ErrShortWrite
+	}
+	n, er := l.W.Write(p)
+	if er != nil {
+		err = er
+	}
+	l.N -= int64(n)
+	return n, err
+}
+
+// LimitWriter works like io.LimitReader. It writes at most n bytes
+// to the underlying Writer. It returns io.ErrShortWrite if more than n
+// bytes are attempted to be written.
+func LimitWriter(w io.Writer, n int64) io.Writer {
+	return &limitedWriter{W: w, N: n}
+}

limio/limit_test.go

@@ -0,0 +1,62 @@
+package limio
+
+import (
+	"bytes"
+	"crypto/rand"
+	"io"
+	"testing"
+)
+
+func TestLimits(t *testing.T) {
+	LimitTestN(0, t)
+	LimitTestN(1, t)
+	LimitTestN(2, t)
+	LimitTestN(1024, t)
+
+	OverLimitTestN(0, t)
+	OverLimitTestN(1, t)
+	OverLimitTestN(2, t)
+	OverLimitTestN(1024, t)
+}
+
+func OverLimitTestN(N int64, t *testing.T) {
+	got := bytes.NewBuffer(nil)
+	exp := bytes.NewBuffer(nil)
+
+	lw := LimitWriter(got, N)
+	w := io.MultiWriter(exp, lw)
+
+	_, err := io.CopyN(w, rand.Reader, 2*N)
+	if err != io.ErrShortWrite && N != 2*N {
+		t.Error("did not throw io.ErrShortWrite!")
+		if err != nil {
+			t.Error(err.Error())
+		}
+	}
+
+	if !bytes.Equal(got.Bytes(), exp.Bytes()[:N]) {
+		t.Errorf("Exp: %s\n Got: %s\n",
+			string(exp.Bytes()[:N]),
+			string(got.Bytes()))
+	}
+
+	if int64(got.Len()) != N {
+		t.Errorf("LimitWriter did not cap at %d", N)
+	}
+}
+
+func LimitTestN(N int64, t *testing.T) {
+	got := bytes.NewBuffer(nil)
+	exp := bytes.NewBuffer(nil)
+
+	lw := LimitWriter(got, N)
+	w := io.MultiWriter(lw, exp)
+
+	io.CopyN(w, rand.Reader, N)
+
+	if !bytes.Equal(got.Bytes(), exp.Bytes()) {
+		t.Errorf("Exp: %s\n Got: %s\n",
+			string(exp.Bytes()[:N]),
+			string(got.Bytes()))
+	}
+}