TFTP客户端搭建

#include <stdio.h>
#include <sys/socket.h>
#include <arpa/inet.h>
#include <netinet/in.h>
#include <netinet/ip.h>
#include <unistd.h>
#include <string.h>
#include <sys/stat.h>
#include <sys/types.h>
#include <fcntl.h>
#include <stdlib.h>
#include <errno.h>
#define ERR_MSG(msg) do{\
    fprintf(stderr, "line:%d ", __LINE__);\
    perror(msg);\
}while(0)
            
#define PORT 69
int do_download(int sfd, struct sockaddr_in sin);
int do_upload(int sfd, struct sockaddr_in sin);
int main(int argc, const char *argv[])
{
	if(argc < 2)
	{
		fprintf(stderr, "请输入服务器的IP\n");
		return -1;
	}
	//创建套接字
	int sfd = socket(AF_INET, SOCK_DGRAM, 0);
	if(sfd < 0)
	{
		ERR_MSG("socket");
		return -1;
	}
	//绑定客户端的ip和端口，非必须的
	//填充服务器的ip地址和端口号
	struct sockaddr_in sin;
	sin.sin_family = AF_INET;
	sin.sin_port = htons(PORT);
	sin.sin_addr.s_addr = inet_addr(argv[1]);
	socklen_t slen = sizeof(sin);
	while(1)
	{
		printf("**********************\n");
		printf("*******1.下载*********\n");
		printf("*******2.上传*********\n");
		printf("*******3.退出*********\n");
		printf("**********************\n");
		int choose = 0;
		scanf("%d", &choose);
		while(getchar()!=10);
		switch(choose)
		{
		case 1:
			do_download(sfd, sin);
			break;
		case 2:
			do_upload(sfd, sin);
			break;
		case 3:
			close(sfd);
			exit(1);
			break;
		default:
			printf("输入错误\n");
			break;
		}
	}
END:
	//关闭所有文件描述符
	close(sfd);
	return 0;
}
int do_upload(int sfd, struct sockaddr_in sin)
{
	char filename[20] = "";
	printf("请输入要上传的文件");
	fgets(filename, sizeof(filename), stdin);
	filename[strlen(filename)-1] = 0;
	//判断该文件是否存在
	int fd = open(filename, O_RDONLY);
	if(fd < 0)
	{
		if(errno == ENOENT)
		{
			printf("文件不存在，请重新输入文件名\n");
			return -1;
		}
		else
		{
			ERR_MSG("open");
			return -2;
		}
	}
	//发送上传请求
	char buf[516] = "";
	int n = sprintf(buf, "%c%c%s%c%s%c", 0, 2, filename, 0, "octet", 0);
	if(sendto(sfd, buf, n, 0, (struct sockaddr*)&sin, sizeof(sin)) < 0)
	{
		ERR_MSG("sendto");
		return -1;
	}
	socklen_t slen = sizeof(sin);
	ssize_t recv_len;
	unsigned short num = 0;
	while(1)
	{
		bzero(buf, sizeof(buf));
		recv_len = recvfrom(sfd, buf, sizeof(buf), 0, (struct sockaddr*)&sin,&slen);
		if(recv_len < 0)
		{
			ERR_MSG("recvfrom");
			return -2;
		}
		//打印操作
		if(4 == buf[1]) //如果回复的是ACK，则客户端发送数据包
		{
			if(num == ntohs(*(unsigned short*)(buf+2)))
			{
				//修改操作码为数据包
				buf[1] = 3;
				//填充块编号
				num++;
				*(unsigned short*)(buf+2) = htons(num);
				//读取文件中的数据
				int res = read(fd, buf+4, 512);
				if(res < 0)
				{
					ERR_MSG("read");
					return -2;
				}
				else if(0 == res)
				{
					printf("上传成功\n");
					break;
				}
				//发送数据包
				if(sendto(sfd, buf, res+4, 0, (struct sockaddr*)&sin,sizeof(sin)) <0)
				{
					ERR_MSG("sendto");
					return -1;
				}
			}
			else
			{
				printf("文件上传失败\n");
				break;
			}
		}
		else if(5 == buf[1]) //错误包
		{
			printf("ERROR\n");
			break;
		}
	}
	close(fd);
	return 0;
}
int do_download(int sfd, struct sockaddr_in sin)
{
	char filename[20] = "";
	printf("请输入要下载的文件");
	fgets(filename, sizeof(filename), stdin);
	filename[strlen(filename)-1] = 0;
	//打开文件
	int fd = open(filename, O_WRONLY|O_CREAT|O_TRUNC, 0664);
	if(fd < 0)
	{
		ERR_MSG("open");
		return -1;
	}
	//发送读请求
	char buf[516] = "";
	int n = sprintf(buf, "%c%c%s%c%s%c", 0, 1, filename, 0, "octet", 0);
	if(sendto(sfd, buf, n, 0, (struct sockaddr*)&sin, sizeof(sin)) < 0)
	{ 
		ERR_MSG("sendto");
		return -1;
	}
	//循环接收数据包，发送ack
	int recv_len = 0;
	unsigned int num = 0;
	socklen_t slen = sizeof(sin);
	while(1)
	{
		bzero(buf, sizeof(buf));
		recv_len = recvfrom(sfd, buf, sizeof(buf), 0, (struct sockaddr*)&sin,&slen);
		if(recv_len < 0)
		{
			ERR_MSG("recvfrom");
			return -1;
		}
		if(3 == buf[1]) //数据包
		{
			//保证收到的块编号是上一次块编号+1
			if(num+1 == ntohs(*(unsigned short*)(buf+2)))
			{
					write(fd, buf+4, recv_len-4);
					buf[1] = 0x04;
					if(sendto(sfd, buf, 4, 0, (struct sockaddr*)&sin, sizeof(sin)) <0)
					{
						ERR_MSG("sendto");
						return -1;
					}
					//下载完毕
					if(recv_len-4<512)
					{
						printf("下载完毕\n");
						break;
					}
					//当前包编号更新
					num++;
				}
			}
			else if(5 == buf[1]) //错误包
			{
				printf("ERROR:%s\n", buf+4);
				break;
			}
	}
	close(fd);
	return 0;
}