/*++

   Copyright (c) 2000 Microsoft Corporation.  All rights reserved.

--*/

// EBScannerDataFilter.cpp : Implementation of CEBScannerDataFilter
#include "stdafx.h"
#include "exeblock.h"
#include "EBScannerDataFilter.h"

/////////////////////////////////////////////////////////////////////////////
// CEBScannerDataFilter



/////////////////////////////////////////////////////////////////////////////
// CEBScannerDataFilter::SetSockets
//
// Implements IFWXDataFilter::SetSockets
//
// Initialize the data-pump for scanned files.
//
// We create a temporary file that will hold the downloaded file, the file
// will be scanned, and only if it is ok, the file will be passed to the
// client.
//
/////////////////////////////////////////////////////////////////////////////
STDMETHODIMP
CEBScannerDataFilter::SetSockets(
    IN   IFWXSocket *piInternalSocket,
    IN   IFWXSocket *piExternalSocket,
    IN   IFWXConnection *piConnection,
    IN   IUnknown *punkFilterContext
    )
{
    UNREFERENCED_PARAMETER(punkFilterContext);
    UNREFERENCED_PARAMETER(piConnection);

    Lock();
    m_spInternalSocket = piInternalSocket;
    m_spExternalSocket = piExternalSocket;
    Unlock();

    TCHAR TmpFileName[MAX_PATH];
    TCHAR TmpPath[MAX_PATH];

    if (GetTempPath(sizeof(TmpPath)/sizeof(TCHAR), TmpPath) == 0 ||
        GetTempFileName(TmpPath, _T("FtP"), 0, TmpFileName) == 0) 
    {
        _AbortWithConstantMessage("451 exeblock - Can't create temp filename\r\n");
        return HRESULT_FROM_WIN32(GetLastError());
    }

    m_FileHandle = CreateFile(
        TmpFileName,
        GENERIC_READ | GENERIC_WRITE,
        FILE_SHARE_READ | FILE_SHARE_DELETE,
        NULL,
        OPEN_ALWAYS,
        FILE_ATTRIBUTE_TEMPORARY | FILE_FLAG_DELETE_ON_CLOSE | FILE_FLAG_OVERLAPPED,
        NULL
        );

    if (m_FileHandle == INVALID_HANDLE_VALUE) 
    {
        _AbortWithConstantMessage("451 exeblock - Can't create temp file\r\n");
        return HRESULT_FROM_WIN32(GetLastError());
    }

    HRESULT hr;
    hr = m_spCallBackInterface->HookToCompletionPort((DWORD)m_FileHandle);

    if (FAILED(hr)) 
    {
        return hr;
    }

    return piExternalSocket->Recv(NULL, this, 0);
}


/////////////////////////////////////////////////////////////////////////////
// CEBScannerDataFilter::_AbortWithConstantMessage
//
// Abort the data transfer and return the given FTP reply. This reply will
// override the success reply from the FTP server.
//
/////////////////////////////////////////////////////////////////////////////
void
CEBScannerDataFilter::_AbortWithConstantMessage(LPSTR FtpReply)
{
//
// NOTE: This functions assumes that FtpReply is a constant string
//       hence the call to CreateConstBuffer.  If the string is dynamically
//       created, CreateBuffer should be used instead, and the string has
//       to be copied to the buffer.
//

    IFWXIOBuffer *iobuf;
    HRESULT hr = m_spCallBackInterface->CreateConstBuffer(
        (PBYTE)FtpReply,
        strlen(FtpReply),
        &iobuf);
    if (SUCCEEDED(hr)) 
    {
        m_pControlDataFilter->OverrideReply(iobuf);
        iobuf->Release();
    }

    _Abort();
}

/////////////////////////////////////////////////////////////////////////////
// CEBScannerDataFilter::_Abort
//
// Abort the data connection by closeing both sockets.
//
/////////////////////////////////////////////////////////////////////////////
void
CEBScannerDataFilter::_Abort()
{
    IFWXSocket *pIFWXSocket;

    //
    // Note that the shutdown here is graceful.  If we do an abortive
    // shutdown FTP.EXE assume that the connection was lost, and
    // aborts also the control connection and we loose the
    // option to display the reason for the failure.
    //

    pIFWXSocket = _GetInternalSocket();
    if (pIFWXSocket) 
    {
        pIFWXSocket->Close(FALSE);
        pIFWXSocket->Release();
    }

    pIFWXSocket = _GetExternalSocket();
    if (pIFWXSocket) 
    {
        pIFWXSocket->Close(FALSE);
        pIFWXSocket->Release();
    }
}



/////////////////////////////////////////////////////////////////////////////
// CEBScannerDataFilter::CompleteAsyncIO
//
// Implements IFWXExternalIOCompletion::CompleteAsyncIO
//
// This methods is called when the overlapped write of a buffer to the
// temp file is completed.  This is the time to receive the next buffer.
//
/////////////////////////////////////////////////////////////////////////////
HRESULT
STDMETHODCALLTYPE
CEBScannerDataFilter::CompleteAsyncIO(
    /* [in] */ BOOL fSuccess,
    /* [in] */ DWORD NumberOfBytesTransfered,
    /* [in] */ DWORD Win32ErrorCode,
    /* [in] */ IFWXOverlapped __RPC_FAR *pOverlapped,
    /* [in] */ UserContextType UserData
    )
{  
    UNREFERENCED_PARAMETER(UserData);
    UNREFERENCED_PARAMETER(Win32ErrorCode);
    UNREFERENCED_PARAMETER(NumberOfBytesTransfered);

    if (!fSuccess) _Abort();
    pOverlapped->Release();
    IFWXSocket *pIFWXSocket = _GetExternalSocket();

    if (pIFWXSocket) 
    {
        pIFWXSocket->Recv(NULL, this, 0);
        pIFWXSocket->Release();
    }

    return S_OK;
}



/////////////////////////////////////////////////////////////////////////////
// CEBScannerDataFilter::CompleteAsyncIO
//
// Implements IFWXIOCompletion::CompleteAsyncIO
//
// This methods is called when a Recv() call completed (all Recv calls
// are made with UserData = 0), or when a Send() call is completed
// (dwUserData != 0)
/////////////////////////////////////////////////////////////////////////////
HRESULT
CEBScannerDataFilter::CompleteAsyncIO(
    BOOL      fSuccess,
    DWORD     Win32ErrorCode,
    IFWXIOBuffer *pIOBuffer,
    UserContextType  UserData,
    PSOCKADDR From,
    INT       FromLen
    )
{
    UNREFERENCED_PARAMETER(FromLen);
    UNREFERENCED_PARAMETER(From);
    UNREFERENCED_PARAMETER(Win32ErrorCode);

    DWORD DataLength = 0;
    HRESULT hr;
    PBYTE pbBuffer = NULL;


    if (pIOBuffer) 
    {
        pIOBuffer->GetBufferAndSize(&pbBuffer, &DataLength);
    }

    if (UserData == 0) 
    {
        if (!fSuccess) 
        {
            _Abort();
            return S_OK;
        }

        if (DataLength == 0) 
        {
            _ScanFile();
            return S_OK;
        }

        IFWXOverlapped * pOverlapped = NULL;

        hr = m_spCallBackInterface->CreateOverlapped(&pOverlapped);

        OVERLAPPED *pov = NULL;

        if (SUCCEEDED(hr)) 
        {
            hr = pOverlapped->GetOverlapped(&pov);
        }

        if (SUCCEEDED(hr)) 
        {
	        hr = pOverlapped->SetNotificationInterface(this, 0);
        }

        if (FAILED(hr)) 
        {
            if (pOverlapped) pOverlapped->Release();
            _Abort();
            return hr;
        }

        pov->Offset = m_FileSize;
        m_FileSize += DataLength;
        if (!WriteFile(
            m_FileHandle,
            pbBuffer,
            DataLength,
            NULL,
            pov) && GetLastError() != ERROR_IO_PENDING) 
        {
            pOverlapped->Release();
            _Abort();
        }

    }
	else 
	{
        if (fSuccess) 
        {
            _WriteNextChunk();
        }
        else 
        {
            _Abort();
        }
    }


    return S_OK;
}


//////////////////////////////////////////////////////////////////////////////
//      CEBScannerDataFilter::Detach
//
// Implements: IFWXDataFilter::Detach
//
// Release the reference held to the external and internal sockets.
//////////////////////////////////////////////////////////////////////////////
STDMETHODIMP CEBScannerDataFilter::Detach()
{

    Lock();
    m_spExternalSocket = NULL;
    m_spInternalSocket = NULL;
    Unlock();

    return S_OK;
}

//////////////////////////////////////////////////////////////////////////////
// CEBScannerDataFilter::_ScanFile
//
// Scan the file (this sample blocks all executables). If the file is ok
// start sending it to the client.
//////////////////////////////////////////////////////////////////////////////
void
CEBScannerDataFilter::_ScanFile()
{
    HRESULT hr;
    m_BytesSent = 0;
    const char ExeFileHeader[] = {'M','Z'};
    DWORD ExeFileHeaderSize = sizeof(ExeFileHeader)/sizeof(char);
    BYTE BlockedData;
    DWORD DataSize = 1;
    if (!m_PreparedData) 
    {
         _AbortWithConstantMessage("451 exeblock internal error (m_PreparedData)\r\n");
        return;    
    }
    hr = m_PreparedData->GetData(&DataSize, &BlockedData);
    if (FAILED(hr)) 
    {
        _AbortWithConstantMessage("451 exeblock internal error (m_PreparedData->GetData)\r\n");
        return;
    }

    // Deallocate previously allocated resources.
    if (m_FileView) 
    {
        UnmapViewOfFile(m_FileView);
        m_FileView = NULL;
    }

    if (m_FileMapping) 
    {
        CloseHandle(m_FileMapping);
        m_FileMapping = NULL;
    }

    // The file size can be zero when listing the contents of an empty
    // directory.
    if (m_FileSize) 
    {
        m_FileMapping = CreateFileMapping(
            m_FileHandle,
            NULL,
            PAGE_READONLY,
            0,
            m_FileSize,
            NULL
            );

        if (m_FileMapping == NULL) 
        {
            _AbortWithConstantMessage("451 exeblock internal error (mapping)\r\n");
            return;
        }

        m_FileView = (PBYTE) MapViewOfFile(
            m_FileMapping,
            FILE_MAP_READ,
            0,
            0,
            m_FileSize
            );
    }

    //
    // Actual virus checking should go here.
    // This sample will only scan for 'MZ' or 'ZM' at the beginning of a file
    // and consider such file as infected.  This will make all .EXE files
    // seem infected.
    //

    if ((m_FileView == NULL) && m_FileSize) 
    {
        _AbortWithConstantMessage("451 exeblock internal error (map view)\r\n");
    } 
	else if (m_FileView && (ExeFileHeaderSize < m_FileSize) &&  
                !memcmp(m_FileView, ExeFileHeader, ExeFileHeaderSize)&&
                (BlockedData != 0))
    {
        _AbortWithConstantMessage("451 exeblock: aborting transfer of file with bad data\r\n");
    } 
    else 
    {
        // File is ok - start sending it to the client.
        _WriteNextChunk();
    }
}


//////////////////////////////////////////////////////////////////////////////
// CEBScannerDataFilter::_WriteNextChunk
//
// Send one block of data from the file.
//////////////////////////////////////////////////////////////////////////////
void
CEBScannerDataFilter::_WriteNextChunk()
{
    IFWXSocket *pInternalSocket = _GetInternalSocket();
    if (pInternalSocket == NULL) {
        return;
    }

    // The choice of 8kb here is somewhat arbitrary.
    // 8kb is the default size of the TCP buffers so it seems a reasonable
    // choice.
    DWORD BytesToSend = min(8192, m_FileSize - m_BytesSent);

    if (BytesToSend) 
    {
        IFWXIOBuffer  * cbuf;
        HRESULT hr = m_spCallBackInterface->CreateConstBuffer(
            (PBYTE)m_FileView + m_BytesSent,
            BytesToSend,
            &cbuf
			);
        if (FAILED(hr)) 
        {
            _Abort();
            return;
        }
        m_BytesSent += BytesToSend;
        pInternalSocket->Send(cbuf, this, 1);
        cbuf->Release();
    } 
    else 
    {
        // Nothing more to send.  Shutdown the sockets and send 226 reply
        _AbortWithConstantMessage("226  Transfer complete. (scanned by exeblock)\r\n");
    }

    pInternalSocket->Release();
}