/*++

   Copyright (c) 1998-2002 Microsoft Corporation.  All rights reserved.

--*/
#include "stdafx.h"
#include "WebExeBlockFilterImpl.h"
#include "WebExeBlockFilter.h"
#include "OutputDebugStringF.h"
#include <assert.h>
#include <limits.h>

#define EXECUTABLE_SIGNATURE    "MZ"
#define CONTENT_LENGTH_SZ       "Content-Length:"
#define HTTP11                  "HTTP/1.1"
#define HTTP10                  "HTTP/1.0"
#define STRING_CONST_SIZE(x)    (sizeof(x) - 1)

/////////////////////////////////////////////////////////////////////////////
//  CWebExeBlockFilterImpl()
/////////////////////////////////////////////////////////////////////////////
CWebExeBlockFilterImpl::CWebExeBlockFilterImpl()
{
}

/////////////////////////////////////////////////////////////////////////////
//  ~CWebExeBlockFilterImpl()
/////////////////////////////////////////////////////////////////////////////
CWebExeBlockFilterImpl::~CWebExeBlockFilterImpl()
{
}

/////////////////////////////////////////////////////////////////////////////
//  Init()
/////////////////////////////////////////////////////////////////////////////
HRESULT CWebExeBlockFilterImpl::Init()
{
    //Initialize our stuff
    return S_OK;
}

/////////////////////////////////////////////////////////////////////////////
//  Reload()
/////////////////////////////////////////////////////////////////////////////
HRESULT CWebExeBlockFilterImpl::Reload()
{
    //Reload new configuration stuff
    return S_OK;
}

/////////////////////////////////////////////////////////////////////////////
//  Shutdown()
/////////////////////////////////////////////////////////////////////////////
HRESULT CWebExeBlockFilterImpl::Shutdown()
{
    //Shutdown/cleanup our stuff
    return S_OK;
}

/////////////////////////////////////////////////////////////////////////////
//  OnReceiveRawData()
//  Handle the SF_NOTIFY_RECEIVE_RESPONSE_RAW_DATA notification
/////////////////////////////////////////////////////////////////////////////
DWORD CWebExeBlockFilterImpl::OnReceiveResponseRawData(
                                PHTTP_FILTER_CONTEXT    pfc,
                                PHTTP_FILTER_RAW_DATA   pRawData
                                )
{
    OutputDebugStringF("WEBEXEBLOCK: OnReceiveResponseRawData called\n");

    DWORD dwError = NO_ERROR;
    BOOL  fNeedMoreNotifications = FALSE; //this will be set to TRUE, if we still need to process response in future notifications
    DWORD dwStatus = SF_STATUS_REQ_NEXT_NOTIFICATION;    
    
    CWebExeBlockRequestContext* pContext = (CWebExeBlockRequestContext*)pfc->pFilterContext;
    if (pContext == NULL) //first time we're called, don't have a context yet, allocate one
    {
        pContext = new CWebExeBlockRequestContext();
        if (pContext == NULL)
        {
            OutputDebugStringF("WEBEXEBLOCK: Failed to allocate memory for body\n");
            dwError = ERROR_OUTOFMEMORY;
            dwStatus = SF_STATUS_REQ_ERROR;
            goto Exit;
        }
        pfc->pFilterContext = pContext;
    }
    //else - we already have a context with some raw data saved. 

    //
    // save raw data just received
    //    
    PWPX_FILTER_CONTEXT pfcWpx = TO_WPX_FILTER_CONTEXT(pfc);
    assert(pfcWpx);
    
    char* pszRawData = (char*)pfcWpx->AllocMemoryPerRequest(pfc, pContext->m_dwRawDataLength + pRawData->cbInData, 0);
    if (pszRawData == NULL)
    {
        OutputDebugStringF("WEBEXEBLOCK: Failed to allocate memory for body\n");
        dwError = ERROR_OUTOFMEMORY;
        dwStatus = SF_STATUS_REQ_ERROR;
        goto Exit;
    }

    //
    // copy the old data
    //
    if (pContext->m_dwRawDataLength > 0) //if we have something from before, copy it first
    {
        memcpy(pszRawData, pContext->m_pszRawData, pContext->m_dwRawDataLength);
    }
    
    //
    // copy new data
    //
    memcpy(pszRawData + pContext->m_dwRawDataLength, pRawData->pvInData, pRawData->cbInData);

    //
    // no need to delete the old memory, since it will be freed by the proxy at the
    // end of this request.
    //
    pContext->m_pszRawData = pszRawData;
    pContext->m_dwRawDataLength += pRawData->cbInData;
    
    if (!pContext->m_fFinishedReceivingHeaders)
    {
        DWORD dwHeadersLen = 0;
        
        if (!FindEndOfHeaders(pContext->m_pszRawData, pContext->m_dwRawDataLength, &dwHeadersLen))
        {
            fNeedMoreNotifications = TRUE; //we haven't scanned the body for the signature yet
            dwStatus = SF_STATUS_REQ_NEXT_NOTIFICATION;
            goto Exit;
        }

        //
        // Fix pointer to the begining of the body, and adjust size of body
        //
        pContext->m_fFinishedReceivingHeaders = TRUE;
        assert(pContext->m_dwRawDataLength >= dwHeadersLen);
        pContext->m_dwBodyLength = pContext->m_dwRawDataLength - dwHeadersLen;  
        pContext->m_dwContentLength = GetContentLength(pContext->m_pszRawData, pContext->m_pszRawData + dwHeadersLen);

        //
        // Important Note:
        // --------------
        // In this sample we only handle 200 OK responses. 
        // In real web filter, need to think about 100-continue, 206 responses, etc...
        // In addition, we don't handle chuncked encoding, compression, etc...
        //                
        DWORD dwHttpStatus;
        dwError = GetHttpStatus(pContext->m_pszRawData, pContext->m_pszRawData + dwHeadersLen, &dwHttpStatus);
        if (dwError != ERROR_SUCCESS)
        {
            dwStatus = SF_STATUS_REQ_ERROR;
            goto Exit;
        }

        if (dwHttpStatus != 200)
        {
            OutputDebugStringF("WEBEXEBLOCK: Return code is: %d. This sample only handles 200 OK\n", dwHttpStatus);                    

            fNeedMoreNotifications = FALSE; 
            dwStatus = SF_STATUS_REQ_NEXT_NOTIFICATION; 
            goto Exit;
        }
    }
    else
    {
        pContext->m_dwBodyLength += pRawData->cbInData; //anything received after the end of the headers is part of the body
    }

    if (pContext->m_dwBodyLength > 0)
    {
        pContext->m_pszBody = pContext->m_pszRawData + pContext->m_dwRawDataLength - pContext->m_dwBodyLength;
    }
    
    //
    // If we have in the body more than > length_of(EXECUTABLE_SIGNATURE) - look for EXECUTABLE_SIGNATURE, else continue.
    // NOTE - in this sample, the signature we look for is in the begining and is very small, so there's
    // no problem in the fact that we accumulate the data and perform the check before forwarding the response to the
    // client. In other cases, if your signature is large, or might appear in the middle of a file, then you
    // need to handle responses carefully - especially when content length is unknown. Avoid accumulating too much data!
    //
    if (pContext->m_dwBodyLength >= STRING_CONST_SIZE(EXECUTABLE_SIGNATURE))
    {
        if (strncmp(EXECUTABLE_SIGNATURE, pContext->m_pszBody, STRING_CONST_SIZE(EXECUTABLE_SIGNATURE)) == 0)
        {
            //
            // In case we found the signature in the begining of the body, this is an attachment.
            // We will block it.
            // (In this sample, we just return an error from the filter. The request will be dropped.
            //  Potentially, you could return your own error page to the client).
            //
            dwError = ERROR_REQUEST_ABORTED; //In this sample, we'll use this error to indicate an exe was blocked
            dwStatus = SF_STATUS_REQ_ERROR;
            goto Exit;
        }
        else //this doesn't start with the signature, so this is not an exe
        {
            OutputDebugStringF("WEBEXEBLOCK: Checked response body, this is not an exe\n");
             
            //
            // We know this is not an exe, so we can send all the data we have accumulated until now.
            //
            fNeedMoreNotifications = FALSE; 
            dwStatus = SF_STATUS_REQ_NEXT_NOTIFICATION; 
            goto Exit;
        }
    }
    else if (pContext->m_dwContentLength < STRING_CONST_SIZE(EXECUTABLE_SIGNATURE))
    {
        OutputDebugStringF("WEBEXEBLOCK: The response body is not an exe\n");
             
        //
        // We know this is not an exe, so we can send all the data we have accumulated until now.
        //
        fNeedMoreNotifications = FALSE; 
        dwStatus = SF_STATUS_REQ_NEXT_NOTIFICATION; 
        goto Exit;
    }
    else //We didn't receive yet enough of the response body, so we need future notifications to check for the signature 
    {
        fNeedMoreNotifications = TRUE; 
        dwStatus = SF_STATUS_REQ_READ_NEXT;
    }

Exit:
    //
    // We don't require anymore notifications:
    // - either we already know if this is an exe or not
    // - or we had an error, so no point to continue checking the response
    //
    if (!fNeedMoreNotifications && dwStatus != SF_STATUS_REQ_ERROR)
    {
        DWORD dwLocalError = DisableResponseRawDataNotification(pfc);
        if (dwLocalError != NO_ERROR)
        {
            if (dwError == NO_ERROR) //override the error only if we don't have a previous error
            {   
                dwError = dwLocalError;
            }
            dwStatus = SF_STATUS_REQ_ERROR;
        }

        //
        // set the raw data to point to our accumulated buffer, and pass it back 
        // to the proxy, to send to the client.
        //
        pRawData->pvInData = pContext->m_pszRawData;
        pRawData->cbInBuffer = pContext->m_dwRawDataLength;
        pRawData->cbInData = pContext->m_dwRawDataLength;

        pContext->m_pszRawData = NULL;
        pContext->m_dwRawDataLength = 0;
    }
    else
    {
        //
        // set the raw data to 0, until we determine if this response is ok or not
        //
        pRawData->pvInData = NULL;
        pRawData->cbInBuffer = 0;
        pRawData->cbInData = 0;
    }
    
    if (dwStatus == SF_STATUS_REQ_ERROR)
    {
        assert(!fNeedMoreNotifications);
        SetLastError(dwError);

        if (pContext != NULL)
        {
            delete pContext;
            pfc->pFilterContext = NULL; 
        }
    }
    
    return dwStatus;
}

/////////////////////////////////////////////////////////////////////////////
//  OnEndOfRequest()
//  Handle the SF_NOTIFY_END_OF_REQUEST notification
/////////////////////////////////////////////////////////////////////////////
DWORD CWebExeBlockFilterImpl::OnEndOfRequest(
                            PHTTP_FILTER_CONTEXT pfc)
{
    OutputDebugStringF("WEBEXEBLOCK: OnEndOfRequest called\n");

    DWORD dwStatus = SF_STATUS_REQ_NEXT_NOTIFICATION;
    CWebExeBlockRequestContext* pContext = (CWebExeBlockRequestContext*)pfc->pFilterContext;

    if (pContext != NULL && pContext->m_dwRawDataLength > 0)
    {   
        //
        // send the accumulated data to the client
        //
        if (!pfc->WriteClient(pfc, pContext->m_pszRawData, &pContext->m_dwRawDataLength, 0))
        {
            dwStatus = SF_STATUS_REQ_ERROR;
        }
        
        pContext->m_pszRawData = NULL;
        pContext->m_dwRawDataLength = 0;
    }

    pfc->pFilterContext = NULL; 
    delete pContext;
    
    return dwStatus;
}

//
// Helper function to disable SF_NOTIFY_RECEIVE_RESPONSE_RAW_DATA notification
//
DWORD CWebExeBlockFilterImpl::DisableResponseRawDataNotification(PHTTP_FILTER_CONTEXT pfc)
{
    DWORD dwError = NO_ERROR;
    BOOL  fRet = FALSE;
    //
    // Note that we don't disable SF_NOTIFY_END_OF_REQUEST, since we still need it for cleanup
    //

    PWPX_FILTER_CONTEXT pfcWpx = TO_WPX_FILTER_CONTEXT(pfc);
    assert(pfcWpx);
        
    fRet = pfcWpx->WPXSupportFunction( 
                            pfc, 
                            SF_REQ_DISABLE_WPX_NOTIFICATIONS,
                            NULL,
                            SF_NOTIFY_RECEIVE_RESPONSE_RAW_DATA, 
                            0 );
    if (!fRet)
    {
        dwError = GetLastError();
        OutputDebugStringF("WEBEXEBLOCK: Failed to disable WPX notifications. Error: %d\n", dwError);
        return dwError;
    }
    
    return NO_ERROR;
}

/////////////////////////////////////////////////////////////////////////////
//  FindEndOfHeaders()
//  Helper function which checks if we've received the end of the headers.
//  Returns TRUE if received all the headers, and FALSE otherwise.
//  On return, *pdwHeadersLen is set to be the length of the headers.
/////////////////////////////////////////////////////////////////////////////
BOOL CWebExeBlockFilterImpl::FindEndOfHeaders(
    const char* pszData, 
    const DWORD dwDataLen, 
    DWORD* pdwHeadersLen
    )
{
    for (DWORD i = 0; i < dwDataLen; i++)
    {
        if (pszData[i] == '\n' && i + 1 < dwDataLen)
        {
            if (pszData[i + 1] == '\n')
            {
                if (pdwHeadersLen != NULL)
                {
                    *pdwHeadersLen = i + 2;
                }
            
                return TRUE;
            }
            else if (pszData[i + 1] == '\r' && i + 2 < dwDataLen && pszData[i + 2] == '\n') 
            {
                if (pdwHeadersLen != NULL)
                {
                    *pdwHeadersLen = i + 3;
                }
            
                return TRUE;
            }
        }
    }

    if (pdwHeadersLen != NULL)
    {
        *pdwHeadersLen = 0;
    }

    return FALSE;
}

/////////////////////////////////////////////////////////////////////////////
//  GetContentLength()
//  Helper function which returns the content length value from the headers.
//  If the content length header does no exist, the function returns ULONG_MAX.
/////////////////////////////////////////////////////////////////////////////
DWORD CWebExeBlockFilterImpl::GetContentLength(const char* pszHeaders, char* pszHeadersEnd)
{
    const char cOld = *(pszHeadersEnd - 1);
    *(pszHeadersEnd - 1) = '\0';

    char* pszContentLength = strstr(pszHeaders, CONTENT_LENGTH_SZ);
    *(pszHeadersEnd - 1) = cOld;
    
    if (pszContentLength == NULL)
    {
        return ULONG_MAX;
    }

    pszContentLength += STRING_CONST_SIZE(CONTENT_LENGTH_SZ);

    //
    // skip white spaces
    //
    while (pszContentLength < pszHeadersEnd && (*pszContentLength == ' ' || *pszContentLength == '\t'))
    {
        pszContentLength++;
    }

    char* pszContentLengthEnd = NULL;
    return strtoul(pszContentLength, &pszContentLengthEnd, 10);
}


/////////////////////////////////////////////////////////////////////////////
//  GetHttpStatus()
//  Helper function which returns HTTP status code.
/////////////////////////////////////////////////////////////////////////////
DWORD CWebExeBlockFilterImpl::GetHttpStatus(
    const char* pszHeaders, 
    const char* pszHeadersEnd,
    DWORD* pdwHttpStatus
    )
{
    if (pszHeadersEnd - pszHeaders <= STRING_CONST_SIZE(HTTP11) ||
        pszHeadersEnd - pszHeaders <= STRING_CONST_SIZE(HTTP10) ||
        pdwHttpStatus == NULL)
    {
        return ERROR_REQUEST_ABORTED;
    }

    const char* pszCurr = pszHeaders;

    if (strncmp(pszCurr, HTTP11, STRING_CONST_SIZE(HTTP11)) != 0)
    {
        if (strncmp(pszCurr, HTTP10, STRING_CONST_SIZE(HTTP10)) != 0)
        {
            return ERROR_REQUEST_ABORTED;
        }

        pszCurr += STRING_CONST_SIZE(HTTP10);
    }
    else
    {
        pszCurr += STRING_CONST_SIZE(HTTP11);
    }

    //
    // skip white spaces
    //
    while (pszCurr < pszHeadersEnd && (*pszCurr == ' ' || *pszCurr == '\t'))
    {
        pszCurr++;
    }

    char* pszStatusEnd = NULL;
    *pdwHttpStatus = strtoul(pszCurr, &pszStatusEnd, 10);

    return ERROR_SUCCESS;
}