/*++

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

--*/

// socksfilter.cpp : Implementation of Csocksfilter
#include "stdafx.h"
#include "Socksfltr.h"
#include "socksfilter.h"
#include "outputdebugstringf.h"
#include "winsock2.h"
#include "SocksConnection.h"
#include "FilterDataStructures.h"
#include "impersonator.h"

//
// the following socket option disallows local address reuse
// The definition here is needed for old versions of winsock2.h
//
#ifndef SO_EXCLUSIVEADDRUSE
#define SO_EXCLUSIVEADDRUSE ((int)(~SO_REUSEADDR))
#endif

/////////////////////////////////////////////////////////////////////////////
// Csocksfilter
Csocksfilter::Csocksfilter () 
{
    m_SocksHookEvents.dwGlobalEvents = 0;
    m_UserIPSessionItemMap.clear();
}

HRESULT STDMETHODCALLTYPE Csocksfilter::FilterInit(
                           /* [in] */ IFWXFirewall __RPC_FAR *pIFWXFirewall,
                           /* [out] */ FwxFilterHookEvents __RPC_FAR *pFilterHookEvents)
{
    HRESULT ret;
    SOCKADDR_IN sin;

    m_spIFWXFirewall = pIFWXFirewall;

    // Create a listenning Socket
    ret = m_spIFWXFirewall->CreateNetworkSocket(FWX_PROTOCOL_TCP,&m_spListenSocket);
    if (FAILED(ret))
        return ret;

    BOOL val = TRUE;
    ret = m_spListenSocket->SetSockOpt(SOL_SOCKET, SO_EXCLUSIVEADDRUSE, (char*)&val, sizeof(val));
    if (FAILED(ret))
    {
        return ret;
    }

    sin.sin_family = AF_INET;
    sin.sin_port   = htons(SOCKS_PORT);
    sin.sin_addr.s_addr   = INADDR_ANY;

    ret =  m_spListenSocket->Bind(LPSOCKADDR(&sin), sizeof(sin), NULL);
    if (FAILED(ret)) {
        return ret;
    }

    ret =  m_spListenSocket->Listen(5);
    if (FAILED(ret)) {
        return ret;
    }

    ret =  m_spListenSocket->Accept(this,0);
    if (FAILED(ret)) {
        return ret;
    }

    FwxFilterHookEvents SocksHookEvents ;
    SocksHookEvents.dwGlobalEvents = 0;

    *pFilterHookEvents = SocksHookEvents;

    return S_OK;
}

HRESULT STDMETHODCALLTYPE Csocksfilter::ReloadConfiguration( void)
{
    return S_OK;
}

HRESULT STDMETHODCALLTYPE Csocksfilter::FilterShutdown( void)
{
    if ( m_spListenSocket)
    {
         m_spListenSocket->Close(FALSE);
    }

    return S_OK;
}

HRESULT STDMETHODCALLTYPE Csocksfilter::AttachToSession(
                        /* [in] */ IFWXSession __RPC_FAR *piSession,
                        /* [out] */ IFWXSessionFilter __RPC_FAR *__RPC_FAR *piSessionFilter,
                        /* [out] */ FwxFilterHookEvents __RPC_FAR *pFilterHookEvents)
{

    UNREFERENCED_PARAMETER(piSession);
    UNREFERENCED_PARAMETER(piSessionFilter);
    UNREFERENCED_PARAMETER(pFilterHookEvents);

    // Not register to any event
    return S_OK;
}




HRESULT
Csocksfilter::CompleteAsyncAccept(
    BOOL                fSuccess,
    DWORD               Win32ErrorCode,
    IFWXNetworkSocket*  /*pListeningSocket */,
    IFWXNetworkSocket*  pAcceptSocket,
    PSOCKADDR           RemoteAddress,
    DWORD               /* RemoteAddressLength*/,
    PSOCKADDR           /* LocalAddress */,
    DWORD               /* LocalAddressLength */,
    UserContextType     /* UserData */
    )
{
    HRESULT hr;
    CComObject<CSocksConnection>  *pSocksConnection;

    if (!fSuccess) {
        // This error is expected after FilterShutdown is called
        if (Win32ErrorCode == ERROR_OPERATION_ABORTED &&
            m_spListenSocket == 0)
        {
            return S_OK;
        }

    // AcceptEx sometimes fails with the following errors:
    // ERROR_NETNAME_DELETED, ERROR_HOST_UNREACHABLE,
    // possibly other codes.  These errors should be ignored.
        IssueNextAccept();

        return S_OK;
    }

    IssueNextAccept();

    hr = CComObject<CSocksConnection>::CreateInstance(&pSocksConnection);
    if (FAILED(hr)) return hr;

    hr = pSocksConnection->Initialize(this,RemoteAddress,pAcceptSocket,m_spIFWXFirewall);

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



    return S_OK;
}


HRESULT Csocksfilter::InsertSocket(u_long ulIp,LPSTR user,CComPtr<IFWXNetworkSocket> &socket)
{
    UserIP newuser;
    newuser.IPAddr = ulIp;
    newuser.user   = user;

    Lock();
        m_UserIPSessionItemMap[newuser].socklist.push_front (socket);
    Unlock();
    return S_OK;
}
HRESULT Csocksfilter::InsertSession(u_long ulIp,LPSTR user, SessionItem &newSession)
{
    UserIP newuser;
    newuser.IPAddr = ulIp;
    newuser.user   = user;

    Lock();
        m_UserIPSessionItemMap[newuser] = newSession;
    Unlock();
    return S_OK;
}

IT Csocksfilter::findInMap(u_long ulIp,LPSTR user)
{

   UserIP thiskey;
   thiskey.IPAddr = ulIp;
   thiskey.user   = user;

   Lock();
       IT it = m_UserIPSessionItemMap.find(thiskey);
       if (it == m_UserIPSessionItemMap.end() )
       {
           Unlock();
           return NULL;
       }

   Unlock();


   return  it;
}


HRESULT Csocksfilter::DeleteFromMap(IT it)
{
    Lock();
        m_UserIPSessionItemMap.erase(it);
    Unlock();
    return S_OK;
}

HRESULT Csocksfilter::DeleteFromList(SOCKLIST *currList,SOCKLIST::iterator i)
{
    Lock();
        (*currList).erase(i);
    Unlock();
    return S_OK;
}


// Find a socket in the list of a certain Session
HRESULT Csocksfilter::findSocket(u_long ulIP,LPSTR user,DWORD ip,USHORT port,IFWXNetworkSocket** socket)
{
    SOCKADDR Addr;
    int len;
    HRESULT hr=S_OK;

    UserIP thiskey;
    thiskey.IPAddr = ulIP;
    thiskey.user   = user;


    Lock();

    IT it2;

    it2 =  m_UserIPSessionItemMap.find(thiskey);

    if (it2 ==  m_UserIPSessionItemMap.end() )
    // Not in map of sessions - create new enter
        DEBUGPRINT(("===>findSocket::The session NOT found in map\n"));
    else
    {
        SessionItem oldSession;
        oldSession = (*it2).second;
        SOCKLIST socklist;
        socklist = oldSession.socklist;
        if (socklist.empty() )
        {
            DEBUGPRINT(("Error - There are no sockets in the list!! \n"));
        }
        else
        {

            // Find in the socket list
            SOCKLIST::iterator i;

            for (i = socklist.begin(); i != socklist.end(); ++i)
            {
                len = sizeof(Addr);
                (*i).m_T->GetPeerName(&Addr,&len);

                // Compare the IP and PORT number
                if ( (LPSOCKADDR_IN(&Addr)->sin_port == port) && (LPSOCKADDR_IN(&Addr)->sin_addr.s_addr == (u_long)ip ) )
                {
                    Unlock();
            hr = ((CComPtr<IFWXNetworkSocket>)(*i)).CopyTo(socket);
            return hr;
                }

            }



        } // size not 0
    } // found in session's map


    Unlock();
    return hr;
}



HRESULT Csocksfilter::SubCount(SessionItem *session)
{

    Lock();
        session->Count--;
    Unlock();
    return S_OK;
}


HRESULT Csocksfilter::FindOrCreateSession(
    u_long ulIp,
    LPSTR user,
    LPSTR domain,
    SOCKADDR ClientAddr,
    SOCKADDR ProxyAddr,
    IFWXSession **spSession,
    HANDLE token
    )
{

    HRESULT hr = S_OK;

    Lock();

    IT it;
    it = findInMap(ulIp, user);

    /* Try to find this IP in the map*/
    // If not found in map
    if (it ==  NULL )
    {
        CComObject<CImpersonator> *pImper;
        CComPtr<IFWXImpersonator>    spImpersonate;

        hr = CComObject<CImpersonator>::CreateInstance(&pImper);
        if (FAILED(hr)) {
            return hr;
        }

        hr = pImper->Initialize(token);
        if (FAILED(hr)) {
            return hr;
        }

        int UserLen = strlen(user) + 1;
        BSTR UniUser = new WCHAR[UserLen];
        if (UniUser == NULL) {
            return E_OUTOFMEMORY;
        }
        if (!MultiByteToWideChar(CP_ACP, 0, user, UserLen, UniUser, UserLen)) {
            delete[] UniUser;
            return HRESULT_FROM_WIN32(GetLastError());
        }


        hr = pImper->put_Username(UniUser); 
        if (FAILED(hr)) {
            delete[] UniUser;
            return hr;
        }
        delete[] UniUser;


        int DomainLen = strlen(domain) + 1;
        BSTR UniDomain = new WCHAR[DomainLen];
        if (UniDomain == NULL) {
            return E_OUTOFMEMORY;
        }
        if (!MultiByteToWideChar(CP_ACP, 0, domain, DomainLen, UniDomain, DomainLen)) {
            delete[] UniDomain;
            return HRESULT_FROM_WIN32(GetLastError());
        }

        hr = pImper->put_UserGroups(UniDomain); 
        if (FAILED(hr)) {
            delete[] UniDomain;
            return hr;
        }
        delete[] UniDomain;

        hr = pImper->put_Namespace(AUTHENTICATED_AD_USER_NAMESPACE_W);
        if (FAILED(hr)) {
            return hr;
        }

        pImper->Initialize(token);

        // Create new seesion
        hr = m_spIFWXFirewall->CreatePrivateSession(
            &ClientAddr,
            sizeof(ClientAddr),
            &ProxyAddr,
            sizeof(ProxyAddr),
            pImper, //impersonate
            spSession);


        if (SUCCEEDED(hr))
        {
            // add this IP and Session to map.
            SessionItem newSession;
            newSession.Count =1 ;
            newSession.spSession = *spSession;
            InsertSession(ulIp,user,newSession);
        }
        else
        {
            Unlock();
            return hr;
        }


    }
    else // Found in the map
    {
            // no need in new session- Add one to the ref count
        (*it).second.Count++;;
        hr = (*it).second.spSession.CopyTo(spSession);
    }

    Unlock();

    return hr;
}


HRESULT Csocksfilter::FindSession(u_long ulIp,LPSTR user,IFWXSession **spSession)
{

    HRESULT hr=S_OK;

    Lock();

    IT it;
    it = findInMap(ulIp,user);


    /* Find this IP in the IP map - must have done a Connect first*/
    // If not found in map - error
    if (it == NULL )
    {
            // if the session was not found - it is an error
        Unlock();
        return E_FAIL;
    }
    else
    {
        // Add one to the ref count of this entry
        //            SessionItem oldSession = (*it).second;
        //            ++oldSession.Count;
        (*it).second.Count++;
        hr = (*it).second.spSession.CopyTo(spSession);
    }

    Unlock();
    return hr;
}