tag.diff

diff --git a/agent/boundarylogproxy/codec/codec.go b/agent/boundarylogproxy/codec/codec.go
index 0c3d715c0..d93000e30 100644
--- a/agent/boundarylogproxy/codec/codec.go
+++ b/agent/boundarylogproxy/codec/codec.go
@@ -16,14 +16,6 @@ import (
        "golang.org/x/xerrors"
 )
 
-type Tag uint8
-
-const (
-       // TagV1 identifies the first revision of the protocol. This version has a maximum
-       // data length of MaxMessageSizeV1.
-       TagV1 Tag = 1
-)
-
 const (
        // DataLength is the number of bits used for the length of encoded protobuf data.
        DataLength = 24
@@ -42,22 +34,14 @@ var ErrMessageTooLarge = xerrors.New("message too large")
 // WriteFrame writes a framed message with the given tag and data. The data
 // must not exceed 2^DataLength in length.
 func WriteFrame(w io.Writer, tag Tag, data []byte) error {
-       var maxSize uint32
-       switch tag {
-       case TagV1:
-               maxSize = MaxMessageSizeV1
-       default:
-               return xerrors.Errorf("unsupported tag: %d", tag)
-       }
-
-       if len(data) > int(maxSize) {
+       if maxSize := tag.MaxMessageSize(); len(data) > int(maxSize) {
                return xerrors.Errorf("data too large for tag %d: %d > %d", tag, len(data), maxSize)
        }
 
        var header uint32
        //nolint:gosec // The length check above ensures there's no overflow.
        header |= uint32(len(data))
-       header |= uint32(tag) << DataLength
+       header |= uint32(tag.ID()) << DataLength
 
        if err := binary.Write(w, binary.BigEndian, header); err != nil {
                return xerrors.Errorf("write header error: %w", err)
@@ -81,7 +65,7 @@ func WriteFrame(w io.Writer, tag Tag, data []byte) error {
 func ReadFrame(r io.Reader, buf []byte) (Tag, []byte, error) {
        var header uint32
        if err := binary.Read(r, binary.BigEndian, &header); err != nil {
-               return 0, nil, xerrors.Errorf("read header error: %w", err)
+               return nil, nil, xerrors.Errorf("read header error: %w", err)
        }
 
        const lengthMask = (1 << DataLength) - 1
@@ -91,20 +75,15 @@ func ReadFrame(r io.Reader, buf []byte) (Tag, []byte, error) {
        if shifted > tagMask {
                // This is really only here to satisfy the gosec linter. We know from above that
                // shifted <= tagMask.
-               return 0, nil, xerrors.Errorf("invalid tag: %d", shifted)
+               return nil, nil, xerrors.Errorf("invalid tag: %d", shifted)
        }
-       tag := Tag(shifted)
-
-       var maxSize uint32
-       switch tag {
-       case TagV1:
-               maxSize = MaxMessageSizeV1
-       default:
-               return 0, nil, xerrors.Errorf("unsupported tag: %d", tag)
+       tag, err := TagFromID(shifted)
+       if err != nil {
+               return nil, nil, err
        }
 
-       if length > maxSize {
-               return 0, nil, ErrMessageTooLarge
+       if maxSize := tag.MaxMessageSize(); length > maxSize {
+               return nil, nil, ErrMessageTooLarge
        }
 
        if cap(buf) < int(length) {
@@ -114,7 +93,7 @@ func ReadFrame(r io.Reader, buf []byte) (Tag, []byte, error) {
        }
 
        if _, err := io.ReadFull(r, buf[:length]); err != nil {
-               return 0, nil, xerrors.Errorf("read full error: %w", err)
+               return nil, nil, xerrors.Errorf("read full error: %w", err)
        }
 
        return tag, buf[:length], nil
diff --git a/agent/boundarylogproxy/codec/codec_test.go b/agent/boundarylogproxy/codec/codec_test.go
index d09a5e3aa..022af6cf9 100644
--- a/agent/boundarylogproxy/codec/codec_test.go
+++ b/agent/boundarylogproxy/codec/codec_test.go
@@ -15,22 +15,22 @@ func TestRoundTrip(t *testing.T) {
 
        tests := []struct {
                name string
-               tag  codec.Tag
+               tag  uint32
                data []byte
        }{
                {
                        name: "empty data",
-                       tag:  codec.TagV1,
+                       tag:  1,
                        data: []byte{},
                },
                {
                        name: "simple data",
-                       tag:  codec.TagV1,
+                       tag:  1,
                        data: []byte("hello world"),
                },
                {
                        name: "binary data",
-                       tag:  codec.TagV1,
+                       tag:  1,
                        data: []byte{0x00, 0x01, 0x02, 0xff, 0xfe},
                },
        }
@@ -39,14 +39,17 @@ func TestRoundTrip(t *testing.T) {
                t.Run(tt.name, func(t *testing.T) {
                        t.Parallel()
 
+                       tag, err := codec.TagFromID(tt.tag)
+                       require.NoError(t, err)
+
                        var buf bytes.Buffer
-                       err := codec.WriteFrame(&buf, tt.tag, tt.data)
+                       err = codec.WriteFrame(&buf, tag, tt.data)
                        require.NoError(t, err)
 
                        readBuf := make([]byte, codec.MaxMessageSizeV1)
                        tag, data, err := codec.ReadFrame(&buf, readBuf)
                        require.NoError(t, err)
-                       require.Equal(t, tt.tag, tag)
+                       require.IsType(t, &codec.TagV1{}, tag)
                        require.Equal(t, tt.data, data)
                })
        }
@@ -55,10 +58,12 @@ func TestRoundTrip(t *testing.T) {
 func TestReadFrameTooLarge(t *testing.T) {
        t.Parallel()
 
+       tag := &codec.TagV1{}
+
        // Hand construct a header that indicates the message size exceeds the maximum
        // message size for codec.TagV1 by one. We just write the header to buf because
        // we expect codec.ReadFrame to bail out when reading the invalid length.
-       header := uint32(codec.TagV1)<<codec.DataLength | (codec.MaxMessageSizeV1 + 1)
+       header := uint32(tag.ID())<<codec.DataLength | (codec.MaxMessageSizeV1 + 1)
        data := make([]byte, 4)
        binary.BigEndian.PutUint32(data, header)
 
@@ -109,14 +114,14 @@ func TestReadFrameAllocatesWhenNeeded(t *testing.T) {
 
        var buf bytes.Buffer
        data := []byte("this message is longer than the buffer")
-       err := codec.WriteFrame(&buf, codec.TagV1, data)
+       err := codec.WriteFrame(&buf, &codec.TagV1{}, data)
        require.NoError(t, err)
 
        // Buffer with insufficient capacity triggers allocation.
        readBuf := make([]byte, 4)
        tag, got, err := codec.ReadFrame(&buf, readBuf)
        require.NoError(t, err)
-       require.Equal(t, codec.TagV1, tag)
+       require.IsType(t, &codec.TagV1{}, tag)
        require.Equal(t, data, got)
 }
 
@@ -125,21 +130,11 @@ func TestWriteFrameDataSize(t *testing.T) {
 
        var buf bytes.Buffer
        data := make([]byte, codec.MaxMessageSizeV1)
-       err := codec.WriteFrame(&buf, codec.TagV1, data)
+       err := codec.WriteFrame(&buf, &codec.TagV1{}, data)
        require.NoError(t, err)
 
        //nolint: makezero // This intentionally increases the slice length.
        data = append(data, 0) // One byte over the maximum
-       err = codec.WriteFrame(&buf, codec.TagV1, data)
-       require.Error(t, err)
-}
-
-func TestWriteFrameInvalidTag(t *testing.T) {
-       t.Parallel()
-
-       var buf bytes.Buffer
-       readBuf := make([]byte, 1)
-       const bogusTag = 2
-       err := codec.WriteFrame(&buf, codec.Tag(bogusTag), readBuf)
+       err = codec.WriteFrame(&buf, &codec.TagV1{}, data)
        require.Error(t, err)
 }
diff --git a/agent/boundarylogproxy/codec/tag.go b/agent/boundarylogproxy/codec/tag.go
new file mode 100644
index 000000000..4baecb770
--- /dev/null
+++ b/agent/boundarylogproxy/codec/tag.go
@@ -0,0 +1,24 @@
+package codec
+
+import "golang.org/x/xerrors"
+
+type Tag interface {
+       ID() uint8
+       MaxMessageSize() uint32
+}
+
+type TagV1 struct{}
+
+var _ Tag = &TagV1{}
+
+func (t TagV1) ID() uint8              { return 1 }
+func (t TagV1) MaxMessageSize() uint32 { return 1 << 15 }
+
+func TagFromID(id uint32) (Tag, error) {
+       switch id {
+       case 1:
+               return &TagV1{}, nil
+       default:
+               return nil, xerrors.Errorf("unknown tag: %d", id)
+       }
+}