From 5367a32ee9ff484acf942bdfd58b7fcdd42046c1 Mon Sep 17 00:00:00 2001
From: Abhishek Shah <abshah@google.com>
Date: Thu, 29 Oct 2015 14:45:29 -0700
Subject: [PATCH] Read Iptables-save output in a more-memory-efficient way

---
 pkg/proxy/iptables/proxier.go      | 71 +++++++++++++++++++++++++-----
 pkg/proxy/iptables/proxier_test.go | 33 ++++++++++++++
 2 files changed, 94 insertions(+), 10 deletions(-)

diff --git a/pkg/proxy/iptables/proxier.go b/pkg/proxy/iptables/proxier.go
index a45f2dcd441..b7dfe8705d4 100644
--- a/pkg/proxy/iptables/proxier.go
+++ b/pkg/proxy/iptables/proxier.go
@@ -765,33 +765,84 @@ func makeChainLine(chain utiliptables.Chain) string {
 // getChainLines parses a table's iptables-save data to find chains in the table.
 // It returns a map of iptables.Chain to string where the string is the chain line from the save (with counters etc).
 func getChainLines(table utiliptables.Table, save []byte) map[utiliptables.Chain]string {
-	// get lines
-	lines := strings.Split(string(save), "\n")
 	chainsMap := make(map[utiliptables.Chain]string)
 	tablePrefix := "*" + string(table)
-	lineNum := 0
+	readIndex := 0
 	// find beginning of table
-	for ; lineNum < len(lines); lineNum++ {
-		if strings.HasPrefix(strings.TrimSpace(lines[lineNum]), tablePrefix) {
-			lineNum++
+	for readIndex < len(save) {
+		line, n := readLine(readIndex, save)
+		readIndex = n
+		if strings.HasPrefix(line, tablePrefix) {
 			break
 		}
 	}
 	// parse table lines
-	for ; lineNum < len(lines); lineNum++ {
-		line := strings.TrimSpace(lines[lineNum])
+	for readIndex < len(save) {
+		line, n := readLine(readIndex, save)
+		readIndex = n
+		if len(line) == 0 {
+			continue
+		}
 		if strings.HasPrefix(line, "COMMIT") || strings.HasPrefix(line, "*") {
 			break
-		} else if len(line) == 0 || strings.HasPrefix(line, "#") {
+		} else if strings.HasPrefix(line, "#") {
 			continue
 		} else if strings.HasPrefix(line, ":") && len(line) > 1 {
 			chain := utiliptables.Chain(strings.SplitN(line[1:], " ", 2)[0])
-			chainsMap[chain] = lines[lineNum]
+			chainsMap[chain] = line
 		}
 	}
 	return chainsMap
 }
 
+func readLine(readIndex int, byteArray []byte) (string, int) {
+	currentReadIndex := readIndex
+
+	// consume left spaces
+	for currentReadIndex < len(byteArray) {
+		if byteArray[currentReadIndex] == ' ' {
+			currentReadIndex++
+		} else {
+			break
+		}
+	}
+
+	// leftTrimIndex stores the left index of the line after the line is left-trimmed
+	leftTrimIndex := currentReadIndex
+
+	// rightTrimIndex stores the right index of the line after the line is right-trimmed
+	// it is set to -1 since the correct value has not yet been determined.
+	rightTrimIndex := -1
+
+	for ; currentReadIndex < len(byteArray); currentReadIndex++ {
+		if byteArray[currentReadIndex] == ' ' {
+			// set rightTrimIndex
+			if rightTrimIndex == -1 {
+				rightTrimIndex = currentReadIndex
+			}
+		} else if (byteArray[currentReadIndex] == '\n') || (currentReadIndex == (len(byteArray) - 1)) {
+			// end of line or byte buffer is reached
+			if currentReadIndex <= leftTrimIndex {
+				return "", currentReadIndex + 1
+			}
+			// set the rightTrimIndex
+			if rightTrimIndex == -1 {
+				rightTrimIndex = currentReadIndex
+				if currentReadIndex == (len(byteArray)-1) && (byteArray[currentReadIndex] != '\n') {
+					// ensure that the last character is part of the returned string,
+					// unless the last character is '\n'
+					rightTrimIndex = currentReadIndex + 1
+				}
+			}
+			return string(byteArray[leftTrimIndex:rightTrimIndex]), currentReadIndex + 1
+		} else {
+			// unset rightTrimIndex
+			rightTrimIndex = -1
+		}
+	}
+	return "", currentReadIndex
+}
+
 func isLocalIP(ip string) (bool, error) {
 	addrs, err := net.InterfaceAddrs()
 	if err != nil {
diff --git a/pkg/proxy/iptables/proxier_test.go b/pkg/proxy/iptables/proxier_test.go
index d0ccc07f7e9..3db6cb53140 100644
--- a/pkg/proxy/iptables/proxier_test.go
+++ b/pkg/proxy/iptables/proxier_test.go
@@ -35,6 +35,39 @@ func checkAllLines(t *testing.T, table utiliptables.Table, save []byte, expected
 	}
 }
 
+func TestReadLinesFromByteBuffer(t *testing.T) {
+
+	testFn := func(byteArray []byte, expected []string) {
+		index := 0
+		readIndex := 0
+		for ; readIndex < len(byteArray); index++ {
+			line, n := readLine(readIndex, byteArray)
+			readIndex = n
+			if expected[index] != line {
+				t.Errorf("expected:%q, actual:%q", expected[index], line)
+			}
+		} // for
+		if readIndex < len(byteArray) {
+			t.Errorf("Byte buffer was only partially read. Buffer length is:%d, readIndex is:%d", len(byteArray), readIndex)
+		}
+		if index < len(expected) {
+			t.Errorf("All expected strings were not compared. expected arr length:%d, matched count:%d", len(expected), index-1)
+		}
+	}
+
+	byteArray1 := []byte("\n  Line 1  \n\n\n L ine4  \nLine 5 \n \n")
+	expected1 := []string{"", "Line 1", "", "", "L ine4", "Line 5", ""}
+	testFn(byteArray1, expected1)
+
+	byteArray1 = []byte("")
+	expected1 = []string{}
+	testFn(byteArray1, expected1)
+
+	byteArray1 = []byte("\n\n")
+	expected1 = []string{"", ""}
+	testFn(byteArray1, expected1)
+}
+
 func TestgetChainLines(t *testing.T) {
 	iptables_save := `# Generated by iptables-save v1.4.7 on Wed Oct 29 14:56:01 2014
 	*nat