AsyncSecureSocket.h

Go to the documentation of this file.

/*  AsyncSecureSocket.h - SSL communication through the Schannel API
    Copyright (C) 2001-2004 Mark Weaver
    Written by Mark Weaver <mark@npsl.co.uk>

    Part of the Open-Win32 library.
    This library is free software; you can redistribute it and/or
    modify it under the terms of the GNU Library General Public
    License as published by the Free Software Foundation; either
    version 2 of the License, or (at your option) any later version.

    This library is distributed in the hope that it will be useful,
    but WITHOUT ANY WARRANTY; without even the implied warranty of
    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
    Library General Public License for more details.

    You should have received a copy of the GNU Library General Public
    License along with this library; if not, write to the
    Free Software Foundation, Inc., 59 Temple Place - Suite 330,
    Boston, MA  02111-1307, USA.
*/

#ifndef OW32_AsyncSecureSocket_h
#define OW32_AsyncSecureSocket_h

#include <OW32/windows.h>
#include <OW32/AsyncSocket.h>
#include <OW32/SSLInit.h>
#include <OW32/auto_array_ptr.h>
#include <tchar.h>
#include <cassert>

// Open Win32 namespace
namespace OW32
{

template <class T>
class CAsyncSecureSocket : public T
{
private:
    CAsyncSecureSocket(const CAsyncSecureSocket& );
    CAsyncSecureSocket& operator= (const CAsyncSecureSocket& );
    void initialise();

public:
    void SetCallback(CAsyncSocketCallback* pCallback) { m_pCallback = pCallback; }
    CAsyncSocketCallback* GetCallback() const { return m_pCallback; }

    virtual void onConnectCompletion(BOOL /*bRet*/) {}
    virtual void onCloseCompletion(BOOL /*bRet*/) {}
    virtual void onSendCompletion(BOOL bRet, DWORD cbSent);
    virtual void onReadCompletion(BOOL bRet, DWORD cbReceived);
    virtual void onTimeout() = 0;

    CAsyncSecureSocket(CAsyncSocketCallback* pCallback, SOCKET s) :
        T(&m_callback,s),
        m_pCallback(pCallback)
    {
        initialise();
    }

    CAsyncSecureSocket(CAsyncSocketCallback* pCallback) :
        T(&m_callback,s),
        m_pCallback(pCallback)
    {
        initialise();
    }

    void negotiate() { m_State = State_Negotiate; negotiateLoop();}

    ~CAsyncSecureSocket();

    virtual int send(const char* buf, int len);
    virtual int recv(char* buf, int len);

    void setCredentials(CredHandle hCreds) { m_hCreds = hCreds; }

    void setRequireClientAuth(bool bRequireClientAuth) { m_bRequireClientAuth = bRequireClientAuth; }

    static SECURITY_STATUS createCredentialsFromCertificate(CredHandle* phCreds, PCCERT_CONTEXT pCertContext,
            DWORD dwDirection, DWORD dwEnabledProtocols = 0);

    SECURITY_STATUS getRemoteCert(PCCERT_CONTEXT* pRemoteCertContext);

    void freeCredentials();

protected:
    virtual void negotiateLoop()=0;

    // interception layer
    class CAsyncSecureSocketCallback : public CAsyncSocketCallback
    {
        virtual void onReadCompletion(BOOL bRet, DWORD cbReceived) {
            CAsyncSecureSocket* pParent = (CAsyncSecureSocket*)((char*)this - offsetof(CAsyncSecureSocket,m_callback));
            pParent->onReadCompletion(bRet,cbReceived); 
        }
        virtual void onSendCompletion(BOOL bRet, DWORD cbSent) {
            CAsyncSecureSocket* pParent = (CAsyncSecureSocket*)((char*)this - offsetof(CAsyncSecureSocket,m_callback));
            pParent->onSendCompletion(bRet,cbSent); 
        }
        virtual void onConnectCompletion(BOOL bRet) { 
            CAsyncSecureSocket* pParent = (CAsyncSecureSocket*)((char*)this - offsetof(CAsyncSecureSocket,m_callback));
            pParent->onConnectCompletion(bRet); 
        }
        virtual void onCloseCompletion(BOOL bRet) { 
            CAsyncSecureSocket* pParent = (CAsyncSecureSocket*)((char*)this - offsetof(CAsyncSecureSocket,m_callback));
            pParent->onCloseCompletion(bRet); 
        }
        virtual void onTimeout() {
            CAsyncSecureSocket* pParent = (CAsyncSecureSocket*)((char *)this - offsetof(CAsyncSecureSocket,m_callback));
            pParent->GetCallback()->onTimeout();
        }
        virtual void onAllNotificationsProcessed() {
            CAsyncSecureSocket* pParent = (CAsyncSecureSocket*)((char *)this - offsetof(CAsyncSecureSocket,m_callback));
            pParent->GetCallback()->onAllNotificationsProcessed();
        }
    };
    friend CAsyncSecureSocketCallback;
    CAsyncSecureSocketCallback m_callback;

    void ReadCompletion(BOOL bRet, DWORD cbReceived);
    void SendCompletion(BOOL bRet, DWORD cbSent);
    SECURITY_STATUS querySizes();

    SecPkgContext_StreamSizes m_Sizes;
    CtxtHandle m_hContext;
    CredHandle m_hCreds;

    auto_array_ptr<char> m_Extra;
    int m_ExtraCount;
    int m_ExtraDecrypted,m_ExtraDecryptedPos;
    int m_recvLen;
    char* m_recvBuf;

    auto_array_ptr<char> m_SendBuf;
    const char* m_sendData;
    int m_bufferSent,m_totalBuffered;
    int m_sendProcessed,m_sendLength;

    bool m_bRequireClientAuth;
    bool m_ownCredentials;
    bool m_fUseIoCompletion;

    enum ConnectionState
    {
        State_Initial,
        State_Negotiate,
        State_Connected,
        State_Renegotiate,
        State_Shutdown
    };
    ConnectionState m_State;

    CAsyncSocketCallback* m_pCallback;
};

template <class T>
void CAsyncSecureSocket<T>::initialise()
{
    m_bRequireClientAuth = false;
    m_ownCredentials = false;
    m_State = State_Initial;

    m_ExtraCount = 0;
    m_ExtraDecrypted = 0;
    m_ExtraDecryptedPos = 0;
    
    m_sendData = 0;
    m_bufferSent = m_totalBuffered = 0;
    m_sendProcessed = m_sendLength = 0;
    
    SecInvalidateHandle(&m_hContext);
    SecInvalidateHandle(&m_hCreds); 

    // TODO:
    m_Sizes.cbMaximumMessage = (DWORD) -1;
    m_Sizes.cbHeader = (DWORD)-1;
    m_Sizes.cbTrailer = (DWORD)-1;
    m_Sizes.cbBlockSize = (DWORD)-1;
    m_Sizes.cBuffers = (DWORD)-1;
}

template <class T>
CAsyncSecureSocket<T>::~CAsyncSecureSocket()
{
    if (SecIsValidHandle(&m_hContext))
        g_SecurityFunc.DeleteSecurityContext(&m_hContext);
    freeCredentials();
}

template <class T>
void CAsyncSecureSocket<T>::freeCredentials()
{
    // free credentials if (a) we own them and (b) they are valid
    if (m_ownCredentials && SecIsValidHandle(&m_hCreds)) {
        g_SecurityFunc.DeleteSecurityContext(&m_hContext);
        SecInvalidateHandle(&m_hCreds);
        m_ownCredentials = false;
    }
}

template <class T>
SECURITY_STATUS CAsyncSecureSocket<T>::getRemoteCert(PCCERT_CONTEXT* pRemoteCertContext)
{
    return g_SecurityFunc.QueryContextAttributes(&m_hContext,
                                    SECPKG_ATTR_REMOTE_CERT_CONTEXT,
                                    (PVOID)pRemoteCertContext);
}

template <class T>
int CAsyncSecureSocket<T>::send(const char* buf, int len)
{
    if (m_State == State_Initial)
    {
        return T::send(buf,len);
    }
    m_bufferSent = m_totalBuffered = m_sendProcessed = 0;
    m_sendLength = len;
    m_sendData = buf;
    SendCompletion(TRUE, 0);
    return 0;
}

template <class T>
void CAsyncSecureSocket<T>::onSendCompletion(BOOL bRet, DWORD cbSent)
{
    if (!bRet || cbSent == 0) {
        m_pCallback->onSendCompletion(bRet, cbSent);
        return;
    }
    SendCompletion(bRet, cbSent);
}

template <class T>
void CAsyncSecureSocket<T>::SendCompletion(BOOL /*bRet*/, DWORD cbSent)
{
    // bump up the "amount of buffered data that we have to send that
    // has been sent" count.
    m_bufferSent += cbSent;

    SecBuffer       Buffers[4];
    SecBufferDesc   Message;

    Message.ulVersion = SECBUFFER_VERSION;
    Message.cBuffers = 4;
    Message.pBuffers = Buffers;

    // Create a single buffer for the header, trailer & message data
    // We will line these up so data can be sent in a single send() call.
    // TODO::
        assert(m_Sizes.cbMaximumMessage != -1 && m_Sizes.cbHeader != -1 &&
            m_Sizes.cbTrailer != -1 && m_Sizes.cbBlockSize != -1 && m_Sizes.cBuffers != -1);
    if (!m_SendBuf.get())
        m_SendBuf.reset(new char[m_Sizes.cbMaximumMessage+m_Sizes.cbHeader+m_Sizes.cbTrailer]);

    for (;;)
    {
        // Send the data
        if (m_bufferSent < m_totalBuffered) 
        {
            int ret = T::send(m_SendBuf.get() + m_bufferSent, m_totalBuffered - m_bufferSent);
            if (ret != 0)
                m_pCallback->onSendCompletion(FALSE, 0);
            return;
        }

        // Ok, we've now sent all of the buffered encrypted data; reset it
        m_totalBuffered = m_bufferSent = 0;

        // See if we have sent everything that we were supposed to, quit if so
        if (m_sendProcessed >= m_sendLength)
        {
            m_pCallback->onSendCompletion(TRUE, m_sendProcessed);
            return;
        }

        // Otherwise encrypt another chunk of data
        Buffers[0].pvBuffer = m_SendBuf.get();
        Buffers[0].cbBuffer = m_Sizes.cbHeader;
        Buffers[0].BufferType = SECBUFFER_STREAM_HEADER;

        // extract a chunk from the message
        DWORD chunksize = (m_sendLength-m_sendProcessed);
        if (chunksize > m_Sizes.cbMaximumMessage)
            chunksize = m_Sizes.cbMaximumMessage;
        CopyMemory(m_SendBuf.get()+m_Sizes.cbHeader, 
            m_sendData+m_sendProcessed, chunksize);

        Buffers[1].pvBuffer = m_SendBuf.get()+m_Sizes.cbHeader;
        Buffers[1].cbBuffer = chunksize;
        Buffers[1].BufferType = SECBUFFER_DATA;

        Buffers[2].pvBuffer = m_SendBuf.get()+m_Sizes.cbHeader+chunksize;
        Buffers[2].cbBuffer = m_Sizes.cbTrailer;
        Buffers[2].BufferType = SECBUFFER_STREAM_TRAILER;

        Buffers[3].BufferType = SECBUFFER_EMPTY;

        SECURITY_STATUS scRet = g_SecurityFunc.EncryptMessage(&m_hContext, 0, &Message, 0);
        if (FAILED(scRet))
        {
            SetLastError(scRet);
            m_pCallback->onSendCompletion(FALSE, 0);
            return;
        }

        m_sendProcessed += chunksize;
        m_totalBuffered = m_Sizes.cbHeader+chunksize+m_Sizes.cbTrailer;

    // loop around and send more data if required
    }
}

template <class T>
int CAsyncSecureSocket<T>::recv(char* buf, int len)
{
    if (m_State == State_Initial)
    {
        return T::recv(buf,len);
    }
    // A zero length read just returns 0
    // This is indicative of a weird client application, as 0
    // also indicates graceful close, but it is what winsock does.
    // TODO: broken, as read completion is not called
    if (len == 0) {
        assert(0);
        return 0;
    }
    m_recvBuf = buf;
    m_recvLen = len;
    ReadCompletion(TRUE, 0);
    return 0;
}

template <class T>
void CAsyncSecureSocket<T>::onReadCompletion(BOOL bRet, DWORD cbIoBuffer)
{
    // Pass on some kind of error or a connection closed 
    // notification  directly to our client
    if (!bRet || cbIoBuffer == 0) {
        m_pCallback->onReadCompletion(bRet, cbIoBuffer);
        return;
    }
    ReadCompletion(bRet, cbIoBuffer);
}

template <class T>
void CAsyncSecureSocket<T>::ReadCompletion(BOOL /*bRet*/, DWORD cbIoBuffer)
{
    // Set up some SecBuffers
    // The schannel API specifies that we need 4
    SecBuffer       Buffers[4];
    SecBufferDesc   Message;

    Message.ulVersion = SECBUFFER_VERSION;
    Message.cBuffers = 4;
    Message.pBuffers = Buffers;

    // Start with no data, and no problems
    SECURITY_STATUS scRet = SEC_E_OK;

    // We will "cheat" here and always try to read up to the maximum SSL 
    // message size.  This is because sometimes, in order to read the amount 
    // of data requested (e.g. 1 byte), we need more data to actually
    // decrypt.  We also always try and decrypt the maximum amount of data
    // possible (again to avoid SEC_E_INCOMPLETE_MESSAGE).  Any decrypted
    // data in excess of that which the user desires is buffered.
    // There is a further buffering requirement if we end up with `extra'
    // data.

    // suck off any pre-decrypted data
    int DecryptedCount = m_ExtraDecrypted-m_ExtraDecryptedPos;
    if (DecryptedCount > 0) {
        m_recvLen = min(DecryptedCount, m_recvLen);

        CopyMemory(m_recvBuf, m_Extra.get()+m_ExtraDecryptedPos, m_recvLen);
        m_ExtraDecryptedPos += m_recvLen;
        m_pCallback->onReadCompletion(TRUE, m_recvLen);
        return;
    }

    if (!m_Extra.get())
        m_Extra.reset(new char[m_Sizes.cbMaximumMessage+m_Sizes.cbHeader+m_Sizes.cbTrailer]);

    // Add any data from a previous handshake or read into the decryption buffer
    if (m_ExtraCount) 
    {
        if (m_ExtraDecrypted)
            MoveMemory(m_Extra.get(), m_Extra.get()+m_ExtraDecrypted, m_ExtraCount);
        cbIoBuffer += m_ExtraCount;
    }

    // Clear the extra data as we are now finished with it; either it is
    // queued for decryption or it has been stuffed into the output buffer
    m_ExtraCount = 0;
    m_ExtraDecryptedPos = 0;
    m_ExtraDecrypted = 0;

    do
    {
        // TODO: check cbIOBuffer <= m_Sizes.cbMax?
        // Fetch more data if we need to do so (remember, we might be decrypting 
        // the 'extra' data)
        if (cbIoBuffer == 0 || scRet == SEC_E_INCOMPLETE_MESSAGE) 
        {
            // Trap needing more data than the maximum message size
            int get = m_Sizes.cbHeader + m_Sizes.cbMaximumMessage + m_Sizes.cbTrailer - cbIoBuffer;
            if (get <= 0) {
                SetLastError((DWORD)E_UNEXPECTED); // TODO::
                m_pCallback->onReadCompletion(FALSE, 0);
                return;
            }

            int ret = T::recv(m_Extra.get() + cbIoBuffer, get);
            m_ExtraCount = cbIoBuffer;
            if (ret != 0)
                m_pCallback->onReadCompletion(FALSE, 0);
            return;
        }

        Buffers[0].pvBuffer = m_Extra.get();
        Buffers[0].cbBuffer = cbIoBuffer;
        Buffers[0].BufferType = SECBUFFER_DATA;

        Buffers[1].BufferType = SECBUFFER_EMPTY;
        Buffers[2].BufferType = SECBUFFER_EMPTY;
        Buffers[3].BufferType = SECBUFFER_EMPTY;

        scRet = g_SecurityFunc.DecryptMessage(&m_hContext, &Message, 0, NULL);
    }
    while (scRet == SEC_E_INCOMPLETE_MESSAGE);

    // docs inconsistent on which we get
    if (scRet == SEC_I_CONTEXT_EXPIRED || scRet == SEC_E_CONTEXT_EXPIRED) {
/*      Buffers[0].pvBuffer = NULL;
        Buffers[0].cbBuffer = 0;
        Buffers[0].BufferType = SECBUFFER_EMPTY;
        Message.cBuffers = 1;

        scRet = EncryptMessage(&m_hContext, 0, &Message, 0);
        if (scRet != SEC_E_OK) {
            SetLastError(scRet);
            return SOCKET_ERROR;
        }

        if (Buffers[0].cbBuffer > 0 && Buffers[0].pvBuffer != NULL) {
            CSocket::send(Buffers[0].pvBuffer, Buffers[0].cbBuffer);
            // TODO:: ? g_SecurityFunc.FreeContextBuffer( OutBuffers[0].pvBuffer );
        }*/
        // No sample, docs unclear, fuck it, just return zero
        // TODO::
        m_pCallback->onReadCompletion(TRUE, 0);
        return ;
    }

    // check for an error
    if (scRet != SEC_E_OK && scRet != SEC_I_RENEGOTIATE)
    {
        SetLastError(scRet);
        m_pCallback->onReadCompletion(FALSE, m_recvLen);
        return;
    }

    // ( Just assume the data buffer is #1 - it always is )
    // Copy an appropriate amount of data to the user provided buffer
    DWORD cbGot = 0;
    if (Buffers[1].BufferType == SECBUFFER_DATA) {
        cbGot = min(m_recvLen, (int)Buffers[1].cbBuffer);
        CopyMemory(m_recvBuf, Buffers[1].pvBuffer, cbGot);
    
        // Store extra decrypted data for later
        if ((int)Buffers[1].cbBuffer > cbGot) {
            MoveMemory(m_Extra.get(), (BYTE*)Buffers[1].pvBuffer + cbGot, Buffers[1].cbBuffer - cbGot);
            m_ExtraDecrypted = Buffers[1].cbBuffer - cbGot;
        }
    }

    // ( Just assume the extra buffer is #3 - it always is )
    // See if extra data is returned
    if (Buffers[3].BufferType == SECBUFFER_EXTRA && Buffers[3].cbBuffer > 0) {
        m_ExtraCount = Buffers[3].cbBuffer;
        MoveMemory(m_Extra.get() + m_ExtraDecrypted, 
            m_Extra.get() + cbIoBuffer - m_ExtraCount, m_ExtraCount);
    }

    // Handle a renegotiate request
    if (scRet == SEC_I_RENEGOTIATE) {
        // TODO : what if cbGot != 0 ?
        // It shouldn't happen, AFAICT schannel should only decrypt up
        // to the renegotiate token
        m_State = State_Renegotiate;
        negotiateLoop();
        return;
    }
    m_pCallback->onReadCompletion(TRUE, cbGot);
}

template <class T>
SECURITY_STATUS CAsyncSecureSocket<T>::querySizes()
{
    //return 
#if 0
    // REF: http://support.microsoft.com/default.aspx?kbid=300562
    // Basically, this doesn't report the correct sizes.  Despite the "fixed" notification,
    // this still isn't giving the right answer, so we just hard-code the values.
    return g_SecurityFunc.QueryContextAttributes(&m_hContext, SECPKG_ATTR_STREAM_SIZES, &m_Sizes);
#else
    m_Sizes.cbMaximumMessage = 16384;
    m_Sizes.cbHeader         = 5;
    m_Sizes.cbTrailer        = 16;
    m_Sizes.cBuffers         = 4;
    m_Sizes.cbBlockSize      = 1;
    return S_OK;
#endif
}

template <class T>
SECURITY_STATUS CAsyncSecureSocket<T>::createCredentialsFromCertificate(
    CredHandle* phCredHandle, 
    PCCERT_CONTEXT pCertContext,
    DWORD dwDirection,
    DWORD dwEnabledProtocols)
{
    // Build Schannel credential structure, with the desired protocols
    SCHANNEL_CRED SchannelCred;
    ZeroMemory(&SchannelCred, sizeof(SchannelCred));

    SchannelCred.dwVersion  = SCHANNEL_CRED_VERSION;
    SchannelCred.grbitEnabledProtocols = dwEnabledProtocols;

    OSVERSIONINFO osVer;
    ZeroMemory(&osVer, sizeof(osVer));
    osVer.dwOSVersionInfoSize = sizeof(osVer);
    if (!GetVersionEx(&osVer))
        return HRESULT_FROM_WIN32(::GetLastError());

    if (osVer.dwMajorVersion >= 5) {
        if (dwDirection == SECPKG_CRED_INBOUND) {
            SchannelCred.dwFlags |= SCH_CRED_NO_SYSTEM_MAPPER;
        } else {
            SchannelCred.dwFlags |= SCH_CRED_NO_DEFAULT_CREDS;
            SchannelCred.dwFlags |= SCH_CRED_MANUAL_CRED_VALIDATION;
        }
    }

    if (pCertContext) {
        SchannelCred.cCreds = 1;
        SchannelCred.paCred = &pCertContext;
    }

    // Now create an SSPI credential.
    TimeStamp tsExpiry;

    SECURITY_STATUS Status = g_SecurityFunc.AcquireCredentialsHandle(
                        NULL,                   // Name of principal    
                        UNISP_NAME,             // Name of package
                        dwDirection,            // Flags indicating use
                        NULL,                   // Pointer to logon ID
                        &SchannelCred,          // Package specific data
                        NULL,                   // Pointer to GetKey() func
                        NULL,                   // Value to pass to GetKey()
                        phCredHandle,           // (out) Cred Handle
                        &tsExpiry);             // (out) Lifetime (optional)
    return Status;
}

} // namespace OW32

#endif // OW32_AsyncSecureSocket_h

---