[1/2] iphlpapi: Add support for the listener and connection classes in GetExtendedTcpTable.

[Hans Leidekker]

---
 dlls/iphlpapi/iphlpapi_main.c  |    4 +++-
 dlls/iphlpapi/ipstats.c        |   38 ++++++++++++++++++++++++++++++++++++--
 dlls/iphlpapi/tests/iphlpapi.c |   36 ++++++++++++++++++++++++++++++++++++
 3 files changed, 75 insertions(+), 3 deletions(-)

diff --git a/dlls/iphlpapi/iphlpapi_main.c b/dlls/iphlpapi/iphlpapi_main.c
index aeb91c1..e970162 100644
--- a/dlls/iphlpapi/iphlpapi_main.c
+++ b/dlls/iphlpapi/iphlpapi_main.c
@@ -1860,7 +1860,9 @@ DWORD WINAPI GetExtendedTcpTable(PVOID pTcpTable, PDWORD pdwSize, BOOL bOrder,
     if (!pdwSize) return ERROR_INVALID_PARAMETER;
 
     if (ulAf != AF_INET ||
-        (TableClass != TCP_TABLE_BASIC_ALL && TableClass != TCP_TABLE_OWNER_PID_ALL))
+        TableClass == TCP_TABLE_OWNER_MODULE_LISTENER ||
+        TableClass == TCP_TABLE_OWNER_MODULE_CONNECTIONS ||
+        TableClass == TCP_TABLE_OWNER_MODULE_ALL)
     {
         FIXME("ulAf = %u, TableClass = %u not supported\n", ulAf, TableClass);
         return ERROR_NOT_SUPPORTED;
diff --git a/dlls/iphlpapi/ipstats.c b/dlls/iphlpapi/ipstats.c
index 1f52fcc..87651d3 100644
--- a/dlls/iphlpapi/ipstats.c
+++ b/dlls/iphlpapi/ipstats.c
@@ -1981,6 +1981,33 @@ static unsigned int find_owning_pid( struct pid_map *map, unsigned int num_entri
 #endif
 }
 
+static BOOL match_class( TCP_TABLE_CLASS class, MIB_TCP_STATE state )
+{
+    switch (class)
+    {
+    case TCP_TABLE_BASIC_ALL:
+    case TCP_TABLE_OWNER_PID_ALL:
+    case TCP_TABLE_OWNER_MODULE_ALL:
+        return TRUE;
+
+    case TCP_TABLE_BASIC_LISTENER:
+    case TCP_TABLE_OWNER_PID_LISTENER:
+    case TCP_TABLE_OWNER_MODULE_LISTENER:
+        if (state == MIB_TCP_STATE_LISTEN) return TRUE;
+        return FALSE;
+
+    case TCP_TABLE_BASIC_CONNECTIONS:
+    case TCP_TABLE_OWNER_PID_CONNECTIONS:
+    case TCP_TABLE_OWNER_MODULE_CONNECTIONS:
+        if (state == MIB_TCP_STATE_ESTAB) return TRUE;
+        return FALSE;
+
+    default:
+        ERR( "unhandled class %u\n", class );
+        return FALSE;
+    }
+}
+
 DWORD build_tcp_table( TCP_TABLE_CLASS class, void **tablep, BOOL order, HANDLE heap, DWORD flags,
                        DWORD *size )
 {
@@ -2007,7 +2034,9 @@ DWORD build_tcp_table( TCP_TABLE_CLASS class, void **tablep, BOOL order, HANDLE
             unsigned int dummy, num_entries = 0;
             int inode;
 
-            if (class == TCP_TABLE_OWNER_PID_ALL) map = get_pid_map( &num_entries );
+            if (class == TCP_TABLE_OWNER_PID_ALL ||
+                class == TCP_TABLE_OWNER_PID_LISTENER ||
+                class == TCP_TABLE_OWNER_PID_CONNECTIONS) map = get_pid_map( &num_entries );
 
             /* skip header line */
             ptr = fgets(buf, sizeof(buf), fp);
@@ -2020,7 +2049,10 @@ DWORD build_tcp_table( TCP_TABLE_CLASS class, void **tablep, BOOL order, HANDLE
                 row.dwLocalPort = htons( row.dwLocalPort );
                 row.dwRemotePort = htons( row.dwRemotePort );
                 row.dwState = TCPStateToMIBState( row.dwState );
-                if (class == TCP_TABLE_OWNER_PID_ALL)
+                if (!match_class( class, row.dwState )) continue;
+                if (class == TCP_TABLE_OWNER_PID_ALL ||
+                    class == TCP_TABLE_OWNER_PID_LISTENER ||
+                    class == TCP_TABLE_OWNER_PID_CONNECTIONS)
                     row.dwOwningPid = find_owning_pid( map, num_entries, inode );
 
                 if (!(table = append_tcp_row( class, heap, flags, table, &count, &row, row_size )))
@@ -2048,6 +2080,7 @@ DWORD build_tcp_table( TCP_TABLE_CLASS class, void **tablep, BOOL order, HANDLE
                     row.dwRemoteAddr = entry->tcpConnRemAddress;
                     row.dwRemotePort = htons( entry->tcpConnRemPort );
                     row.dwState = entry->tcpConnState;
+                    if (!match_class( class, row.dwState )) continue;
                     if (!(table = append_tcp_row( class, heap, flags, table, &count, &row, row_size )))
                         break;
                 }
@@ -2128,6 +2161,7 @@ DWORD build_tcp_table( TCP_TABLE_CLASS class, void **tablep, BOOL order, HANDLE
             row.dwRemoteAddr = pINData->inp_faddr.s_addr;
             row.dwRemotePort = pINData->inp_fport;
             row.dwState = TCPStateToMIBState (pTCPData->t_state);
+            if (!match_class( class, row.dwState )) continue;
             if (!(table = append_tcp_row( class, heap, flags, table, &count, &row, row_size )))
                 break;
         }
diff --git a/dlls/iphlpapi/tests/iphlpapi.c b/dlls/iphlpapi/tests/iphlpapi.c
index 26bf26e..60616e1 100644
--- a/dlls/iphlpapi/tests/iphlpapi.c
+++ b/dlls/iphlpapi/tests/iphlpapi.c
@@ -1212,6 +1212,24 @@ static void test_GetExtendedTcpTable(void)
     HeapFree( GetProcessHeap(), 0, table );
 
     size = 0;
+    ret = pGetExtendedTcpTable( NULL, &size, TRUE, AF_INET, TCP_TABLE_BASIC_LISTENER, 0 );
+    ok( ret == ERROR_INSUFFICIENT_BUFFER, "got %u\n", ret );
+
+    table = HeapAlloc( GetProcessHeap(), 0, size );
+    ret = pGetExtendedTcpTable( table, &size, TRUE, AF_INET, TCP_TABLE_BASIC_LISTENER, 0 );
+    ok( ret == ERROR_SUCCESS, "got %u\n", ret );
+    HeapFree( GetProcessHeap(), 0, table );
+
+    size = 0;
+    ret = pGetExtendedTcpTable( NULL, &size, TRUE, AF_INET, TCP_TABLE_BASIC_CONNECTIONS, 0 );
+    ok( ret == ERROR_INSUFFICIENT_BUFFER, "got %u\n", ret );
+
+    table = HeapAlloc( GetProcessHeap(), 0, size );
+    ret = pGetExtendedTcpTable( table, &size, TRUE, AF_INET, TCP_TABLE_BASIC_CONNECTIONS, 0 );
+    ok( ret == ERROR_SUCCESS, "got %u\n", ret );
+    HeapFree( GetProcessHeap(), 0, table );
+
+    size = 0;
     ret = pGetExtendedTcpTable( NULL, &size, TRUE, AF_INET, TCP_TABLE_OWNER_PID_ALL, 0 );
     ok( ret == ERROR_INSUFFICIENT_BUFFER, "got %u\n", ret );
 
@@ -1219,6 +1237,24 @@ static void test_GetExtendedTcpTable(void)
     ret = pGetExtendedTcpTable( table_pid, &size, TRUE, AF_INET, TCP_TABLE_OWNER_PID_ALL, 0 );
     ok( ret == ERROR_SUCCESS, "got %u\n", ret );
     HeapFree( GetProcessHeap(), 0, table_pid );
+
+    size = 0;
+    ret = pGetExtendedTcpTable( NULL, &size, TRUE, AF_INET, TCP_TABLE_OWNER_PID_LISTENER, 0 );
+    ok( ret == ERROR_INSUFFICIENT_BUFFER, "got %u\n", ret );
+
+    table_pid = HeapAlloc( GetProcessHeap(), 0, size );
+    ret = pGetExtendedTcpTable( table_pid, &size, TRUE, AF_INET, TCP_TABLE_OWNER_PID_LISTENER, 0 );
+    ok( ret == ERROR_SUCCESS, "got %u\n", ret );
+    HeapFree( GetProcessHeap(), 0, table_pid );
+
+    size = 0;
+    ret = pGetExtendedTcpTable( NULL, &size, TRUE, AF_INET, TCP_TABLE_OWNER_PID_CONNECTIONS, 0 );
+    ok( ret == ERROR_INSUFFICIENT_BUFFER, "got %u\n", ret );
+
+    table_pid = HeapAlloc( GetProcessHeap(), 0, size );
+    ret = pGetExtendedTcpTable( table_pid, &size, TRUE, AF_INET, TCP_TABLE_OWNER_PID_CONNECTIONS, 0 );
+    ok( ret == ERROR_SUCCESS, "got %u\n", ret );
+    HeapFree( GetProcessHeap(), 0, table_pid );
 }
 
 static void test_GetExtendedUdpTable(void)
-- 
1.7.10.4