Move packet size overflow checking around

1 files changed, 10 insertions, 12 deletions

diff --git a/jawol.c b/jawol.c
index b301153..d604455 100644
--- a/jawol.c
+++ b/jawol.c
@@ -70,6 +70,11 @@ static const char usage_msg[] =
 #include <sys/ioctl.h>
 #include <getopt.h>
 
+#define PACKET_SIZE (200)
+#define CHECK_PACKET_SIZE(x) do { \
+		if ( (x) >= PACKET_SIZE ) { fprintf( stderr, "Packet size overflow.\n" ); abort(); } \
+	} while(0)
+
 static struct {
 	const char *ifname;
 	struct in_addr dest;
@@ -160,7 +165,7 @@ int main(int argc, char **argv)
 
 static int etherwake(int count, char **maclist)
 {
-	unsigned char outpack[200];
+	unsigned char outpack[PACKET_SIZE];
 	int s = rawsocket();
 	if ( s == -1 )
 		return 2;
@@ -200,6 +205,7 @@ static int etherwake(int count, char **maclist)
 			memcpy( outpack + packet_size, options.wolpw, options.wolpwlen );
 			packet_size += options.wolpwlen;
 		}
+		CHECK_PACKET_SIZE( packet_size );
 		if ( options.debug ) {
 			printf( "The final packet for %s is:", maclist[idx] );
 			for ( int i = 0; i < packet_size; i++ ) printf( " %2.2x", outpack[i] );
@@ -278,7 +284,7 @@ int udpsocket()
 static int wakeonlan(int count, char **maclist)
 {
 	struct ether_addr dest_ether;
-	unsigned char outpack[200];
+	unsigned char outpack[PACKET_SIZE];
 	int sock = udpsocket();
 	if ( sock == -1 )
 		return 2;
@@ -297,6 +303,7 @@ static int wakeonlan(int count, char **maclist)
 			memcpy( outpack + packet_size, options.wolpw, options.wolpwlen );
 			packet_size += options.wolpwlen;
 		}
+		CHECK_PACKET_SIZE( packet_size );
 		if ( options.debug ) {
 			printf( "The final packet for %s is:", maclist[idx] );
 			for ( int i = 0; i < packet_size; i++ ) printf( " %2.2x", outpack[i] );
@@ -343,16 +350,7 @@ static int get_fill(unsigned char *pkt, const uint8_t *bindestmac, bool with_hea
 		memcpy(pkt + offset, bindestmac, 6);
 		offset += 6;
 	}
-	if ( offset > 180 ) {
-		fprintf( stderr, "Packet buffer overflow.\n" );
-		abort();
-	}
-	if (options.debug) {
-		printf("Packet is ");
-		for (i = 0; i < offset; i++)
-			printf(" %2.2x", pkt[i]);
-		printf(".\n");
-	}
+	CHECK_PACKET_SIZE( offset );
 	return offset;
 }