sourcemod-central: tools/daemon/smud_connections.cpp@f3b8bbe17f85

     1 #include <assert.h>

     2 #include "smud_connections.h"

     3 #include "smud.h"

     5 ConnectionPool::ConnectionPool()

6 {

     7 	pthread_mutex_init(&m_AddLock, NULL);

     8 	m_timeOut = 1000;

9 }

    11 ConnectionPool::~ConnectionPool()

    12 {

    13 	pthread_mutex_destroy(&m_AddLock);

    14 }

    16 void ConnectionPool::AddConnection( int fd )

    17 {

    18 	smud_connection *connection = new smud_connection(fd);

    20 	pthread_mutex_lock(&m_AddLock);

    21 	m_AddQueue.push_back(connection);

    22 	pthread_mutex_unlock(&m_AddLock);

    23 }

    25 void ConnectionPool::Process( bool *terminate )

    26 {

    27 	struct timespec ts_wait;

    29 	ts_wait.tv_sec = 0;

    30 	ts_wait.tv_nsec = 50000000; /* 50ms */

    32 	std::list<smud_connection *>::iterator iter;

    33 	smud_connection *con = NULL;

    35 	while (1)

    36 	{

    37 		if (*terminate)

    38 		{

    39 			return;

    40 		}

    42 		iter = m_Links.begin();

    44 		QueryResult result = QueryResult_Continue;

    45 		int pollReturn = 0;

    47 		/* Add all connections that want processing to the sets */

    48 		while (iter != m_Links.end())

    49 		{

    50 			con = *iter;

    52 			pollReturn = poll(&(con->pollData), 1, 0);

    53 			assert(pollReturn <= 1);

    55 			if (pollReturn == -1)

    56 			{

    57 				//Something went badly wrong or the connection closed.

    58 				result = QueryResult_Complete;

    59 			}

    60 			else if (pollReturn == 1)

    61 			{

    62 				//Poll returns the number of sockets available (which can only ever be 1)

    63 				result = ProcessConnection(con);

    64 			}

    66 			if (result == QueryResult_Complete)

    67 			{

    68 				iter = m_Links.erase(iter);

    69 #if defined DEBUG

    70 				fprintf(stdout, "Closing socket %d\n", con->fd);

    71 #endif

    72 				closesocket(con->fd);

    73 				delete con;

    74 				continue;

    75 			}

    77 			iter++;

    78 		}

    80 		/* Add new items to process */

    82 		iter = m_Links.end();

    83 		pthread_mutex_lock(&m_AddLock);

    84 		m_Links.splice(iter, m_AddQueue);

    85 		pthread_mutex_unlock(&m_AddLock);

    87 		nanosleep(&ts_wait, NULL);

    88 	}

    89 }

    91 QueryResult ConnectionPool::ProcessConnection( smud_connection *con )

    92 {

    93 	switch (con->state)

    94 	{

    95 		case ConnectionState_ReadQueryHeader:

    96 		{

    97 			ReadQueryHeader(con);

    98 			break;

    99 		}

   101 		case ConnectionState_ReadQueryData:

   102 		{

   103 			ReadQueryContent(con);

   104 			break;

   105 		}

   107 		case ConnectionState_ReplyQuery:

   108 		{

   109 			ReplyQuery(con);

   110 			break;

   111 		}

   113 		case ConnectionState_SendingFiles:

   114 		{

   115 			SendFile(con);

   116 			break;

   117 		}

   119 		case ConnectionState_SendUnknownList:

   120 		{

   121 			SendUnknownList(con);

   122 			break;

   123 		}

   125 		case ConnectionState_Complete:

   126 		{

   127 			break;

   128 		}

   129 	}

   131 	if (con->state == ConnectionState_Complete)

   132 	{

   133 		printf("Ending connection because it marked itself as finished\n");

   134 		return QueryResult_Complete;

   135 	}

   137 	if (con->start + m_timeOut < time(NULL))

   138 	{

   139 		printf("Ending connection because it has passed maximum allowed time\n");

   140 		return QueryResult_Complete;

   141 	}

   143 	return QueryResult_Continue;

   144 }

   146 void ConnectionPool::ReadQueryHeader( smud_connection *con )

   147 {

   148 	if (con->buffer == NULL)

   149 	{

   150 		con->buffer = new char[QUERY_HEADER_SIZE];

   151 		con->writtenCount = 0;

   152 	}

   154 	int bytesReceived = 0;

   156 	bytesReceived = recv(con->fd, con->buffer+con->writtenCount, QUERY_HEADER_SIZE-con->writtenCount, 0);

   158 	if (bytesReceived == -1)

   159 	{

   160 		if (errno != EAGAIN && errno != EWOULDBLOCK)

   161 		{

   162 			con->state = ConnectionState_Complete;

   163 			delete [] con->buffer;

   164 			con->buffer = NULL;

   165 			con->writtenCount = 0;

   166 		}

   168 		return;

   169 	}

   171 	con->writtenCount += bytesReceived;

   173 	assert(con->writtenCount <= QUERY_HEADER_SIZE);

   175 	if (con->writtenCount < QUERY_HEADER_SIZE)

   176 	{

   177 		/* Don't change the connection status, so next cycle we will come back to here and continue receiving data */

   178 		return;

   179 	}

   181 	if (con->buffer[0] != 'A' || con->buffer[1] != 'G')

   182 	{

   183 		con->state = ConnectionState_Complete;

   184 		delete [] con->buffer;

   185 		con->buffer = NULL;

   186 		con->writtenCount = 0;

   187 		return;

   188 	}

   190 	//Ignore the next 8 bytes for the moment. Versioning data is currently unused

   191 	// uint16[4] - source version major/minor/something/rev

   193 	con->sentSums = con->buffer[10];

   195 	con->state = ConnectionState_ReadQueryData;

   196 #if defined DEBUG

   197 	fprintf(stdout, "Query Header Read Complete, %i md5's expected\n", con->sentSums);

   198 #endif

   200 	delete [] con->buffer;

   201 	con->buffer = NULL;

   202 	con->writtenCount = 0;

   203 }

   205 void ConnectionPool::ReplyQuery(smud_connection *con)

   206 {

   207 	char data[12];

   208 	data[0] = 'A';

   209 	data[1] = 'G';

   211 	data[2] = (char)Update_Unknown; //unused versioning crap

   212 	*(short *)&data[3] = 1;

   213 	*(short *)&data[5] = 0;

   214 	*(short *)&data[7] = 0;

   215 	*(short *)&data[9] = 3;

   217 	data[11] = (char)con->sendCount;

   219 	int bytesSent = send(con->fd, data+con->writtenCount, sizeof(data)-con->writtenCount, 0);

   221 	if (bytesSent == -1)

   222 	{

   223 		if (errno != EAGAIN && errno != EWOULDBLOCK)

   224 		{

   225 			con->state = ConnectionState_Complete;

   226 			con->writtenCount = 0;

   227 		}

   229 		return;

   230 	}

   232 	con->writtenCount += bytesSent;

   234 	assert(con->writtenCount <= 12);

   236 	if (con->writtenCount < 12)

   237 	{

   238 		/** Still more data needs to be sent - Return so we come back here next cycle */

   239 		return;

   240 	}

   242 	con->state = ConnectionState_SendingFiles;

   243 	con->writtenCount = 0;

   244 #if defined DEBUG

   245 	printf("Query Reply Header Complete\n");

   246 #endif

   247 }

   249 void ConnectionPool::ReadQueryContent( smud_connection *con )

   250 {

   251 	if (con->buffer == NULL)

   252 	{

   253 		con->buffer = new char[QUERY_CONTENT_SIZE*con->sentSums];

   254 		con->writtenCount = 0;

   255 	}

   257 	int bytesReceived = 0;

   259 	bytesReceived = recv(con->fd, con->buffer+con->writtenCount, (QUERY_CONTENT_SIZE*con->sentSums)-con->writtenCount, 0);

   261 	if (bytesReceived == -1)

   262 	{

   263 		if (errno != EAGAIN && errno != EWOULDBLOCK)

   264 		{

   265 			con->state = ConnectionState_Complete;

   266 			delete [] con->buffer;

   267 			con->buffer = NULL;

   268 			con->writtenCount = 0;

   269 		}

   271 		return;

   272 	}

   274 	con->writtenCount += bytesReceived;

   276 	assert(con->writtenCount <= (QUERY_CONTENT_SIZE*con->sentSums));

   278 	if (con->writtenCount < (QUERY_CONTENT_SIZE*con->sentSums))

   279 	{

   280 		/* Don't change the connection status, so next cycle we will come back to here and continue receiving data */

   281 		return;

   282 	}

   284 	con->shouldSend = new MD5Status[con->sentSums]();

   285 	con->fileLocation = new int[con->sentSums]();

   286 	con->headerSent = new bool[con->sentSums]();

   288 	for (int i=0; i<con->sentSums; i++)

   289 	{

   290 		con->fileLocation[i] = -1;

   291 		con->shouldSend[i] = GetMD5UpdateStatus(con->buffer + (QUERY_CONTENT_SIZE*i), con, i);

   293 		if (con->shouldSend[i] == MD5Status_NeedsUpdate)

   294 		{

   295 #if defined DEBUG

   296 			fprintf(stdout, "File %i needs updating\n", i);

   297 #endif

   298 			con->sendCount++;

   299 			con->headerSent[i] = false;

   300 			continue;

   301 		}

   303 		if (con->shouldSend[i] == MD5Status_Unknown)

   304 		{

   305 #if defined DEBUG

   306 			fprintf(stdout, "File %i is unknown\n", i);

   307 #endif

   308 			con->unknownCount++;

   309 		}

   310 	}

   312 	con->state = ConnectionState_ReplyQuery;

   313 	con->pollData.events = POLLOUT;

   314 	delete [] con->buffer;

   315 	con->buffer = NULL;

   316 	con->writtenCount = 0;

   317 #if defined DEBUG

   318 	fprintf(stdout, "Query Data Read Complete\n");

   319 #endif

   320 }

   322 MD5Status ConnectionPool::GetMD5UpdateStatus( const char *md5 , smud_connection *con, int fileNum)

   323 {

   324 	//Try find a file with this name in some directory.

   325 	char path[100] = "./md5/";

   326 	char temp[4];

   327 	char md5String[33] = "";

   329 	for (int i=0; i<16; i++)

   330 	{

   331 		snprintf(temp, sizeof(temp), "%02x", (unsigned char)md5[i]);

   332 		strcat(md5String, temp);

   333 	}

   335 	strcat(path, md5String);

   337 #if defined DEBUG

   338 	fprintf(stdout, "checking for file \"%s\"\n", path);

   339 #endif

   341 	FILE *file = fopen(path, "r");

   343 	if (file == NULL)

   344 	{

   345 		printf("Couldn't find file!\n");

   346 		return MD5Status_Unknown;

   347 	}

   349 	char latestMD5[33];

   350 	fgets(latestMD5, 33, file);

   351 #if defined DEBUG

   352 	fprintf(stdout, "Latest md5 is: %s\n", latestMD5);

   353 #endif

   355 	if (strcmp(latestMD5, md5String) == 0)

   356 	{

   357 		fclose(file);

   358 		return MD5Status_Current;

   359 	}

   361 	char filename[100];

   362 	filename[0] = '\n';

   364 	while (filename[0] == '\n')

   365 	{

   366 		fgets(filename, sizeof(filename), file);

   367 	}

   369 	if (filename[strlen(filename)-1] == '\n')

   370 	{

   371 		filename[strlen(filename)-1] = '\0';

   372 	}

   374 	fclose(file);

   376 #if defined DEBUG

   377 	fprintf(stdout, "Filename is %s\n", filename);

   378 #endif

   380 	//We now need to match this filename with one of our mmap'd files in memory and store it until send gets called.

   381 	for (int i=0; i<NUM_FILES; i++)

   382 	{

   383 		if (strcmp(fileNames[i], filename) == 0)

   384 		{

   385 			con->fileLocation[fileNum] = i;

   386 #if defined DEBUG

   387 			fprintf(stdout, "File %i mapped to local file %i\n", fileNum, i);

   388 #endif

   389 			return MD5Status_NeedsUpdate;

   390 		}

   391 	}

   393 	return MD5Status_Unknown;

   394 }

   396 void ConnectionPool::SendFile( smud_connection *con )

   397 {

   398 	//Find the next file to send.

   399 	while (con->writtenCount == 0 &&

   400 			con->currentFile < con->sentSums &&

   401 			con->shouldSend[con->currentFile] != MD5Status_NeedsUpdate)

   402 	{

   403 		con->currentFile++;

   404 	}

   406 	//All files have been sent.

   407 	if (con->currentFile >= con->sentSums)

   408 	{

   409 #if defined DEBUG

   410 		fprintf(stdout, "All files sent!\n");

   411 #endif

   412 		con->state = ConnectionState_SendUnknownList;

   413 		con->writtenCount = 0;

   414 		return;

   415 	}

   417 	void *file = fileLocations[con->fileLocation[con->currentFile]];

   418 	int filelength = fileLength[con->fileLocation[con->currentFile]];

   420 #if defined DEBUG

   421 	fprintf(stdout, "Sending file of length %i\n", filelength);

   422 	fprintf(stdout,

   423 			"Current file index is: %i, maps to file index: %i\n",

   424 			con->currentFile,

   425 			con->fileLocation[con->currentFile]);

   426 #endif

   428 	int sentBytes = 0;

   430 	if (!con->headerSent[con->currentFile])

   431 	{

   432 		char buffer[5];

   433 		buffer[0] = con->currentFile;

   434 		*((int *)&buffer[1]) = filelength;

   436 		sentBytes = send(con->fd, buffer+con->writtenCount, 5-con->writtenCount, 0);

   438 		if (sentBytes == -1)

   439 		{

   440 			if (errno != EAGAIN && errno != EWOULDBLOCK)

   441 			{

   442 				con->state = ConnectionState_Complete;

   443 				con->writtenCount = 0;

   444 			}

   446 			return;

   447 		}

   449 		con->writtenCount += sentBytes;

   451 		assert(con->writtenCount <= 5);

   453 		if (con->writtenCount < 5)

   454 		{

   455 			return;

   456 		}

   458 		con->headerSent[con->currentFile] = true;

   459 		con->writtenCount = 0;

   460 	}

   462 	sentBytes = send(con->fd, (unsigned char *)file+con->writtenCount, filelength-con->writtenCount, 0);

   464 	if (sentBytes == -1)

   465 	{

   466 		if (errno != EAGAIN && errno != EWOULDBLOCK)

   467 		{

   468 			con->state = ConnectionState_Complete;

   469 			con->writtenCount = 0;

   470 		}

   472 		return;

   473 	}

   475 	con->writtenCount += sentBytes;

   477 	assert(con->writtenCount <= filelength);

   479 	if (con->writtenCount < filelength)

   480 	{

   481 		return;

   482 	}

   484 	con->currentFile++;

   485 	con->writtenCount = 0;

   486 #if defined DEBUG

   487 	fprintf(stdout, "Sent a file!: %s\n", fileNames[con->fileLocation[con->currentFile-1]]);

   488 #endif

   489 }

   491 void ConnectionPool::SendUnknownList( smud_connection *con )

   492 {

   493 	int size = con->unknownCount+1;

   494 	char *packet = new char[size]();

   496 	packet[0] = con->unknownCount;

   498 #if defined DEBUG

   499 	fprintf(stdout, "%i Files are unknown\n", con->unknownCount);

   500 #endif

   502 	int i=1;

   504 	for (int j=0; j<con->sentSums; j++)

   505 	{

   506 		if (con->shouldSend[j] == MD5Status_Unknown)

   507 		{

   508 			packet[i] = j;

   509 			i++;

   510 		}

   511 	}

   513 	int sentBytes = send(con->fd, packet+con->writtenCount, size-con->writtenCount, 0);

   515 	if (sentBytes == -1)

   516 	{

   517 		if (errno != EAGAIN && errno != EWOULDBLOCK)

   518 		{

   519 			con->state = ConnectionState_Complete;

   520 			con->writtenCount = 0;

   521 		}

   523 		return;

   524 	}

   526 	con->writtenCount += sentBytes;

   528 	assert(con->writtenCount <= size);

   530 	if (con->writtenCount < size)

   531 	{

   532 		return;

   533 	}

   535 	con->state = ConnectionState_Complete;

   536 	con->writtenCount = 0;

   537 #if defined DEBUG

   538 	fprintf(stdout, "Unknowns Sent\n");

   539 #endif

   540 }

author	David Anderson <dvander@alliedmods.net>
	Wed Feb 25 20:07:12 2009 -0500 (2009-02-25)
changeset 2607	f3b8bbe17f85
parent 2462	6040047a9da1
permissions	-rw-r--r--