tcptransport.cpp

Windows CE 6.0 Server 源码

            }

            TRACEMSG(ZONE_TCPIP, (L"SMBSRV: Added TCP Connection -- Total Connections: %d", InterlockedIncrement(&TCP_TRANSPORT::g_lAliveSockets)));
            InitializeCriticalSection(&pNewConn->csSockLock);
            pNewConn->refCnt = 1; //NOTE: there is only one ref and 2 copies
                                  //  this is because ProcessingThread will call
                                  // DecrementConnectionCounter which removes
                                  // from the list
            pNewConn->sock = newSock;
            pNewConn->hHandle = CreateThread(NULL, 0, SMBSRV_TCPProcessingThread, (VOID *)pNewConn, CREATE_SUSPENDED, NULL);
            pNewConn->hOverlappedHandle = CreateEvent(NULL, TRUE, FALSE, NULL);

            if(NULL != pNewConn->hHandle) {
                EnterCriticalSection(&g_csLockTCPTransportGlobals);
                    if(!g_ConnectionList.push_front(pNewConn)) {
                        hr = E_OUTOFMEMORY;
                    }
                    TRACEMSG(ZONE_TCPIP, (L"SMBSRV: TCP CONNECTIONLIST-push: %d", g_ConnectionList.size()));
                LeaveCriticalSection(&g_csLockTCPTransportGlobals);
                ResumeThread(pNewConn->hHandle);
            }
        }
    }

    DoneSocketInList:
        {
        IFDBG(BOOL fFound = FALSE;);
        csLock.Lock();
        //
        // Remove ourself from the listen list.
        TRACEMSG(ZONE_TCPIP, (TEXT("SMBSRV-LSTNTHREAD:Exiting thread!")));

        ce::list<LISTEN_NODE, TCP_LISTEN_NODE_ALLOC >::iterator itConn;
        for(itConn = g_SocketListenList.begin(); itConn!=g_SocketListenList.end();++itConn) {
            if(itConn->s == s) {
                g_SocketListenList.erase(itConn);
                IFDBG(fFound = TRUE;);
                break;
            }
        }
        IFDBG(ASSERT(fFound));
        }

     Done:
        if (INVALID_SOCKET != s)
            closesocket(s);

        if(FAILED(hr))
            return -1;
        else 
            return 0;
}


int timed_recv(SOCKET s,char *buf,int len,int flags, UINT seconds, BOOL *pfTimedOut)
{
    int retValue = 0;
    int recvd = 0;

    if(INVALID_SOCKET == s) {
        return SOCKET_ERROR;
    }

    *pfTimedOut = FALSE;

    while(SOCKET_ERROR != retValue && len)
    {
        //
        // If we need to block for only a certain period of time do so
        if(0xFFFFFFFF != seconds) {
            FD_SET fd;
            FD_ZERO (&fd);
            FD_SET (s, &fd);
            timeval tv;
            tv.tv_sec = seconds;
            tv.tv_usec = 0;
            if (1 != select (0, &fd, NULL, NULL, &tv)) {    
                TRACEMSG(ZONE_TCPIP, (L"SMB_SRV: TCP transport -- timed_recv failed on select() with %d", WSAGetLastError()));
                *pfTimedOut = TRUE;
                retValue = SOCKET_ERROR;
                goto Done;
            }
            //
            // If there is something available, fetch it
            if(FALSE == FD_ISSET(s, &fd)) {
                continue;
            }

        }

        int iJustRead = recv(s, (char *)buf, len, flags);
        TRACEMSG(ZONE_TCPIP, (L"SMB_SRV: TCP transport -- just recv'ed %d bytes", iJustRead));
        if(SOCKET_ERROR == iJustRead) {
            TRACEMSG(ZONE_TCPIP, (L"SMB_SRV: TCP transport -- timed_recv failed on recv() with %d", WSAGetLastError()));
            retValue = iJustRead;
            goto Done;
        }

        //
        // If our socket has been closed, return with error
        if(0 == iJustRead) {
            TRACEMSG(ZONE_TCPIP, (L"SMB_SRV: TCP transport -- timed_recv failed on recv() with 0 bytes %d -- sock prob closed", WSAGetLastError()));
            retValue = SOCKET_ERROR;
            goto Done;
        }
        ASSERT(iJustRead <= len);

        len -= iJustRead;
        recvd += iJustRead;
        buf += iJustRead;
    }

    retValue = recvd;
    Done:
        return retValue;
}

HRESULT writeAll(SOCKET s, char *buffer, int len) 
{
    char *pPacketTemp = (char *)buffer;
    HRESULT hRet = S_OK;

    if(SOCKET_ERROR == s) {
        return E_INVALIDARG;
    }

    while(len)
    {
        int retCode = send(s, pPacketTemp, len, 0);

        if(retCode == 0 || retCode == SOCKET_ERROR) {
            TRACEMSG(ZONE_ERROR, (L"SMB_SRV: TCP transport -- Error(0x%08x)--send failed", WSAGetLastError()));
            hRet = E_FAIL;
            break;
        }

        pPacketTemp += retCode;
        len -= retCode;
    }
    return hRet;
}


HRESULT writeAll_BlockedOverLapped(SOCKET s, HANDLE h, WSABUF *pBufs, UINT uiNumBufs) 
{
    HRESULT hRet = E_FAIL;
    DWORD dwSent = 0;
    WSAOVERLAPPED overlapped;
    DWORD dwLeft = 0;

    if(SOCKET_ERROR == s) {
        return E_INVALIDARG;
    }

    for(UINT i=0; i<uiNumBufs; i++) {
        dwLeft += pBufs[i].len;
    }

    while(dwLeft) {
        overlapped.hEvent = h;
        int retCode = WSASend(s, pBufs, uiNumBufs, &dwSent, 0, &overlapped, NULL);

        if(0 == retCode) {
            TRACEMSG(ZONE_TCPIP, (L"SMB_SRV: TCP transport -- no error on send, data went w/o overlapped call needed"));
        } else if(WSA_IO_PENDING == WSAGetLastError()) {
            DWORD dwFlags;
            TRACEMSG(ZONE_TCPIP, (L"SMB_SRV: TCP transport -- needed to block on overlapped WSASend()"));
            if(WAIT_OBJECT_0 != WaitForSingleObject(h, INFINITE)) {
                TRACEMSG(ZONE_TCPIP, (L"SMB_SRV: TCP transport -- error waiting on return event for WSASend()"));
                ASSERT(FALSE);
                goto Done;
            }

            //
            // See how much data we sent
            ResetEvent(overlapped.hEvent);
            if(TRUE != WSAGetOverlappedResult(s, &overlapped, &dwSent, TRUE, &dwFlags)) {
                TRACEMSG(ZONE_ERROR, (L"SMB_SRV: TCP transport -- couldnt get overlapped results for WSASend() = GLE: %d", WSAGetLastError()));    
                goto Done;
            }
        } else {
            TRACEMSG(ZONE_ERROR, (L"SMB_SRV: TCP transport -- error code %d from WSASend()", WSAGetLastError()));
            goto Done;
        }

        if(dwSent > dwLeft) {
            TRACEMSG(ZONE_ERROR, (L"SMB_SRV: TCP transport -- more data sent that requested?! WSASend()"));
            ASSERT(FALSE);
            goto Done;
        }

        //
        // Adjust buffers to represent what was sent
        dwLeft -= dwSent;
        while(uiNumBufs && dwSent >= pBufs[0].len) {
            dwSent -= pBufs[0].len;
            pBufs++;
            uiNumBufs --;
        }

        if(uiNumBufs && dwSent) {
            ASSERT(dwSent < pBufs[0].len);
            pBufs[0].len -= dwSent;
            pBufs[0].buf += dwSent;
        }
    }

   hRet = S_OK;
   Done:
        return hRet;
}


VOID IncrementConnectionCounter(CONNECTION_HOLDER *pMyConnection)
{
    InterlockedIncrement(&pMyConnection->refCnt);
}

VOID DecrementConnectionCounter(CONNECTION_HOLDER *pMyConnection) 
{
    BOOL fFound = FALSE; 

    if(0 != InterlockedDecrement(&pMyConnection->refCnt)) {
        return;
    }

    CCritSection csLock(&TCP_TRANSPORT::g_csLockTCPTransportGlobals);
    csLock.Lock();
        ce::list<CONNECTION_HOLDER *, TCP_CONNECTION_HLDR_ALLOC>::iterator it;
        ce::list<CONNECTION_HOLDER *, TCP_CONNECTION_HLDR_ALLOC>::iterator itEnd = g_ConnectionList.end();

        for(it = g_ConnectionList.begin(); it != itEnd; ++it) {
            if(pMyConnection == *it) {
                g_ConnectionList.erase(it++);
                TRACEMSG(ZONE_TCPIP, (L"SMBSRV: TCP CONNECTIONLIST - DecCntr: %d", g_ConnectionList.size()));
                fFound = TRUE;
                break;
            }
       }
    csLock.UnLock();

    //
    //if we dont find it, thats bad -- leak memory and quit!!
    if(!fFound) {
        ASSERT(FALSE);
        return;
    }

    ASSERT(FALSE == csLock.IsLocked());
    ASSERT(0 == pMyConnection->refCnt);

    DeleteCriticalSection(&pMyConnection->csSockLock);   

    if(INVALID_SOCKET != pMyConnection->sock) {
        closesocket(pMyConnection->sock); 
    }

    CloseHandle(pMyConnection->hHandle); 
    CloseHandle(pMyConnection->hOverlappedHandle);
    pMyConnection->hHandle = INVALID_HANDLE_VALUE;
    pMyConnection->hOverlappedHandle = INVALID_HANDLE_VALUE;
    delete pMyConnection;
}


HRESULT 
TCP_TRANSPORT::DeleteTCPTransportToken(VOID *pToken)
{
    CONNECTION_HOLDER *pNew = (CONNECTION_HOLDER *)pToken;
    DecrementConnectionCounter(pNew);  
    return S_OK;
}


HRESULT
TCP_TRANSPORT::CopyTCPTransportToken(VOID *pToken, VOID **pNewToken)
{
    CONNECTION_HOLDER *pNew = (CONNECTION_HOLDER *)pToken;
    IncrementConnectionCounter(pNew);  
    *pNewToken = (VOID *)pNew;
    return S_OK;
}


//PERFPERF: we could have one thread for all tcp instead of one per connection
//   there isnt any real reason not to do that (just batch them up in a big array and block
//   on all of them)
//
//   As for the refcount (IncrementConnectionCounter) -- this thread proc
//      owns one by default (from pre CreateThread call).  then we inc this
//      everytime we send a connection to the cracker.  its dec'ed in our
//      own memory deconstruction
DWORD 
TCP_TRANSPORT::SMBSRV_TCPProcessingThread(LPVOID _pAdapter)
{
    CONNECTION_HOLDER *pMyConnection = (CONNECTION_HOLDER *)_pAdapter;
    SOCKET sSock = pMyConnection->sock;
    HRESULT hr = E_FAIL;    
    USHORT  sConnectionID;
    ULONG   ConnectionID=0xFFFF;
    SMB_PACKET *pNewPacket = NULL;

    //
    // Build up a unique connection ID
    if(FAILED(g_ConnectionID.GetID(&sConnectionID))) {
        sConnectionID = 0xFFFF;
        ASSERT(FALSE);
        hr = E_FAIL;
        goto Done;
    }
    ConnectionID = (ULONG)(SMB_Globals::TCP_TRANSPORT << 16) | sConnectionID;
    pMyConnection->ulConnectionID = ConnectionID;

    while(FALSE == g_fStopped)
    {
        DWORD dwHeader;
        CHAR *pHeader = (CHAR *)&dwHeader;
        BOOL fTimedOut;

        //
        // Get some memory out of our free pool
        ASSERT(NULL == pNewPacket);
        if(NULL == (pNewPacket = SMB_Globals::g_SMB_Pool.Alloc())) {
            hr = E_OUTOFMEMORY;
            goto Done;
        }

        //
        // Block for the packet header, on timeout inform the packet cracker
        if(4 != timed_recv(sSock, pHeader, 4, 0, g_uiTCPTimeoutInSeconds, &fTimedOut)) {
            if(fTimedOut) {
                TRACEMSG(ZONE_ERROR, (TEXT("SMBSRV-TCP -- reading SMB TCP header timed out!! telling the cracker")));
                pNewPacket->uiPacketType = SMB_CLIENT_IDLE_TIMEOUT;
                pNewPacket->pInSMB = 0;
                pNewPacket->uiInSize = 0;