/*++

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

--*/

/*
 *
 *  This filter is a flavor of the filter IIS uses for compression.
 * (Original filter written by David Treadwell on July 1997.)
 *
 * ISA and IIS let you accumulate Requests chunks into a complete
 * Request.
 * The following filter is an example to a filter that collects the
 * response chunks and then allows you to change them depending on the
 * complete response.

 * This filter works with 2 notifications.
 * In Send Raw Data it collects the response's chunks, sends 0 bytes
 * instead of them (i.e. sends nothing).

 * Then, when all the chunks of this response passed Send Raw Data
 * notification, ISA thinks the complete response was sent. So
 * it calls End Of Request Notification. End Of Request Notification
 * will be the place where we will send the complete response.
 *
 */


#include <windows.h>
#include <httpfilt.h>
#include <stdio.h>
#include <stdlib.h>
#include <errno.h>

const char LINE_TO_INSERT[] = "<h3>Brought to you by an ISA Web Filter</h3>\r\n";
const char CONTENT_LENGTH[] = "Content-Length:";

static DWORD OnSendRawData(PHTTP_FILTER_CONTEXT pfc, PHTTP_FILTER_RAW_DATA pRawData);
static DWORD OnEndOfRequest(PHTTP_FILTER_CONTEXT pfc);
static void DisableNotifications(PHTTP_FILTER_CONTEXT pfc, DWORD   flags);
static int strnstr(const char *string, const char *strCharSet , int n);
static bool InsertOurHTMLHeader(PHTTP_FILTER_CONTEXT pfc,LPBYTE *ppBuffer,LPDWORD lpdwLen);
static DWORD dwfnContentLen(LPSTR ContentLenLine,int iLineLen);


BOOL WINAPI TerminateFilter ( DWORD dwFlags )
{
    UNREFERENCED_PARAMETER(dwFlags);
    return TRUE;
}


BOOL WINAPI GetFilterVersion ( PHTTP_FILTER_VERSION pVer )
{
    if (pVer == NULL) 
    {
        SetLastError( ERROR_INVALID_PARAMETER); 
        return FALSE;
    }
    
    pVer->dwFilterVersion = HTTP_FILTER_REVISION;

    pVer->dwFlags =  SF_NOTIFY_ORDER_HIGH | SF_NOTIFY_SECURE_PORT | SF_NOTIFY_NONSECURE_PORT
                    | SF_NOTIFY_END_OF_REQUEST | SF_NOTIFY_SEND_RAW_DATA ;

    return TRUE;
}


DWORD WINAPI HttpFilterProc (
    PHTTP_FILTER_CONTEXT pfc,
    DWORD NotificationType,
    LPVOID pvNotification
)
{

    DWORD dwRet = SF_STATUS_REQ_NEXT_NOTIFICATION;

    switch (NotificationType)
    {
        case SF_NOTIFY_SEND_RAW_DATA:
            dwRet = OnSendRawData(pfc, (PHTTP_FILTER_RAW_DATA)pvNotification);
            break;
        case SF_NOTIFY_END_OF_REQUEST:
            dwRet = OnEndOfRequest(pfc);
            break;
        default:
            // We cannot reach here, unless Web Filter support has a BAD ERROR.
            SetLastError( ERROR_INVALID_PARAMETER); 
            dwRet = SF_STATUS_REQ_ERROR;
            break;
    }

    return dwRet;
}

/*
 * OnSendRawData():
 * During Send Raw Data Notification we do the following:
 * 1) Append each chunk to an accumulation buffer (pRawData->cvInData).
 * 2) Resize the Current chunk to 0 ( don't send anything.)
 */
static DWORD OnSendRawData ( PHTTP_FILTER_CONTEXT pfc, PHTTP_FILTER_RAW_DATA pInRawData )
{

    if ( NULL == pfc || NULL == pInRawData)
    {
        SetLastError( ERROR_INVALID_PARAMETER); 
        return SF_STATUS_REQ_ERROR;
    }
  
    /*
     * Called first time for this request - then allocate pRawData.
     */
    DWORD dwReserved = 0;

    if ( NULL == pfc->pFilterContext )
    {
        pfc->pFilterContext = (LPVOID)pfc->AllocMem( pfc, sizeof(HTTP_FILTER_RAW_DATA), dwReserved);
        if ( NULL == pfc->pFilterContext )
        {
            DisableNotifications( pfc, SF_NOTIFY_END_OF_REQUEST | SF_NOTIFY_SEND_RAW_DATA);
            SetLastError( ERROR_NOT_ENOUGH_MEMORY );
            return SF_STATUS_REQ_ERROR;
        }

        PHTTP_FILTER_RAW_DATA pRawData = (PHTTP_FILTER_RAW_DATA)pfc->pFilterContext;

        pRawData->cbInBuffer = pInRawData->cbInBuffer;
        pRawData->pvInData = (LPVOID)pfc->AllocMem( pfc, pRawData->cbInBuffer, dwReserved);
        if ( NULL == pRawData->pvInData )
        {
            DisableNotifications(pfc,SF_NOTIFY_END_OF_REQUEST | SF_NOTIFY_SEND_RAW_DATA);
            SetLastError( ERROR_NOT_ENOUGH_MEMORY );
            return SF_STATUS_REQ_ERROR;
        }

        pRawData->cbInData = 0;
        pRawData->dwReserved = pInRawData->dwReserved ;
    }


    /*
     * Get the pRawData from the Request Context.
     */
    PHTTP_FILTER_RAW_DATA pRawData = (PHTTP_FILTER_RAW_DATA)pfc->pFilterContext;

    /*
     * If Not enough buffer in pRawData -> increase buffer.
     */
    if ( pInRawData->cbInData + pRawData->cbInData > pRawData->cbInBuffer)
    {
        pRawData->cbInBuffer = pInRawData->cbInData + pRawData->cbInBuffer ;
        LPBYTE lpBuffer = (LPBYTE)pfc->AllocMem(pfc,pRawData->cbInBuffer, dwReserved);
        if ( NULL == lpBuffer )
        {
            DisableNotifications(pfc,SF_NOTIFY_END_OF_REQUEST | SF_NOTIFY_SEND_RAW_DATA);
            SetLastError( ERROR_NOT_ENOUGH_MEMORY );
            return SF_STATUS_REQ_ERROR;
        }

        memcpy((LPVOID)lpBuffer, pRawData->pvInData, pRawData->cbInData);
        pRawData->pvInData = (void *)lpBuffer;
    }

    /*
     * Append InRawData ( new chunk )  to accumulation buffer.
     */
    LPBYTE lpBuffer = (LPBYTE)pRawData->pvInData;
    memcpy((LPVOID)(&lpBuffer[pRawData->cbInData]), pInRawData->pvInData, pInRawData->cbInData);
    pRawData->cbInData = pRawData->cbInData + pInRawData->cbInData;

    /*
     * Mark current chunk as size 0 ( i.e. Don't send enything. )
     */
    pInRawData->cbInData = 0;


    return SF_STATUS_REQ_NEXT_NOTIFICATION;
}



/*
 * OnEndOfRequest():
 * During End Of Request Notification we:
 * 1) Get the complete response from the filter Request Specific Data.
 * 2) May manipulate that response.
 * 3) Send that response.
 */
static DWORD OnEndOfRequest ( PHTTP_FILTER_CONTEXT pfc )
{

    if ( NULL == pfc)
    {
        SetLastError( ERROR_INVALID_PARAMETER); 
        return SF_STATUS_REQ_ERROR;
    }
  
    /*
     * Block the following SEND_RESPONSE and END_OF_REQUEST notifications for
     * this request. ( WriteClient() bellow will generate them.)
     *
     */
    DisableNotifications( pfc, SF_NOTIFY_END_OF_REQUEST | SF_NOTIFY_SEND_RAW_DATA);


    /*
     * OK all data was accumulated. Now extract from Request Data
     * the response buffer. (which now contains the complete response).
     */
    PHTTP_FILTER_RAW_DATA pRawData = (PHTTP_FILTER_RAW_DATA)pfc->pFilterContext;
    if ( NULL == pRawData)
    {
        return SF_STATUS_REQ_NEXT_NOTIFICATION;
    }

    LPBYTE lpBuffer = (LPBYTE)pRawData->pvInData;
    DWORD bytesToSend = pRawData->cbInData;
    /*
     * Empty Request data to make it ready for next response
     * in case the connection is kept alive
     */
    pfc->pFilterContext = NULL;

    /*
     * Add here code to modify the response. It is now the complete response.
     * any decisions and changes that demand the complete Response in hand
     * should be done here.
     */
     if (!InsertOurHTMLHeader( pfc, &lpBuffer, &bytesToSend))
     {
        SetLastError( ERROR_NOT_ENOUGH_MEMORY );
        return SF_STATUS_REQ_ERROR;
     }


    /*
     * Send the complete response in one chunk.
     */
    DWORD dwReserved = 0 ;
    if ( pfc->WriteClient(pfc,(LPVOID)lpBuffer,&bytesToSend,dwReserved) )
    {
        return SF_STATUS_REQ_NEXT_NOTIFICATION;
    }
    else
    {
        return SF_STATUS_REQ_ERROR;
    }


}


#define STRING_SIZE(str)    (sizeof(str) - 1)
/*
 *
 * InsertOurHTMLHeader()
 * Given the full response chunk, test if the response is in HTML format (has
 * <HTML>...</HTML> in it). If so, insert an HTML header line into the response
 * HTML. To do that we look at the input HTML as made of 3 chunks:
 * Chunk1 from start till "Content-Length: ..." header ( if exist.)
 * Chunk2 from "\r\n" after "Content-Length:" header untill last character of
 * "<HTML>".
 * Chunk3 from first character after "<HTML>" till the end of the response.
 * We then calculate the new Content-Length value by adding the length of the
 * line we will insert to the original content length, and finally allocate a new
 * buffer and paste into it the chunks in the following order:
 * Chunk1, Updateded-Content-Length-Header, Chunk2, Line-we-wish-to-insert, chunk3.
 *
 */
static bool InsertOurHTMLHeader ( PHTTP_FILTER_CONTEXT pfc, LPBYTE *ppBuffer, LPDWORD lpdwLen )
{
    LPSTR lpBuffer = (LPSTR)*ppBuffer;
    DWORD dwLen = *lpdwLen;
    bool bStatus = true;

    if ( 0 < strnstr( lpBuffer, "<HTML>", dwLen)  &&
        0 < strnstr( lpBuffer, "</HTML>", dwLen) )
    {
        /*
         *
         * Find first chunk. If no Content-Length header then first chunk will be
         * zero sized.
         *
         */
        int iStartChunk1 = 0;
        int iEndChunk1 = strnstr( lpBuffer, CONTENT_LENGTH, dwLen);
        if ( 0 >= iEndChunk1 )
            iEndChunk1 = 0;


        int iTmp = 0;

        if ( 0 < iEndChunk1 )
        {
            /*
             *
             * If we are here then there is a Content-Length header, so we
             * find the first "\r\n" after the the Content-Length header.
             *
             */
            iTmp = strnstr( lpBuffer + iEndChunk1, "\r\n" , dwLen - iEndChunk1);
            if ( 0 >= iTmp ) 
                goto ReturnBuffer;
        }

        /*
         *
         * Find the second chunk. It starts on the "\r\n" after the Content-Length
         * Header, and ends after the first "<HTML>".
         *
         */
        int iStartChunk2 = iEndChunk1 + iTmp;
        iTmp = strnstr( lpBuffer + iStartChunk2, "<HTML>", dwLen - iStartChunk2);
        if ( 0 >= iTmp ) 
            goto ReturnBuffer;
        int iEndChunk2 = iStartChunk2 + iTmp + STRING_SIZE("<HTML>");

        /*
         *
         * The third and last chunk starts after the first "<HTML>" and goes till
         * the end of the input response.
         *
         */
        int iStartChunk3 = iEndChunk2;
        int iEndChunk3 = (int)dwLen;


        char szContentLen[100];
        if ( 0 < iEndChunk1 )
        {
            /*
             *
             * If there is Content-Length Header then find the Input Content Length,
             * and add to it the length of the inserted line.
             *
             */
            DWORD dwContentLen =
                dwfnContentLen( lpBuffer + iEndChunk1, iStartChunk2 - iEndChunk1) +
                STRING_SIZE(LINE_TO_INSERT);
            char *p = (char *) memcpy( szContentLen, CONTENT_LENGTH, STRING_SIZE(CONTENT_LENGTH)) + 
                STRING_SIZE(CONTENT_LENGTH);
            *p++ = ' ';
            _ultoa( dwContentLen, p, 10);
        }
        else
        {
            szContentLen[0] = 0;
        }

        /*
         *
         * Allocate a buffer for the new response.
         *
         */
        iTmp = iEndChunk1 - iStartChunk1 + strlen(szContentLen) + iEndChunk2 - iStartChunk2 +
               STRING_SIZE(LINE_TO_INSERT) + iEndChunk3 - iStartChunk3 + 1;
        DWORD dwReserved = 0;
        LPSTR lpNewBuffer = (LPSTR)pfc->AllocMem(pfc,iTmp, dwReserved);
        if ( !lpNewBuffer )
        {
            bStatus = false;
            goto ReturnBuffer;
        }

#define MEMCAT(lpTarget,lpSource,SourceLen,TargetLen) \
        memcpy(lpTarget,lpSource,SourceLen); TargetLen += SourceLen;


        DWORD dwNewLen = 0;

        MEMCAT( lpNewBuffer + dwNewLen, lpBuffer + iStartChunk1, iEndChunk1 - iStartChunk1, dwNewLen);

        MEMCAT( lpNewBuffer + dwNewLen, szContentLen, strlen(szContentLen), dwNewLen);

        MEMCAT( lpNewBuffer + dwNewLen, lpBuffer + iStartChunk2, iEndChunk2 - iStartChunk2, dwNewLen);

        MEMCAT( lpNewBuffer + dwNewLen, LINE_TO_INSERT, strlen(LINE_TO_INSERT), dwNewLen);

        MEMCAT( lpNewBuffer + dwNewLen, lpBuffer + iStartChunk3, iEndChunk3 - iStartChunk3, dwNewLen);


        /*
         * We don't have to do that, but making it a string is good for debugging
         */
        lpNewBuffer[dwNewLen] = 0;

        lpBuffer = lpNewBuffer;
        dwLen = dwNewLen;
    }

ReturnBuffer:
    *ppBuffer = (LPBYTE)lpBuffer;
    *lpdwLen = dwLen;
    return bStatus;
}

/*
 *
 * Extract the input Content-Length value.
 *
 */
static ULONG dwfnContentLen(LPSTR ContentLenLine, int iLineLen)
{
    int iColumn = strnstr(ContentLenLine,":",iLineLen) + strlen(":");
    char *s = new char[iLineLen - iColumn + 1];
    if (!s)
        return 0;
    strncpy(s, &ContentLenLine[iColumn],iLineLen - iColumn);
    s[iLineLen - iColumn] = 0;
    LPSTR stopstring;
    DWORD dwCL =  strtoul( s, &stopstring, 10 );       
    if (errno == ERANGE) // overflow
    {
        delete [] s;     
        return 0;
    }

    delete [] s;
    return dwCL;
}

/*
 *
 * Utility function to notify that a filter is not to be called
 * to a notification throughout the lifetime of the current Request.
 *
 */
static void DisableNotifications ( PHTTP_FILTER_CONTEXT pfc, DWORD   flags )
{
    pfc->ServerSupportFunction(
                      pfc,
                      SF_REQ_DISABLE_NOTIFICATIONS,
                      NULL,
                      flags,
                      0
                      );
}

/*
 * strnstr()
 * finds first appearance of strCharSet in string ignoring
 * letters case.
 */
static int strnstr ( const char *string, const char *strCharSet , int n)
{
    int len = (strCharSet  != NULL ) ? ((int)strlen(strCharSet )) : 0 ;

    if ( 0 == n || 0 == len )
    {
        return -1;
    }

    int ret = -1;
    BOOLEAN found = FALSE;
    for (int I = 0 ; I <= n - len && !(found) ; I++)
    {
        int J = 0 ;
        for ( ; J < len ; J++ )
        {
            if (toupper(string[I + J]) != toupper(strCharSet [J]))
            {
                break; // Exit For(J)
            }
        }

        if ( J == len)
        {
            found = TRUE;
            ret = I;
        }
    }

    return ret;
}