header tests for encrypted packet

2 files changed, 61 insertions, 3 deletions

diff --git a/satp/packet.go b/satp/packet.go
index 22cf4ba..b52b0e1 100644
--- a/satp/packet.go
+++ b/satp/packet.go
@@ -114,7 +114,7 @@ type EncryptedPacket struct {
 func NewEncryptedPacket() (ep *EncryptedPacket) {
 	ep = &EncryptedPacket{}
 	ep.header = ep.buffer[:8:8]
-	ep.payload = ep.buffer[8:]
+	ep.payload = ep.buffer[8:8]
 	ep.authTag = nil
 	return
 }
@@ -177,7 +177,11 @@ func (ep *EncryptedPacket) ReadFrom(r io.Reader) (int64, error) {
 	return int64(n), err
 }
 
+func (ep *EncryptedPacket) getPacket() (data []byte) {
+	return ep.buffer[:len(ep.header)+len(ep.payload)+len(ep.authTag)]
+}
+
 func (ep *EncryptedPacket) WriteTo(w io.Writer) (int64, error) {
-	n, err := w.Write(ep.buffer[:len(ep.header)+len(ep.payload)+len(ep.authTag)])
+	n, err := w.Write(ep.getPacket())
 	return int64(n), err
 }
diff --git a/satp/packet_test.go b/satp/packet_test.go
index 1284880..2ae7aad 100644
--- a/satp/packet_test.go
+++ b/satp/packet_test.go
@@ -69,7 +69,7 @@ var (
 		0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0xC0, 0xA8, 0x2A, 0x01}
 )
 
-func TestPlainPacketPayloadType(t *testing.T) {
+func TestPlainPacketHeader(t *testing.T) {
 	testvectors := []uint16{0, IPv4Type, IPv6Type, EthernetType, 0xAA55, 0xF00F}
 
 	for _, vector := range testvectors {
@@ -195,6 +195,60 @@ func TestPlainPacketReadWrite(t *testing.T) {
 	}
 }
 
+func TestEncryptedPacketHeader(t *testing.T) {
+	testvectors := []struct {
+		SequenceNumber uint32
+		SenderID       uint16
+		Mux            uint16
+	}{
+		{0, 0, 0},
+		{23, 0, 0},
+		{0, 12, 0},
+		{0, 0, 42},
+		{1, 2, 3},
+		{4294967295, 0, 0},
+		{0, 65535, 0},
+		{0, 0, 65535},
+		{1254234, 2, 23},
+	}
+
+	for _, vector := range testvectors {
+		pkt := NewEncryptedPacket()
+		if vector.SequenceNumber > 0 {
+			pkt.SetSequenceNumber(vector.SequenceNumber)
+		}
+		if vector.SenderID > 0 {
+			pkt.SetSenderID(vector.SenderID)
+		}
+		if vector.Mux > 0 {
+			pkt.SetMux(vector.Mux)
+		}
+		result := pkt.getPacket()
+		expected := []byte{
+			byte(vector.SequenceNumber >> 24), byte(vector.SequenceNumber >> 16),
+			byte(vector.SequenceNumber >> 8), byte(vector.SequenceNumber),
+			byte(vector.SenderID >> 8), byte(vector.SenderID),
+			byte(vector.Mux >> 8), byte(vector.Mux)}
+
+		if bytes.Compare(result, expected) != 0 {
+			t.Fatalf("resulting packet is invalid, is: '%v', should be '%v'", result, expected)
+		}
+
+		sequenceNumber := pkt.GetSequenceNumber()
+		if vector.SequenceNumber != sequenceNumber {
+			t.Fatalf("resulting sequence number is invalid, is: 0x%04X, should be 0x%04X", sequenceNumber, vector.SequenceNumber)
+		}
+		senderID := pkt.GetSenderID()
+		if vector.SenderID != senderID {
+			t.Fatalf("resulting sequence number is invalid, is: 0x%04X, should be 0x%04X", senderID, vector.SenderID)
+		}
+		Mux := pkt.GetMux()
+		if vector.Mux != Mux {
+			t.Fatalf("resulting sequence number is invalid, is: 0x%04X, should be 0x%04X", Mux, vector.Mux)
+		}
+	}
+}
+
 // func TestEncryptedPacketMarshal(t *testing.T) {
 // 	testvectors := []struct {
 // 		packet   EncryptedPacket