From 0962b92c623b0cb018c07c674a5e2a57384f2541 Mon Sep 17 00:00:00 2001
From: Corbin Phelps <corbin.phelps@bluemedora.com>
Date: Mon, 1 Oct 2018 09:07:42 -0400
Subject: [PATCH] Added hard coded topic limit (#34)

---
 CHANGELOG.md      |  4 ++++
 src/kafka.go      | 19 ++++++++++++++++++-
 src/kafka_test.go | 38 ++++++++++++++++++++++++++++++++++++++
 3 files changed, 60 insertions(+), 1 deletion(-)
 create mode 100644 src/kafka_test.go

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 757e5103..4b1d504b 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -5,6 +5,10 @@ All notable changes to this project will be documented in this file.
 The format is based on [Keep a Changelog](http://keepachangelog.com/)
 and this project adheres to [Semantic Versioning](http://semver.org/).
 
+## 0.2.2 - 2018-10-01
+### Added
+- Hardcoded limit for topic collection
+
 ## 0.2.1 - 2018-09-25
 ### Changed
 - Source Type on all metrics with `attr=Count` to `Rate`
diff --git a/src/kafka.go b/src/kafka.go
index 87b915ef..da54bbc5 100644
--- a/src/kafka.go
+++ b/src/kafka.go
@@ -16,7 +16,7 @@ import (
 
 const (
 	integrationName    = "com.newrelic.kafka"
-	integrationVersion = "0.2.1"
+	integrationVersion = "0.2.2"
 )
 
 func main() {
@@ -57,6 +57,9 @@ func coreCollection(zkConn zookeeper.Connection, kafkaIntegration *integration.I
 	collectedTopics, err := tc.GetTopics(zkConn)
 	ExitOnErr(err)
 
+	// Enforce hard limits on Topics
+	collectedTopics = enforceTopicLimit(collectedTopics)
+
 	// Setup wait group
 	var wg sync.WaitGroup
 
@@ -90,3 +93,17 @@ func ExitOnErr(err error) {
 		os.Exit(1)
 	}
 }
+
+// maxTopics is the maximum amount of Topics that can be collect.
+// If there are more than this number of Topics then collection of
+// Topics will fail.
+const maxTopics = 300
+
+func enforceTopicLimit(collectedTopics []string) []string {
+	if length := len(collectedTopics); length > maxTopics {
+		log.Error("There are %d topics in collection, the maximum amount of topics to collect is %d. Use the topic whitelist configuration parameter to limit collection size.", length, maxTopics)
+		return make([]string, 0)
+	}
+
+	return collectedTopics
+}
diff --git a/src/kafka_test.go b/src/kafka_test.go
new file mode 100644
index 00000000..327a37b3
--- /dev/null
+++ b/src/kafka_test.go
@@ -0,0 +1,38 @@
+package main
+
+import (
+	"reflect"
+	"testing"
+)
+
+func Test_enforceTopicLimit(t *testing.T) {
+	overLimit := make([]string, maxTopics+1)
+
+	for i := 0; i < maxTopics+1; i++ {
+		overLimit[i] = "topic"
+	}
+
+	testCases := []struct {
+		name   string
+		topics []string
+		want   []string
+	}{
+		{
+			"Over Limit",
+			overLimit,
+			[]string{},
+		},
+		{
+			"Under Limit",
+			[]string{"topic"},
+			[]string{"topic"},
+		},
+	}
+
+	for _, tc := range testCases {
+		out := enforceTopicLimit(tc.topics)
+		if !reflect.DeepEqual(out, tc.want) {
+			t.Errorf("Test Case %s Failed: Expected '%+v' got '%+v'", tc.name, tc.want, out)
+		}
+	}
+}