add a way to set flags and vendor code

MATRIXXSoftware · Nov 10, 2023 · 7b6a247 · 7b6a247
1 parent e0632c5
commit 7b6a247
Show file tree

Hide file tree

Showing 2 changed files with 44 additions and 14 deletions.
diff --git a/diameter.go b/diameter.go
@@ -107,25 +107,55 @@ func (*Diameter) NewMessage(name string) *DiameterMessage {
 	}
 }
 
-func (m *DiameterMessage) AddAVP(code uint32) *DiameterAVP {
+func (m *DiameterMessage) AddAVP() *DiameterAVP {
+	// populate later
 	avp := DiameterAVP{
-		code:   code,
-		flags:  avp.Mbit,
+		code:   0,
+		flags:  0,
 		vendor: 0,
-		data:   nil, // populate later
+		data:   nil,
 	}
 	m.avps = append(m.avps, &avp)
 	return &avp
 }
 
-func (avp *DiameterAVP) XUTF8String(value string) {
-	avp.data = datatype.UTF8String(value)
+func (a *DiameterAVP) XCode(code uint32) *DiameterAVP {
+	a.code = code
+	return a
 }
 
-func (avp *DiameterAVP) XDiameterIdentity(value string) {
-	avp.data = datatype.DiameterIdentity(value)
+func (a *DiameterAVP) XMbit() *DiameterAVP {
+	a.flags = a.flags | avp.Mbit
+	return a
 }
 
+func (a *DiameterAVP) XPbit() *DiameterAVP {
+	a.flags = a.flags | avp.Pbit
+	return a
+}
+
+func (a *DiameterAVP) XVbit() *DiameterAVP {
+	a.flags = a.flags | avp.Vbit
+	return a
+}
+
+func (a *DiameterAVP) XVendor(vendor uint32) *DiameterAVP {
+	a.vendor = vendor
+	return a
+}
+
+func (a *DiameterAVP) XUTF8String(value string) *DiameterAVP {
+	a.data = datatype.UTF8String(value)
+	return a
+}
+
+func (a *DiameterAVP) XDiameterIdentity(value string) *DiameterAVP {
+	a.data = datatype.DiameterIdentity(value)
+	return a
+}
+
+// TODO add more data type
+
 func (d *Diameter) Send(client *DiameterClient, msg *DiameterMessage) (uint32, error) {
 
 	req := msg.diamMsg

diff --git a/example.js b/example.js
@@ -10,12 +10,12 @@ let client = diameter.newClient();
 
 export default function () {
     let msg = diameter.newMessage("CCR");
-    msg.addAVP(263).UTF8String("Session-8888");         // Session ID
-    msg.addAVP(264).DiameterIdentity("origin.host")     // Origin-Host
-    msg.addAVP(296).DiameterIdentity("origin.realm")    // Origin-Realm
-    msg.addAVP(283).DiameterIdentity("dest.host")       // Destination-Host
-    msg.addAVP(293).DiameterIdentity("dest.realm")      // Destination-Realm
-    msg.addAVP(1).UTF8String("ValueFooBar");            // User-Name
+    msg.addAVP().Code(263).Mbit().UTF8String("Session-8888");      // Session ID
+    msg.addAVP().Code(264).DiameterIdentity("origin.host")         // Origin-Host
+    msg.addAVP().Code(296).DiameterIdentity("origin.realm")        // Origin-Realm
+    msg.addAVP().Code(283).DiameterIdentity("dest.host")           // Destination-Host
+    msg.addAVP().Code(293).DiameterIdentity("dest.realm")          // Destination-Realm
+    msg.addAVP().Code(1).Vendor(10415).UTF8String("ValueFooBar");  // User-Name
 
     const response = diameter.send(client, msg);
     console.log("Result Code:", response);