some more tests

2 files changed, 84 insertions, 8 deletions

diff --git a/satp/packet.go b/satp/packet.go
index c7751ce..259c396 100644
--- a/satp/packet.go
+++ b/satp/packet.go
@@ -126,7 +126,7 @@ func (ep *EncryptedPacket) SetAuthTagLength(length int) error {
 		ep.authTag = nil
 		return nil
 	}
-	if length > total {
+	if length > (total - 2) { // minimal payload length = 2 (Payload Type)
 		return ErrTooShort
 	}
 	ep.payload = ep.payload[:total-length]
@@ -170,7 +170,7 @@ func (ep *EncryptedPacket) ReadFrom(r io.Reader) (int64, error) {
 	if err != nil && err != io.EOF {
 		return 0, err
 	}
-	if n < len(ep.header)+2+len(ep.authTag) {
+	if n < len(ep.header)+2+len(ep.authTag) { // minimal payload length = 2 (Payload Type)
 		return 0, ErrTooShort
 	}
 	ep.payload = ep.payload[:n-len(ep.header)]
diff --git a/satp/packet_test.go b/satp/packet_test.go
index ffff08e..e35eb17 100644
--- a/satp/packet_test.go
+++ b/satp/packet_test.go
@@ -254,9 +254,52 @@ func TestEncryptedPacketHeader(t *testing.T) {
 	}
 }
 
-// func TestEncryptedPacketSetAuthTagLen(t *testing.T) {
-// TODO: implement this
-// }
+func TestEncryptedPacketSetAuthTagLen(t *testing.T) {
+	var buffer [PACKET_BUFFER_SIZE]byte
+	tRand.Read(buffer[:])
+
+	testvectors := []struct {
+		Data       []byte
+		AuthTagLen int
+		valid      bool
+	}{
+		{buffer[:10], 0, true},
+		{buffer[:10], 1, false},
+		{buffer[:11], 1, true},
+		{buffer[:200], 20, true},
+		{buffer[:200], 190, true},
+		{buffer[:200], 191, false},
+		{buffer[:1440], 20, true},
+	}
+
+	for _, vector := range testvectors {
+		pkt := NewEncryptedPacket()
+		n, err := pkt.ReadFrom(bytes.NewReader(vector.Data))
+		if err != nil {
+			t.Fatal("unexpected error:", err)
+		}
+
+		err = pkt.SetAuthTagLength(vector.AuthTagLen)
+		if vector.valid {
+			if err != nil {
+				t.Fatal("unexpected error:", err)
+			}
+			if n != int64(len(vector.Data)) {
+				t.Fatalf("returned length is invalid, is: %d, should be %d", n, len(vector.Data))
+			}
+
+			out := pkt.GetAuthTagLength()
+			if out != vector.AuthTagLen {
+				t.Fatalf("resulting auth-tag length is wrong, is: %d, should be %d", out, vector.AuthTagLen)
+			}
+		} else {
+			if err == nil {
+				t.Fatalf("setting auth-tag length to %d on packet with length %d should give an error", vector.AuthTagLen, len(vector.Data))
+			}
+		}
+	}
+
+}
 
 func TestEncryptedPacketReadFrom(t *testing.T) {
 	testvectors := []struct {
@@ -310,9 +353,42 @@ func TestEncryptedPacketReadFrom(t *testing.T) {
 	}
 }
 
-// func TestEncryptedPacketWriteTo(t *testing.T) {
-// TODO: implement this
-// }
+func TestEncryptedPacketWriteTo(t *testing.T) {
+	var buffer [PACKET_BUFFER_SIZE]byte
+	tRand.Read(buffer[:])
+
+	testvectors := []struct {
+		Data           []byte
+		SequenceNumber uint32
+		SenderID       uint16
+		Mux            uint16
+	}{
+		{[]byte{0, 0, 0, 0, 0, 0, 0, 0, 0, 0}, 0, 0, 0},
+		{[]byte{1, 2, 3, 4, 5, 6, 7, 8, 17, 32}, 0x01020304, 0x0506, 0x0708},
+	}
+
+	for _, vector := range testvectors {
+		pkt := NewEncryptedPacket()
+		pkt.SetSequenceNumber(vector.SequenceNumber)
+		pkt.SetSenderID(vector.SenderID)
+		pkt.SetMux(vector.Mux)
+		pkt.payload = pkt.payload[:len(vector.Data)-8]
+		copy(pkt.payload, vector.Data[8:])
+
+		w := &bytes.Buffer{}
+		n, err := pkt.WriteTo(w)
+		if err != nil {
+			t.Fatal("unexpected error:", err)
+		}
+		if n != int64(len(vector.Data)) {
+			t.Fatalf("returned length is invalid, is: %d, should be %d", n, len(vector.Data))
+		}
+		out := w.Bytes()
+		if bytes.Compare(out, vector.Data) != 0 {
+			t.Fatalf("resulting payload is invalid, is: '%v', should be '%v'", out, vector.Data)
+		}
+	}
+}
 
 func generateRandomTestDataEncryptedPacket() (authTagLen int, packetData []byte) {
 	minLen := int32(8 + 2)