From d80c76f7464148a22fb56ce894e46530b697213a Mon Sep 17 00:00:00 2001
From: Alberto Bertogli <albertito@blitiri.com.ar>
Date: Mon, 19 Feb 2018 22:33:06 +0000
Subject: [PATCH] envelope: Handle zero-length keys in AddHeader

We don't expect the AddHeader function to be given an empty key;
however, if that were to happen, it currently crashes.

This patch fixes the bug, while also adding tests for that and other
similar cases.
---
 internal/envelope/envelope.go      | 10 ++++++----
 internal/envelope/envelope_test.go | 26 ++++++++++++++++++++++++++
 2 files changed, 32 insertions(+), 4 deletions(-)

diff --git a/internal/envelope/envelope.go b/internal/envelope/envelope.go
index d0eb1d2..c504ad0 100644
--- a/internal/envelope/envelope.go
+++ b/internal/envelope/envelope.go
@@ -39,11 +39,13 @@ func DomainIn(addr string, locals *set.String) bool {
 }
 
 func AddHeader(data []byte, k, v string) []byte {
-	// If the value contains newlines, indent them properly.
-	if v[len(v)-1] == '\n' {
-		v = v[:len(v)-1]
+	if len(v) > 0 {
+		// If the value contains newlines, indent them properly.
+		if v[len(v)-1] == '\n' {
+			v = v[:len(v)-1]
+		}
+		v = strings.Replace(v, "\n", "\n\t", -1)
 	}
-	v = strings.Replace(v, "\n", "\n\t", -1)
 
 	header := []byte(fmt.Sprintf("%s: %s\n", k, v))
 	return append(header, data...)
diff --git a/internal/envelope/envelope_test.go b/internal/envelope/envelope_test.go
index e8746a7..661548d 100644
--- a/internal/envelope/envelope_test.go
+++ b/internal/envelope/envelope_test.go
@@ -41,3 +41,29 @@ func TestDomainIn(t *testing.T) {
 		}
 	}
 }
+
+func TestAddHeader(t *testing.T) {
+	cases := []struct {
+		data, k, v, expected string
+	}{
+		{"", "Key", "Value", "Key: Value\n"},
+		{"data", "Key", "Value", "Key: Value\ndata"},
+		{"data", "Key", "Value\n", "Key: Value\ndata"},
+		{"data", "Key", "L1\nL2", "Key: L1\n\tL2\ndata"},
+		{"data", "Key", "L1\nL2\n", "Key: L1\n\tL2\ndata"},
+
+		// Degenerate cases: we don't expect to ever produce these, and the
+		// output is admittedly not nice, but they should at least not cause
+		// chasquid to crash.
+		{"data", "Key", "", "Key: \ndata"},
+		{"data", "", "", ": \ndata"},
+		{"", "", "", ": \n"},
+	}
+	for i, c := range cases {
+		got := string(AddHeader([]byte(c.data), c.k, c.v))
+		if got != c.expected {
+			t.Errorf("%d (%q -> %q): expected %q, got %q",
+				i, c.k, c.v, c.expected, got)
+		}
+	}
+}