gsecuresocket.cpp

一个非常有用的开源代码

/*
	Copyright (C) 2006, Mike Gashler

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

	see http://www.gnu.org/copyleft/lesser.html
*/

#include "GSecureSocket.h"
#include "GMacros.h"
#include "GKeyPair.h"
#include "GBigNumber.h"
#include "GSpinLock.h"
#include "GPointerQueue.h"
#include "GRand.h"
#include "GArray.h"
#include "GCrypto.h"
#include "GXML.h"
#include "GThread.h"


// Handshake Protocol
// ------------------
// 1- Client: (unencrypted) My public key is ____.
// 2- Server: (unencrypted) My public key is ____.
// 3- Server: (enc w/ client public key) My passphrase is ____.
// 4- Client: (enc w/ server public key) My passphrase is ____.
// 5- Anyone: (enc w/ other passphrase) ...

#define HANDSHAKE_MAGIC 0x5ec50ce7

struct HandshakeHeader
{
	unsigned long nMagic; // HANDSHAKE_MAGIC
	unsigned short nMessage; // 0 = My public key is..., 1 = My passphrase is...
	// Data
};

GSecureSocketServer::GSecureSocketServer(int nMaxPacketSize, GRand* pRand, int nPort) : GSocketServer(false, nMaxPacketSize, nPort, 1)
{
	m_pRand = pRand;
	m_pKeyPair = new GKeyPair();
	m_pKeyPair->GenerateKeyPair(pRand);
	m_nPassphraseSize = pRand->GetRandByteCount();
	m_pPassphrases = new GPointerArray(16);
	m_pClientPassphrases = new GPointerArray(16);
	m_nClientPassphraseSizes = new GIntArray(16);
}

GSecureSocketServer::~GSecureSocketServer()
{
	int n;
	for(n = 0; n < m_pPassphrases->GetSize(); n++)
		delete((unsigned char*)m_pPassphrases->GetPointer(n));
	delete(m_pPassphrases);

	for(n = 0; n < m_pClientPassphrases->GetSize(); n++)
		delete((unsigned char*)m_pClientPassphrases->GetPointer(n));
	delete(m_pClientPassphrases);
	delete(m_nClientPassphraseSizes);
	delete(m_pKeyPair);
}

/*static*/ GSecureSocketServer* GSecureSocketServer::HostSecureSocket(u_short nPort, int nMaxPacketSize, GRand* pRand)
{
	// Check input
	GAssert(pRand->GetRandByteCount() >= (int)sizeof(int), "pRand not seeded big enough");
	GAssert(nMaxPacketSize > 0, "Bad max packet size");

	// Make the socket
	GSecureSocketServer* pSocket = new GSecureSocketServer(nMaxPacketSize, pRand, nPort);
	if(!pSocket)
		return NULL;
	return pSocket;
}

int GSecureSocketServer::GetMessageCount()
{
	while(true)
	{
		int nMessageCount = m_pMessageQueue->GetSize();
		if(nMessageCount > 0)
		{
			// See if we are done handshaking for this client
			GSocketMessage* pMessage = (GSocketMessage*)m_pMessageQueue->Peek();
			int nConnection = pMessage->GetConnection();

			// Make sure we have slots for the passphrases
			while(m_pPassphrases->GetSize() <= nConnection)
				m_pPassphrases->AddPointer(NULL);
			while(m_pClientPassphrases->GetSize() <= nConnection)
				m_pClientPassphrases->AddPointer(NULL);
			while(m_nClientPassphraseSizes->GetSize() <= nConnection)
				m_nClientPassphraseSizes->AddInt(0);

			// See if it's a handshake message or a real message
			const unsigned char* pClientPassphrase = (const unsigned char*)m_pClientPassphrases->GetPointer(nConnection);
			if(pClientPassphrase)
				return nMessageCount;
			else
				Handshake();
		}
		else
			return 0;
	}
}

void GSecureSocketServer::Handshake()
{
	// Check the header
	int nSize;
	int nConnection;
	unsigned char* pBuf = GSocketServer::GetNextMessage(&nSize, &nConnection);
	struct HandshakeHeader* pHeader = (struct HandshakeHeader*)pBuf;
	if(nSize <= (int)sizeof(struct HandshakeHeader) || pHeader->nMagic != HANDSHAKE_MAGIC)
	{
		GAssert(false, "Bad header packet");
		memset(pBuf, '\0', nSize);
		delete(pBuf);
		return;
	}

	// Process the payload
	unsigned char* pPayload = pBuf + sizeof(struct HandshakeHeader);
	int nPayloadSize = nSize - sizeof(struct HandshakeHeader);
	if(pHeader->nMessage == 0)
	{
		bool bRet = OnReceivePublicKey((const char*)pPayload, nPayloadSize, nConnection);
		if(!bRet)
			GAssert(false, "error processing public key");
	}
	else
	{
		GAssert(pHeader->nMessage == 1, "unexpected handshake packet");
		bool bRet = OnReceivePassphrase((const unsigned char*)pPayload, nPayloadSize, nConnection);
		if(!bRet)
			GAssert(false, "error processing passphrase");
	}
	memset(pBuf, '\0',nSize);
	delete(pBuf);
}

unsigned char* GSecureSocketServer::GetNextMessage(int* pnSize, int* pnOutConnectionNumber)
{
	// Get the encrypted message
	if(GetMessageCount() <= 0)
	{
		*pnOutConnectionNumber = -1;
		*pnSize = 0;
		return NULL;
	}
	unsigned char* pBuf = GSocketServer::GetNextMessage(pnSize, pnOutConnectionNumber);
	const unsigned char* pPassphrase = (const unsigned char*)m_pPassphrases->GetPointer(*pnOutConnectionNumber);
	GAssert(pPassphrase, "should have passphrase by now");

	// Decrypt the message
	GCrypto::Decrypt(pBuf, *pnSize, pPassphrase, m_nPassphraseSize);

	return pBuf;
}

bool GSecureSocketServer::SendPublicKey(int nConnection)
{
	bool bRet = true;
	GXMLTag* pTag = m_pKeyPair->ToXML(false);
	char* pMyPublicKey = pTag->ToString();
	int nPublicKeyBlobSize = strlen(pMyPublicKey);
	struct HandshakeHeader* pHeader = (struct HandshakeHeader*)alloca(sizeof(struct HandshakeHeader) + nPublicKeyBlobSize);
	memcpy(((unsigned char*)pHeader) + sizeof(struct HandshakeHeader), pMyPublicKey, nPublicKeyBlobSize);
	pHeader->nMagic = HANDSHAKE_MAGIC;
	pHeader->nMessage = 0;
	if(!GSocketServer::Send(pHeader, sizeof(struct HandshakeHeader) + nPublicKeyBlobSize, nConnection))
		bRet = false;
	delete(pMyPublicKey);
	delete(pTag);
	return bRet;
}

bool GSecureSocketServer::SendPassphrase(int nConnection, GKeyPair* pKeyPair)
{
	bool bRet = true;

	// Make a passphrase for this client
	unsigned char* pPassphrase = new unsigned char[m_nPassphraseSize];
	memcpy(pPassphrase, m_pRand->GetRand(), m_nPassphraseSize);
	if(pPassphrase[m_nPassphraseSize - 1] == 0)
		pPassphrase[m_nPassphraseSize - 1] = 1;
	m_pPassphrases->SetPointer(nConnection, pPassphrase);

	// Encrypt my passphrase with client's public key
	int nCypherSize;
	unsigned char* pCypher = pKeyPair->PowerMod(pPassphrase, m_nPassphraseSize, true, &nCypherSize);

	// Send it to the client
	struct HandshakeHeader* pHeader = (struct HandshakeHeader*)alloca(sizeof(struct HandshakeHeader) + nCypherSize);
	memcpy(((unsigned char*)pHeader) + sizeof(struct HandshakeHeader), pCypher, nCypherSize);
	pHeader->nMagic = HANDSHAKE_MAGIC;
	pHeader->nMessage = 1;
	if(!GSocketServer::Send(pHeader, sizeof(struct HandshakeHeader) + nCypherSize, nConnection))
	{
		GAssert(false, "Error sending passphrase");
		bRet = false;
	}
	delete(pCypher);
	return bRet;
}

bool GSecureSocketServer::OnReceivePublicKey(const char* pXML, int nXMLSize, int nConnectionNumber)
{
	// Decode client's public key
	GKeyPair kp;
	if(!kp.FromXML(pXML, nXMLSize))
	{
		GAssert(false, "Bad keypair");
		return false;
	}

	// Send my public key
	if(!SendPublicKey(nConnectionNumber))
	{
		GAssert(false, "Error sending my public key");
	}

	// Send my passphrase for this client
	if(!SendPassphrase(nConnectionNumber, &kp))
	{
		GAssert(false, "Error sending passphrase");
	}
	
	return true;
}

bool GSecureSocketServer::OnReceivePassphrase(const unsigned char* pCypher, int nCypherSize, int nConnectionNumber)
{
	// Decrypt it with my private key
	int nPassphraseSize;
	unsigned char* pPassphrase = m_pKeyPair->PowerMod(pCypher, nCypherSize, false, &nPassphraseSize);
	m_pClientPassphrases->SetPointer(nConnectionNumber, pPassphrase);
	m_nClientPassphraseSizes->SetInt(nConnectionNumber, nPassphraseSize);
	return true;
}

bool GSecureSocketServer::Send(const void* pBuf, int nLen, int nConnectionNumber)
{
	if(m_pClientPassphrases->GetSize() <= nConnectionNumber)
	{
		GAssert(false, "not done handshaking yet");
		return false;
	}

	// Encrypt the message with the client's passphrase
	unsigned char* pClientPassphrase = (unsigned char*)m_pClientPassphrases->GetPointer(nConnectionNumber);
	if(!pClientPassphrase)
	{
		GAssert(false, "not done handshaking yet");
		return false;
	}
	int nClientPassphraseSize = m_nClientPassphraseSizes->GetInt(nConnectionNumber);
	unsigned char* pCypher = (unsigned char*)alloca(nLen);
	memcpy(pCypher, pBuf, nLen);
	GCrypto::Encrypt(pCypher, nLen, pClientPassphrase, nClientPassphraseSize);
	
	// Send it
	return GSocketServer::Send(pCypher, nLen, nConnectionNumber);
}

// --------------------------------------------------------------------------

GSecureSocketClient::GSecureSocketClient(int nMaxPacketSize, GRand* pRand) : GSocketClient(false, nMaxPacketSize)
{
	m_pRand = pRand;
	m_pKeyPair = new GKeyPair();
	m_pKeyPair->GenerateKeyPair(pRand);
	m_nPassphraseSize = pRand->GetRandByteCount();
	m_pPassphrase = new unsigned char[m_nPassphraseSize];
	memcpy(m_pPassphrase, pRand->GetRand(), m_nPassphraseSize);
	if(m_pPassphrase[m_nPassphraseSize - 1] == 0)
		m_pPassphrase[m_nPassphraseSize - 1] = 1;
	m_pServerPassphrase = NULL;
	m_nServerPassphraseSize = 0;
}

GSecureSocketClient::~GSecureSocketClient()
{
	delete(m_pPassphrase);
	delete(m_pServerPassphrase);
	delete(m_pKeyPair);
}

/*static*/ GSecureSocketClient* GSecureSocketClient::ConnectToSecureSocket(const char* szAddress, u_short nPort, int nMaxPacketSize, GRand* pRand)
{
	// Check input
	GAssert(pRand->GetRandByteCount() >= (int)sizeof(int), "pRand not seeded big enough");
	GAssert(nMaxPacketSize > 0, "Bad max packet size");

	// Make the socket
	GSecureSocketClient* pSocket = new GSecureSocketClient(nMaxPacketSize, pRand);
	if(!pSocket)
		return NULL;
	if(!pSocket->Connect(szAddress, nPort))
	{
		delete(pSocket);
		return NULL;
	}

	// Handshake
	if(!pSocket->SendPublicKey())
	{
		GAssert(false, "error sending public key");
		delete(pSocket);
		return NULL;
	}
	int nTimeout = 200;
	while(!pSocket->HandShake())
	{
		if(--nTimeout <= 0)
		{
			delete(pSocket);
			return NULL;
		}
		GThread::sleep(50);
	}

	return pSocket;
}

bool GSecureSocketClient::HandShake()
{
	// See if we already did handshaking
	if(m_pServerPassphrase)
		return true;

	// Make sure there's a message to process
	if(m_pMessageQueue->GetSize() < 1)
		return false;

	// Check the header
	int nSize;
	unsigned char* pBuf = GSocketClient::GetNextMessage(&nSize);
	struct HandshakeHeader* pHeader = (struct HandshakeHeader*)pBuf;
	if(nSize <= (int)sizeof(struct HandshakeHeader) || pHeader->nMagic != HANDSHAKE_MAGIC)
	{
		GAssert(false, "Bad header packet");
		memset(pBuf, '\0', nSize);
		delete(pBuf);
		return false;
	}

	// Process the payload
	unsigned char* pPayload = pBuf + sizeof(struct HandshakeHeader);
	int nPayloadSize = nSize - sizeof(struct HandshakeHeader);
	if(pHeader->nMessage == 0)
	{
		bool bRet = OnReceivePublicKey((const char*)pPayload, nPayloadSize);
		if(!bRet)
			GAssert(false, "error processing public key");
	}
	else
	{
		GAssert(pHeader->nMessage == 1, "unexpected handshake packet");
		bool bRet = OnReceivePassphrase((const unsigned char*)pPayload, nPayloadSize);
		if(!bRet)
			GAssert(false, "error processing passphrase");
	}
	memset(pBuf, '\0', nSize);
	delete(pBuf);
	if(m_pServerPassphrase)
		return true;
	else
		return false;
}

bool GSecureSocketClient::SendPublicKey()
{
	bool bRet = true;
	GXMLTag* pTag = m_pKeyPair->ToXML(false);
	char* pMyPublicKey = pTag->ToString();
	int nPublicKeyBlobSize = strlen(pMyPublicKey);
	struct HandshakeHeader* pHeader = (struct HandshakeHeader*)alloca(sizeof(struct HandshakeHeader) + nPublicKeyBlobSize);
	memcpy(((unsigned char*)pHeader) + sizeof(struct HandshakeHeader), pMyPublicKey, nPublicKeyBlobSize);
	pHeader->nMagic = HANDSHAKE_MAGIC;
	pHeader->nMessage = 0;
	if(!GSocketClient::Send(pHeader, sizeof(struct HandshakeHeader) + nPublicKeyBlobSize))
		bRet = false;
	delete(pMyPublicKey);
	delete(pTag);
	return bRet;
}

bool GSecureSocketClient::SendPassphrase(GKeyPair* pKeyPair)
{
	bool bRet = true;

	// Encrypt my passphrase with client's public key
	int nCypherSize;
	unsigned char* pCypher = pKeyPair->PowerMod(m_pPassphrase, m_nPassphraseSize, true, &nCypherSize);

	// Send it to the client
	struct HandshakeHeader* pHeader = (struct HandshakeHeader*)alloca(sizeof(struct HandshakeHeader) + nCypherSize);
	memcpy(((unsigned char*)pHeader) + sizeof(struct HandshakeHeader), pCypher, nCypherSize);
	pHeader->nMagic = HANDSHAKE_MAGIC;
	pHeader->nMessage = 1;
	if(!GSocketClient::Send(pHeader, sizeof(struct HandshakeHeader) + nCypherSize))
	{
		GAssert(false, "Error sending passphrase");
		bRet = false;
	}
	delete(pCypher);
	return bRet;
}

bool GSecureSocketClient::OnReceivePublicKey(const char* pXML, int nXMLSize)
{
	// Decode the server's public key
	GKeyPair kp;
	if(!kp.FromXML(pXML, nXMLSize))
	{
		GAssert(false, "Bad keypair");
		return false;
	}

	// Send my passphrase
	if(!SendPassphrase(&kp))
	{
		GAssert(false, "Error sending passphrase");
	}
	return true;
}

bool GSecureSocketClient::OnReceivePassphrase(const unsigned char* pCypher, int nCypherSize)
{
	// Decrypt it with my private key
	m_pServerPassphrase = m_pKeyPair->PowerMod(pCypher, nCypherSize, false, &m_nServerPassphraseSize);
	return true;
}

bool GSecureSocketClient::Send(const void* pBuf, int nLen)
{
	if(!m_pServerPassphrase)
	{
		GAssert(false, "not done handshaking yet");
		return false;
	}

	// Encrypt the message with the server's passphrase
	unsigned char* pCypher = (unsigned char*)alloca(nLen);
	memcpy(pCypher, pBuf, nLen);
	GCrypto::Encrypt(pCypher, nLen, m_pServerPassphrase, m_nServerPassphraseSize);

	// Send it
	return GSocketClient::Send(pCypher, nLen);
}

int GSecureSocketClient::GetMessageCount()
{
	return GSocketClient::GetMessageCount();
}

unsigned char* GSecureSocketClient::GetNextMessage(int* pnSize)
{
	// Get the encrypted message
	unsigned char* pBuf = GSocketClient::GetNextMessage(pnSize);
	if(!pBuf)
		return NULL;

	// Decrypt the message
	GCrypto::Decrypt(pBuf, *pnSize, m_pPassphrase, m_nPassphraseSize);

	return pBuf;
}