rfc1035label: using a structure to hold original data

2 files changed, 105 insertions, 37 deletions

diff --git a/rfc1035label/label.go b/rfc1035label/label.go
index 26d8a49..5093de8 100644
--- a/rfc1035label/label.go
+++ b/rfc1035label/label.go
@@ -8,11 +8,78 @@ import (
 // This implements RFC 1035 labels, including compression.
 // https://tools.ietf.org/html/rfc1035#section-4.1.4
 
-// LabelsFromBytes decodes a serialized stream and returns a list of labels
-func LabelsFromBytes(buf []byte) ([]string, error) {
+// Labels represents RFC1035 labels
+type Labels struct {
+	// original contains the original bytes if the object was parsed from a byte
+	// sequence, or nil otherwise. The `original` field is necessary to deal
+	// with compressed labels. If the labels are further modified, the original
+	// content is invalidated and no compression will be used.
+	original []byte
+	// Labels contains the parsed labels. A change here invalidates the
+	// `original` object.
+	Labels []string
+}
+
+// same compares two string arrays
+func same(a, b []string) bool {
+	if len(a) != len(b) {
+		return false
+	}
+	for i := 0; i < len(a); i++ {
+		if a[i] != b[i] {
+			return false
+		}
+	}
+	return true
+}
+
+// ToBytes returns a byte sequence representing the labels. If the original
+// sequence is modified, the labels are parsed again, otherwise the original
+// byte sequence is returned.
+func (l *Labels) ToBytes() []byte {
+	// if the original byte sequence has been modified, invalidate it and
+	// serialize again.
+	// NOTE: this function is not thread-safe. If multiple threads modify
+	// the `Labels` field, the result may be wrong.
+	originalLabels, err := labelsFromBytes(l.original)
+	// if the original object has not been modified, or we cannot parse it,
+	// return the original bytes.
+	if err != nil || (l.original != nil && same(originalLabels, l.Labels)) {
+		return l.original
+	}
+	return labelsToBytes(l.Labels)
+}
+
+// Length returns the length in bytes of the serialized labels
+func (l *Labels) Length() int {
+	return len(l.ToBytes())
+}
+
+// NewLabels returns an initialized Labels object.
+func NewLabels() *Labels {
+	return &Labels{
+		Labels: make([]string, 0),
+	}
+}
+
+// FromBytes returns a Labels object from the given byte sequence, or an error if
+// any.
+func FromBytes(data []byte) (*Labels, error) {
+	lab := NewLabels()
+	l, err := labelsFromBytes(data)
+	if err != nil {
+		return nil, err
+	}
+	lab.original = data
+	lab.Labels = l
+	return lab, nil
+}
+
+// fromBytes decodes a serialized stream and returns a list of labels
+func labelsFromBytes(buf []byte) ([]string, error) {
 	var (
-		pos, oldPos     int
 		labels          = make([]string, 0)
+		pos, oldPos     int
 		label           string
 		handlingPointer bool
 	)
@@ -58,8 +125,8 @@ func LabelsFromBytes(buf []byte) ([]string, error) {
 	return labels, nil
 }
 
-// LabelToBytes encodes a label and returns a serialized stream of bytes
-func LabelToBytes(label string) []byte {
+// labelToBytes encodes a label and returns a serialized stream of bytes
+func labelToBytes(label string) []byte {
 	var encodedLabel []byte
 	if len(label) == 0 {
 		return []byte{0}
@@ -71,12 +138,12 @@ func LabelToBytes(label string) []byte {
 	return append(encodedLabel, 0)
 }
 
-// LabelsToBytes encodes a list of labels and returns a serialized stream of
+// labelsToBytes encodes a list of labels and returns a serialized stream of
 // bytes
-func LabelsToBytes(labels []string) []byte {
+func labelsToBytes(labels []string) []byte {
 	var encodedLabels []byte
 	for _, label := range labels {
-		encodedLabels = append(encodedLabels, LabelToBytes(label)...)
+		encodedLabels = append(encodedLabels, labelToBytes(label)...)
 	}
 	return encodedLabels
 }
diff --git a/rfc1035label/label_test.go b/rfc1035label/label_test.go
index 3a69f2b..6098e44 100644
--- a/rfc1035label/label_test.go
+++ b/rfc1035label/label_test.go
@@ -7,42 +7,35 @@ import (
 )
 
 func TestLabelsFromBytes(t *testing.T) {
-	labels, err := LabelsFromBytes([]byte{
+	expected := []byte{
 		0x9, 's', 'l', 'a', 'c', 'k', 'w', 'a', 'r', 'e',
 		0x2, 'i', 't',
 		0x0,
-	})
+	}
+	labels, err := FromBytes(expected)
 	require.NoError(t, err)
-	require.Equal(t, 1, len(labels))
-	require.Equal(t, "slackware.it", labels[0])
+	require.Equal(t, 1, len(labels.Labels))
+	require.Equal(t, len(expected), labels.Length())
+	require.Equal(t, expected, labels.ToBytes())
+	require.Equal(t, "slackware.it", labels.Labels[0])
 }
 
 func TestLabelsFromBytesZeroLength(t *testing.T) {
-	labels, err := LabelsFromBytes([]byte{})
+	labels, err := FromBytes([]byte{})
 	require.NoError(t, err)
-	require.Equal(t, 0, len(labels))
+	require.Equal(t, 0, len(labels.Labels))
+	require.Equal(t, 0, labels.Length())
+	require.Equal(t, []byte{}, labels.ToBytes())
 }
 
 func TestLabelsFromBytesInvalidLength(t *testing.T) {
-	labels, err := LabelsFromBytes([]byte{0x5, 0xaa, 0xbb}) // short length
+	_, err := FromBytes([]byte{0x5, 0xaa, 0xbb}) // short length
 	require.Error(t, err)
-	require.Equal(t, 0, len(labels))
 }
 
 func TestLabelsFromBytesInvalidLengthOffByOne(t *testing.T) {
-	labels, err := LabelsFromBytes([]byte{0x3, 0xaa, 0xbb}) // short length
+	_, err := FromBytes([]byte{0x3, 0xaa, 0xbb}) // short length
 	require.Error(t, err)
-	require.Equal(t, 0, len(labels))
-}
-
-func TestLabelToBytes(t *testing.T) {
-	encodedLabel := LabelToBytes("slackware.it")
-	expected := []byte{
-		0x9, 's', 'l', 'a', 'c', 'k', 'w', 'a', 'r', 'e',
-		0x2, 'i', 't',
-		0x0,
-	}
-	require.Equal(t, expected, encodedLabel)
 }
 
 func TestLabelsToBytes(t *testing.T) {
@@ -55,14 +48,20 @@ func TestLabelsToBytes(t *testing.T) {
 		2, 'i', 't',
 		0,
 	}
-	encodedLabels := LabelsToBytes([]string{"slackware.it", "insomniac.slackware.it"})
-	require.Equal(t, expected, encodedLabels)
+	labels := Labels{
+		Labels: []string{
+			"slackware.it",
+			"insomniac.slackware.it",
+		},
+	}
+	require.Equal(t, expected, labels.ToBytes())
 }
 
 func TestLabelToBytesZeroLength(t *testing.T) {
-	encodedLabel := LabelToBytes("")
-	expected := []byte{0}
-	require.Equal(t, expected, encodedLabel)
+	labels := Labels{
+		Labels: []string{""},
+	}
+	require.Equal(t, []byte{0}, labels.ToBytes())
 }
 
 func TestCompressedLabel(t *testing.T) {
@@ -89,9 +88,11 @@ func TestCompressedLabel(t *testing.T) {
 		"systemboot.org",
 	}
 
-	labels, err := LabelsFromBytes(data)
+	labels, err := FromBytes(data)
 	require.NoError(t, err)
-	require.Equal(t, expected, labels)
+	require.Equal(t, 4, len(labels.Labels))
+	require.Equal(t, expected, labels.Labels)
+	require.Equal(t, len(data), labels.Length())
 }
 
 func TestShortCompressedLabel(t *testing.T) {
@@ -105,7 +106,7 @@ func TestShortCompressedLabel(t *testing.T) {
 		192,
 	}
 
-	_, err := LabelsFromBytes(data)
+	_, err := FromBytes(data)
 	require.Error(t, err)
 }
 
@@ -121,6 +122,6 @@ func TestNestedCompressedLabel(t *testing.T) {
 		9, 'i', 'n', 's', 'o', 'm', 'n', 'i', 'a', 'c',
 		192, 5,
 	}
-	_, err := LabelsFromBytes(data)
+	_, err := FromBytes(data)
 	require.Error(t, err)
 }