tftpserver.cpp

一个简单的tftp服务器

/**************************************************************************
*   Copyright (C) 2005 by Achal Dhir                                      *
*   [e-mail]                                                   *
*                                                                         *
*   This program is free software; you can redistribute it and/or modify  *
*   it under the terms of the GNU General Public License as published by  *
*   the Free Software Foundation; either version 2 of the License, or     *
*   (at your option) any later version.                                   *
*                                                                         *
*   This program 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 General Public License for more details.                          *
*                                                                         *
*   You should have received a copy of the GNU General Public License     *
*   along with this program; if not, write to the                         *
*   Free Software Foundation, Inc.,                                       *
*   59 Temple Place - Suite 330, Boston, MA  02111-1307, USA.             *
***************************************************************************/

// tftpserver.cpp
#include <sys/types.h>
#include <limits.h>
#include <sys/socket.h>
#include <netinet/in.h>
#include <arpa/inet.h>
#include <netdb.h>
#include <unistd.h>
#include <signal.h>
#include <stdio.h>
#include <fcntl.h>
#include <errno.h>
#include <sys/time.h>
#include <time.h>
#include <stdlib.h>
#include <limits.h>
#include <memory.h>
#include <sys/stat.h>
#include <stdio.h>
#include <syslog.h>
#include <string>
#include <pthread.h>
#include <sys/socket.h>
#include <sys/ioctl.h>
#include <map>
using namespace std;
#include "tftpserver.h"

//Global Variables
bool kRunning = true;
myMap tftpCache;
myMultiMap tftpAge;
bool verbatim = false;
char iniFile[256]="";
char logFile[256]="";
WORD blksize = 65464;
WORD timeout = 3;
data2 cfig;
char tempbuff[256];
char logBuff[512];
char sVersion[] = "TFTP Server SinglePort Version 1.55 Unix Built 1551";
packet* datain;

int main(int argc, char **argv)
{
	signal(SIGINT, catch_int);
	signal(SIGABRT, catch_int);
	signal(SIGTERM, catch_int);
	signal(SIGQUIT, catch_int);
	signal(SIGTSTP, catch_int);
	signal(SIGHUP, catch_int);

    logBuff[0] = 0;

    for (int i = 1; i < argc; i++)
    {
        if (!strcasecmp(argv[i], "-v"))
            verbatim = true;
        else if (!strcmp(argv[i], "-i") && argc > i + 1 && argv[i + 1][0] != '-' )
        {
            myTrim(iniFile, argv[i + 1]);
            i++;
        }
        else if (!strcmp(argv[i], "-l") && argc > i + 1 && argv[i + 1][0] != '-' )
        {
            myTrim(logFile, argv[i + 1]);
            i++;
        }
        else if (!strncasecmp(argv[i], "-i", 2))
            myTrim(iniFile, argv[i] + 2);
        else if (!strncasecmp(argv[i], "-l", 2))
            myTrim(logFile, argv[i] + 2);
        else
            sprintf(logBuff, "Error: Invalid Argument %s", argv[i]);
    }

	if (!iniFile[0])
		strcpy(iniFile,"/etc/tftpserver.ini");

	if (verbatim)
	{
		if (logBuff[0])
		{
			printf("%s\n", logBuff);
			exit(EXIT_FAILURE);
		}

		init();
		timeval tv;
		fd_set readfds;
		request req;
		datain = (packet*)calloc(1, blksize + 4);
		int fdsReady = 0;

		if (!datain)
		{
			sprintf(logBuff,"Memory Error");
			logMess(logBuff, 0);
			exit(1);
		}

		if (cfig.tftpConn[0].server)
		{
			printf("\nAccepting requests..\n");

			do
			{
				FD_ZERO(&readfds);
				tv.tv_sec = 1;
				tv.tv_usec = 0;

				for (int i = 0; i < MAX_SERVERS && cfig.tftpConn[i].server; i++)
					FD_SET(cfig.tftpConn[i].sock, &readfds);

				fdsReady = select(cfig.maxFD, &readfds, NULL, NULL, &tv);

				//if (errno)
				//	printf("%s\n", strerror(errno));

				for (int i = 0; fdsReady > 0 && i < MAX_SERVERS && cfig.tftpConn[i].server; i++)
				{
					if (FD_ISSET(cfig.tftpConn[i].sock, &readfds))
					{
						fdsReady--;
						memset(&req, 0, sizeof(request));
						memset(datain, 0, blksize + 4);
						req.clientsize = sizeof(req.client);
						req.sockInd = i;
						errno = 0;
						req.bytesRecd = recvfrom(cfig.tftpConn[req.sockInd].sock, (char*)datain, blksize + 4, 0, (sockaddr*)&req.client, &req.clientsize);

						if (req.bytesRecd < 4 || errno)
							continue;

						//printf("%u=%u\n", req.bytesRecd, blksize + 4);
						sprintf(req.mapname, "%s:%u", inet_ntoa(req.client.sin_addr), ntohs(req.client.sin_port));
						request *req1 = tftpCache[req.mapname];

						if (!req1)
							tftpCache.erase(req.mapname);

						//printf("%u\n",req1);

						//printf("Here\n");

						if (req1)
						{
							req1->bytesRecd = req.bytesRecd;

							if (ntohs(datain->opcode) == 3 && req1->opcode == 2)
							{
								if ((WORD)req1->bytesRecd  <= req1->blksize + 4)
								{
									req1->tblock = req1->block + 1;
									if (req1->attempt <= 3 && ntohs(datain->block) == req1->tblock)
									{
										req1->block = req1->tblock;
										req1->fblock++;
										req1->attempt = 0;
										req1->acout.opcode = htons(4);
										req1->acout.block = ntohs(req1->block);
										processRecv(req1);
									}
								}
								else
								{
									req1->serverError.opcode = htons(5);
									req1->serverError.errorcode = htons(4);
									req1->bytesSent = sendto(cfig.tftpConn[req1->sockInd].sock, (const char*)&req1->serverError, strlen(req1->serverError.errormessage) + 5, 0, (sockaddr*)&req1->client, req1->clientsize);
									sprintf(req1->serverError.errormessage, "Error: Incoming Packet too large");
									logMess(req1, 1);
									req1->attempt = UCHAR_MAX;
								}
							}
							else if (ntohs(datain->opcode) == 4 && req1->opcode == 1)
							{
								if (req1->bytesRecd >= 4)
								{
									if (req1->attempt <= 3 && ntohs(datain->block) == req1->block)
									{
										req1->block++;
										req1->fblock++;
										req1->attempt = 0;
										processSend(req1);
									}
								}
							}
							else if (req1->bytesRecd > 512)
							{
								req1->serverError.opcode = htons(5);
								req1->serverError.errorcode = htons(4);
								req1->bytesSent = sendto(cfig.tftpConn[req1->sockInd].sock, (const char*)&req1->serverError, strlen(req1->serverError.errormessage) + 5, 0, (sockaddr*)&req1->client, req1->clientsize);
								sprintf(req1->serverError.errormessage, "Error: Incoming Packet too large");
								logMess(req1, 1);
								req1->attempt = UCHAR_MAX;
							}
							else if (ntohs(datain->opcode) == 1 || ntohs(datain->opcode) == 2)
							{
								if (req1->file)
								{
									fclose(req1->file);
									req1->file = 0;
								}

								if (!processNew(&req))
								{
									memcpy(req1, &req, sizeof(request));
								}
							}
							else if (ntohs(datain->opcode) == 5)
							{
								sprintf(req1->serverError.errormessage, "Error %i at Client, %s", ntohs(datain->block), &datain->buffer);
								logMess(req1, 1);
								req1->attempt = UCHAR_MAX;
								if (req1->file)
								{
									fclose(req1->file);
									req1->file = 0;
								}
							}
							else
							{
								req1->serverError.opcode = htons(5);
								req1->serverError.errorcode = htons(4);
								sprintf(req1->serverError.errormessage, "Unexpected Option Code %u", ntohs(datain->opcode));
								req1->bytesSent = sendto(cfig.tftpConn[req1->sockInd].sock, (const char*)&req1->serverError, strlen(req1->serverError.errormessage) + 5, 0, (sockaddr*)&req1->client, req1->clientsize);
								logMess(req1, 1);
								req1->attempt = UCHAR_MAX;
								if (req1->file)
								{
									fclose(req1->file);
									req1->file = 0;
								}
							}
						}
						else if (req.bytesRecd < 4 || errno)
						{
							req.serverError.opcode = htons(5);
							req.serverError.errorcode = htons(0);
							sprintf(req.serverError.errormessage, "Communication Error");
							logMess(&req, 1);
						}
						else if (ntohs(datain->opcode) != 1 && ntohs(datain->opcode) != 2)
						{
							req.serverError.opcode = htons(5);
							req.serverError.errorcode = htons(5);
							sprintf(req.serverError.errormessage, "Unknown transfer ID");
							req.bytesSent = sendto(cfig.tftpConn[i].sock, (const char*)&req.serverError, strlen(req.serverError.errormessage) + 5, 0, (sockaddr*)&req.client, req.clientsize);
							logMess(&req, 1);
						}
						else if (req.bytesRecd > 512)
						{
							req.serverError.opcode = htons(5);
							req.serverError.errorcode = htons(4);
							sprintf(req.serverError.errormessage, "Error: Incoming Packet too large");
							req.bytesSent = sendto(cfig.tftpConn[i].sock, (const char*)&req.serverError, strlen(req.serverError.errormessage) + 5, 0, (sockaddr*)&req.client, req.clientsize);
							logMess(&req, 1);
						}
						else if (ntohs(datain->opcode) == 5)
						{
							sprintf(req.serverError.errormessage, "Error %i at Client, %s", ntohs(datain->block), &datain->buffer);
							logMess(&req, 1);
						}
						else
						{
							if (cfig.hostRanges[0].rangeStart)
							{
								DWORD iip = ntohl(req.client.sin_addr.s_addr);
								BYTE allowed = false;

								for (WORD j = 0; j <= sizeof(cfig.hostRanges) && cfig.hostRanges[j].rangeStart; j++)
								{
									if (iip >= cfig.hostRanges[j].rangeStart && iip <= cfig.hostRanges[j].rangeEnd)
									{
										allowed = true;
										break;
									}
								}

								if (!allowed)
								{
									req.serverError.opcode = htons(5);
									req.serverError.errorcode = htons(2);
									strcpy(req.serverError.errormessage, "Access Denied");
									logMess(&req, 1);
									req.bytesSent = sendto(cfig.tftpConn[i].sock, (const char*)&req.serverError, strlen(req.serverError.errormessage) + 5, 0, (sockaddr*)&req.client, req.clientsize);
									continue;
								}
							}

							if (!processNew(&req))
							{
								request *req1 = (request*)calloc(1, sizeof(request));

								if (!req1)
								{
									sprintf(logBuff,"Memory Error");
									logMess(logBuff, 1);
									continue;
								}

								memcpy(req1, &req, sizeof(request));
								tftpCache[req1->mapname] = req1;
								tftpAge.insert(pair<long, request*>(req1->expiry, req1));
							}
						}
					}
				}

				myMultiMap::iterator p = tftpAge.begin();
				myMultiMap::iterator q;
				time_t currentTime = time(NULL);

				while (p != tftpAge.end())
				{
					if (!tftpAge.size())
						break;

					request *req = p->second;

					if (p->first > currentTime)
					{
						break;
					}
					else if (p->first < req->expiry && req->expiry > currentTime)
					{
						q = p;
						p++;
						tftpAge.erase(q);
						tftpAge.insert(pair<long, request*>(req->expiry, req));
					}
					else if (req->expiry <= currentTime && req->attempt >= 3)
					{
						if (req->attempt < UCHAR_MAX)
						{
							req->serverError.opcode = htons(5);
							req->serverError.errorcode = htons(0);

							if (req->fblock && !req->block)
								strcpy(req->serverError.errormessage, "Large File, Block# Rollover not supported by Client");
							else
								strcpy(req->serverError.errormessage, "Timeout");

							req->bytesSent = sendto(cfig.tftpConn[req->sockInd].sock, (const char*) &req->serverError, strlen(req->serverError.errormessage) + 5, 0, (sockaddr*)&req->client, req->clientsize);
							logMess(req, 1);
						}

						q = p;
						p++;
						tftpAge.erase(q);
						tftpCache.erase(req->mapname);
						clean(req);
					}
					else if (req->expiry <= currentTime)
					{
						if (ntohs(req->acout.opcode) == 3)
						{
							if (processSend(req))
								req->attempt = 255;
							else
							{
								req->attempt++;
								req->expiry = currentTime + req->timeout;
							}
						}
						else
						{
							errno = 0;
							req->bytesSent = sendto(cfig.tftpConn[req->sockInd].sock, (const char*)&req->acout, req->bytesSent, 0, (sockaddr*)&req->client, req->clientsize);

							if (errno)
								req->attempt = 255;
							else
							{
								req->attempt++;
								req->expiry = currentTime + req->timeout;
							}
						}
						p++;
					}
					else
						p++;
				}
			}
			while (kRunning);
		}
	}
	else
	{
		if(logBuff[0])
		{
			syslog(LOG_MAKEPRI(LOG_LOCAL1, LOG_CRIT), logBuff);
			exit(EXIT_FAILURE);
		}

		/* Our process ID and Session ID */
		pid_t pid, sid;

		/* Fork off the parent process */
		pid = fork();
		if (pid < 0)
		{
			exit(EXIT_FAILURE);
		}
		/* If we got a good PID, then
		we can exit the parent process. */
		if (pid > 0)
		{
			exit(EXIT_SUCCESS);
		}

		/* Change the file mode mask */
		umask(0);

		/* Open any logs here */

		/* Create a new SID for the child process */
		sid = setsid();
		if (sid < 0)
		{
			/* Log the failure */
			exit(EXIT_FAILURE);
		}

		/* Change the current working directory */
		if ((chdir("/")) < 0)
		{
			/* Log the failure */
			exit(EXIT_FAILURE);
		}

		/* Close out the standard file descriptors */
		close(STDIN_FILENO);
		close(STDOUT_FILENO);
		close(STDERR_FILENO);

		/* Daemon-specific initialization goes here */
		//Initialize
		init();
		timeval tv;
		fd_set readfds;
		request req;
		datain = (packet*)calloc(1, blksize + 4);
		int fdsReady = 0;

		if (!datain)
		{
			sprintf(logBuff,"Memory Error");
			logMess(logBuff, 0);
			exit(1);
		}

		if (cfig.tftpConn[0].server)
		{
			do
			{
				FD_ZERO(&readfds);
				tv.tv_sec = 1;
				tv.tv_usec = 0;

				for (int i = 0; i < MAX_SERVERS && cfig.tftpConn[i].server; i++)
					FD_SET(cfig.tftpConn[i].sock, &readfds);

				fdsReady = select(cfig.maxFD, &readfds, NULL, NULL, &tv);

				//if (errno)
				//	printf("%s\n", strerror(errno));

				for (int i = 0; fdsReady > 0 && i < MAX_SERVERS && cfig.tftpConn[i].server; i++)
				{
					if (FD_ISSET(cfig.tftpConn[i].sock, &readfds))
					{
						fdsReady--;
						memset(&req, 0, sizeof(request));
						memset(datain, 0, blksize + 4);
						req.clientsize = sizeof(req.client);
						req.sockInd = i;
						errno = 0;
						req.bytesRecd = recvfrom(cfig.tftpConn[req.sockInd].sock, (char*)datain, blksize + 4, 0, (sockaddr*)&req.client, &req.clientsize);

						if (req.bytesRecd < 4 || errno)
							continue;

						//printf("%u=%u\n", req.bytesRecd, blksize + 4);
						sprintf(req.mapname, "%s:%u", inet_ntoa(req.client.sin_addr), ntohs(req.client.sin_port));