added forgotten set functions to packets

2 files changed, 16 insertions, 0 deletions

diff --git a/src/encrypted_packet.c b/src/encrypted_packet.c
index 0f76589..176241b 100644
--- a/src/encrypted_packet.c
+++ b/src/encrypted_packet.c
@@ -96,7 +96,13 @@ u_int32_t encrypted_packet_get_payload_length(encrypted_packet_t* packet)
 
 void encrypted_packet_set_payload_length(encrypted_packet_t* packet, u_int32_t len)
 {
+  if(!packet)
+    return;
 
+  if(len > ENCRYPTED_PACKET_SIZE_MAX || (len + sizeof(encrypted_packet_header_t)) > ENCRYPTED_PACKET_SIZE_MAX)
+    len = ENCRYPTED_PACKET_SIZE_MAX - sizeof(encrypted_packet_header_t);
+
+  packet->payload_length_ = len;
 }
 
 seq_nr_t encrypted_packet_get_seq_nr(encrypted_packet_t* packet)
diff --git a/src/plain_packet.c b/src/plain_packet.c
index 9d31467..69338c1 100644
--- a/src/plain_packet.c
+++ b/src/plain_packet.c
@@ -65,7 +65,17 @@ u_int32_t plain_packet_get_length(plain_packet_t* packet)
 
 void plain_packet_set_length(plain_packet_t* packet, u_int32_t len)
 {
+  if(!packet)
+    return;
   
+  if(len > PLAIN_PACKET_SIZE_MAX)
+    len = PLAIN_PACKET_SIZE_MAX - sizeof(payload_type_t);
+  else if(len < sizeof(payload_type_t))
+    len = 0;
+  else
+    len -= sizeof(payload_type_t);
+
+  packet->payload_length_ = len;  
 }
 
 u_int8_t* plain_packet_get_payload(plain_packet_t* packet)