check_tcp: Properly deal will partial recv(3)s

The np_expect_match() function now returns one of three possible states
instead of just TRUE or FALSE:

- NP_MATCH_SUCCESS
- NP_MATCH_FAILURE
- NP_MATCH_RETRY

The NP_MATCH_RETRY state indicates that matching might succeed if
np_expect_match() is called with a longer input string.  This allows
check_tcp to decide whether it makes sense to wait for additional data
from the server.

4 files changed, 74 insertions, 38 deletions

diff --git a/lib/tests/test_tcp.c b/lib/tests/test_tcp.c
index 8e9d43c8..ae6dc1f4 100644
--- a/lib/tests/test_tcp.c
+++ b/lib/tests/test_tcp.c
@@ -25,7 +25,7 @@ main (int argc, char **argv)
 {
 	char** server_expect;
 	int server_expect_count = 3;
-	plan_tests(8);
+	plan_tests(9);
 
 	server_expect = malloc(sizeof(char*) * server_expect_count);
 
@@ -33,21 +33,23 @@ main (int argc, char **argv)
 	server_expect[1] = strdup("bb");
 	server_expect[2] = strdup("CC");
 	
-	ok(np_expect_match("AA bb CC XX", server_expect, server_expect_count, NP_MATCH_EXACT) == TRUE,
+	ok(np_expect_match("AA bb CC XX", server_expect, server_expect_count, NP_MATCH_EXACT) == NP_MATCH_SUCCESS,
 	   "Test matching any string at the beginning (first expect string)");
-	ok(np_expect_match("bb AA CC XX", server_expect, server_expect_count, NP_MATCH_EXACT) == TRUE,
+	ok(np_expect_match("bb AA CC XX", server_expect, server_expect_count, NP_MATCH_EXACT) == NP_MATCH_SUCCESS,
 	   "Test matching any string at the beginning (second expect string)");
-	ok(np_expect_match("XX bb AA CC XX", server_expect, server_expect_count, NP_MATCH_EXACT) == FALSE,
+	ok(np_expect_match("b", server_expect, server_expect_count, NP_MATCH_EXACT) == NP_MATCH_RETRY,
+	   "Test matching any string at the beginning (substring match)");
+	ok(np_expect_match("XX bb AA CC XX", server_expect, server_expect_count, NP_MATCH_EXACT) == NP_MATCH_FAILURE,
 	   "Test with strings not matching at the beginning");
-	ok(np_expect_match("XX CC XX", server_expect, server_expect_count, NP_MATCH_EXACT) == FALSE,
+	ok(np_expect_match("XX CC XX", server_expect, server_expect_count, NP_MATCH_EXACT) == NP_MATCH_FAILURE,
 	   "Test matching any string");
-	ok(np_expect_match("XX", server_expect, server_expect_count, 0) == FALSE,
+	ok(np_expect_match("XX", server_expect, server_expect_count, 0) == NP_MATCH_RETRY,
 	   "Test not matching any string");
-	ok(np_expect_match("XX AA bb CC XX", server_expect, server_expect_count, NP_MATCH_ALL) == TRUE,
+	ok(np_expect_match("XX AA bb CC XX", server_expect, server_expect_count, NP_MATCH_ALL) == NP_MATCH_SUCCESS,
 	   "Test matching all strings");
-	ok(np_expect_match("XX bb CC XX", server_expect, server_expect_count, NP_MATCH_ALL) == FALSE,
+	ok(np_expect_match("XX bb CC XX", server_expect, server_expect_count, NP_MATCH_ALL) == NP_MATCH_RETRY,
 	   "Test not matching all strings");
-	ok(np_expect_match("XX XX", server_expect, server_expect_count, NP_MATCH_ALL) == FALSE,
+	ok(np_expect_match("XX XX", server_expect, server_expect_count, NP_MATCH_ALL) == NP_MATCH_RETRY,
 	   "Test not matching any string (testing all)");
 	 
 
diff --git a/lib/utils_tcp.c b/lib/utils_tcp.c
index cf67b116..497a1701 100644
--- a/lib/utils_tcp.c
+++ b/lib/utils_tcp.c
@@ -3,7 +3,7 @@
 * Library for check_tcp
 * 
 * License: GPL
-* Copyright (c) 1999-2007 Nagios Plugins Development Team
+* Copyright (c) 1999-2013 Nagios Plugins Development Team
 * 
 * Description:
 * 
@@ -29,29 +29,44 @@
 #include "common.h"
 #include "utils_tcp.h"
 
-int
+#define VERBOSE(message)                        \
+	do {                                    \
+		if (flags & NP_MATCH_VERBOSE)   \
+			puts(message);          \
+	} while (0)
+
+enum np_match_result
 np_expect_match(char* status, char** server_expect, int expect_count, int flags)
 {
-	int match = 0;
-	int i;
+	int i, match = 0, partial = 0;
 	for (i = 0; i < expect_count; i++) {
 		if (flags & NP_MATCH_VERBOSE)
 			printf ("looking for [%s] %s [%s]\n", server_expect[i],
 					(flags & NP_MATCH_EXACT) ? "in beginning of" : "anywhere in",
 					status);
 
-		if ((flags & NP_MATCH_EXACT &&
-			!strncmp(status, server_expect[i], strlen(server_expect[i]))) ||
-			(!(flags & NP_MATCH_EXACT) && strstr(status, server_expect[i])))
-		{
-			if(flags & NP_MATCH_VERBOSE) puts("found it");
-			match += 1;
-		} else
-			if(flags & NP_MATCH_VERBOSE) puts("couldn't find it");
+		if (flags & NP_MATCH_EXACT) {
+			if (strncmp(status, server_expect[i], strlen(server_expect[i])) == 0) {
+				VERBOSE("found it");
+				match++;
+				continue;
+			} else if (strncmp(status, server_expect[i], strlen(status)) == 0) {
+				VERBOSE("found a substring");
+				partial++;
+				continue;
+			}
+		} else if (strstr(status, server_expect[i]) != NULL) {
+				VERBOSE("found it");
+				match++;
+				continue;
+		}
+		VERBOSE("couldn't find it");
 	}
 	if ((flags & NP_MATCH_ALL && match == expect_count) ||
-		(!(flags & NP_MATCH_ALL) && match >= 1)) {
-		return TRUE;
-	} else
-		return FALSE;
+	    (!(flags & NP_MATCH_ALL) && match >= 1))
+		return NP_MATCH_SUCCESS;
+	else if (partial > 0 || !(flags & NP_MATCH_EXACT))
+		return NP_MATCH_RETRY;
+	else
+		return NP_MATCH_FAILURE;
 }
diff --git a/lib/utils_tcp.h b/lib/utils_tcp.h
index 34b771d6..0328a9cf 100644
--- a/lib/utils_tcp.h
+++ b/lib/utils_tcp.h
@@ -4,5 +4,19 @@
 #define NP_MATCH_EXACT          0x2
 #define NP_MATCH_VERBOSE        0x4
 
-int np_expect_match(char* status, char** server_expect, int server_expect_count,
-                    int flags);
+/*
+ * The NP_MATCH_RETRY state indicates that matching might succeed if
+ * np_expect_match() is called with a longer input string.  This allows the
+ * caller to decide whether it makes sense to wait for additional data from the
+ * server.
+ */
+enum np_match_result {
+	NP_MATCH_FAILURE,
+	NP_MATCH_SUCCESS,
+	NP_MATCH_RETRY
+};
+
+enum np_match_result np_expect_match(char *status,
+                                     char **server_expect,
+                                     int server_expect_count,
+                                     int flags);
diff --git a/plugins/check_tcp.c b/plugins/check_tcp.c
index e8d7ec68..517b6b5d 100644
--- a/plugins/check_tcp.c
+++ b/plugins/check_tcp.c
@@ -3,7 +3,7 @@
 * Nagios check_tcp plugin
 *
 * License: GPL
-* Copyright (c) 1999-2008 Nagios Plugins Development Team
+* Copyright (c) 1999-2013 Nagios Plugins Development Team
 *
 * Description:
 *
@@ -277,25 +277,30 @@ main (int argc, char **argv)
 			status = realloc(status, len + i + 1);
 			memcpy(&status[len], buffer, i);
 			len += i;
+			status[len] = '\0';
 
 			/* stop reading if user-forced */
 			if (maxbytes && len >= maxbytes)
 				break;
+
+			if ((match = np_expect_match(status,
+			    server_expect,
+			    server_expect_count,
+			    match_flags)) != NP_MATCH_RETRY)
+				break;
 		}
 
 		/* no data when expected, so return critical */
 		if (len == 0)
 			die (STATE_CRITICAL, _("No data received from host\n"));
 
-		/* force null-termination and strip whitespace from end of output */
-		status[len--] = '\0';
 		/* print raw output if we're debugging */
 		if(flags & FLAG_VERBOSE)
 			printf("received %d bytes from host\n#-raw-recv-------#\n%s\n#-raw-recv-------#\n",
 			       (int)len + 1, status);
-		while(isspace(status[len])) status[len--] = '\0';
-
-		match = np_expect_match(status, server_expect, server_expect_count, match_flags);
+		/* strip whitespace from end of output */
+		while(--len > 0 && isspace(status[len]))
+			status[len] = '\0';
 	}
 
 	if (server_quit != NULL) {
@@ -315,7 +320,7 @@ main (int argc, char **argv)
 		result = STATE_WARNING;
 
 	/* did we get the response we hoped? */
-	if(match == FALSE && result != STATE_CRITICAL)
+	if(match != NP_MATCH_SUCCESS && result != STATE_CRITICAL)
 		result = expect_mismatch_state;
 
 	/* reset the alarm */
@@ -326,10 +331,10 @@ main (int argc, char **argv)
 	 * the response we were looking for. if-else */
 	printf("%s %s - ", SERVICE, state_text(result));
 
-	if(match == FALSE && len && !(flags & FLAG_HIDE_OUTPUT))
+	if(match != NP_MATCH_SUCCESS && len && !(flags & FLAG_HIDE_OUTPUT))
 		printf("Unexpected response from host/socket: %s", status);
 	else {
-		if(match == FALSE)
+		if(match != NP_MATCH_SUCCESS)
 			printf("Unexpected response from host/socket on ");
 		else
 			printf("%.3f second response time on ", elapsed_time);
@@ -339,13 +344,13 @@ main (int argc, char **argv)
 			printf("socket %s", server_address);
 	}
 
-	if (match != FALSE && !(flags & FLAG_HIDE_OUTPUT) && len)
+	if (match == NP_MATCH_SUCCESS && !(flags & FLAG_HIDE_OUTPUT) && len)
 		printf (" [%s]", status);
 
 	/* perf-data doesn't apply when server doesn't talk properly,
 	 * so print all zeroes on warn and crit. Use fperfdata since
 	 * localisation settings can make different outputs */
-	if(match == FALSE)
+	if(match != NP_MATCH_SUCCESS)
 		printf ("|%s",
 				fperfdata ("time", elapsed_time, "s",
 				(flags & FLAG_TIME_WARN ? TRUE : FALSE), 0,