Ability to load additional diameter dictionary definition (#2)

* add dict_generator cli
* add dict.load(dictionary)

cmd/dict_generator/main.go

@@ -0,0 +1,87 @@
+package main
+
+import (
+	"flag"
+	"fmt"
+	"io"
+	"os"
+	"strings"
+
+	"github.com/fiorix/go-diameter/v4/diam/dict"
+	log "github.com/sirupsen/logrus"
+)
+
+func main() {
+	var dictionary string
+	var output string
+
+	flag.StringVar(&dictionary, "dictionary", "dictionary.xml", "Dictionary")
+	flag.StringVar(&output, "output", "const.js", "Output file")
+	flag.Parse()
+
+	file, err := os.Open(dictionary)
+	if err != nil {
+		log.Fatalf("Error opening dictioanry: %s\n", err)
+	}
+	defer file.Close()
+
+	parser := dict.Default
+	parser.Load(file)
+	if err != nil {
+		log.Fatalf("Error parsing dictioanry: %s\n", err)
+	}
+
+	w, err := os.OpenFile(output, os.O_CREATE|os.O_WRONLY|os.O_TRUNC, 0666)
+	if err != nil {
+		log.Fatalf("Error opening file: %s\n", err)
+	}
+	defer w.Close()
+
+	PrintFlags(w)
+	PrintAvpCode(w, parser)
+	PrintVendorId(w, parser)
+}
+
+func PrintFlags(w io.Writer) {
+	fmt.Fprintf(w, "export const flags = {\n")
+	fmt.Fprintf(w, "    Vbit: 0x80,\n")
+	fmt.Fprintf(w, "    Mbit: 0x40,\n")
+	fmt.Fprintf(w, "    Pbit: 0x20,\n")
+	fmt.Fprintf(w, "}\n")
+	fmt.Fprintf(w, "\n")
+}
+
+func PrintAvpCode(w io.Writer, parser *dict.Parser) {
+	fmt.Fprintf(w, "export const avpCode = {\n")
+	for _, app := range parser.Apps() {
+		fmt.Fprintf(w, "    // %s\n", app.Name)
+		for _, avp := range app.AVP {
+			name := strings.ReplaceAll(avp.Name, "-", "")
+			fmt.Fprintf(w, "    %s: %d,\n", name, avp.Code)
+		}
+	}
+	fmt.Fprintf(w, "}\n")
+	fmt.Fprintf(w, "\n")
+
+}
+
+func PrintVendorId(w io.Writer, parser *dict.Parser) {
+	vendorIds := make(map[uint32]struct{})
+	exists := struct{}{}
+
+	fmt.Fprintf(w, "export const vendorId = {\n")
+	for _, app := range parser.Apps() {
+		for _, vendor := range app.Vendor {
+			// Remove duplicate vendorId
+			_, found := vendorIds[vendor.ID]
+			if found {
+				continue
+			}
+			vendorIds[vendor.ID] = exists
+
+			fmt.Fprintf(w, "    %s: %d,\n", vendor.Name, vendor.ID)
+		}
+	}
+	fmt.Fprintf(w, "}\n")
+	fmt.Fprintf(w, "\n")
+}

diameter.go

@@ -3,6 +3,7 @@ package diameter
 import (
 	"errors"
 	"net"
+	"os"
 	"time"
 
 	"github.com/fiorix/go-diameter/diam/avp"
@@ -31,7 +32,9 @@ type DataType struct{}
 
 type AVP struct{}
 
-func (*Diameter) NewClient() (*DiameterClient, error) {
+type Dict struct{}
+
+func (*Diameter) XClient() (*DiameterClient, error) {
 
 	// TODO make all this configurable later
 	cfg := &sm.Settings{
@@ -118,7 +121,7 @@ func (d *Diameter) Send(client *DiameterClient, msg *DiameterMessage) (uint32, e
 
 	// Wait for CCA
 	resp := <-client.hopIds[hopByHopID]
-	//log.Infof("Received CCA \n%s", resp)
+	log.Infof("Received CCA \n%s", resp)
 
 	delete(client.hopIds, hopByHopID)
 
@@ -228,7 +231,22 @@ func (a *AVP) XNew(code uint32, vendor uint32, flags uint8, data datatype.Type)
 	return diam.NewAVP(code, flags, vendor, data)
 }
 
+func (*Dict) Load(dictionary string) error {
+	file, err := os.Open(dictionary)
+	if err != nil {
+		return err
+	}
+	defer file.Close()
+
+	dict.Default.Load(file)
+	if err != nil {
+		return err
+	}
+	return nil
+}
+
 func init() {
 	modules.Register("k6/x/diameter", &Diameter{})
 	modules.Register("k6/x/diameter/avp", &AVP{})
+	modules.Register("k6/x/diameter/dict", &Dict{})
 }

dict/dictionary.xml

@@ -0,0 +1,9 @@
+<?xml version="1.0" encoding="UTF-8"?>
+<diameter>
+	<application id="4" type="auth" name="Extra">
+        <vendor name="Matrixxsoftware" id="35838"/>
+        <avp name="Event-Name-Code" code="13001" vendor-id="35838">
+            <data type="UTF8String"/>
+        </avp> 
+    </application>
+</diameter>