rfc1035label: support compression pointers

2 files changed, 127 insertions, 45 deletions

diff --git a/rfc1035label/label.go b/rfc1035label/label.go
index c63b4af..26d8a49 100644
--- a/rfc1035label/label.go
+++ b/rfc1035label/label.go
@@ -1,39 +1,61 @@
 package rfc1035label
 
 import (
-	"fmt"
+	"errors"
 	"strings"
 )
 
-// This implements the compression from RFC 1035, section 4.1.4
-// https://tools.ietf.org/html/rfc1035
+// 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) {
 	var (
-		pos     = 0
-		domains = make([]string, 0)
-		label   = ""
+		pos, oldPos     int
+		labels          = make([]string, 0)
+		label           string
+		handlingPointer bool
 	)
+
 	for {
 		if pos >= len(buf) {
-			return domains, nil
+			break
 		}
 		length := int(buf[pos])
 		pos++
+		var chunk string
 		if length == 0 {
-			domains = append(domains, label)
+			labels = append(labels, label)
 			label = ""
+			if handlingPointer {
+				pos = oldPos
+				handlingPointer = false
+			}
+		} else if length&0xc0 == 0xc0 {
+			// compression pointer
+			if handlingPointer {
+				return nil, errors.New("rfc1035label: cannot handle nested pointers")
+			}
+			handlingPointer = true
+			if pos+1 > len(buf) {
+				return nil, errors.New("rfc1035label: pointer buffer too short")
+			}
+			off := int(buf[pos-1]&^0xc0)<<8 + int(buf[pos])
+			oldPos = pos + 1
+			pos = off
+		} else {
+			if pos+length > len(buf) {
+				return nil, errors.New("rfc1035label: buffer too short")
+			}
+			chunk = string(buf[pos : pos+length])
+			if label != "" {
+				label += "."
+			}
+			label += chunk
+			pos += length
 		}
-		if len(buf)-pos < length {
-			return nil, fmt.Errorf("DomainNamesFromBytes: invalid short label length")
-		}
-		if label != "" {
-			label += "."
-		}
-		label += string(buf[pos : pos+length])
-		pos += length
 	}
+	return labels, nil
 }
 
 // LabelToBytes encodes a label and returns a serialized stream of bytes
diff --git a/rfc1035label/label_test.go b/rfc1035label/label_test.go
index f91110b..3a69f2b 100644
--- a/rfc1035label/label_test.go
+++ b/rfc1035label/label_test.go
@@ -1,8 +1,9 @@
 package rfc1035label
 
 import (
-	"bytes"
 	"testing"
+
+	"github.com/stretchr/testify/require"
 )
 
 func TestLabelsFromBytes(t *testing.T) {
@@ -11,38 +12,27 @@ func TestLabelsFromBytes(t *testing.T) {
 		0x2, 'i', 't',
 		0x0,
 	})
-	if err != nil {
-		t.Fatal(err)
-	}
-	if len(labels) != 1 {
-		t.Fatalf("Invalid labels length. Expected: 1, got: %v", len(labels))
-	}
-	if labels[0] != "slackware.it" {
-		t.Fatalf("Invalid label. Expected: %v, got: %v'", "slackware.it", labels[0])
-	}
+	require.NoError(t, err)
+	require.Equal(t, 1, len(labels))
+	require.Equal(t, "slackware.it", labels[0])
 }
 
 func TestLabelsFromBytesZeroLength(t *testing.T) {
 	labels, err := LabelsFromBytes([]byte{})
-	if err != nil {
-		t.Fatal(err)
-	}
-	if len(labels) != 0 {
-		t.Fatalf("Invalid labels length. Expected: 0, got: %v", len(labels))
-	}
+	require.NoError(t, err)
+	require.Equal(t, 0, len(labels))
 }
 
 func TestLabelsFromBytesInvalidLength(t *testing.T) {
+	labels, err := LabelsFromBytes([]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
-	if err == nil {
-		t.Fatal("Expected error, got nil")
-	}
-	if len(labels) != 0 {
-		t.Fatalf("Invalid labels length. Expected: 0, got: %v", len(labels))
-	}
-	if labels != nil {
-		t.Fatalf("Invalid label. Expected nil, got %v", labels)
-	}
+	require.Error(t, err)
+	require.Equal(t, 0, len(labels))
 }
 
 func TestLabelToBytes(t *testing.T) {
@@ -52,15 +42,85 @@ func TestLabelToBytes(t *testing.T) {
 		0x2, 'i', 't',
 		0x0,
 	}
-	if !bytes.Equal(encodedLabel, expected) {
-		t.Fatalf("Invalid label. Expected: %v, got: %v", expected, encodedLabel)
+	require.Equal(t, expected, encodedLabel)
+}
+
+func TestLabelsToBytes(t *testing.T) {
+	expected := []byte{
+		9, 's', 'l', 'a', 'c', 'k', 'w', 'a', 'r', 'e',
+		2, 'i', 't',
+		0,
+		9, 'i', 'n', 's', 'o', 'm', 'n', 'i', 'a', 'c',
+		9, 's', 'l', 'a', 'c', 'k', 'w', 'a', 'r', 'e',
+		2, 'i', 't',
+		0,
 	}
+	encodedLabels := LabelsToBytes([]string{"slackware.it", "insomniac.slackware.it"})
+	require.Equal(t, expected, encodedLabels)
 }
 
 func TestLabelToBytesZeroLength(t *testing.T) {
 	encodedLabel := LabelToBytes("")
 	expected := []byte{0}
-	if !bytes.Equal(encodedLabel, expected) {
-		t.Fatalf("Invalid label. Expected: %v, got: %v", expected, encodedLabel)
+	require.Equal(t, expected, encodedLabel)
+}
+
+func TestCompressedLabel(t *testing.T) {
+	data := []byte{
+		// slackware.it
+		9, 's', 'l', 'a', 'c', 'k', 'w', 'a', 'r', 'e',
+		2, 'i', 't',
+		0,
+		// insomniac.slackware.it
+		9, 'i', 'n', 's', 'o', 'm', 'n', 'i', 'a', 'c',
+		192, 0,
+		// mail.systemboot.org
+		4, 'm', 'a', 'i', 'l',
+		10, 's', 'y', 's', 't', 'e', 'm', 'b', 'o', 'o', 't',
+		3, 'o', 'r', 'g',
+		0,
+		// systemboot.org
+		192, 31,
+	}
+	expected := []string{
+		"slackware.it",
+		"insomniac.slackware.it",
+		"mail.systemboot.org",
+		"systemboot.org",
+	}
+
+	labels, err := LabelsFromBytes(data)
+	require.NoError(t, err)
+	require.Equal(t, expected, labels)
+}
+
+func TestShortCompressedLabel(t *testing.T) {
+	data := []byte{
+		// slackware.it
+		9, 's', 'l', 'a', 'c', 'k', 'w', 'a', 'r', 'e',
+		2, 'i', 't',
+		0,
+		// insomniac.slackware.it
+		9, 'i', 'n', 's', 'o', 'm', 'n', 'i', 'a', 'c',
+		192,
+	}
+
+	_, err := LabelsFromBytes(data)
+	require.Error(t, err)
+}
+
+func TestNestedCompressedLabel(t *testing.T) {
+	data := []byte{
+		// it
+		3, 'i', 't',
+		0,
+		// slackware.it
+		9, 's', 'l', 'a', 'c', 'k', 'w', 'a', 'r', 'e',
+		192, 0,
+		// insomniac.slackware.it
+		9, 'i', 'n', 's', 'o', 'm', 'n', 'i', 'a', 'c',
+		192, 5,
 	}
+	_, err := LabelsFromBytes(data)
+	require.Error(t, err)
 }