3 files changed, 119 insertions, 7 deletions
diff --git a/internal/color/brush/brush.go b/internal/color/brush/brush.go
index 63d63d8..1400d4e 100644
--- a/internal/color/brush/brush.go
+++ b/internal/color/brush/brush.go
@@ -37,6 +37,13 @@ func paintSeverity(sb *strings.Builder, text string) bool {
 
 func paintRemote(sb *strings.Builder, line string) {
 	splitted := strings.SplitN(line, protocol.FieldDelimiter, 6)
+	if len(splitted) < 6 {
+		// Malformed or short frame (e.g. from an older/buggy server):
+		// fall back to the plain-text default branch instead of
+		// indexing out of range.
+		paintDefault(sb, line)
+		return
+	}
 
 	color.PaintWithAttr(sb, splitted[0],
 		config.Client.TermColors.Remote.RemoteFg,
@@ -105,6 +112,10 @@ func paintRemote(sb *strings.Builder, line string) {
 
 func paintClient(sb *strings.Builder, line string) {
 	splitted := strings.SplitN(line, protocol.FieldDelimiter, 3)
+	if len(splitted) < 3 {
+		paintDefault(sb, line)
+		return
+	}
 
 	color.PaintWithAttr(sb, splitted[0],
 		config.Client.TermColors.Client.ClientFg,
@@ -138,6 +149,10 @@ func paintClient(sb *strings.Builder, line string) {
 
 func paintServer(sb *strings.Builder, line string) {
 	splitted := strings.SplitN(line, protocol.FieldDelimiter, 3)
+	if len(splitted) < 3 {
+		paintDefault(sb, line)
+		return
+	}
 
 	color.PaintWithAttr(sb, splitted[0],
 		config.Client.TermColors.Server.ServerFg,
@@ -185,10 +200,17 @@ func Colorfy(line string) string {
 		paintServer(sb, line)
 
 	default:
-		color.PaintWithAttr(sb, line,
-			color.FgDefault,
-			color.BgDefault,
-			color.AttrNone)
+		paintDefault(sb, line)
 	}
 	return sb.String()
 }
+
+// paintDefault writes the line using the default (uncoloured) attributes.
+// It is the fallback used both by Colorfy's default branch and by the
+// paint* functions when the protocol frame is too short to decode safely.
+func paintDefault(sb *strings.Builder, line string) {
+	color.PaintWithAttr(sb, line,
+		color.FgDefault,
+		color.BgDefault,
+		color.AttrNone)
+}
diff --git a/internal/color/brush/brush_test.go b/internal/color/brush/brush_test.go
new file mode 100644
index 0000000..52f31c2
--- /dev/null
+++ b/internal/color/brush/brush_test.go
@@ -0,0 +1,90 @@
+package brush
+
+import (
+	"os"
+	"strings"
+	"testing"
+
+	"github.com/mimecast/dtail/internal/config"
+	"github.com/mimecast/dtail/internal/source"
+)
+
+// TestMain ensures config.Client is populated with defaults so Colorfy can
+// reach the colourised branches without nil-dereferencing TermColors fields.
+func TestMain(m *testing.M) {
+	config.Setup(source.Client, &config.Args{ConfigFile: "none"}, nil)
+	os.Exit(m.Run())
+}
+
+// TestColorfy_ShortFramesDoNotPanic feeds Colorfy with malformed or short
+// protocol frames (e.g. "REMOTE", "REMOTE|host", "CLIENT|host") which used
+// to cause an index-out-of-range panic in paintRemote/paintClient/paintServer
+// when they unconditionally indexed splitted[0..N] after SplitN.
+// For any short frame Colorfy must not panic and must return a non-empty
+// string that still contains the original text (either verbatim in the plain
+// fallback branch or embedded inside color escape sequences).
+func TestColorfy_ShortFramesDoNotPanic(t *testing.T) {
+	cases := []struct {
+		name string
+		in   string
+	}{
+		{"empty", ""},
+		{"plain_text", "hello world"},
+		{"remote_prefix_only", "REMOTE"},
+		{"remote_one_field", "REMOTE|host"},
+		{"remote_two_fields", "REMOTE|host|100"},
+		{"remote_three_fields", "REMOTE|host|100|1"},
+		{"remote_four_fields", "REMOTE|host|100|1|id"},
+		{"remote_full_frame", "REMOTE|host|100|1|id|hello from remote"},
+		{"client_prefix_only", "CLIENT"},
+		{"client_one_field", "CLIENT|host"},
+		{"client_full_frame", "CLIENT|host|hello from client"},
+		{"server_prefix_only", "SERVER"},
+		{"server_one_field", "SERVER|host"},
+		{"server_full_frame", "SERVER|host|hello from server"},
+	}
+
+	for _, tc := range cases {
+		tc := tc
+		t.Run(tc.name, func(t *testing.T) {
+			defer func() {
+				if r := recover(); r != nil {
+					t.Fatalf("Colorfy(%q) panicked: %v", tc.in, r)
+				}
+			}()
+
+			got := Colorfy(tc.in)
+
+			if tc.in == "" {
+				// An empty input must not suddenly produce any content.
+				if got != "" && !strings.Contains(got, "") {
+					t.Fatalf("Colorfy(\"\") = %q, want empty or trivial", got)
+				}
+				return
+			}
+
+			// The original payload should still be recoverable somewhere in
+			// the output so users can see what the server actually sent even
+			// when the frame is malformed.
+			if !strings.Contains(got, tc.in) && !containsAllFields(got, tc.in) {
+				t.Fatalf("Colorfy(%q) = %q, original payload not present", tc.in, got)
+			}
+		})
+	}
+}
+
+// containsAllFields checks whether every non-delimiter field of the input is
+// present in the output. This allows the colorised happy path (which
+// interleaves escape sequences between fields) to still satisfy the
+// "payload preserved" expectation.
+func containsAllFields(out, in string) bool {
+	for _, f := range strings.Split(in, "|") {
+		if f == "" {
+			continue
+		}
+		if !strings.Contains(out, f) {
+			return false
+		}
+	}
+	return true
+}
diff --git a/internal/color/color.go b/internal/color/color.go
index 9d0bc2e..aebc9e4 100644
--- a/internal/color/color.go
+++ b/internal/color/color.go
@@ -86,7 +86,7 @@ func ToFgColor(s string) (FgColor, error) {
 	case "default":
 		return FgDefault, nil
 	default:
-		return FgDefault, fmt.Errorf("unknown foreground text color '" + s + "'")
+		return FgDefault, fmt.Errorf("unknown foreground text color '%s'", s)
 	}
 }
 
@@ -113,7 +113,7 @@ func ToBgColor(s string) (BgColor, error) {
 	case "default":
 		return BgDefault, nil
 	default:
-		return BgDefault, fmt.Errorf("unknown background text color '" + s + "'")
+		return BgDefault, fmt.Errorf("unknown background text color '%s'", s)
 	}
 }
 
@@ -143,6 +143,6 @@ func ToAttribute(s string) (Attribute, error) {
 	case "":
 		return AttrNone, nil
 	default:
-		return AttrNone, fmt.Errorf("unknown text attribute '" + s + "'")
+		return AttrNone, fmt.Errorf("unknown text attribute '%s'", s)
 	}
 }