tftp.c

linux下的tftp socket编程样例

#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <sys/types.h>
#include <sys/socket.h>
#include <sys/time.h>
#include <sys/stat.h>
#include <netdb.h>
#include <netinet/in.h>
#include <arpa/inet.h>
#include <unistd.h>
#include <fcntl.h>


#define TFTP_TIMEOUT 5             /* seconds */


#define TFTP_RRQ   1
#define TFTP_WRQ   2
#define TFTP_DATA  3
#define TFTP_ACK   4
#define TFTP_ERROR 5
#define TFTP_OACK  6

static const char *tftp_bb_error_msg[] = {
	"Undefined error",
	"File not found",
	"Access violation",
	"Disk full or allocation error",
	"Illegal TFTP operation",
	"Unknown transfer ID",
	"File already exists",
	"No such user"
};

static inline int tftp(char *host, const char *remotefile, int localfd)
{
	const int bb_tftp_num_retries = 5;

	struct sockaddr_in sa;
	struct sockaddr_in from;
	struct timeval tv;
	socklen_t fromlen;
	fd_set rfds;
	char *cp;
	unsigned short tmp;
	int socketfd;
	int len;
	int opcode = 0;
	int finished = 0;
	int timeout = bb_tftp_num_retries;
	unsigned short block_nr = 1;
	int port = htons (69);

	/* Can't use RESERVE_CONFIG_BUFFER here since the allocation
	 * size varies meaning BUFFERS_GO_ON_STACK would fail */
	char *buf = malloc(516);

	if ((socketfd = socket(PF_INET, SOCK_DGRAM, 0)) < 0) {
		printf ("Create socket error\n");
		exit (0);
	}

	len = sizeof (sa);
	memset(&sa, 0, len);
	bind(socketfd, (struct sockaddr *)&sa, len);

	sa.sin_family = AF_INET;
	sa.sin_port = port;
	sa.sin_addr.s_addr = inet_addr (host);

	/* build opcode */

	opcode = TFTP_RRQ;

	while (1) {

		cp = buf;

		/* first create the opcode part */

		*((unsigned short *) cp) = htons(opcode);

		cp += 2;

		/* add filename and mode */

		if (opcode == TFTP_RRQ) {
                        int too_long = 0;

			/* see if the filename fits into buf */
			/* and fill in packet                */

			len = strlen(remotefile) + 1;

			if ((cp + len) >= &buf[515]) {
			        too_long = 1;
			} else {
			        strncpy(cp, remotefile, len);
				cp += len;
			}

			if (too_long || ((&buf[515] - cp) < 6)) {
				printf ("too long remote-filename");
				break;
			}

			/* add "mode" part of the package */

			memcpy(cp, "octet", 6);
			cp += 6;
		}

		/* add ack and data */

		if (opcode == TFTP_ACK) {
			*((unsigned short *) cp) = htons(block_nr);
			cp += 2;
			block_nr++;
			printf ("block nr : %d\n", block_nr);
		}


		/* send packet */

		timeout = bb_tftp_num_retries;  /* re-initialize */
		do {
			len = cp - buf;
			if (sendto(socketfd, buf, len, 0, (struct sockaddr *) &sa, sizeof(sa)) < 0) {
				printf ("send error\n");
				len = -1;
				break;
			}

			if (finished && (opcode == TFTP_ACK)) {
				break;
			}

			/* receive packet */

			memset(&from, 0, sizeof(from));
			fromlen = sizeof(from);

			tv.tv_sec = TFTP_TIMEOUT;
			tv.tv_usec = 0;

			FD_ZERO(&rfds);
			FD_SET(socketfd, &rfds);

			switch (select(FD_SETSIZE, &rfds, NULL, NULL, &tv)) {
			case 1:
				len = recvfrom(socketfd, buf, 516, 0,
						(struct sockaddr *) &from, &fromlen);

				if (len < 0) {
					printf ("recvfrom\n");
					break;
				}

				timeout = 0;

				if (sa.sin_port == port) {
					sa.sin_port = from.sin_port;
				}
				if (sa.sin_port == from.sin_port) {
					break;
				}

				/* fall-through for bad packets! */
				/* discard the packet - treat as timeout */
				timeout = bb_tftp_num_retries;

			case 0:
				printf ("timeout\n");

				timeout--;
				if (timeout == 0) {
					len = -1;
					printf ("last timeout");
				}
				break;

			default:
				printf ("select error\n");
				len = -1;
			}

		} while (timeout && (len >= 0));

		if ((finished) || (len < 0)) {
			break;
		}

		/* process received packet */


		opcode = ntohs(*((unsigned short *) buf));
		tmp = ntohs(*((unsigned short *) &buf[2]));

		if (opcode == TFTP_ERROR) {
			char *msg = NULL;

			if (buf[4] != '\0') {
				msg = &buf[4];
				buf[515] = '\0';
			} else if (tmp < (sizeof(tftp_bb_error_msg)
					  / sizeof(char *))) {

				msg = (char *) tftp_bb_error_msg[tmp];
			}

			if (msg) {
				printf ("server says: %s\n", msg);
			}

			break;
		}

		if (opcode == TFTP_DATA) {
			if (tmp == block_nr) {			
				len = write(localfd, &buf[4], len - 4);
				if (len < 0) {
					printf ("write local file error\n");
					break;
				}

				if (len != (512)) {
					finished++;
				}

				opcode = TFTP_ACK;
				continue;
			}
		}

	}

	return finished ? EXIT_SUCCESS : EXIT_FAILURE;
}

int
main()
{
	char *host = "192.168.0.254";
	const char *localfile = "localfile";
	const char *updateimage = "NET.C";
	int result;
	int fd;

	/* figure out what to pass to getopt */

	fd = open(localfile, O_WRONLY | O_CREAT);

	if (fd < 0) {
		printf ("local file error\n");
		exit (0);
	}

	result = tftp(host, updateimage, fd);

	return(result);
}