VOID
SrvTreeRundown(
PLWIO_SRV_TREE pTree
)
{
BOOLEAN bInLock = FALSE;
LWIO_LOCK_RWMUTEX_EXCLUSIVE(bInLock, &pTree->mutex);
LwRtlRBTreeTraverse(
pTree->pAsyncStateCollection,
LWRTL_TREE_TRAVERSAL_TYPE_IN_ORDER,
SrvTreeRundownAsyncStatesRbTreeVisit,
NULL);
if (pTree->pAsyncStateCollection)
{
LwRtlRBTreeRemoveAll(pTree->pAsyncStateCollection);
}
LwRtlRBTreeTraverse(
pTree->pFileCollection,
LWRTL_TREE_TRAVERSAL_TYPE_IN_ORDER,
SrvTreeRundownFileRbTreeVisit,
NULL);
LWIO_UNLOCK_RWMUTEX(bInLock, &pTree->mutex);
}
VOID
SrvSession2Rundown(
PLWIO_SRV_SESSION_2 pSession
)
{
BOOLEAN bInLock = FALSE;
LWIO_LOCK_RWMUTEX_EXCLUSIVE(bInLock, &pSession->mutex);
SrvSession2UpdateLastActivityTime_inlock(pSession);
LwRtlRBTreeTraverse(
pSession->pTreeCollection,
LWRTL_TREE_TRAVERSAL_TYPE_IN_ORDER,
SrvSession2RundownTreeRbTreeVisit,
NULL);
LWIO_UNLOCK_RWMUTEX(bInLock, &pSession->mutex);
}
NTSTATUS
SrvElementsEnumResources(
SRV_RESOURCE_TYPE resourceType,
PFN_ENUM_RESOURCES pfnEnumResourcesCB,
PVOID pUserData
)
{
NTSTATUS ntStatus = STATUS_SUCCESS;
BOOLEAN bInLock = FALSE;
SRV_ELEMENTS_ENUM_RESOURCES enumResources =
{
.resourceType = resourceType,
.pfnEnumResourcesCB = pfnEnumResourcesCB,
.pUserData = pUserData,
.bContinue = TRUE
};
if ((resourceType == SRV_RESOURCE_TYPE_UNKNOWN) || !pfnEnumResourcesCB)
{
ntStatus = STATUS_INVALID_PARAMETER;
BAIL_ON_NT_STATUS(ntStatus);
}
LWIO_LOCK_RWMUTEX_SHARED(bInLock, &gSrvElements.resources.mutex);
ntStatus = LwRtlRBTreeTraverse(
gSrvElements.resources.pResources,
LWRTL_TREE_TRAVERSAL_TYPE_IN_ORDER,
&SrvElementsEnumResourcesInternal,
&enumResources);
BAIL_ON_NT_STATUS(ntStatus);
cleanup:
LWIO_UNLOCK_RWMUTEX(bInLock, &gSrvElements.resources.mutex);
return ntStatus;
error:
goto cleanup;
}
NTSTATUS
SrvSession2GetFileCount(
PLWIO_SRV_SESSION_2 pSession,
PULONG64 pullFileCount
)
{
NTSTATUS ntStatus = STATUS_SUCCESS;
BOOLEAN bInLock = FALSE;
ULONG64 ullTotalFileCount = 0;
ntStatus = SrvSession2UpdateLastActivityTime(pSession);
BAIL_ON_NT_STATUS(ntStatus);
LWIO_LOCK_RWMUTEX_SHARED(bInLock, &pSession->mutex);
ntStatus = LwRtlRBTreeTraverse(
pSession->pTreeCollection,
LWRTL_TREE_TRAVERSAL_TYPE_IN_ORDER,
&SrvSession2CountTotalFiles,
&ullTotalFileCount);
BAIL_ON_NT_STATUS(ntStatus);
*pullFileCount = ullTotalFileCount;
cleanup:
LWIO_UNLOCK_RWMUTEX(bInLock, &pSession->mutex);
return ntStatus;
error:
*pullFileCount = 0;
goto cleanup;
}
static
NTSTATUS
SrvProtocolProcessCandidateConnection2(
PVOID pKey,
PVOID pData,
PVOID pUserData,
PBOOLEAN pbContinue
)
{
NTSTATUS ntStatus = STATUS_SUCCESS;
PLWIO_SRV_SESSION_2 pSession = (PLWIO_SRV_SESSION_2)pData;
PSRV_PROTOCOL_CONNECTION_ENUM_QUERY pConnectionEnumQuery =
(PSRV_PROTOCOL_CONNECTION_ENUM_QUERY)pUserData;
BOOLEAN bInLock = FALSE;
LWIO_LOCK_RWMUTEX_SHARED(bInLock, &pSession->mutex);
ntStatus = LwRtlRBTreeTraverse(
pSession->pTreeCollection,
LWRTL_TREE_TRAVERSAL_TYPE_IN_ORDER,
&SrvProtocolProcessCandidateTree2,
pConnectionEnumQuery);
BAIL_ON_NT_STATUS(ntStatus);
*pbContinue = TRUE;
cleanup:
LWIO_UNLOCK_RWMUTEX(bInLock, &pSession->mutex);
return ntStatus;
error:
*pbContinue = FALSE;
goto cleanup;
}
static
NTSTATUS
SrvProtocolEnumCandidateConnections(
PVOID pKey,
PVOID pData,
PVOID pUserData,
PBOOLEAN pbContinue
)
{
NTSTATUS ntStatus = STATUS_SUCCESS;
PLWIO_SRV_CONNECTION pConnection = (PLWIO_SRV_CONNECTION)pData;
PSRV_PROTOCOL_CONNECTION_ENUM_QUERY pConnectionEnumQuery =
(PSRV_PROTOCOL_CONNECTION_ENUM_QUERY)pUserData;
BOOLEAN bInLock = FALSE;
PWSTR pwszClientHost = NULL;
if (pConnectionEnumQuery->pQueryAddress)
{
/*
* Look for connections by computer address first in case
* that was the qualifier string
*/
struct addrinfo* pCursor = pConnectionEnumQuery->pQueryAddress;
BOOLEAN bMatch = FALSE;
for (; !bMatch && (pCursor != NULL); pCursor = pCursor->ai_next)
{
ntStatus = SrvSocketCompareAddress(
pConnection->pClientAddress,
pConnection->clientAddrLen,
pCursor->ai_addr,
pCursor->ai_addrlen,
&bMatch);
BAIL_ON_NT_STATUS(ntStatus);
}
if (!bMatch)
{
pConnection = NULL;
}
}
if (pConnection)
{
LWIO_LOCK_RWMUTEX_SHARED(bInLock, &pConnection->mutex);
ntStatus = SrvSocketAddressToStringW(pConnection->pClientAddress,
&pwszClientHost);
BAIL_ON_NT_STATUS(ntStatus);
pConnectionEnumQuery->pwszClientHost = pwszClientHost;
pConnectionEnumQuery->pClientAddress = pConnection->pClientAddress;
pConnectionEnumQuery->clientAddrLen = pConnection->clientAddrLen;
pConnectionEnumQuery->ulConnectionResId
= pConnection->resource.ulResourceId;
ntStatus = WireGetCurrentNTTime(&pConnectionEnumQuery->llCurTime);
BAIL_ON_NT_STATUS(ntStatus);
switch (SrvConnectionGetProtocolVersion(pConnection))
{
case SMB_PROTOCOL_VERSION_1:
ntStatus = LwRtlRBTreeTraverse(
pConnection->pSessionCollection,
LWRTL_TREE_TRAVERSAL_TYPE_IN_ORDER,
&SrvProtocolProcessCandidateConnection,
pConnectionEnumQuery);
break;
case SMB_PROTOCOL_VERSION_2:
ntStatus = LwRtlRBTreeTraverse(
pConnection->pSessionCollection,
LWRTL_TREE_TRAVERSAL_TYPE_IN_ORDER,
&SrvProtocolProcessCandidateConnection2,
pConnectionEnumQuery);
break;
case SMB_PROTOCOL_VERSION_UNKNOWN:
/* Ignore connections that are still being established */
break;
default:
ntStatus = STATUS_INTERNAL_ERROR;
break;
}
BAIL_ON_NT_STATUS(ntStatus);
}
*pbContinue = TRUE;
cleanup:
pConnectionEnumQuery->pClientAddress = NULL;
pConnectionEnumQuery->clientAddrLen = 0;
pConnectionEnumQuery->pwszClientHost = NULL;
SRV_SAFE_FREE_MEMORY(pwszClientHost);
LWIO_UNLOCK_RWMUTEX(bInLock, &pConnection->mutex);
return ntStatus;
error:
*pbContinue = FALSE;
goto cleanup;
}
NTSTATUS
SrvProtocolEnumerateConnections(
PWSTR pwszComputerName,
PWSTR pwszShareName,
ULONG ulInfoLevel,
ULONG ulPreferredMaxLength,
PBYTE pBuffer,
ULONG ulBufferSize,
PULONG pulBytesUsed,
PULONG pulEntriesRead,
PULONG pulTotalEntries,
PULONG pulResumeHandle
)
{
NTSTATUS ntStatus = STATUS_SUCCESS;
BOOLEAN bInLock = FALSE;
BOOLEAN bMoreData = FALSE;
SRV_PROTOCOL_CONNECTION_ENUM_QUERY ConnectionEnumQuery =
{
.pwszComputerName = pwszComputerName,
.pwszShareName = pwszShareName,
.pClientAddress = NULL,
.clientAddrLen = 0,
.pwszClientHost = NULL,
.ulInfoLevel = ulInfoLevel,
.ulPreferredMaxLength = ulPreferredMaxLength,
.iEntryIndex = 0,
.iResumeIndex = pulResumeHandle ? *pulResumeHandle : 0,
.ulEntriesRead = 0,
.ulTotalEntries = 0,
.pBuffer = pBuffer,
.ulBufferSize = ulBufferSize,
.ulBytesUsed = 0,
.pQueryAddress = NULL
};
if (pwszComputerName)
{
wchar16_t wszPrefix[] = {'\\','\\', 0};
size_t sPrefixLen =
(sizeof(wszPrefix)-sizeof(wszPrefix[0]))/sizeof(wszPrefix[0]);
if (!SMBWc16snCmp(pwszComputerName, &wszPrefix[0], sPrefixLen))
{
pwszComputerName += sPrefixLen;
}
ConnectionEnumQuery.pwszComputerName = pwszComputerName;
ntStatus = SrvSocketGetAddrInfoW(
pwszComputerName,
&ConnectionEnumQuery.pQueryAddress);
BAIL_ON_NT_STATUS(ntStatus);
}
LWIO_LOCK_RWMUTEX_SHARED(bInLock, &gProtocolApiGlobals.mutex);
ntStatus = LwRtlRBTreeTraverse(
gProtocolApiGlobals.pConnections,
LWRTL_TREE_TRAVERSAL_TYPE_IN_ORDER,
&SrvProtocolCountCandidateConnections,
&ConnectionEnumQuery);
BAIL_ON_NT_STATUS(ntStatus);
ntStatus = LwRtlRBTreeTraverse(
gProtocolApiGlobals.pConnections,
LWRTL_TREE_TRAVERSAL_TYPE_IN_ORDER,
&SrvProtocolEnumCandidateConnections,
&ConnectionEnumQuery);
/* If we still have more data to read, then return MORE_ENTRIES */
if (ntStatus == STATUS_END_OF_FILE &&
ConnectionEnumQuery.ulEntriesRead < ConnectionEnumQuery.ulTotalEntries)
{
bMoreData = TRUE;
ntStatus = STATUS_SUCCESS;
}
BAIL_ON_NT_STATUS(ntStatus);
*pulBytesUsed = ConnectionEnumQuery.ulBytesUsed;
*pulEntriesRead = ConnectionEnumQuery.ulEntriesRead;
*pulTotalEntries = ConnectionEnumQuery.ulTotalEntries;
if (pulResumeHandle)
{
*pulResumeHandle =
ConnectionEnumQuery.iResumeIndex + ConnectionEnumQuery.ulEntriesRead;
}
cleanup:
LWIO_UNLOCK_RWMUTEX(bInLock, &gProtocolApiGlobals.mutex);
SrvProtocolFreeConnectionEnumQueryContents(&ConnectionEnumQuery);
return (NT_SUCCESS(ntStatus) && bMoreData ? STATUS_MORE_ENTRIES : ntStatus);
error:
*pulBytesUsed = 0;
*pulEntriesRead = 0;
*pulTotalEntries = 0;
if (pBuffer && ulBufferSize)
{
memset(pBuffer, 0, ulBufferSize);
}
goto cleanup;
}
static
NTSTATUS
SrvProtocolCountCandidateConnections(
PVOID pKey,
PVOID pData,
PVOID pUserData,
PBOOLEAN pbContinue
)
{
NTSTATUS ntStatus = STATUS_SUCCESS;
PLWIO_SRV_CONNECTION pConnection = (PLWIO_SRV_CONNECTION)pData;
PSRV_PROTOCOL_CONNECTION_ENUM_QUERY pConnectionEnumQuery =
(PSRV_PROTOCOL_CONNECTION_ENUM_QUERY)pUserData;
ULONG ulTotalConnectionCount = pConnectionEnumQuery->ulTotalEntries;
ULONG ulConnectionCount = 0;
if (pConnectionEnumQuery->pQueryAddress)
{
/*
* Look for connections by computer address first in case
* that was the qualifier string
*/
struct addrinfo* pCursor = pConnectionEnumQuery->pQueryAddress;
BOOLEAN bMatch = FALSE;
for (; !bMatch && (pCursor != NULL); pCursor = pCursor->ai_next)
{
ntStatus = SrvSocketCompareAddress(
pConnection->pClientAddress,
pConnection->clientAddrLen,
pCursor->ai_addr,
pCursor->ai_addrlen,
&bMatch);
BAIL_ON_NT_STATUS(ntStatus);
}
if (!bMatch)
{
pConnection = NULL;
}
}
if (pConnection)
{
switch (SrvConnectionGetProtocolVersion(pConnection))
{
case SMB_PROTOCOL_VERSION_1:
ntStatus = SrvConnectionGetConnectionCount(
pConnection,
pConnectionEnumQuery->pwszShareName,
&ulConnectionCount);
break;
case SMB_PROTOCOL_VERSION_2:
ntStatus = SrvConnectionGetConnectionCount2(
pConnection,
pConnectionEnumQuery->pwszShareName,
&ulConnectionCount);
break;
case SMB_PROTOCOL_VERSION_UNKNOWN:
/* Ignore connections that are still being established */
break;
default:
ntStatus = STATUS_INTERNAL_ERROR;
break;
}
BAIL_ON_NT_STATUS(ntStatus);
ulTotalConnectionCount += ulConnectionCount;
pConnectionEnumQuery->ulTotalEntries = ulTotalConnectionCount;
}
*pbContinue = TRUE;
cleanup:
return ntStatus;
error:
*pbContinue = FALSE;
goto cleanup;
}
