added encrypted packet check with 'random' data

2 files changed, 45 insertions, 3 deletions

diff --git a/satp/packet.go b/satp/packet.go
index b52b0e1..c7751ce 100644
--- a/satp/packet.go
+++ b/satp/packet.go
@@ -124,12 +124,13 @@ func (ep *EncryptedPacket) SetAuthTagLength(length int) error {
 	if length <= 0 {
 		ep.payload = ep.payload[:total]
 		ep.authTag = nil
+		return nil
 	}
 	if length > total {
 		return ErrTooShort
 	}
 	ep.payload = ep.payload[:total-length]
-	ep.authTag = ep.buffer[len(ep.header)+len(ep.payload) : total]
+	ep.authTag = ep.buffer[len(ep.header)+len(ep.payload) : len(ep.header)+total]
 	return nil
 }
 
diff --git a/satp/packet_test.go b/satp/packet_test.go
index 16fcee8..98c1cdb 100644
--- a/satp/packet_test.go
+++ b/satp/packet_test.go
@@ -161,8 +161,8 @@ func TestPlainPacketWriteTo(t *testing.T) {
 func generateRandomTestDataPlainPacket() (payloadType uint16, payload []byte) {
 	payloadType = uint16(rand.Uint32())
 
-	packetlen := uint(2 + rand.Int31n(PACKET_BUFFER_SIZE-2))
-	payload = make([]byte, packetlen)
+	packetLen := uint(2 + rand.Int31n(PACKET_BUFFER_SIZE-2))
+	payload = make([]byte, packetLen)
 	rand.Read(payload)
 	return
 }
@@ -249,6 +249,47 @@ func TestEncryptedPacketHeader(t *testing.T) {
 	}
 }
 
+func generateRandomTestDataEncryptedPacket() (authTagLen int, packetData []byte) {
+	minLen := int32(8 + 2)
+	packetLen := uint(minLen + rand.Int31n(PACKET_BUFFER_SIZE-minLen))
+	authTagLen = int(rand.Int31n(int32(packetLen) - minLen))
+	packetData = make([]byte, packetLen)
+	rand.Read(packetData)
+	return
+}
+
+func TestEncryptedPacketReadWrite(t *testing.T) {
+	for i := 0; i < NUM_RANDOM_DATASETS; i++ {
+		pkt := NewEncryptedPacket()
+		authTagLen, in := generateRandomTestDataEncryptedPacket()
+
+		n, err := pkt.ReadFrom(bytes.NewReader(in))
+		if err != nil {
+			t.Fatal("unexpected error:", err)
+		}
+		if n != int64(len(in)) {
+			t.Fatalf("returned length is invalid, is: %d, should be %d", n, len(in))
+		}
+
+		err = pkt.SetAuthTagLength(authTagLen)
+		if err != nil {
+			t.Fatal("unexpected error:", err)
+		}
+
+		out := &bytes.Buffer{}
+		if n, err = pkt.WriteTo(out); err != nil {
+			t.Fatal("unexpected error:", err)
+		}
+		if n != int64(len(in)) {
+			t.Fatalf("returned length is invalid, is: %d, should be %d ... authTagLength is %d", n, len(in), authTagLen)
+		}
+
+		if bytes.Compare(in, out.Bytes()) != 0 {
+			t.Fatalf("resulting packet is invalid, is: '%v', should be '%v'", out.Bytes(), in)
+		}
+	}
+}
+
 // func TestEncryptedPacketMarshal(t *testing.T) {
 // 	testvectors := []struct {
 // 		packet   EncryptedPacket