mirror of https://github.com/tongzx/nt5src
You can not select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
424 lines
10 KiB
424 lines
10 KiB
#include "stdafx.h"
|
|
#include "gkwsock.h"
|
|
|
|
|
|
// ASYNC_ACCEPT --------------------------------------------------------------------------
|
|
|
|
|
|
ASYNC_ACCEPT::ASYNC_ACCEPT (void)
|
|
{
|
|
AcceptSocket = INVALID_SOCKET;
|
|
ClientSocket = INVALID_SOCKET;
|
|
AcceptFunc = NULL;
|
|
AcceptFuncContext = NULL;
|
|
ReferenceCount = 0L;
|
|
StopNotifyEvent = NULL;
|
|
}
|
|
|
|
ASYNC_ACCEPT::~ASYNC_ACCEPT (void)
|
|
{
|
|
assert (AcceptSocket == INVALID_SOCKET);
|
|
assert (ClientSocket == INVALID_SOCKET);
|
|
assert (ReferenceCount == 0L);
|
|
assert (!StopNotifyEvent);
|
|
}
|
|
|
|
HRESULT ASYNC_ACCEPT::StartIo (
|
|
IN SOCKADDR_IN * SocketAddress,
|
|
IN ASYNC_ACCEPT_FUNC ArgAcceptFunc,
|
|
IN PVOID ArgAcceptContext)
|
|
{
|
|
HRESULT Result;
|
|
|
|
assert (SocketAddress);
|
|
assert (ArgAcceptFunc);
|
|
|
|
Lock ();
|
|
|
|
if (AcceptSocket == INVALID_SOCKET && ReferenceCount == 0L) {
|
|
// this object is not currently in use
|
|
// so, it's acceptable to use it
|
|
|
|
assert (!AcceptFunc);
|
|
assert (!AcceptFuncContext);
|
|
assert (!StopNotifyEvent);
|
|
|
|
// This increase in reference count is needed
|
|
// to shut down the service gracefully
|
|
// Reference count on ASYNC_ACCEPT objects
|
|
// will never drop to zero unless StopWait is called.
|
|
// StopWait will call matching Release, which will
|
|
// bring the reference count to the expected value of 0.
|
|
AddRef ();
|
|
|
|
Result = StartIoLocked (SocketAddress);
|
|
|
|
if (Result == S_OK) {
|
|
|
|
assert (AcceptSocket != INVALID_SOCKET);
|
|
|
|
AcceptFunc = ArgAcceptFunc;
|
|
AcceptFuncContext = ArgAcceptContext;
|
|
}
|
|
else {
|
|
|
|
Release ();
|
|
}
|
|
}
|
|
else {
|
|
Debug (_T("ASYNC_ACCEPT::StartIo: this object is already in use, must first call Stop and wait for sync counter\n"));
|
|
Result = E_FAIL;
|
|
}
|
|
|
|
Unlock();
|
|
|
|
return Result;
|
|
}
|
|
|
|
|
|
HRESULT ASYNC_ACCEPT::StartIoLocked (
|
|
IN SOCKADDR_IN * SocketAddress)
|
|
{
|
|
HRESULT Result;
|
|
BOOL KeepaliveOption;
|
|
|
|
assert (SocketAddress);
|
|
assert (AcceptSocket == INVALID_SOCKET);
|
|
assert (ClientSocket == INVALID_SOCKET);
|
|
assert (ReferenceCount == 1);
|
|
assert (!StopNotifyEvent);
|
|
|
|
StopNotifyEvent = CreateEvent (NULL, TRUE, FALSE, NULL);
|
|
|
|
if (!StopNotifyEvent) {
|
|
|
|
Result = GetLastErrorAsResult ();
|
|
DebugLastError (_T("ASYNC_ACCEPT::StartIoLocked: failed to create stop notify event\n"));
|
|
|
|
} else {
|
|
|
|
AcceptSocket = WSASocket (AF_INET, SOCK_STREAM, IPPROTO_TCP, NULL, 0, WSA_FLAG_OVERLAPPED);
|
|
|
|
if (AcceptSocket == INVALID_SOCKET) {
|
|
|
|
Result = GetLastErrorAsResult ();
|
|
DebugLastError (_T("ASYNC_ACCEPT::StartIoLocked: failed to create accept socket\n"));
|
|
|
|
} else {
|
|
|
|
if (bind (AcceptSocket, (SOCKADDR *) SocketAddress, sizeof (SOCKADDR_IN))) {
|
|
|
|
Result = GetLastErrorAsResult ();
|
|
DebugLastErrorF (_T("ASYNC_ACCEPT::StartIoLocked: failed to bind accept socket to address %08X:%04X\n"),
|
|
ntohl (SocketAddress -> sin_addr.s_addr),
|
|
ntohs (SocketAddress -> sin_port));
|
|
|
|
} else {
|
|
|
|
// Set keepalive on the socket
|
|
KeepaliveOption = TRUE;
|
|
if (SOCKET_ERROR == setsockopt (AcceptSocket, SOL_SOCKET,
|
|
SO_KEEPALIVE, (PCHAR) &KeepaliveOption, sizeof (KeepaliveOption)))
|
|
{
|
|
Result = GetLastErrorAsResult ();
|
|
DebugLastError (_T("ASYNC_ACCEPT: Failed to set keepalive on accept socket.\n"));
|
|
|
|
} else {
|
|
|
|
if (listen (AcceptSocket, 10)) {
|
|
|
|
Result = GetLastErrorAsResult ();
|
|
DebugLastError (_T("ASYNC_ACCEPT::StartIoLocked: failed to listen on accept socket\n"));
|
|
|
|
} else {
|
|
|
|
if (!BindIoCompletionCallback ((HANDLE) AcceptSocket, ASYNC_ACCEPT::IoCompletionCallback, 0)) {
|
|
|
|
Result = GetLastErrorAsResult ();
|
|
DebugLastError (_T("ASYNC_ACCEPT::StartIoLocked: failed to bind i/o completion callback\n"));
|
|
|
|
} else {
|
|
|
|
Result = IssueAccept ();
|
|
|
|
if (Result == S_OK) {
|
|
|
|
return Result;
|
|
|
|
}
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
closesocket (AcceptSocket);
|
|
AcceptSocket = INVALID_SOCKET;
|
|
}
|
|
|
|
CloseHandle (StopNotifyEvent);
|
|
StopNotifyEvent = NULL;
|
|
}
|
|
|
|
return Result;
|
|
}
|
|
|
|
|
|
HRESULT ASYNC_ACCEPT::GetListenSocketAddress (
|
|
OUT SOCKADDR_IN * ReturnSocketAddress)
|
|
{
|
|
HRESULT Result;
|
|
INT SocketAddressLength;
|
|
|
|
Lock();
|
|
|
|
if (AcceptSocket != INVALID_SOCKET) {
|
|
|
|
SocketAddressLength = sizeof (SOCKADDR_IN);
|
|
if (getsockname (AcceptSocket, (SOCKADDR *) ReturnSocketAddress, &SocketAddressLength) == SOCKET_ERROR) {
|
|
Result = GetLastErrorAsResult();
|
|
}
|
|
else {
|
|
Result = S_OK;
|
|
}
|
|
}
|
|
else {
|
|
Result = E_INVALIDARG;
|
|
}
|
|
|
|
Unlock();
|
|
|
|
return Result;
|
|
}
|
|
|
|
HRESULT ASYNC_ACCEPT::IssueAccept (void)
|
|
{
|
|
HRESULT Result;
|
|
BOOL KeepaliveOption;
|
|
|
|
AssertLocked();
|
|
assert (ClientSocket == INVALID_SOCKET);
|
|
// assert (ReferenceCount == 0);
|
|
|
|
ClientSocket = WSASocket (AF_INET, SOCK_STREAM, IPPROTO_TCP, NULL, 0, WSA_FLAG_OVERLAPPED);
|
|
if (ClientSocket == INVALID_SOCKET) {
|
|
Result = GetLastErrorAsResult ();
|
|
DebugLastError (_T("ASYNC_ACCEPT::IssueAccept: failed to create client socket.\n"));
|
|
return Result;
|
|
}
|
|
|
|
ZeroMemory (&Overlapped, sizeof (OVERLAPPED));
|
|
|
|
AddRef();
|
|
|
|
if (!AcceptEx (AcceptSocket,
|
|
ClientSocket,
|
|
ClientInfoBuffer,
|
|
0,
|
|
sizeof (SOCKADDR_IN) + 0x10,
|
|
sizeof (SOCKADDR_IN) + 0x10,
|
|
&ClientInfoBufferLength,
|
|
&Overlapped)) {
|
|
|
|
if (WSAGetLastError() != WSA_IO_PENDING) {
|
|
// an error occurred
|
|
Release ();
|
|
Result = GetLastErrorAsResult ();
|
|
DebugLastError (_T("ASYNC_ACCEPT::IssueAccept: failed to issue accept.\n"));
|
|
return Result;
|
|
}
|
|
|
|
// Set keepalive on the socket
|
|
KeepaliveOption = TRUE;
|
|
if (SOCKET_ERROR == setsockopt (ClientSocket, SOL_SOCKET,
|
|
SO_KEEPALIVE, (PCHAR) &KeepaliveOption, sizeof (KeepaliveOption)))
|
|
{
|
|
Release ();
|
|
Result = GetLastErrorAsResult ();
|
|
DebugLastError (_T("ASYNC_ACCEPT: IssueAccept: Failed to set keepalive on client socket.\n"));
|
|
return Result;
|
|
|
|
}
|
|
}
|
|
|
|
return S_OK;
|
|
}
|
|
|
|
// static
|
|
void ASYNC_ACCEPT::IoCompletionCallback (DWORD Status, DWORD BytesTransferred, LPOVERLAPPED Overlapped)
|
|
{
|
|
ASYNC_ACCEPT * AsyncAccept;
|
|
|
|
AsyncAccept = CONTAINING_RECORD (Overlapped, ASYNC_ACCEPT, Overlapped);
|
|
|
|
AsyncAccept -> IoComplete (Status, BytesTransferred);
|
|
|
|
AsyncAccept -> Release ();
|
|
}
|
|
|
|
void ASYNC_ACCEPT::IoComplete (DWORD Status, DWORD BytesTransferred)
|
|
{
|
|
ASYNC_ACCEPT_FUNC LocalAcceptFunc;
|
|
PVOID LocalAcceptFuncContext;
|
|
SOCKADDR_IN LocalAddressCopy;
|
|
SOCKADDR_IN RemoteAddressCopy;
|
|
SOCKET LocalClientSocket;
|
|
SOCKADDR * LocalAddress;
|
|
INT LocalAddressLength;
|
|
SOCKADDR * RemoteAddress;
|
|
INT RemoteAddressLength;
|
|
INT Result;
|
|
|
|
Lock();
|
|
|
|
assert (ClientSocket != INVALID_SOCKET);
|
|
assert (ReferenceCount > 0);
|
|
|
|
if (AcceptSocket == INVALID_SOCKET) {
|
|
// Stop has been called
|
|
// just immediately disconnect the client
|
|
// we'll deal with object lifetime below
|
|
|
|
closesocket (ClientSocket);
|
|
ClientSocket = INVALID_SOCKET;
|
|
}
|
|
else {
|
|
// the context is in the normal state
|
|
// continue processing
|
|
|
|
if (Status == ERROR_SUCCESS) {
|
|
// a client has successfully connected
|
|
|
|
GetAcceptExSockaddrs (
|
|
ClientInfoBuffer,
|
|
0, // no initial recv
|
|
sizeof (SOCKADDR_IN) + 0x10,
|
|
sizeof (SOCKADDR_IN) + 0x10,
|
|
&LocalAddress,
|
|
&LocalAddressLength,
|
|
&RemoteAddress,
|
|
&RemoteAddressLength);
|
|
|
|
// copy information out of the context
|
|
// so that it will be valid after we issue a new accept and unlock
|
|
LocalAddressCopy = *(SOCKADDR_IN *) LocalAddress;
|
|
RemoteAddressCopy = *(SOCKADDR_IN *) RemoteAddress;
|
|
LocalClientSocket = ClientSocket;
|
|
LocalAcceptFunc = AcceptFunc;
|
|
LocalAcceptFuncContext = AcceptFuncContext;
|
|
|
|
ClientSocket = INVALID_SOCKET;
|
|
|
|
// update the accept context
|
|
Result = setsockopt (ClientSocket, SOL_SOCKET, SO_UPDATE_ACCEPT_CONTEXT,
|
|
reinterpret_cast <char *> (&AcceptSocket), sizeof (SOCKET));
|
|
|
|
// issue a new accept
|
|
IssueAccept();
|
|
|
|
Unlock();
|
|
|
|
(*LocalAcceptFunc) (LocalAcceptFuncContext, LocalClientSocket, &LocalAddressCopy, &RemoteAddressCopy);
|
|
|
|
Lock();
|
|
}
|
|
else {
|
|
// some error has occurred
|
|
// this is usually (but not always) fatal
|
|
|
|
assert (ClientSocket != INVALID_SOCKET);
|
|
|
|
closesocket (ClientSocket);
|
|
ClientSocket = INVALID_SOCKET;
|
|
|
|
switch (Status) {
|
|
case STATUS_CANCELLED:
|
|
Debug (_T("ASYNC_ACCEPT::IoComplete: accept failed, STATUS_CANCELED, original thread probably exited, resubmitting request...\n"));
|
|
break;
|
|
|
|
default:
|
|
DebugError (Status, _T("AsyncAccept: async accept FAILED, sleeping 2000ms and retrying...\n"));
|
|
Sleep (2000);
|
|
break;
|
|
}
|
|
|
|
IssueAccept();
|
|
}
|
|
}
|
|
|
|
Unlock();
|
|
}
|
|
|
|
void ASYNC_ACCEPT::StopWait (void)
|
|
{
|
|
DWORD Status;
|
|
|
|
Lock();
|
|
|
|
if (AcceptSocket != INVALID_SOCKET) {
|
|
|
|
// closing the socket cancels all pending i/o
|
|
// we do NOT close the ClientSocket
|
|
// only the i/o completion callback path may do that
|
|
closesocket (AcceptSocket);
|
|
AcceptSocket = INVALID_SOCKET;
|
|
AcceptFunc = NULL;
|
|
AcceptFuncContext = NULL;
|
|
|
|
if (ClientSocket != INVALID_SOCKET) {
|
|
// an accept is still pending. it may complete successfully,
|
|
// or it may complete with STATUS_CANCELED (since we just closed AcceptSocket)
|
|
// in either case, we must wait for the i/o complete callback to run.
|
|
// AcceptSocket = INVALID_SOCKET is an indicator to the completion callback
|
|
// that it should abort / return immediately.
|
|
|
|
assert (StopNotifyEvent);
|
|
|
|
Unlock ();
|
|
|
|
// This is the counterpart to the AddRef called in
|
|
// StartIoLocked (see comment there)
|
|
Release ();
|
|
|
|
DebugF (_T("ASYNC_ACCEPT::StopWait: waiting for i/o completion thread...\n"));
|
|
|
|
Status = WaitForSingleObject (StopNotifyEvent, INFINITE);
|
|
|
|
assert (Status == WAIT_OBJECT_0);
|
|
|
|
Lock ();
|
|
}
|
|
}
|
|
else {
|
|
assert (!AcceptFunc);
|
|
assert (!AcceptFuncContext);
|
|
}
|
|
|
|
if (StopNotifyEvent) {
|
|
CloseHandle (StopNotifyEvent);
|
|
StopNotifyEvent = NULL;
|
|
}
|
|
|
|
Unlock();
|
|
}
|
|
|
|
|
|
void ASYNC_ACCEPT::AddRef (void) {
|
|
assert (ReferenceCount >= 0L);
|
|
InterlockedIncrement (&ReferenceCount);
|
|
}
|
|
|
|
|
|
void ASYNC_ACCEPT::Release (void) {
|
|
LONG Count;
|
|
|
|
assert (ReferenceCount >= 0L);
|
|
|
|
Count = InterlockedDecrement (&ReferenceCount);
|
|
|
|
if (Count == 0L) {
|
|
DebugF (_T("ASYNC_ACCEPT::Release -- Reference count dropped to zero. (this is %x)\n"), this);
|
|
|
|
assert (StopNotifyEvent);
|
|
SetEvent (StopNotifyEvent);
|
|
}
|
|
}
|