LogDog: Add butler protocol reader/write library.

client/logdog/butlerproto/proto_test.go

Index: client/logdog/butlerproto/proto_test.go
diff --git a/client/logdog/butlerproto/proto_test.go b/client/logdog/butlerproto/proto_test.go
new file mode 100644
index 0000000000000000000000000000000000000000..b70edd444899eb9392d0a1e2cf43968f5e00aa18
--- /dev/null
+++ b/client/logdog/butlerproto/proto_test.go
@@ -0,0 +1,210 @@
+// Copyright 2015 The Chromium Authors. All rights reserved.
+// Use of this source code is governed by a BSD-style license that can be
+// found in the LICENSE file.
+
+package butlerproto
+
+import (
+	"bytes"
+	"io"
+	"strings"
+	"testing"
+
+	"github.com/golang/protobuf/proto"
+	"github.com/luci/luci-go/common/logdog/protocol"
+	"github.com/luci/luci-go/common/recordio"
+	. "github.com/smartystreets/goconvey/convey"
+)
+
+func read(ir io.Reader) (*Reader, error) {
+	r := Reader{}
+	if err := r.Read(ir); err != nil {
+		return nil, err
+	}
+	return &r, nil
+}
+
+func TestReader(t *testing.T) {
+	Convey(`A Reader instance`, t, func() {
+		r := Reader{}
+		buf := bytes.Buffer{}
+		fw := recordio.NewWriter(&buf)
+
+		writeFrame := func(data []byte) error {
+			if _, err := fw.Write(data); err != nil {
+				return err
+			}
+			if err := fw.Flush(); err != nil {
+				return err
+			}
+			return nil
+		}
+
+		push := func(m proto.Message) error {
+			data, err := proto.Marshal(m)
+			if err != nil {
+				return err
+			}
+			return writeFrame(data)
+		}
+
+		Convey(`Can read a ButlerLogBundle entry.`, func() {
+			md := protocol.ButlerMetadata{
+				Type: protocol.ButlerMetadata_ButlerLogBundle,
+			}
+			bundle := protocol.ButlerLogBundle{
+				Source: "test source",
+			}
+
+			So(push(&md), ShouldBeNil)
+			So(push(&bundle), ShouldBeNil)
+
+			So(r.Read(&buf), ShouldBeNil)
+			So(r.Bundle, ShouldNotBeNil)
+			So(r.Bundle.Source, ShouldEqual, "test source")
+		})
+
+		Convey(`Will fail to Read an unknown type.`, func() {
+			// Assert that we are testing an unknown type.
+			unknownType := protocol.ButlerMetadata_ContentType(-1)
+			So(protocol.ButlerMetadata_ContentType_name[int32(unknownType)], ShouldEqual, "")
+
+			md := protocol.ButlerMetadata{
+				Type: unknownType,
+			}
+			So(push(&md), ShouldBeNil)
+
+			err := r.Read(&buf)
+			So(err, ShouldNotBeNil)
+			So(err.Error(), ShouldContainSubstring, "unknown data type")
+		})
+
+		Convey(`Will fail to read junk metadata.`, func() {
+			So(writeFrame([]byte{0xd0, 0x6f, 0x00, 0xd5}), ShouldBeNil)
+
+			err := r.Read(&buf)
+			So(err, ShouldNotBeNil)
+			So(err.Error(), ShouldContainSubstring, "failed to unmarshal Metadata frame")
+		})
+
+		Convey(`With a proper Metadata frame`, func() {
+			md := protocol.ButlerMetadata{
+				Type: protocol.ButlerMetadata_ButlerLogBundle,
+			}
+			So(push(&md), ShouldBeNil)
+
+			Convey(`Will fail if the bundle data is junk.`, func() {
+				So(writeFrame([]byte{0xd0, 0x6f, 0x00, 0xd5}), ShouldBeNil)
+
+				err := r.Read(&buf)
+				So(err, ShouldNotBeNil)
+				So(err.Error(), ShouldContainSubstring, "failed to unmarshal Bundle frame")
+			})
+		})
+
+		Convey(`With a proper compressed Metadata frame`, func() {
+			md := protocol.ButlerMetadata{
+				Type:        protocol.ButlerMetadata_ButlerLogBundle,
+				Compression: protocol.ButlerMetadata_ZLIB,
+			}
+			So(push(&md), ShouldBeNil)
+
+			Convey(`Will fail if the data frame is missing.`, func() {
+				err := r.Read(&buf)
+				So(err, ShouldNotBeNil)
+				So(err.Error(), ShouldContainSubstring, "failed to read Bundle data")
+			})
+
+			Convey(`Will fail if there is junk compressed data.`, func() {
+				So(writeFrame(bytes.Repeat([]byte{0x55, 0xAA}, 16)), ShouldBeNil)
+
+				err := r.Read(&buf)
+				So(err, ShouldNotBeNil)
+				So(err.Error(), ShouldContainSubstring, "failed to initialize zlib reader")
+			})
+		})
+
+		Convey(`Will refuse to read a frame larger than our maximum size.`, func() {
+			r.maxSize = 16
+			So(writeFrame(bytes.Repeat([]byte{0x55}, 17)), ShouldBeNil)
+
+			err := r.Read(&buf)
+			So(err, ShouldEqual, recordio.ErrFrameTooLarge)
+		})
+
+		Convey(`Will refuse to read a compressed protobuf larger than our maximum size.`, func() {
+			// We are crafting this data such that its compressed (frame) size is
+			// below our threshold (16), but its compressed size exceeds it. Repeated
+			// bytes compress very well :)
+			//
+			// Because the frame it smaller than our threshold, our FrameReader will
+			// not outright reject the frame. However, the data is still larger than
+			// we're allowed, and we must reject it.
+			r.maxSize = 16
+			w := Writer{
+				Compress:          true,
+				CompressThreshold: 0,
+			}
+			So(w.writeData(recordio.NewWriter(&buf), protocol.ButlerMetadata_ButlerLogBundle,
+				bytes.Repeat([]byte{0x55}, 64)), ShouldBeNil)
+
+			err := r.Read(&buf)
+			So(err, ShouldNotBeNil)
+			So(err.Error(), ShouldContainSubstring, "limit exceeded")
+		})
+	})
+}
+
+func TestWriter(t *testing.T) {
+	Convey(`A Writer instance outputting to a Buffer`, t, func() {
+		buf := bytes.Buffer{}
+		w := Writer{}
+		bundle := protocol.ButlerLogBundle{
+			Source: "test source",
+			Entries: []*protocol.ButlerLogBundle_Entry{
+				{},
+			},
+		}
+
+		Convey(`When configured to compress with a threshold of 64`, func() {
+			w.Compress = true
+			w.CompressThreshold = 64
+
+			Convey(`Will not compress if below the compression threshold.`, func() {
+				So(w.Write(&buf, &bundle), ShouldBeNil)
+
+				r, err := read(&buf)
+				So(err, ShouldBeNil)
+				So(r.Metadata.Compression, ShouldEqual, protocol.ButlerMetadata_NONE)
+			})
+
+			Convey(`Will not write data larger than the maximum bundle size.`, func() {
+				w.maxSize = 16
+				bundle.Source = strings.Repeat("A", 17)
+				err := w.Write(&buf, &bundle)
+				So(err, ShouldNotBeNil)
+				So(err.Error(), ShouldContainSubstring, "exceeds soft cap")
+			})
+
+			Convey(`Will compress data >= the threshold.`, func() {
+				bundle.Source = strings.Repeat("A", 64)
+				So(w.Write(&buf, &bundle), ShouldBeNil)
+
+				r, err := read(&buf)
+				So(err, ShouldBeNil)
+				So(r.Metadata.Compression, ShouldEqual, protocol.ButlerMetadata_ZLIB)
+				So(r.Bundle.Source, ShouldResemble, bundle.Source)
+
+				Convey(`And can be reused.`, func() {
+					bundle.Source = strings.Repeat("A", 64)
+					So(w.Write(&buf, &bundle), ShouldBeNil)
+
+					r, err := read(&buf)
+					So(err, ShouldBeNil)
+					So(r.Metadata.Compression, ShouldEqual, protocol.ButlerMetadata_ZLIB)
+					So(r.Bundle.Source, ShouldEqual, bundle.Source)
+				})
+			})
+		})
+	})
+}