recvfile.c

开源的组播程序

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

#include <rs.h>

#ifdef DMALLOC
#include "dmalloc.h"
#endif

#include "pmtp.h"

struct pmtp_download {
	struct pmsp_header session;
	uint32_t bytes_received;
	uint32_t packets_received;
	uint32_t segments;
	uint32_t segments_received;
	int fd;
};

uint16_t get_udpsocket (struct pmsp_header *session,
		const unsigned char *source, const unsigned char *group,
		const uint16_t port);
int open_outputfile (const unsigned char *dir, const unsigned char *filename);

uint32_t add_packet_to_segment (struct pmtp_segment *segment,
		uint32_t packet_index, uint16_t length, const unsigned char *payload);

int16_t receive_pmsp_packet (struct pmtp_download *download,
		unsigned char *download_directory, 
		unsigned char *packet, uint16_t length);

int main (int argc, char *argv[]) {
	uint32_t packet_index, session_index;
	uint32_t segmentid;
	unsigned char packet[PMTP_HEADER_SIZE + PMTP_PAYLOAD];
	uint8_t download_completed = 0;
	void *codec;
	unsigned char sha_printable[60];
	unsigned char ftt[24], ftt_str[50], dirname[256];

	uint32_t bytes_left = 0;

	struct sockaddr_in s_in;
	int s, n, len;

	struct pmtp_segment current, new;
	uint32_t new_segment_packets = 0;

	struct pmtp_download download;
	download.session.sessionid = 0;
	download.session.status = PMSP_STATE_INIT;

	if (argc != 2) {
		fprintf (stderr, "Usage: %s <download directory>\n", argv[0]);
		exit (1);
	}
	if (strlen (argv[1]) > 127) {
		fprintf (stderr, "%s: Name of directory too long.\n", argv[0]);
		exit (1);
	}
	strncpy (dirname, argv[1], 127);
	strcat (dirname, "/");
	dirname[127] = 0;

	s = get_udpsocket (&download.session, "0.0.0.0", PMTP_GROUP, PMTP_PORT);

	if (DEBUG & DEBUG_CODING) 
		fprintf (stderr, "Code (%d / %d)\n", NN, NROOTS);

	if((codec = init_rs_char(SYMSIZE,GENPOLY,FCS,PRIM,NROOTS)) == NULL) {
		fprintf(stderr, "init_rs_char failed!\n");
		exit (1);
	}
	len = sizeof (s_in);

	init_segment (&current, codec);
	init_segment (&new, codec);
	do {
		n = recvfrom (s, packet, sizeof(packet), 0,
			(struct sockaddr *) &s_in, &len);
		if ((decode_header_8 (packet) >> 4) != 0) {
			receive_pmsp_packet (&download, dirname, packet, n);
			bytes_left = download.session.partlength;
			download_completed = 0;
		} else {
			packet_index = decode_header_16 (&packet[PMTP_HEADER_PACKETID])
				 & 0x0fff;
			session_index = decode_header_16 (&packet[PMTP_HEADER_SESSIONID]);
			segmentid = decode_header_32 (&packet[PMTP_HEADER_SEGMENTID]);
			if (DEBUG & DEBUG_PACKETS_RECEIVED) 
				fprintf (stderr,
					"Received (%d) with size %d from segment %d, index %d\n",
					current.packets_received + 1, n, segmentid, packet_index);
			if (DEBUG & DEBUG_PACKETS_SHA)
				fprintf (stderr, "Packet %d, size %d, segment %d, SHA 0x%s\n",
					packet_index, n, segmentid,
						get_printable_sha (sha_printable, packet, n));
	
			if (current.packets == NULL) {
				current.segmentid = segmentid;
				current.payload_size = n - PMTP_HEADER_SIZE;
			}
			if (segmentid == current.segmentid) {
				add_packet_to_segment (&current, packet_index,
					n - PMTP_HEADER_SIZE, packet + PMTP_HEADER_SIZE);
			} else {
				if (new.packets == NULL) {
					new.segmentid = segmentid;
					new.payload_size = n - PMTP_HEADER_SIZE;
				}
				if (segmentid == new.segmentid) {
					add_packet_to_segment (&new, packet_index,
						n - PMTP_HEADER_SIZE, packet + PMTP_HEADER_SIZE);
					new_segment_packets++;
				} else {
					if (DEBUG & DEBUG_SEGMENTS) 
						fprintf (stderr,
							"Segment %d is not current %d and not new %d.\n", 
								segmentid, current.segmentid,
								new.segmentid);
					continue;
				}
			}
			if (current.packets_received == NN ||
					new_segment_packets >= MAX_PACKETS_LOST ||
					(download.segments_received == download.segments - 1 &&
						current.packets_received == (NN - NROOTS / 2))) {
				if (DEBUG & DEBUG_SEGMENTS) 
					fprintf (stderr, "Swapping current segment %d for %d\n",
						current.segmentid, new.segmentid);
				bytes_left = decode_segment
					(&current, download.fd, NROOTS, bytes_left);
				download.segments_received++;
				free_segment (&current);
				init_segment (&current, codec);
				current = new;
				init_segment (&new, codec);
				new_segment_packets = 0;
				if (bytes_left == 0) 
					download_completed = 1;
				printf ("Bytes left: %d\n", bytes_left);
			}
		}
	} while (! download_completed);
	free_segment (&current);
	free_segment (&new);
	close (download.fd);
	close (s);
	strcat (dirname, download.session.filename);
	get_tigertree (ftt, dirname,0, 0);
	printf ("File: %s, TigerTree: %s\n", download.session.filename,
		get_printable_hex (ftt_str, ftt, 24));
	return 0;
}

int16_t receive_pmsp_packet (struct pmtp_download *download,
		unsigned char *download_directory, 
		unsigned char *packet, uint16_t length) {
	unsigned char *bname;
	unsigned char filename[128];
	uint8_t version;
	uint16_t codingid, sessionid, payload, start, port;
	uint32_t partoffset, partlength, bitrate, source, group;
	uint64_t filesize;
	struct in_addr in_addr;
	uint8_t filenamelength = length - PMSP_HEADER_FILENAME;
	unsigned char file_tigertree[24], part_tigertree[24],
		ftt_str[50], ptt_str[50];
	
	version = decode_header_8 (packet) >> 4;
	codingid = decode_header_16 (packet) & 0x0fff;
	sessionid = decode_header_16 (packet + PMSP_HEADER_SESSIONID);
	payload = decode_header_16 (packet + PMSP_HEADER_PAYLOAD);
	start = decode_header_16 (packet + PMSP_HEADER_START)>>4;

	source =  *((uint32_t *) (packet + PMSP_HEADER_SOURCE));
	group = *((uint32_t *) (packet + PMSP_HEADER_GROUP));
	port = *((uint16_t *) (packet + PMSP_HEADER_PORT));
	bitrate = decode_header_16 (packet + PMSP_HEADER_BITRATE);

	filesize = decode_header_32 (packet + PMSP_HEADER_FILESIZE + 4);
	memcpy (file_tigertree, packet + PMSP_HEADER_FILETIGERTREE, 24);
	
	partoffset = decode_header_32 (packet + PMSP_HEADER_PARTOFFSET);
	partlength = decode_header_32 (packet + PMSP_HEADER_PARTLENGTH);
	memcpy (part_tigertree, packet + PMSP_HEADER_PARTTIGERTREE, 24);

	strncpy (filename, packet + PMSP_HEADER_FILENAME, filenamelength);
	filename[filenamelength] = 0;
	bname = basename (filename);

	if (download->session.status == PMSP_STATE_LEARNED
			 && sessionid == download->session.sessionid) {
		if (strncmp (bname, download->session.filename, filenamelength)) {
			fprintf (stderr, "%s doesn't match %s.\n", bname,
				download->session.filename);
			exit(1);
		}
		if (version == download->session.version &&
				codingid == download->session.codingid &&
				payload == download->session.payload &&
				source == download->session.source &&
				group == download->session.group &&
				port == download->session.port &&
				bitrate == download->session.bitrate && 
				filesize == download->session.filesize &&
				!memcmp (file_tigertree, download->session.file_tigertree,24) &&
				partoffset == download->session.partoffset &&
				partlength == download->session.partlength &&
				!memcmp (part_tigertree, download->session.part_tigertree,24) &&
				(strncmp (bname, download->session.filename, filenamelength)
				== 0)) {
			download->session.start = start;
			if (DEBUG & DEBUG_PMSP) 
				fprintf (stderr, "PMSP session %u confirmed, starting in %u seconds\n", sessionid, start);
		} else {
			download->session.status = PMSP_STATE_CONFLICT;
			if (DEBUG & DEBUG_PMSP) {
				fprintf (stderr, "Received conflicting PMSP announcement\nStarting in %u seconds\nVersion: %u\nCodingID: %u\nSessionID: %u\nPayload: %u\nFilesize: %u\nFile TigerTree: %s\nPart Offset: %u\nPart Length: %u\nPart TigerTree: %s\nBitrate: %u Kbit/s\nFilename: %s.\n", 
					download->session.start, download->session.version,
					download->session.codingid, download->session.sessionid,
					download->session.payload,
					(uint32_t) download->session.filesize,
					get_printable_hex (ftt_str, file_tigertree, 24),
					download->session.partoffset, download->session.partlength,
					get_printable_hex (ptt_str, part_tigertree, 24),
					download->session.bitrate * 8, download->session.filename);
				in_addr.s_addr = download->session.source;
				fprintf (stderr, "Source: %s\n",
					inet_ntoa (in_addr));
				in_addr.s_addr = download->session.group;
				fprintf (stderr, "Group: %s\nPort: %u\n",
					inet_ntoa (in_addr),
					ntohs(download->session.port));
				
			}
		}
	}
	if (download->session.status == PMSP_STATE_INIT) {
		download->session.status = PMSP_STATE_LEARNED;
		download->session.sessionid = sessionid;
		download->session.version = version;
		download->session.codingid = codingid;
		download->session.payload = payload;
		download->session.start = start;
		download->session.family = AF_INET;
		download->session.protocol = SOCK_DGRAM;
		download->session.source = source;
		download->session.group = group;
		download->session.port = port;
		download->session.filesize = filesize;
		memcpy (download->session.file_tigertree, file_tigertree, 24);
		download->session.partoffset = partoffset;
		download->session.partlength = partlength;
		memcpy (download->session.part_tigertree, part_tigertree, 24);
		download->session.bitrate = bitrate;
		strncpy (download->session.filename, bname, filenamelength + 1);
		download->fd = open_outputfile (download_directory, 
			download->session.filename);
		download->bytes_received = download->packets_received =
			download->segments_received = 0;
		download->segments = partlength / ((NN - NROOTS) * payload);
		if (partlength != download->segments * ((NN - NROOTS) * payload)) {
			download->segments++;
		}

		if (DEBUG & DEBUG_PMSP) {
			fprintf (stderr, "Received new PMSP announcement of size %d\nStarting in %u seconds\nVersion: %u\nCodingID: %u\nSessionID: %u\nPayload: %u\nFilesize: %u\nFile TigerTree: %s\nPart Offset: %u\nPart Length: %u\nPart TigerTree: %s\nBitrate: %u Kbit/s\nFilename: %s.\n", 
				length, download->session.start, download->session.version,
				download->session.codingid, download->session.sessionid,
				download->session.payload,
				(uint32_t) download->session.filesize,
				get_printable_hex (ftt_str, file_tigertree, 24),
				download->session.partoffset, download->session.partlength,
				get_printable_hex (ptt_str, part_tigertree, 24),
				download->session.bitrate * 8, download->session.filename);
				in_addr.s_addr = download->session.source;
				fprintf (stderr, "Source: %s\n",
					inet_ntoa (in_addr));
				in_addr.s_addr = download->session.group;
				fprintf (stderr, "Group: %s\nPort: %u\n",
					inet_ntoa (in_addr),
					ntohs(download->session.port));
		}
	}
	return PMSP_HEADER_FILENAME + strlen (download->session.filename);
}
	
uint32_t add_packet_to_segment (struct pmtp_segment *segment,
		uint32_t packet_index, uint16_t length, const unsigned char *payload) {
	uint32_t segmentsize = NN * segment->payload_size;
	if (segment == NULL) {
		fprintf (stderr, "No segment structure pointer\n");
		exit (1);
	}
	if (payload == NULL) {
		fprintf (stderr, "No payload pointer\n");
		exit (1);
	}
	if (length > segment->payload_size) {
		fprintf (stderr,
			"Packet with length %d is too big for segment %d with payload %d\n",
			length, segment->segmentid, segment->payload_size);
		exit (1);
	}

	if (segment->packets == NULL) {
		if ((segment->packets = (unsigned char*) malloc (segmentsize))==NULL) {
			fprintf (stderr, "Couldn't malloc %d bytes for segment\n",
				segmentsize);
			exit (1);
		}
	}
	if (packet_index > sizeof (segment->received_packets)) {
		fprintf (stderr, "Packet index %d is out of range\n", packet_index);
		exit (1);
	}
	if (segment->received_packets[packet_index] == 1) {
		segment->duplicate_packets++;
	} else {
		segment->packets_received++;
		segment->received_packets[packet_index] = 1;
	}
	memcpy (segment->packets + packet_index * segment->payload_size, payload,
		segment->payload_size);
	return segment->packets_received;
}

uint16_t get_udpsocket (struct pmsp_header *session,
		const unsigned char *source, const unsigned char *group,
		const uint16_t port) {
	int s;
	struct sockaddr_in s_in;
	struct ip_mreqn m_in;

	if ((s = socket (AF_INET, SOCK_DGRAM, 0)) < 0) {
		perror ("socket");
		exit (1);
	}
	memset(&m_in, 0, sizeof (m_in));
	if (!inet_aton (group, &(m_in.imr_multiaddr))) {
		fprintf (stderr, "Couldn't convert %s\n", group);
	}
	if (setsockopt (s, SOL_IP, IP_ADD_MEMBERSHIP, &m_in, sizeof (m_in))) {
		perror ("setsockopt");
		exit (1);
	}

	/* `source' is currently unused because we're not using IGMPv3 yet */
	session->source = 0;

	session->protocol = SOCK_DGRAM;
	session->group = m_in.imr_multiaddr.s_addr;

	memset ((char *) &s_in, 0, sizeof (s_in));
	session->family = s_in.sin_family = AF_INET;
	session->port = s_in.sin_port = htons (PMTP_PORT);
	if (bind(s, (struct sockaddr *) &s_in, sizeof (s_in)) < 0) {
		perror ("bind");
		exit (1);
	}
	return s;
}

int open_outputfile (const unsigned char *dir, const unsigned char *filename) {
	unsigned char full [256];
	strcpy (full, dir);
	int fd;
	if (filename == NULL) {
		fprintf (stderr, "No filename specified\n");
		exit (1);
	}
	strncat (full, filename, 128);
	if ((fd = open (full, O_WRONLY | O_CREAT | O_EXCL, 00755 )) < 0) {
		fprintf (stderr, "Error opening %s for writing\n", full);
		exit (1);
	}
	return fd;
}