1 files changed, 18 insertions, 16 deletions

diff --git a/src/net.c b/src/net.c
index d7222e1..ed7e69f 100644
--- a/src/net.c
+++ b/src/net.c
@@ -5,40 +5,42 @@
 #include <assert.h>
 #include "net.h"
 
-void* net_recv(int sock, msg_type_t wanted)
+void* net_recv(int sock, msg_type_t type)
 {
-	void* result = NULL;
-	uint8_t peekbuf[2], type, payload_len, *packet, *payload;
-	ssize_t len = recv(sock, peekbuf, 2, MSG_PEEK); // just peek into packet to determine type
+	msg_t m;
+	void *result;
+	ssize_t len = recv(sock, &m.hdr, sizeof(msg_header_t), MSG_PEEK); // just peek into packet to determine message header
 
 	assert(len != -1);
 
-	type = peekbuf[NET_MSG_OFFSET_TYPE];
-	payload_len = peekbuf[NET_MSG_OFFSET_PAYLOAD_LENGTH];
-
-	if(type != wanted)
+	if(m.hdr.type != type)
 	{
-		printf("client_recv: received type %d instead of %d", type, wanted);
+		printf("net_recv: received message type %d instead of %d", m.hdr.type, type);
 		return NULL;
 	}
 
-	packet = malloc(payload_len+NET_MSG_OFFSET_PAYLOAD);
-	recv(sock, packet, payload_len+NET_MSG_OFFSET_PAYLOAD, 0);
-	payload = &packet[NET_MSG_OFFSET_PAYLOAD];
+	m.payload = malloc(m.hdr.payload_length);	// Allocate space for message payload
+	recv(sock, &m.hdr, sizeof(msg_header_t), 0);	// Remove message header from socket
+	recv(sock, m.payload, m.hdr.payload_length, 0);// And then receive the payload
 
 	switch(type)
 	{
 		case msg_type_hello:
-			result = server_recv_hello(payload, payload_len);
+			result = server_recv_hello(m.payload, m.hdr.payload_length);
 			break;
 		case msg_type_start_game:
-			result = client_recv_player_list(payload, payload_len);
+			result = client_recv_player_list(m.payload, m.hdr.payload_length);
 			break;
 		case msg_type_deal_cards:
-			result = client_recv_deal_cards(payload, payload_len);
+			result = client_recv_deal_cards(m.payload, m.hdr.payload_length);
+			break;
+		default:
+			printf("net_recv: Unknown message type %d received!\n", type);
+			exit(EXIT_FAILURE);
 			break;
 	}
-	free(packet);
+
+	free(m.payload);
 
 	return result;
 }