Changed how packet generation functions get access to the buffer, so they won't let you overflow the buffer, and they will keep track of the packet size "internally".

3 files changed, 71 insertions, 44 deletions

diff --git a/mactelnet.c b/mactelnet.c
index 7486a9b..dce51da 100644
--- a/mactelnet.c
+++ b/mactelnet.c
@@ -18,13 +18,16 @@
 */
 #include <string.h>
 #include <stdio.h>
+#include <stdlib.h>
 #include <linux/if_ether.h>
 #include "mactelnet.h"
 #include "config.h"
 
 unsigned char mt_mactelnet_cpmagic[4] = { 0x56, 0x34, 0x12, 0xff };
 
-int initPacket(unsigned char *data, unsigned char ptype, unsigned char *srcmac, unsigned char *dstmac, unsigned short sessionkey, unsigned short counter) {
+
+int initPacket(struct mt_packet *packet, unsigned char ptype, unsigned char *srcmac, unsigned char *dstmac, unsigned short sessionkey, unsigned short counter) {
+	unsigned char *data = packet->data;
 
 	/* Packet version */
 	data[0] = 1;
@@ -53,10 +56,25 @@ int initPacket(unsigned char *data, unsigned char ptype, unsigned char *srcmac,
 	data[21] = counter & 0xff;
 
 	/* 22 bytes header */
+	packet->size = 22;
 	return 22;
 }
 
-int addControlPacket(unsigned char *data, unsigned char cptype, void *cpdata, int data_len) {
+int addControlPacket(struct mt_packet *packet, char cptype, void *cpdata, int data_len) {
+	unsigned char *data = packet->data + packet->size;
+
+	/* Something is really wrong. Packets should never become over 1500 bytes */
+	if (packet->size + MT_CPHEADER_LEN + data_len > MT_PACKET_LEN) {
+		fprintf(stderr, "addControlPacket: ERROR, too large packet. Exceeds %d bytes\n", MT_PACKET_LEN);
+		exit(1);
+	}
+
+	if (cptype == MT_CPTYPE_PLAINDATA) {
+		memcpy(data, cpdata, data_len);
+		packet->size += data_len;
+		return data_len;
+	}
+
 	/* Control Packet Magic id */
 	memcpy(data,  mt_mactelnet_cpmagic, sizeof(mt_mactelnet_cpmagic));
 
@@ -71,11 +89,12 @@ int addControlPacket(unsigned char *data, unsigned char cptype, void *cpdata, in
 
 	/* Insert data */
 	if (data_len) {
-		memcpy(data+9, cpdata, data_len);
+		memcpy(data + MT_CPHEADER_LEN, cpdata, data_len);
 	}
 
+	packet->size += MT_CPHEADER_LEN + data_len;
 	/* Control packet header length + data length */
-	return 9+data_len;
+	return MT_CPHEADER_LEN + data_len;
 }
 
 void parsePacket(unsigned char *data, struct mt_mactelnet_hdr *pkthdr) {
diff --git a/mactelnet.h b/mactelnet.h
index 2a62353..d992de4 100644
--- a/mactelnet.h
+++ b/mactelnet.h
@@ -22,6 +22,8 @@
 #define MT_HEADER_LEN 22
 #define MT_CPHEADER_LEN 9
 
+#define MT_PACKET_LEN 1500
+
 /* Packet type */
 #define MT_PTYPE_SESSIONSTART 0
 #define MT_PTYPE_DATA 1
@@ -58,8 +60,13 @@ struct mt_mactelnet_control_hdr {
 	unsigned char *data;
 };
 
-extern int initPacket(unsigned char *data, unsigned char ptype, unsigned char *srcmac, unsigned char *dstmac, unsigned short sessionkey, unsigned short counter);
-extern int addControlPacket(unsigned char *data, unsigned char cptype, void *cpdata, int data_len);
+struct mt_packet {
+	int size;
+	unsigned char data[MT_PACKET_LEN];
+};
+
+extern int initPacket(struct mt_packet *packet, unsigned char ptype, unsigned char *srcmac, unsigned char *dstmac, unsigned short sessionkey, unsigned short counter);
+extern int addControlPacket(struct mt_packet *packet, char cptype, void *cpdata, int data_len);
 extern void parsePacket(unsigned char *data, struct mt_mactelnet_hdr *pkthdr);
 extern int parseControlPacket(unsigned char *data, const int data_len, struct mt_mactelnet_control_hdr *cpkthdr);
 
diff --git a/main.c b/main.c
index 4f35662..545c317 100644
--- a/main.c
+++ b/main.c
@@ -40,9 +40,9 @@
 
 int sockfd;
 int deviceIndex;
-int outcounter=0;
-int incounter=0;
-int sessionkey=0;
+int outcounter = 0;
+int incounter = 0;
+int sessionkey = 0;
 int running = 1;
 
 unsigned char terminalMode = 0;
@@ -58,12 +58,12 @@ unsigned char encryptionkey[128];
 unsigned char username[255];
 unsigned char password[255];
 
-int sendUDP(const unsigned char *data, int len) {
-	return sendCustomUDP(sockfd, deviceIndex, srcmac, dstmac, &sourceip,  sourceport, &destip, 20561, data, len);
+int sendUDP(struct mt_packet *packet) {
+	return sendCustomUDP(sockfd, deviceIndex, srcmac, dstmac, &sourceip,  sourceport, &destip, 20561, packet->data, packet->size);
 }
 
 void sendAuthData(unsigned char *username, unsigned char *password) {
-	unsigned char data[1500];
+	struct mt_packet data;
 	unsigned char *terminal = (unsigned char *)getenv("TERM");
 	unsigned short width = 0;
 	unsigned short height = 0;
@@ -86,37 +86,37 @@ void sendAuthData(unsigned char *username, unsigned char *password) {
 	md5sum[0] = 0;
 
 	/* Send combined packet to server */
-	plen = initPacket(data, MT_PTYPE_DATA, srcmac, dstmac, sessionkey, outcounter);
+	plen = initPacket(&data, MT_PTYPE_DATA, srcmac, dstmac, sessionkey, outcounter);
 	databytes = plen;
-	plen += addControlPacket(data + plen, MT_CPTYPE_PASSWORD, md5sum, 17);
-	plen += addControlPacket(data + plen, MT_CPTYPE_USERNAME, username, strlen(username));
-	plen += addControlPacket(data + plen, MT_CPTYPE_TERM_TYPE, terminal, strlen(terminal));
+	plen += addControlPacket(&data, MT_CPTYPE_PASSWORD, md5sum, 17);
+	plen += addControlPacket(&data, MT_CPTYPE_USERNAME, username, strlen(username));
+	plen += addControlPacket(&data, MT_CPTYPE_TERM_TYPE, terminal, strlen(terminal));
 
 	if (getTerminalSize(&width, &height) > 0) {
-		plen += addControlPacket(data + plen, MT_CPTYPE_TERM_WIDTH, &width, 2);
-		plen += addControlPacket(data + plen, MT_CPTYPE_TERM_HEIGHT, &height, 2);
+		plen += addControlPacket(&data, MT_CPTYPE_TERM_WIDTH, &width, 2);
+		plen += addControlPacket(&data, MT_CPTYPE_TERM_HEIGHT, &height, 2);
 	}
 
 	outcounter += plen - databytes;
 
 	/* TODO: handle result */
-	result = sendUDP(data, plen);
+	result = sendUDP(&data);
 }
 
 void sig_winch(int sig) {
 	unsigned short width,height;
-	unsigned char data[1500];
+	struct mt_packet data;
 	int result,plen,databytes;
 
 	/* terminal height/width has changed, inform server */
 	if (getTerminalSize(&width, &height) > 0) {
-		plen = initPacket(data, MT_PTYPE_DATA, srcmac, dstmac, sessionkey, outcounter);
+		plen = initPacket(&data, MT_PTYPE_DATA, srcmac, dstmac, sessionkey, outcounter);
 		databytes = plen;
-		plen += addControlPacket(data + plen, MT_CPTYPE_TERM_WIDTH, &width, 2);
-		plen += addControlPacket(data + plen, MT_CPTYPE_TERM_HEIGHT, &height, 2);
+		plen += addControlPacket(&data, MT_CPTYPE_TERM_WIDTH, &width, 2);
+		plen += addControlPacket(&data, MT_CPTYPE_TERM_HEIGHT, &height, 2);
 		outcounter += plen - databytes;
 
-		result = sendUDP(data, plen);
+		result = sendUDP(&data);
 	}
 
 	/* reinstate signal handler */
@@ -138,14 +138,14 @@ void handlePacket(unsigned char *data, int data_len) {
 
 	/* Handle data packets */
 	if (pkthdr.ptype == MT_PTYPE_DATA) {
-		unsigned char odata[200];
+		struct mt_packet odata;
 		int plen=0,result=0;
 		int rest = 0;
 		unsigned char *p = data;
 
 		/* Always transmit ACKNOWLEDGE packets in response to DATA packets */
-		plen = initPacket(odata, MT_PTYPE_ACK, srcmac, dstmac, sessionkey, pkthdr.counter + (data_len - MT_HEADER_LEN));
-		result = sendUDP(odata, plen);
+		plen = initPacket(&odata, MT_PTYPE_ACK, srcmac, dstmac, sessionkey, pkthdr.counter + (data_len - MT_HEADER_LEN));
+		result = sendUDP(&odata);
 
 		if (DEBUG)
 			printf("ACK: Plen = %d, Send result: %d\n", plen, result);
@@ -209,12 +209,12 @@ void handlePacket(unsigned char *data, int data_len) {
 
 	/* The server wants to terminate the connection, we have to oblige */
 	else if (pkthdr.ptype == MT_PTYPE_END) {
-		unsigned char odata[200];
+		struct mt_packet odata;
 		int plen=0,result=0;
 
 		/* Acknowledge the disconnection by sending a END packet in return */
-		plen = initPacket(odata, MT_PTYPE_END, srcmac, dstmac, pkthdr.seskey, 0);
-		result = sendUDP(odata, plen);
+		plen = initPacket(&odata, MT_PTYPE_END, srcmac, dstmac, pkthdr.seskey, 0);
+		result = sendUDP(&odata);
 
 		fprintf(stderr, "Connection closed.\n");
 
@@ -231,7 +231,7 @@ void handlePacket(unsigned char *data, int data_len) {
 int main (int argc, char **argv) {
 	int insockfd;
 	int result;
-	unsigned char data[1500];
+	struct mt_packet data;
 	struct sockaddr_in si_me;
 	unsigned char buff[1500];
 	int plen = 0;
@@ -338,8 +338,8 @@ int main (int argc, char **argv) {
 
 	printf("Connecting to %s...", ether_ntoa((struct ether_addr *)dstmac));
 
-	plen = initPacket(data, MT_PTYPE_SESSIONSTART, srcmac, dstmac, sessionkey, 0);
-	result = sendUDP(data, plen);
+	plen = initPacket(&data, MT_PTYPE_SESSIONSTART, srcmac, dstmac, sessionkey, 0);
+	result = sendUDP(&data);
 
 	/* Try to connect with a timeout */
 	FD_ZERO(&read_fds);
@@ -367,12 +367,12 @@ int main (int argc, char **argv) {
 	 * TODO: Should resubmit whenever a PTYPE_DATA packet is sent, and an ACK packet with correct datacounter is received
 	 * or time out the connection, in all other cases.
 	*/
-	plen = initPacket(data, MT_PTYPE_DATA, srcmac, dstmac, sessionkey, 0);
-	plen += addControlPacket(data + plen, MT_CPTYPE_BEGINAUTH, NULL, 0);
-	outcounter += plen - MT_HEADER_LEN;
+	plen = initPacket(&data, MT_PTYPE_DATA, srcmac, dstmac, sessionkey, 0);
+	plen = addControlPacket(&data, MT_CPTYPE_BEGINAUTH, NULL, 0);
+	outcounter += plen;
 
 	/* TODO: handle result of sendUDP */
-	result = sendUDP(data, plen);
+	result = sendUDP(&data);
 
 	while (running) {
 		int reads;
@@ -397,22 +397,23 @@ int main (int argc, char **argv) {
 			if (FD_ISSET(0, &read_fds)) {
 				unsigned char keydata[100];
 				int datalen;
+
 				datalen = read(STDIN_FILENO, &keydata, 100);
-				memset(data, 0, sizeof(data));
-				plen = initPacket(data, MT_PTYPE_DATA, srcmac, dstmac, sessionkey, outcounter);
+
+				plen = initPacket(&data, MT_PTYPE_DATA, srcmac, dstmac, sessionkey, outcounter);
+				plen += addControlPacket(&data, MT_CPTYPE_PLAINDATA, &keydata, datalen);
 				outcounter += datalen;
-				memcpy(data + plen, &keydata, datalen);
-				result = sendUDP(data, plen + datalen);
+				result = sendUDP(&data);
 			}
 		/* Handle select() timeout */
 		} else {
 			/* handle keepalive counter, transmit keepalive packet every 10 seconds
 			   of inactivity  */
 			if ((keepalive_counter++ % 10) == 0) {
-				char odata[MT_HEADER_LEN];
+				struct mt_packet odata;
 				int plen=0,result=0;
-				plen = initPacket(odata, MT_PTYPE_ACK, srcmac, dstmac, sessionkey, 0);
-				result = sendUDP(odata, plen);
+				plen = initPacket(&odata, MT_PTYPE_ACK, srcmac, dstmac, sessionkey, 0);
+				result = sendUDP(&odata);
 			}
 		}
 	}