refactor: add NewHeader & NewMessage functions

free5gc · Oct 1, 2024 · 1edde69 · 1edde69
1 parent c86a6ee
commit 1edde69
Show file tree

Hide file tree

Showing 7 changed files with 339 additions and 114 deletions.
diff --git a/ike_test.go b/ike_test.go
@@ -62,7 +62,7 @@ func TestEncodeDecode(t *testing.T) {
 			MajorVersion: 2,
 			MinorVersion: 0,
 			ExchangeType: message.IKE_AUTH,
-			Initiator:    true,
+			Flags:        message.InitiatorBitCheck,
 			MessageID:    0x03,
 			NextPayload:  uint8(message.TypeEAP),
 		},
@@ -92,7 +92,7 @@ func TestEncodeDecode(t *testing.T) {
 	b, err := EncodeEncrypt(expIkeMsg, ikeSAKey, message.Role_Initiator)
 	require.NoError(t, err)
 
-	ikehdr, err := message.ParseIkeHeader(b)
+	ikehdr, err := message.ParseHeader(b)
 	require.NoError(t, err)
 
 	ikeMsg, err := DecodeDecrypt(b, ikehdr, ikeSAKey, message.Role_Responder)
@@ -162,9 +162,21 @@ func TestDecodeDecrypt(t *testing.T) {
 					MajorVersion: 2,
 					MinorVersion: 0,
 					ExchangeType: message.IKE_AUTH,
-					Initiator:    true,
+					Flags:        message.InitiatorBitCheck,
 					MessageID:    0x03,
 					NextPayload:  uint8(message.TypeSK),
+					PayloadBytes: []byte{
+						0x30, 0x00, 0x00, 0x50, 0xec, 0x50, 0x31, 0x16,
+						0x2c, 0x69, 0x2f, 0xbb, 0xfc, 0x4d, 0x20, 0x64,
+						0x0c, 0x91, 0x21, 0xeb, 0xe9, 0x47, 0x5e, 0xf9,
+						0x4f, 0x9b, 0x02, 0x95, 0x9d, 0x31, 0x24, 0x2e,
+						0x53, 0x5e, 0x9c, 0x3c, 0x4d, 0xca, 0xec, 0xd1,
+						0xbf, 0xd6, 0xdd, 0x80, 0xaa, 0x81, 0x2b, 0x07,
+						0xde, 0x36, 0xde, 0xe9, 0xb7, 0x50, 0x94, 0x35,
+						0xf6, 0x35, 0xe1, 0xaa, 0xae, 0x1c, 0x38, 0x25,
+						0xf4, 0xea, 0xe3, 0x38, 0x49, 0x03, 0xf7, 0x24,
+						0xf4, 0x44, 0x17, 0x0c, 0x68, 0x45, 0xca, 0x80,
+					},
 				},
 				Payloads: message.IKEPayloadContainer{
 					&message.EAP{
@@ -415,7 +427,7 @@ func TestEncryptMsg(t *testing.T) {
 			MajorVersion: 2,
 			MinorVersion: 0,
 			ExchangeType: message.IKE_AUTH,
-			Initiator:    true,
+			Flags:        message.InitiatorBitCheck,
 			MessageID:    0x03,
 		},
 	}

diff --git a/message/build.go b/message/build.go
@@ -7,25 +7,6 @@ import (
 	"github.com/pkg/errors"
 )
 
-func (ikeMessage *IKEMessage) BuildIKEHeader(
-	initiatorSPI uint64,
-	responderSPI uint64,
-	exchangeType uint8,
-	initiator bool,
-	response bool,
-	messageID uint32,
-) {
-	ikeMessage.IKEHeader = new(IKEHeader)
-	ikeMessage.InitiatorSPI = initiatorSPI
-	ikeMessage.ResponderSPI = responderSPI
-	ikeMessage.MajorVersion = 2
-	ikeMessage.MinorVersion = 0
-	ikeMessage.Initiator = initiator
-	ikeMessage.Response = response
-	ikeMessage.ExchangeType = exchangeType
-	ikeMessage.MessageID = messageID
-}
-
 func (container *IKEPayloadContainer) Reset() {
 	*container = nil
 }

diff --git a/message/header.go b/message/header.go
@@ -0,0 +1,112 @@
+package message
+
+import (
+	"encoding/binary"
+
+	"github.com/pkg/errors"
+)
+
+// Flags
+const (
+	ResponseBitCheck  = 0x20
+	VersionBitCheck   = 0x10
+	InitiatorBitCheck = 0x08
+)
+
+const IKE_HEADER_LEN int = 28
+
+type IKEHeader struct {
+	InitiatorSPI uint64
+	ResponderSPI uint64
+	MajorVersion uint8
+	MinorVersion uint8
+	ExchangeType uint8
+	Flags        uint8
+	MessageID    uint32
+	NextPayload  uint8
+	PayloadBytes []byte
+}
+
+func NewHeader(
+	iSPI, rSPI uint64, exchgType uint8,
+	response, initiator bool, mId uint32,
+	nextPayload uint8, payloadBytes []byte,
+) *IKEHeader {
+	h := &IKEHeader{
+		InitiatorSPI: iSPI,
+		ResponderSPI: rSPI,
+		ExchangeType: exchgType,
+		MajorVersion: 2,
+		MinorVersion: 0,
+		MessageID:    mId,
+		NextPayload:  nextPayload,
+		PayloadBytes: payloadBytes,
+	}
+	if response {
+		h.Flags |= ResponseBitCheck
+	}
+	if initiator {
+		h.Flags |= InitiatorBitCheck
+	}
+	return h
+}
+
+func (h *IKEHeader) Marshal() ([]byte, error) {
+	b := make([]byte, IKE_HEADER_LEN)
+
+	binary.BigEndian.PutUint64(b[0:8], h.InitiatorSPI)
+	binary.BigEndian.PutUint64(b[8:16], h.ResponderSPI)
+	b[16] = h.NextPayload
+	b[17] = (h.MajorVersion << 4) | (h.MinorVersion & 0x0F)
+	b[18] = h.ExchangeType
+	b[19] = h.Flags
+	binary.BigEndian.PutUint32(b[20:24], h.MessageID)
+
+	totalLen := IKE_HEADER_LEN + len(h.PayloadBytes)
+	if totalLen > 0xFFFFFFFF {
+		return nil, errors.Errorf("length exceeds uint32 limit: %d", totalLen)
+	}
+
+	binary.BigEndian.PutUint32(b[24:IKE_HEADER_LEN], uint32(totalLen))
+	if len(h.PayloadBytes) > 0 {
+		b = append(b, h.PayloadBytes...)
+	}
+	return b, nil
+}
+
+func (h *IKEHeader) IsResponse() bool {
+	return (h.Flags & ResponseBitCheck) != 0
+}
+
+func (h *IKEHeader) IsInitiator() bool {
+	return (h.Flags & InitiatorBitCheck) != 0
+}
+
+func ParseHeader(b []byte) (*IKEHeader, error) {
+	// IKE message packet format this implementation referenced is
+	// defined in RFC 7296, Section 3.1
+	// bounds checking
+	if len(b) < IKE_HEADER_LEN {
+		return nil, errors.Errorf("ParseHeader(): Received broken IKE header")
+	}
+
+	totalLen := binary.BigEndian.Uint32(b[24:IKE_HEADER_LEN])
+	if totalLen < uint32(IKE_HEADER_LEN) {
+		return nil, errors.Errorf("ParseHeader(): Illegal IKE message length %d < header length %d",
+			totalLen, IKE_HEADER_LEN)
+	}
+
+	h := &IKEHeader{
+		InitiatorSPI: binary.BigEndian.Uint64(b[:8]),
+		ResponderSPI: binary.BigEndian.Uint64(b[8:16]),
+		NextPayload:  b[16],
+		MajorVersion: b[17] >> 4,
+		MinorVersion: b[17] & 0x0F,
+		ExchangeType: b[18],
+		Flags:        b[19],
+		MessageID:    binary.BigEndian.Uint32(b[20:24]),
+		PayloadBytes: b[IKE_HEADER_LEN:],
+	}
+
+	return h, nil
+}
diff --git a/message/header_test.go b/message/header_test.go
@@ -0,0 +1,61 @@
+package message
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/require"
+)
+
+func TestNewAndMarshalIKEHeader(t *testing.T) {
+	ikeHdr := NewHeader(
+		0x000000000006f708, 0xc9e2e31f8b64053d, IKE_AUTH,
+		false, true, 0x03, uint8(NoNext), nil,
+	)
+	require.Equal(t, &IKEHeader{
+		InitiatorSPI: 0x000000000006f708,
+		ResponderSPI: 0xc9e2e31f8b64053d,
+		MajorVersion: 2,
+		MinorVersion: 0,
+		ExchangeType: IKE_AUTH,
+		Flags:        InitiatorBitCheck,
+		MessageID:    0x03,
+		NextPayload:  uint8(NoNext),
+	}, ikeHdr)
+	require.True(t, ikeHdr.IsInitiator())
+	require.False(t, ikeHdr.IsResponse())
+
+	b, err := ikeHdr.Marshal()
+	require.NoError(t, err)
+	require.Equal(t, []byte{
+		0x00, 0x00, 0x00, 0x00, 0x00, 0x06, 0xf7, 0x08,
+		0xc9, 0xe2, 0xe3, 0x1f, 0x8b, 0x64, 0x05, 0x3d,
+		0x00, 0x20, 0x23, 0x08, 0x00, 0x00, 0x00, 0x03,
+		0x00, 0x00, 0x00, 0x1c,
+	}, b)
+}
+
+func TestParseIKEHeader(t *testing.T) {
+	ikeHdr, err := ParseHeader(
+		[]byte{
+			0x00, 0x00, 0x00, 0x00, 0x00, 0x06, 0xf7, 0x08,
+			0xc9, 0xe2, 0xe3, 0x1f, 0x8b, 0x64, 0x05, 0x3d,
+			0x00, 0x20, 0x23, 0x08, 0x00, 0x00, 0x00, 0x03,
+			0x00, 0x00, 0x00, 0x1c,
+		},
+	)
+	require.NoError(t, err)
+
+	require.Equal(t, &IKEHeader{
+		InitiatorSPI: 0x000000000006f708,
+		ResponderSPI: 0xc9e2e31f8b64053d,
+		MajorVersion: 2,
+		MinorVersion: 0,
+		ExchangeType: IKE_AUTH,
+		Flags:        InitiatorBitCheck,
+		MessageID:    0x03,
+		NextPayload:  uint8(NoNext),
+		PayloadBytes: []byte{},
+	}, ikeHdr)
+	require.True(t, ikeHdr.IsInitiator())
+	require.False(t, ikeHdr.IsResponse())
+}
diff --git a/message/message.go b/message/message.go
@@ -6,114 +6,58 @@ import (
 	"github.com/pkg/errors"
 )
 
-const IKE_HEADER_LEN int = 28
-
-type IKEHeader struct {
-	InitiatorSPI uint64
-	ResponderSPI uint64
-	MajorVersion uint8
-	MinorVersion uint8
-	ExchangeType uint8
-	Initiator    bool
-	Response     bool
-	MessageID    uint32
-	NextPayload  uint8
-}
-
 type IKEMessage struct {
 	*IKEHeader
 	Payloads IKEPayloadContainer
 }
 
-func ParseIkeHeader(b []byte) (*IKEHeader, error) {
-	// IKE message packet format this implementation referenced is
-	// defined in RFC 7296, Section 3.1
-	// bounds checking
-	if len(b) < IKE_HEADER_LEN {
-		return nil, errors.Errorf("ParseIkeHeader(): Received broken IKE header")
-	}
-	ikeMessageLength := binary.BigEndian.Uint32(b[24:IKE_HEADER_LEN])
-	if ikeMessageLength < uint32(IKE_HEADER_LEN) {
-		return nil, errors.Errorf("ParseIkeHeader(): Illegal IKE message length %d < header length %d",
-			ikeMessageLength, IKE_HEADER_LEN)
-	}
-	// len() return int, which is 64 bit on 64-bit host and 32 bit
-	// on 32-bit host, so this implementation may potentially cause
-	// problem on 32-bit machine
-	if len(b) != int(ikeMessageLength) {
-		return nil, errors.Errorf("ParseIkeHeader(): The length of received message " +
-			"not matchs the length specified in header")
+func NewMessage(
+	iSPI, rSPI uint64, exchgType uint8,
+	response, initiator bool, mId uint32,
+	payloads IKEPayloadContainer,
+) *IKEMessage {
+	m := &IKEMessage{
+		IKEHeader: NewHeader(iSPI, rSPI, exchgType,
+			response, initiator, mId, uint8(NoNext), nil),
+		Payloads: payloads,
 	}
-
-	ikeHeader := new(IKEHeader)
-
-	ikeHeader.InitiatorSPI = binary.BigEndian.Uint64(b[:8])
-	ikeHeader.ResponderSPI = binary.BigEndian.Uint64(b[8:16])
-	ikeHeader.MajorVersion = b[17] >> 4
-	ikeHeader.MinorVersion = b[17] & 0x0F
-	ikeHeader.ExchangeType = b[18]
-	ikeHeader.Initiator = (b[19] & InitiatorBitCheck) == InitiatorBitCheck
-	ikeHeader.Response = (b[19] & ResponseBitCheck) == ResponseBitCheck
-	ikeHeader.MessageID = binary.BigEndian.Uint32(b[20:24])
-	ikeHeader.NextPayload = b[16]
-
-	return ikeHeader, nil
+	return m
 }
 
-func (ikeMessage *IKEMessage) Encode() ([]byte, error) {
-	b := make([]byte, IKE_HEADER_LEN)
-
-	binary.BigEndian.PutUint64(b[0:8], ikeMessage.InitiatorSPI)
-	binary.BigEndian.PutUint64(b[8:16], ikeMessage.ResponderSPI)
-	b[17] = (ikeMessage.MajorVersion << 4) | (ikeMessage.MinorVersion & 0x0F)
-	b[18] = ikeMessage.ExchangeType
-	if ikeMessage.Initiator {
-		b[19] |= InitiatorBitCheck
-	}
-	if ikeMessage.Response {
-		b[19] |= ResponseBitCheck
-	}
-	binary.BigEndian.PutUint32(b[20:24], ikeMessage.MessageID)
-
-	if len(ikeMessage.Payloads) > 0 {
-		b[16] = byte(ikeMessage.Payloads[0].Type())
+func (m *IKEMessage) Encode() ([]byte, error) {
+	if len(m.Payloads) > 0 {
+		m.IKEHeader.NextPayload = uint8(m.Payloads[0].Type())
 	} else {
-		b[16] = byte(NoNext)
+		m.IKEHeader.NextPayload = uint8(NoNext)
 	}
 
-	ikeMessagePayloadData, err := ikeMessage.Payloads.Encode()
+	var err error
+	m.IKEHeader.PayloadBytes, err = m.Payloads.Encode()
 	if err != nil {
 		return nil, errors.Errorf("Encode(): EncodePayload failed: %+v", err)
 	}
-
-	b = append(b, ikeMessagePayloadData...)
-	ikeMsgDataLen := len(b)
-	if ikeMsgDataLen > 0xFFFFFFFF {
-		return nil, errors.Errorf("Encode(): ikeMessageData length exceeds uint32 limit: %d", ikeMsgDataLen)
-	}
-	binary.BigEndian.PutUint32(b[24:IKE_HEADER_LEN], uint32(ikeMsgDataLen))
-	return b, nil
+	return m.IKEHeader.Marshal()
 }
 
-func (ikeMessage *IKEMessage) Decode(b []byte) error {
+func (m *IKEMessage) Decode(b []byte) error {
 	var err error
-	ikeMessage.IKEHeader, err = ParseIkeHeader(b)
+	m.IKEHeader, err = ParseHeader(b)
 	if err != nil {
 		return errors.Wrapf(err, "Decode()")
 	}
 
-	err = ikeMessage.DecodePayload(b[IKE_HEADER_LEN:])
+	err = m.DecodePayload(m.PayloadBytes)
 	if err != nil {
-		return errors.Errorf("Decode(): DecodePayload failed: %+v", err)
+		return errors.Errorf("Decode(): DecodePayload failed: %v", err)
 	}
 
 	return nil
 }
 
-func (ikeMessage *IKEMessage) DecodePayload(b []byte) error {
-	err := ikeMessage.Payloads.Decode(ikeMessage.NextPayload, b)
+func (m *IKEMessage) DecodePayload(b []byte) error {
+	err := m.Payloads.Decode(m.NextPayload, b)
 	if err != nil {
-		return errors.Errorf("DecodePayload(): DecodePayload failed: %+v", err)
+		return errors.Errorf("DecodePayload(): DecodePayload failed: %v", err)
 	}
 
 	return nil