windows-nt/Source/XPSP1/NT/shell/ext/hnw/wizard/netenum.cpp
2020-09-26 16:20:57 +08:00

337 lines
11 KiB
C++

//
// NetEnum.cpp
//
// Functions to enumerate computers and/or shares on the network
//
#include "stdafx.h"
#include "NetEnum.h"
#include "NetUtil.h"
#include "Util.h"
static CNetEnum* g_pNetEnum = NULL;
//////////////////////////////////////////////////////////////////////////////
void InitNetEnum()
{
ASSERT(g_pNetEnum == NULL);
g_pNetEnum = new CNetEnum;
}
void TermNetEnum()
{
delete g_pNetEnum;
}
void EnumComputers(NETENUMCALLBACK pfnCallback, LPVOID pvCallbackParam)
{
if (g_pNetEnum)
g_pNetEnum->EnumComputers(pfnCallback, pvCallbackParam);
}
//////////////////////////////////////////////////////////////////////////////
CNetEnum::CNetEnum()
{
m_hThread = NULL;
InitializeCriticalSection(&m_cs);
m_bAbort = FALSE;
}
CNetEnum::~CNetEnum()
{
m_bAbort = TRUE;
// Wait for the thread to die
EnterCriticalSection(&m_cs);
HANDLE hThread = m_hThread;
m_hThread = NULL;
LeaveCriticalSection(&m_cs);
if (hThread != NULL)
{
WaitForSingleObject(hThread, INFINITE);
CloseHandle(hThread);
}
DeleteCriticalSection(&m_cs);
}
void CNetEnum::Abort()
{
EnterCriticalSection(&m_cs);
m_bAbort = TRUE;
LeaveCriticalSection(&m_cs);
}
void CNetEnum::EnumComputers(NETENUMCALLBACK pfnCallback, LPVOID pvCallbackParam)
{
EnumHelper(jtEnumComputers, pfnCallback, pvCallbackParam);
}
void CNetEnum::EnumNetPrinters(NETENUMCALLBACK pfnCallback, LPVOID pvCallbackParam)
{
EnumHelper(jtEnumPrinters, pfnCallback, pvCallbackParam);
}
void CNetEnum::EnumHelper(JOBTYPE eJobType, NETENUMCALLBACK pfnCallback, LPVOID pvCallbackParam)
{
EnterCriticalSection(&m_cs);
HANDLE hThread = m_hThread;
m_pfnCallback = pfnCallback;
m_pvCallbackParam = pvCallbackParam;
m_eJobType = eJobType;
m_bAbort = FALSE;
m_bNewJob = TRUE; // if thread in progress, tell it to start a new job
LeaveCriticalSection(&m_cs);
if (hThread == NULL)
{
DWORD dwThreadId;
m_hThread = CreateThread(NULL, 0, EnumThreadProc, this, CREATE_SUSPENDED, &dwThreadId);
if (m_hThread)
{
ResumeThread(m_hThread);
}
}
}
DWORD WINAPI CNetEnum::EnumThreadProc(LPVOID pvParam)
{
((CNetEnum*)pvParam)->EnumThreadProc();
return 0;
}
#ifdef _DEBUG
void TraceNetResource(const NETRESOURCE* pNetRes)
{
DWORD dwScope = pNetRes->dwScope;
DWORD dwType = pNetRes->dwType;
DWORD dwDisplayType = pNetRes->dwDisplayType;
DWORD dwUsage = pNetRes->dwUsage;
TRACE("NETRESOURCE (0x%08X):\n\tdwScope = %s\n\tdwType = %s\n\tdwDisplayType = %s\n\tdwUsage = %s\n\tlpLocalName = %s\n\tlpRemoteName = %s\n\tlpComment = %s\n\tlpProvider = %s\n",
(DWORD_PTR)pNetRes,
(dwScope == RESOURCE_CONNECTED) ? "RESOURCE_CONNECTED" : (dwScope == RESOURCE_GLOBALNET) ? "RESOURCE_GLOBALNET" : (dwScope == RESOURCE_REMEMBERED) ? "RESOURCE_REMEMBERED" : "(unknown)",
(dwType == RESOURCETYPE_ANY) ? "RESOURCETYPE_ANY" : (dwType == RESOURCETYPE_DISK) ? "RESOURCETYPE_DISK" : (dwType == RESOURCETYPE_PRINT) ? "RESOURCETYPE_PRINT" : "(unknown)",
(dwDisplayType == RESOURCEDISPLAYTYPE_DOMAIN) ? "RESOURCEDISPLAYTYPE_DOMAIN" : (dwDisplayType == RESOURCEDISPLAYTYPE_GENERIC) ? "RESOURCEDISPLAYTYPE_GENERIC" : (dwDisplayType == RESOURCEDISPLAYTYPE_SERVER) ? "RESOURCEDISPLAYTYPE_SERVER" : (dwDisplayType == RESOURCEDISPLAYTYPE_SHARE) ? "RESOURCEDISPLAYTYPE_SHARE" : "(unknown)",
(dwUsage == RESOURCEUSAGE_CONNECTABLE) ? "RESOURCEUSAGE_CONNECTABLE" : (dwUsage == RESOURCEUSAGE_CONTAINER) ? "RESOURCEUSAGE_CONTAINER" : "(unknown)",
pNetRes->lpLocalName == NULL ? L"(null)" : pNetRes->lpLocalName,
pNetRes->lpRemoteName == NULL ? L"(null)" : pNetRes->lpRemoteName,
pNetRes->lpComment == NULL ? L"(null)" : pNetRes->lpComment,
pNetRes->lpProvider == NULL ? L"(null)" : pNetRes->lpProvider);
}
#endif
void CNetEnum::EnumThreadProc()
{
// Init stuff we don't want to do more than once
NETRESOURCE* prgNetResOuter = (NETRESOURCE*)malloc(1024);
NETRESOURCE* prgNetResInnerT = (NETRESOURCE*)malloc(1024);
TCHAR szComputerName[MAX_COMPUTERNAME_LENGTH+1];
DWORD cch = _countof(szComputerName);
GetComputerName(szComputerName, &cch);
HANDLE hEnumOuter = NULL;
begin:
// If the job ID changes out from under us, that means we need to stop
// the current task and jump back to the beginning.
EnterCriticalSection(&m_cs);
JOBTYPE eJobType = m_eJobType;
m_bNewJob = FALSE;
LeaveCriticalSection(&m_cs);
// Close enumeration left open by a previous job
if (hEnumOuter != NULL)
{
WNetCloseEnum(hEnumOuter);
hEnumOuter = NULL;
}
#ifdef _DEBUG
// Sleep(eJobType == jtEnumComputers ? 6000 : 12000); // simulate WNetOpenEnum taking a long time
#endif
// REVIEW: should we look for computers outside the current workgroup?
DWORD dwResult;
if (eJobType == jtEnumComputers)
{
dwResult = WNetOpenEnum(RESOURCE_CONTEXT, RESOURCETYPE_ANY, RESOURCEUSAGE_CONTAINER,
NULL, &hEnumOuter);
}
else
{
ASSERT(eJobType == jtEnumPrinters);
dwResult = WNetOpenEnum(RESOURCE_CONTEXT, RESOURCETYPE_PRINT, RESOURCEUSAGE_CONNECTABLE,
NULL, &hEnumOuter);
}
if (dwResult == NO_ERROR)
{
if (m_bAbort) goto cleanup;
if (m_bNewJob) goto begin;
BOOL bCallbackResult = TRUE;
// Keep looping until no more items
for (;;)
{
DWORD cOuterEntries = 20;
DWORD cbBuffer = 1024;
dwResult = WNetEnumResource(hEnumOuter, &cOuterEntries, prgNetResOuter, &cbBuffer);
if (dwResult == ERROR_NO_MORE_ITEMS)
break;
for (DWORD iOuter = 0; iOuter < cOuterEntries; iOuter++)
{
NETRESOURCE* pNetResOuter = &prgNetResOuter[iOuter];
BOOL bDoCallback = FALSE;
#ifdef _DEBUG
if (eJobType == jtEnumPrinters)
TraceNetResource(pNetResOuter);
#endif
if (pNetResOuter->dwDisplayType != RESOURCEDISPLAYTYPE_SERVER)
continue;
if (DoComputerNamesMatch(pNetResOuter->lpRemoteName, szComputerName))
continue;
HANDLE hEnumInner = NULL;
if (eJobType == jtEnumPrinters)
{
#ifdef _DEBUG
DWORD dwTicksBefore = GetTickCount();
#endif
dwResult = WNetOpenEnum(RESOURCE_GLOBALNET, RESOURCETYPE_PRINT, RESOURCEUSAGE_CONNECTABLE,
pNetResOuter, &hEnumInner);
#ifdef _DEBUG
DWORD dwTicks = GetTickCount() - dwTicksBefore;
if (dwTicks > 100)
{
TRACE("PERFORMANCE NOTE - took %d.%d sec to look for printers on %s\r\n", dwTicks / 1000, (dwTicks % 1000) - (dwTicks % 100), pNetResOuter->lpRemoteName);
}
#endif
if (dwResult != NO_ERROR)
continue;
}
for (;;)
{
DWORD cInnerEntries;
const NETRESOURCE* prgNetResInner = NULL;
if (eJobType == jtEnumPrinters)
{
cInnerEntries = 20;
cbBuffer = 1024;
dwResult = WNetEnumResource(hEnumInner, &cInnerEntries, prgNetResInnerT, &cbBuffer);
if (dwResult == ERROR_NO_MORE_ITEMS)
break;
prgNetResInner = prgNetResInnerT;
}
else
{
cInnerEntries = 1;
prgNetResInner = prgNetResOuter + iOuter;
}
for (DWORD iInner = 0; iInner < cInnerEntries; iInner++)
{
const NETRESOURCE* pNetResInner = &prgNetResInner[iInner];
LPCTSTR pszShareName;
#ifdef _DEBUG
if (eJobType == jtEnumPrinters)
TraceNetResource(pNetResInner);
#endif
if (eJobType == jtEnumComputers)
{
if (pNetResInner->dwDisplayType == RESOURCEDISPLAYTYPE_SERVER)
{
bDoCallback = TRUE;
pszShareName = NULL;
}
}
else // eJobType == jtEnumPrinters
{
bDoCallback = TRUE;
pszShareName = FindFileTitle(pNetResInner->lpRemoteName);
}
// We must call the callback inside the same critical section where
// we check if we should stop or restart, otherwise we might call
// the wrong callback!
// TODO: Get the real printer share name!!
EnterCriticalSection(&m_cs);
if (m_bAbort || m_bNewJob)
bCallbackResult = FALSE;
else if (bDoCallback)
bCallbackResult = (*m_pfnCallback)(m_pvCallbackParam, pNetResOuter->lpRemoteName, pszShareName);
LeaveCriticalSection(&m_cs);
if (!bCallbackResult)
break;
}
if (eJobType == jtEnumComputers)
break;
}
if (eJobType == jtEnumPrinters)
{
WNetCloseEnum(hEnumInner);
}
}
if (m_bAbort) goto cleanup;
if (m_bNewJob) goto begin;
if (!bCallbackResult)
break;
}
}
cleanup:
if (hEnumOuter != NULL)
{
WNetCloseEnum(hEnumOuter);
hEnumOuter = NULL;
}
// Be careful to close m_hThread only if we don't need to start another job
{
EnterCriticalSection(&m_cs);
BOOL bThreadDone = (m_bAbort || !m_bNewJob);
if (bThreadDone)
{
// Call callback function one more time
if (!m_bAbort)
{
(*m_pfnCallback)(m_pvCallbackParam, NULL, NULL);
}
CloseHandle(m_hThread);
m_hThread = NULL;
}
LeaveCriticalSection(&m_cs);
// Check if another job has been requested
if (!bThreadDone)
goto begin;
}
free(prgNetResInnerT);
free(prgNetResOuter);
}