3 files changed, 58 insertions, 27 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);