windows-nt/Source/XPSP1/NT/net/rras/ip/nath323/gkwsock.cpp
2020-09-26 16:20:57 +08:00

425 lines
10 KiB
C++

#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);
}
}