Modify delete payload struct (#11)

Co-authored-by: Allen00991 <allen.chen@saviah.com>
free5gc · Oct 1, 2024 · 7474f6e · 7474f6e
1 parent 1edde69
commit 7474f6e
Show file tree

Hide file tree

Showing 3 changed files with 27 additions and 14 deletions.
diff --git a/message/build.go b/message/build.go
@@ -151,7 +151,7 @@ func (container *ProposalContainer) BuildProposal(proposalNumber uint8, protocol
 }
 
 func (container *IKEPayloadContainer) BuildDeletePayload(
-	protocolID uint8, spiSize uint8, numberOfSPI uint16, spis []byte,
+	protocolID uint8, spiSize uint8, numberOfSPI uint16, spis []uint32,
 ) {
 	deletePayload := new(Delete)
 	deletePayload.ProtocolID = protocolID

diff --git a/message/payload_delete.go b/message/payload_delete.go
@@ -12,14 +12,14 @@ type Delete struct {
 	ProtocolID  uint8
 	SPISize     uint8
 	NumberOfSPI uint16
-	SPIs        []byte
+	SPIs        []uint32
 }
 
 func (d *Delete) Type() IKEPayloadType { return TypeD }
 
 func (d *Delete) marshal() ([]byte, error) {
-	if len(d.SPIs) != (int(d.SPISize) * int(d.NumberOfSPI)) {
-		return nil, errors.Errorf("Total bytes of all SPIs not correct")
+	if len(d.SPIs) != int(d.NumberOfSPI) {
+		return nil, errors.Errorf("Number of SPI not correct")
 	}
 
 	deleteData := make([]byte, 4)
@@ -29,7 +29,11 @@ func (d *Delete) marshal() ([]byte, error) {
 	binary.BigEndian.PutUint16(deleteData[2:4], d.NumberOfSPI)
 
 	if int(d.NumberOfSPI) > 0 {
-		deleteData = append(deleteData, d.SPIs...)
+		byteSlice := make([]byte, d.SPISize)
+		for _, v := range d.SPIs {
+			binary.BigEndian.PutUint32(byteSlice, v)
+			deleteData = append(deleteData, byteSlice...)
+		}
 	}
 
 	return deleteData, nil
@@ -51,7 +55,12 @@ func (d *Delete) unmarshal(b []byte) error {
 		d.SPISize = spiSize
 		d.NumberOfSPI = numberOfSPI
 
-		d.SPIs = append(d.SPIs, b[4:]...)
+		b = b[4:]
+		var spi uint32
+		for i := 0; i < len(b); i += 4 {
+			spi = binary.BigEndian.Uint32(b[i : i+4])
+			d.SPIs = append(d.SPIs, spi)
+		}
 	}
 
 	return nil

diff --git a/message/payload_delete_test.go b/message/payload_delete_test.go
@@ -14,12 +14,12 @@ func TestDeleteMarshal(t *testing.T) {
 		expErr      bool
 	}{
 		{
-			description: "Total bytes of all SPIs not correct",
+			description: "Number of SPI not correct",
 			delete: Delete{
 				ProtocolID:  TypeESP,
 				SPISize:     4,
 				NumberOfSPI: 1,
-				SPIs:        []byte{0x01, 0x02, 0x03},
+				SPIs:        []uint32{0x01, 0x02, 0x03},
 			},
 			expErr: true,
 		},
@@ -41,11 +41,13 @@ func TestDeleteMarshal(t *testing.T) {
 			delete: Delete{
 				ProtocolID:  TypeESP,
 				SPISize:     4,
-				NumberOfSPI: 1,
-				SPIs:        []byte{0x01, 0x02, 0x03, 0x04},
+				NumberOfSPI: 4,
+				SPIs:        []uint32{0x01, 0x02, 0x03, 0x04},
 			},
 			expMarshal: []byte{
-				0x03, 0x04, 0x00, 0x01, 0x01, 0x02, 0x03, 0x04,
+				0x03, 0x04, 0x00, 0x04, 0x00, 0x00, 0x00, 0x01,
+				0x00, 0x00, 0x00, 0x02, 0x00, 0x00, 0x00, 0x03,
+				0x00, 0x00, 0x00, 0x04,
 			},
 			expErr: false,
 		},
@@ -84,13 +86,15 @@ func TestDeleteUnmarshal(t *testing.T) {
 		{
 			description: "Delete Unmarshal",
 			b: []byte{
-				0x03, 0x04, 0x00, 0x01, 0x01, 0x02, 0x03, 0x04,
+				0x03, 0x04, 0x00, 0x04, 0x00, 0x00, 0x00, 0x01,
+				0x00, 0x00, 0x00, 0x02, 0x00, 0x00, 0x00, 0x03,
+				0x00, 0x00, 0x00, 0x04,
 			},
 			expMarshal: Delete{
 				ProtocolID:  TypeESP,
 				SPISize:     4,
-				NumberOfSPI: 1,
-				SPIs:        []byte{0x01, 0x02, 0x03, 0x04},
+				NumberOfSPI: 4,
+				SPIs:        []uint32{0x01, 0x02, 0x03, 0x04},
 			},
 			expErr: false,
 		},