refactor: Improve code readability

bbsakura · Aug 5, 2024 · ade4f81 · ade4f81
1 parent 270d7a2
commit ade4f81
Show file tree

Hide file tree

Showing 3 changed files with 28 additions and 23 deletions.
diff --git a/pkg/diameter/avp.go b/pkg/diameter/avp.go
@@ -15,7 +15,14 @@ type AVP struct {
 	Value interface{}
 }
 
-type AvpMeta struct {
+type AVPMeta struct {
+	code   uint32
+	flag   uint8
+	vendor uint32
+	value  datatype.Type
+}
+
+type AVPMetaC struct {
 	code      uint32
 	flag      uint8
 	vendor    uint32

diff --git a/pkg/diameter/avp_dict.go b/pkg/diameter/avp_dict.go
@@ -4,13 +4,13 @@ import (
 	"github.com/fiorix/go-diameter/v4/diam/avp"
 )
 
-var avpDict map[string]*AvpMeta
+var avpDict map[string]*AVPMetaC
 
 const vendorId3GPP = 10415
 
 func init() {
 	// TODO: generate this part from XML dictionary files
-	avpDict = map[string]*AvpMeta{
+	avpDict = map[string]*AVPMetaC{
 		"User-Name":                            {code: avp.UserName, flag: avp.Mbit, vendor: 0, converter: toUTF8String},
 		"Auth-Session-State":                   {code: avp.AuthSessionState, flag: avp.Mbit, vendor: 0, converter: toEnumerated},
 		"Visited-PLMN-Id":                      {code: avp.VisitedPLMNID, flag: avp.Vbit | avp.Mbit, vendor: vendorId3GPP, converter: toOctetString},

diff --git a/pkg/diameter/diameter.go b/pkg/diameter/diameter.go
@@ -306,17 +306,16 @@ func (c *K6DiameterClient) SendAIR(options ConnectionOptions) (bool, error) {
 		sid = c.generateSessionID()
 	}
 	m := diam.NewRequest(diam.AuthenticationInformation, diam.TGPP_S6A_APP_ID, dict.Default)
-	_, err = m.NewAVP(avp.SessionID, avp.Mbit, 0, datatype.UTF8String(sid))
-	if err != nil {
-		return false, errors.WithMessage(err, "NewAVP SessionID failed")
-	}
-	_, err = m.NewAVP(avp.OriginHost, avp.Mbit, 0, c.cfg.OriginHost)
-	if err != nil {
-		return false, errors.WithMessage(err, "NewAVP OriginHost failed")
+	avps := []AVPMeta{
+		{code: avp.SessionID, flag: avp.Mbit, vendor: 0, value: datatype.UTF8String(sid)},
+		{code: avp.OriginHost, flag: avp.Mbit, vendor: 0, value: c.cfg.OriginHost},
+		{code: avp.OriginRealm, flag: avp.Mbit, vendor: 0, value: c.cfg.OriginRealm},
 	}
-	_, err = m.NewAVP(avp.OriginRealm, avp.Mbit, 0, c.cfg.OriginRealm)
-	if err != nil {
-		return false, errors.WithMessage(err, "NewAVP OriginRealm failed")
+	for _, avp := range avps {
+		_, err = m.NewAVP(avp.code, avp.flag, avp.vendor, avp.value)
+		if err != nil {
+			return false, errors.WithMessage(err, "NewAVP failed")
+		}
 	}
 	if options.ProxiableFlag {
 		m.Header.CommandFlags |= diam.ProxiableFlag
@@ -365,17 +364,16 @@ func (c *K6DiameterClient) SendULR(options ConnectionOptions) (bool, error) {
 		sid = c.generateSessionID()
 	}
 	m := diam.NewRequest(diam.UpdateLocation, diam.TGPP_S6A_APP_ID, dict.Default)
-	_, err = m.NewAVP(avp.SessionID, avp.Mbit, 0, datatype.UTF8String(sid))
-	if err != nil {
-		return false, errors.WithMessage(err, "NewAVP SessionID failed")
-	}
-	_, err = m.NewAVP(avp.OriginHost, avp.Mbit, 0, c.cfg.OriginHost)
-	if err != nil {
-		return false, errors.WithMessage(err, "NewAVP OriginHost failed")
+	avps := []AVPMeta{
+		{code: avp.SessionID, flag: avp.Mbit, vendor: 0, value: datatype.UTF8String(sid)},
+		{code: avp.OriginHost, flag: avp.Mbit, vendor: 0, value: c.cfg.OriginHost},
+		{code: avp.OriginRealm, flag: avp.Mbit, vendor: 0, value: c.cfg.OriginRealm},
 	}
-	_, err = m.NewAVP(avp.OriginRealm, avp.Mbit, 0, c.cfg.OriginRealm)
-	if err != nil {
-		return false, errors.WithMessage(err, "NewAVP OriginRealm failed")
+	for _, avp := range avps {
+		_, err = m.NewAVP(avp.code, avp.flag, avp.vendor, avp.value)
+		if err != nil {
+			return false, errors.WithMessage(err, "NewAVP failed")
+		}
 	}
 	if options.ProxiableFlag {
 		m.Header.CommandFlags |= diam.ProxiableFlag