gistfile1.cpp

#include "stdafx.h"
#include "SOCKSRelay.h"
#include <thread>
#include <chrono>
#include <fstream>
#include <sstream>
#include <memory.h>
#include <string.h>
#include <array>
#include "ClientSocketStream.h"

using namespace std;

#ifndef min
#define min(X,Y) ((X)<(Y)?(X):(Y))
#endif


#if LOGLEVEL >= LOG_STATUS
std::string SOCKSRelay::appname;
atomic<int> SOCKSRelay::activeWorkers;
#endif

#if LOGLEVEL >= LOG_STATUS
SOCKSRelay::SOCKSRelay(SocketStream &A, SocketStream &B, int32_t queueSize, ProxyMode mode, int conID)
#else
SOCKSRelay::SOCKSRelay(SocketStream &A, SocketStream &B, int32_t queueSize, ProxyMode mode)
#endif
{  
	if (mode == PM_socks_uninitilized)
	{
		SOCKSRelayWorker *worker = new SOCKSRelayWorker(A,B,queueSize,PM_socks_uninitilized);
#if LOGLEVEL >= LOG_STATUS
		worker->mPartner = nullptr;
		worker->setWorkerID(conID * 2);	
#endif		
		worker->operator()();
	}
	else
	{
		SOCKSRelayWorker* workerA = new SOCKSRelayWorker(A,B,queueSize,PM_socks_initilized);
		SOCKSRelayWorker* workerB = new SOCKSRelayWorker(B,A,queueSize,PM_socks_initilized);
		workerA->mPartner = workerB;
		workerB->mPartner = workerA;
#if LOGLEVEL >= LOG_STATUS
		workerA->setWorkerID(conID * 2);
		workerB->setWorkerID(conID * 2 + 1);
#endif
		thread([workerA](){workerA->operator()();}).detach();
		workerB->operator()();
	}	
}

SOCKSRelay::SOCKSRelayWorker::SOCKSRelayWorker(SocketStream &in, SocketStream &out, int32_t queueSize, ProxyMode mode) : 
	mIn(in), 
	mOut(out),
	mQueueSize(queueSize),
	mProxyMode(mode),
	mPartner(nullptr)
{
	mBuffer = new char[mQueueSize];	
#if LOGLEVEL >= LOG_STATUS
	SOCKSRelay::activeWorkers ++;
#endif
}

SOCKSRelay::SOCKSRelayWorker::~SOCKSRelayWorker()
{	
	if (mPartner)
		mPartner->mPartner = nullptr;
	else
		mIn.disconnect();
	mOut.disconnect();
#if LOGLEVEL >= LOG_STATUS
	SOCKSRelay::activeWorkers --;
#endif
	delete []mBuffer;	
#if LOGLEVEL >= LOG_PACKAGES
	if (LOGIn.is_open())	
		LOGIn.close();			
	if (LOGOut.is_open())
		LOGOut.close();
#endif
}

#if LOGLEVEL >= LOG_STATUS
void SOCKSRelay::SOCKSRelayWorker::setWorkerID(int32_t ID)
{
	mWorkerID = ID;
#if LOGLEVEL >= LOG_PACKAGES
#ifdef WIN32
	LARGE_INTEGER t1, t2, p1, p2, f;
	unsigned long long r1,r2;

	QueryPerformanceFrequency(&f);
	QueryPerformanceCounter(&p1);
	r1 = __rdtsc();
	
	Sleep(1);
	GetSystemTimeAsFileTime((FILETIME*) (&t1.u));
	do
	{
		Sleep(0);
		GetSystemTimeAsFileTime((FILETIME*) (&t2.u));		
	} while(t2.QuadPart == t1.QuadPart);

	QueryPerformanceCounter(&p2);
	r2 = __rdtsc();

	CPUPerformance = f.QuadPart * (r2 - r1) / (p2.QuadPart-p1.QuadPart);
	CPUStartClock = __rdtsc();
#endif
	char fName[50];
	if (LOGIn.is_open())	
		LOGIn.close();
	if (LOGOut.is_open())
		LOGOut.close();
	
	if (ID % 2 == 0)
	{
		sprintf(fName,"../Log/%03dABIn%s.txt",ID / 2,SOCKSRelay::appname.c_str());LOGIn.open(fName,ios::binary|ios::app);	
		sprintf(fName,"../Log/%03dABOu%s.txt",ID / 2,SOCKSRelay::appname.c_str());LOGOut.open(fName,ios::binary|ios::app);
	}
	else
	{
		sprintf(fName,"../Log/%03dBAIn%s.txt",ID / 2,SOCKSRelay::appname.c_str());LOGIn.open(fName,ios::binary|ios::app);	
		sprintf(fName,"../Log/%03dBAOu%s.txt",ID / 2,SOCKSRelay::appname.c_str());LOGOut.open(fName,ios::binary|ios::app);
	}
#endif
}

template <> void SOCKSRelay::SOCKSRelayWorker::writeLog(LOGType type, const char* data)
{
	writeLog(type, data,strlen(data));
}


void SOCKSRelay::SOCKSRelayWorker::writeLog(LOGType type, const char* data, int32_t length)
{
	ostream* temp = nullptr;
	switch (type)
	{
	case SOCKSRelay::SOCKSRelayWorker::LT_In:
		temp = &LOGIn;
		break;
	case SOCKSRelay::SOCKSRelayWorker::LT_Out:
		temp = &LOGOut;
		break;
	case SOCKSRelay::SOCKSRelayWorker::LT_Partner_In:
		if (mPartner)
			temp = &mPartner->LOGIn;
		break;
	case SOCKSRelay::SOCKSRelayWorker::LT_Partner_Out:
		if (mPartner)
			temp = &mPartner->LOGOut;
		break;
	default:
		break;
	}
	if (temp)
	{
		ostream& out = *temp;
#if LOGLEVEL >= LOG_PACKAGES
#if LOGLEVEL > 3
		out << "--LOG BEGIN--";
		writeTime(type);
#endif
		out.write(data,length);
		out.flush();
#if LOGLEVEL > 3
		out << "--LOG END-- ";
		writeTime(type);
#endif		
#endif
	}
}
#ifdef WIN32
void SOCKSRelay::SOCKSRelayWorker::writeTime(LOGType type)
{
#if LOGLEVEL >= LOG_ALL
	ostream* temp = nullptr;
	switch (type)
	{
	case SOCKSRelay::SOCKSRelayWorker::LT_In:
		temp = &LOGIn;
		break;
	case SOCKSRelay::SOCKSRelayWorker::LT_Out:
		temp = &LOGOut;
		break;
	case SOCKSRelay::SOCKSRelayWorker::LT_Partner_In:
		if (mPartner)
			temp = &mPartner->LOGIn;
		break;
	case SOCKSRelay::SOCKSRelayWorker::LT_Partner_Out:
		if (mPartner)
			temp = &mPartner->LOGOut;
		break;
	default:
		break;
	}
	if (temp)
	{
		ostream& out = *temp;

		out << "-- time = (" << ((__rdtsc() - CPUStartClock) * 1000000ULL) / CPUPerformance << "µs)\n";		
		out.flush();
	}
	#endif
}
#endif

#endif


void SOCKSRelay::SOCKSRelayWorker::operator()()
{		
	int32_t packetLength = 0;
	if (mProxyMode == PM_socks_uninitilized)
		if (!parseHeader())
			goto endThread;

	while (mIn.isConnected() && mOut.isConnected())
	{
		packetLength = mIn.receive(mBuffer, mQueueSize);
//		writeLog(LT_In, mBuffer, packetLength);
		if (packetLength > 0)
		{
			mOut.send(mBuffer, packetLength);
	//		writeLog(LT_Out, mBuffer, packetLength);
		}
	}

endThread:
	delete this;
}

bool SOCKSRelay::SOCKSRelayWorker::parseHeader()
{	
	int32_t matchedLength = 0;
	int32_t packetLength = 0;
	while(packetLength < 8)
	{
		if(mIn.isConnected() && packetLength < mQueueSize)
			packetLength += mIn.receive(mBuffer + packetLength, mQueueSize - packetLength);
		else			
			return false;
	}

	if (mBuffer[0] == 4)
	{
		if (parseV4Header(packetLength,matchedLength))
		{
			mOut.send(mBuffer + matchedLength , packetLength - matchedLength);
			return true;
		}
		else
			return false;
	}
	else
		return false;
}

bool SOCKSRelay::SOCKSRelayWorker::parseV4Header(int32_t& packetLength,int32_t& matchedLength)
{	
	swap(mBuffer [2],mBuffer[3]);
	int32_t lastMatch = 0;
	array<char,10> reply;
	const char* hostName = nullptr;
	memcpy(reply.data(),"\0\x5bZYXWVU",8);						
	if (mBuffer[1] != 1)
	{
		mIn.send(reply.data(),8);
		mIn.disconnect();
#if LOGLEVEL >= LOG_STATUS
		cout << mWorkerID << ": bind request not supported\n";
#endif		
		return false;
	}

	matchedLength = 8;			
	ClientSocketStream host;
	const char* userID = mBuffer + matchedLength;	
	if (!parseV4Header_FindString(packetLength, matchedLength))
	{
		mIn.send(reply.data(),8);
		mIn.disconnect();
		return false;
	}

	if (mBuffer[4] == 0 && mBuffer[5] == 0 && mBuffer[6] == 0 && mBuffer[7] != 0)
	{
		hostName = mBuffer + matchedLength;
		if (!parseV4Header_FindString(packetLength, matchedLength))
		{
			mIn.send(reply.data(),8);
			mIn.disconnect();
			return false;
		}
		host.connectToHost(hostName, *reinterpret_cast<int16_t*>(mBuffer + 2));
	}	
	else
		host.connect(*reinterpret_cast<int32_t*>(mBuffer + 4), *reinterpret_cast<int16_t*>(mBuffer + 2));

	if (!host.isConnected())
	{
		mIn.send(reply.data(),8);
		mIn.disconnect();
#if LOGLEVEL >= LOG_STATUS
		if (mBuffer[4] == 0 && mBuffer[5] == 0 && mBuffer[6] == 0 && mBuffer[7] != 0)
			cout << mWorkerID << ": could not connect to host " 
				<< hostName << ":" 
				<< *reinterpret_cast<int16_t*>(mBuffer + 2) <<"\n";
		else
			cout << mWorkerID << ": could not connect to host " << 
			(int)*reinterpret_cast<int8_t*>(mBuffer + 4) << "." << 
			(int)*reinterpret_cast<int8_t*>(mBuffer + 5) << "." << 
			(int)*reinterpret_cast<int8_t*>(mBuffer + 6) << "." << 
			(int)*reinterpret_cast<int8_t*>(mBuffer + 7) << ":" <<
			(int)*reinterpret_cast<int16_t*>(mBuffer + 2) <<"\n";
#endif		
		return false;
	}
					
	reply[1] = 0x5a;
	mIn.send(reply.data(),8);
	mOut = host;
	mPartner = new SOCKSRelayWorker(mOut, mIn, mQueueSize, PM_socks_initilized);
	mPartner->mPartner = this;			
	thread([this](){this->mPartner->operator()();}).detach();							
	mPartner->setWorkerID(mWorkerID + 1);	
	return true;
}

bool SOCKSRelay::SOCKSRelayWorker::parseV4Header_FindString(int32_t& packetLength,int32_t& matchedLength)
{
	while (mQueueSize <= packetLength)
	{
		while (matchedLength < packetLength)
			if (mBuffer[matchedLength++] == 0)
				return true;							
		packetLength += mIn.receive(mBuffer + packetLength, mQueueSize - packetLength);		
	}
	return false;
}