1 files changed, 91 insertions, 57 deletions

diff --git a/package.cpp b/package.cpp
index fe042d1..4cdba50 100644
--- a/package.cpp
+++ b/package.cpp
@@ -37,202 +37,236 @@
 
 Package::Package()
 {
-  header_ = 0;
-  payload_type_ = 0;
-  auth_tag_ = 0;                
+  has_header_ = false;
+  has_payload_type_ = false;
+  has_auth_tag_ = false;                
 }
 
 Package::Package(u_int32_t length) : Buffer(length)
 {
-  header_ = 0;
-  payload_type_ = 0;
-  auth_tag_ = 0;                
+  has_header_ = false;
+  has_payload_type_ = false;
+  has_auth_tag_ = false;                
 }
 
 Package::Package(const Buffer &src) : Buffer(src)
 {
-  header_ = 0;
-  payload_type_ = 0;
-  auth_tag_ = 0;                
+  has_header_ = false;
+  has_payload_type_ = false;
+  has_auth_tag_ = false;                
 }
   
 bool Package::hasHeader() const
 {
-  return header_;
+  return has_header_;
 }
 
 Package& Package::withHeader(bool b)
 {
   if(b && length_ >= sizeof(struct HeaderStruct))
-    header_ = reinterpret_cast<struct HeaderStruct*>(buf_);
+    has_header_ = true;
   else
-    header_ = 0;
+    has_header_ = false;
 
   return *this;
 }
 
 seq_nr_t Package::getSeqNr() const
 {
-  if(!header_)
+  if(!has_header_)
     return 0;
 
-  return SEQ_NR_T_NTOH(header_->seq_nr);
+  struct HeaderStruct* header;
+  header = reinterpret_cast<struct HeaderStruct*>(buf_);
+  return SEQ_NR_T_NTOH(header->seq_nr);
 }
 
 sender_id_t Package::getSenderId() const
 {
-  if(!header_)
+  if(!has_header_)
     return 0;
 
-  return SENDER_ID_T_NTOH(header_->sender_id);
+  struct HeaderStruct* header;
+  header = reinterpret_cast<struct HeaderStruct*>(buf_);
+  return SENDER_ID_T_NTOH(header->sender_id);
 }
 
 Package& Package::addHeader(seq_nr_t seq_nr, sender_id_t sender_id)
 {
-  if(!header_)
+  if(!has_header_)
   {
     if(sizeof(struct HeaderStruct) > resizeFront(length_ + sizeof(struct HeaderStruct)))
       return *this;
 
-    header_ = reinterpret_cast<struct HeaderStruct*>(buf_);
+    has_header_ = true;
   }
-  header_->seq_nr = SEQ_NR_T_HTON(seq_nr);
-  header_->sender_id = SENDER_ID_T_HTON(sender_id);
+  struct HeaderStruct* header;
+  header = reinterpret_cast<struct HeaderStruct*>(buf_);
+  header->seq_nr = SEQ_NR_T_HTON(seq_nr);
+  header->sender_id = SENDER_ID_T_HTON(sender_id);
   return *this;
 }
 
 Package& Package::removeHeader()
 {
-  if(!header_)
+  if(!has_header_)
     return *this;
 
   if(length_ >= sizeof(struct HeaderStruct))
     resizeFront(length_ - sizeof(struct HeaderStruct));
 
-  header_ = 0;
+  has_header_ = false;
   
   return *this;
 }
 
 Package& Package::setSeqNr(seq_nr_t seq_nr)
 {
-  if(header_)
-    header_->seq_nr = SEQ_NR_T_HTON(seq_nr);
-
+  if(has_header_)
+  {
+    struct HeaderStruct* header;
+    header = reinterpret_cast<struct HeaderStruct*>(buf_);
+    header->seq_nr = SEQ_NR_T_HTON(seq_nr);
+  }
   return *this;
 }
 
 Package& Package::setSenderId(sender_id_t sender_id)
 {
-  if(header_)
-    header_->sender_id = SENDER_ID_T_HTON(sender_id);
-
+  if(has_header_)
+  {
+    struct HeaderStruct* header;
+    header = reinterpret_cast<struct HeaderStruct*>(buf_);
+    header->sender_id = SENDER_ID_T_HTON(sender_id);
+  }
   return *this;
 }
 
+
+
 bool Package::hasPayloadType() const
 {
-  return payload_type_;
+  return has_payload_type_;
 }
 
 Package& Package::withPayloadType(bool b)
 {
-  if(auth_tag_)
-    throw std::runtime_error("can't change payload_type state with existing auth_tag");
-
   if(b && length_ >= sizeof(payload_type_t))
-    payload_type_ = reinterpret_cast<payload_type_t*>(&buf_[length_ - sizeof(payload_type_t)]);
+    has_payload_type_ = true;
   else
-    payload_type_ = 0;
+    has_payload_type_ = false;
 
   return *this;
 }
 
 payload_type_t Package::getPayloadType() const
 {
-  if(!payload_type_)
+  if(!has_payload_type_)
+    return 0;
+
+  if((!has_auth_tag_ && length_ < sizeof(payload_type_t)) ||
+     (has_auth_tag_ && length_ < (sizeof(payload_type_t) + sizeof(auth_tag_t))))
     return 0;
 
-  return PAYLOAD_TYPE_T_NTOH(*payload_type_);
+  payload_type_t* payload_type;
+
+  if(!has_auth_tag_)  
+    payload_type = reinterpret_cast<payload_type_t*>(buf_ + length_ - sizeof(payload_type_t));
+  else
+    payload_type = reinterpret_cast<payload_type_t*>(buf_ + length_ - sizeof(payload_type_t) - sizeof(auth_tag_t));
+  return PAYLOAD_TYPE_T_NTOH(*payload_type);
 }
 
 Package& Package::addPayloadType(payload_type_t payload_type)
 {
-  if(auth_tag_)
+  if(has_auth_tag_)
     throw std::runtime_error("can't add payload_type with existing auth_tag");
 
-  if(!payload_type_)
+  if(!has_payload_type_)
   {
-    if(sizeof(payload_type_t) > resizeBack(length_ + sizeof(payload_type_t)))
+    u_int32_t new_length = length_ + sizeof(payload_type_t);
+    if(new_length > resizeBack(new_length))
       return *this;
 
-    payload_type_ = reinterpret_cast<payload_type_t*>(&buf_[length_ - sizeof(payload_type_t)]);
+    has_payload_type_ = true;
   }
-  *payload_type_ = PAYLOAD_TYPE_T_HTON(payload_type);
+  payload_type_t* payload_type_ptr;
+  payload_type_ptr = reinterpret_cast<payload_type_t*>(buf_ + length_ - sizeof(payload_type_t));
+  *payload_type_ptr = PAYLOAD_TYPE_T_HTON(payload_type);
   return *this;
 }
 
 Package& Package::removePayloadType()
 {
-  if(auth_tag_)
+  if(has_auth_tag_)
     throw std::runtime_error("can't remove payload_type with existing auth_tag");
 
-  if(!payload_type_)
+  if(!has_payload_type_)
     return *this;
 
   if(length_ >= sizeof(payload_type_t))
     resizeBack(length_ - sizeof(payload_type_t));
 
-  payload_type_ = 0;
+  has_payload_type_ = false;
   
   return *this;
 }
 
+
+
 bool Package::hasAuthTag() const
 {
-  return auth_tag_;
+  return has_auth_tag_;
 }
 
 Package& Package::withAuthTag(bool b)
 {
   if(b && length_ >= sizeof(auth_tag_t))
-    auth_tag_ = reinterpret_cast<auth_tag_t*>(&buf_[length_ - sizeof(auth_tag_t)]);
+    has_auth_tag_ = true;
   else
-    auth_tag_ = 0;
-
+    has_auth_tag_ = false;
+  
   return *this;
-}
+}  
 
 auth_tag_t Package::getAuthTag() const
 {
-  if(!auth_tag_)
+  if(!has_auth_tag_)
+    return 0;
+
+  if(length_ < sizeof(auth_tag_t))
     return 0;
 
-  return AUTH_TAG_T_NTOH(*auth_tag_);
+  auth_tag_t* auth_tag;
+  auth_tag = reinterpret_cast<auth_tag_t*>(buf_ + length_ - sizeof(auth_tag_t));
+  return AUTH_TAG_T_NTOH(*auth_tag);
 }
 
 Package& Package::addAuthTag(auth_tag_t auth_tag)
 {
-  if(!auth_tag_)
+  if(!has_auth_tag_)
   {
-    if(sizeof(auth_tag_t) > resizeBack(length_ + sizeof(auth_tag_t)))
+    u_int32_t new_length = length_ + sizeof(auth_tag_t);
+    if(new_length > resizeBack(new_length))
       return *this;
 
-    auth_tag_ = reinterpret_cast<auth_tag_t*>(&buf_[length_ - sizeof(auth_tag_t)]);
+    has_auth_tag_ = true;
   }
-  *auth_tag_ = AUTH_TAG_T_HTON(auth_tag);
+  auth_tag_t* auth_tag_ptr;
+  auth_tag_ptr = reinterpret_cast<auth_tag_t*>(buf_ + length_ - sizeof(auth_tag_t));
+  *auth_tag_ptr = AUTH_TAG_T_HTON(auth_tag);
   return *this;
 }
 
 Package& Package::removeAuthTag()
 {
-  if(!auth_tag_)
+  if(!has_auth_tag_)
     return *this;
 
   if(length_ >= sizeof(auth_tag_t))
     resizeBack(length_ - sizeof(auth_tag_t));
 
-  auth_tag_ = 0;
+  has_auth_tag_ = false;
   
   return *this;
 }