From 3bc8108907b459fef55836df549340cce25f1585 Mon Sep 17 00:00:00 2001 From: Jason Lee Date: Fri, 17 Nov 2023 14:42:34 +0800 Subject: [PATCH 1/4] add module.go --- Makefile | 2 +- diameter.go | 9 --------- module.go | 23 +++++++++++++++++++++++ 3 files changed, 24 insertions(+), 10 deletions(-) create mode 100644 module.go diff --git a/Makefile b/Makefile index 0de722d..e197d1e 100644 --- a/Makefile +++ b/Makefile @@ -8,7 +8,7 @@ xk6: go install go.k6.io/xk6/cmd/xk6@latest generator: - $(GO) build -o bin/dict_generator cmd/dict_generator/main.go + $(GO) build -o bin/dict_generator cmd/dict_generator/*.go build: xk6 xk6 build v0.37.0 --with github.com/matrixxsoftware/xk6-diameter=. --output bin/k6 diff --git a/diameter.go b/diameter.go index 34501b6..9485466 100644 --- a/diameter.go +++ b/diameter.go @@ -12,11 +12,8 @@ import ( "github.com/fiorix/go-diameter/v4/diam/dict" "github.com/fiorix/go-diameter/v4/diam/sm" log "github.com/sirupsen/logrus" - "go.k6.io/k6/js/modules" ) -type Diameter struct{} - type DiameterClient struct { client *sm.Client conn diam.Conn @@ -274,9 +271,3 @@ func (*Dict) Load(dictionary string) error { } return nil } - -func init() { - modules.Register("k6/x/diameter", &Diameter{}) - modules.Register("k6/x/diameter/avp", &AVP{}) - modules.Register("k6/x/diameter/dict", &Dict{}) -} diff --git a/module.go b/module.go new file mode 100644 index 0000000..ed999be --- /dev/null +++ b/module.go @@ -0,0 +1,23 @@ +package diameter + +import ( + "github.com/dop251/goja" + "go.k6.io/k6/js/modules" +) + +type ( + Diameter struct { + vu modules.VU + exports *goja.Object + } + RootModule struct{} + Module struct { + *Diameter + } +) + +func init() { + modules.Register("k6/x/diameter", &Diameter{}) + modules.Register("k6/x/diameter/avp", &AVP{}) + modules.Register("k6/x/diameter/dict", &Dict{}) +} From e1b1082ed5dcc8a9bddf8cee23294c112e3d431e Mon Sep 17 00:00:00 2001 From: Jason Lee Date: Sun, 19 Nov 2023 13:29:45 +0800 Subject: [PATCH 2/4] add config.go --- config.go | 91 ++++++++++++++++++++++++++++++++++++++++++++++ diameter.go | 42 +++++++-------------- example/example.js | 6 ++- 3 files changed, 110 insertions(+), 29 deletions(-) create mode 100644 config.go diff --git a/config.go b/config.go new file mode 100644 index 0000000..05d1b74 --- /dev/null +++ b/config.go @@ -0,0 +1,91 @@ +package diameter + +import ( + "encoding/json" + "time" + + log "github.com/sirupsen/logrus" +) + +type CapacityExchangeConfig struct { + VendorID *uint32 `json:"vendorID"` + ProductName *string `json:"productName,omitempty"` + OriginHost *string `json:"originHost,omitempty"` + OriginRealm *string `json:"originRealm,omitempty"` +} + +type DiameterConfig struct { + MaxRetransmits *uint `json:"maxRetransmits,omitempty"` + RetransmitInterval *time.Duration `json:"retransmitInterval,omitempty"` + EnableWatchdog *bool `json:"enableWatchdog,omitempty"` + WatchdogInterval *time.Duration `json:"watchdogInterval,omitempty"` + WatchdogStream *uint `json:"watchdogStream,omitempty"` + CapacityExchange *CapacityExchangeConfig `json:"capacityExchange,omitempty"` +} + +func processConfig(arg map[string]interface{}) (*DiameterConfig, error) { + + var config DiameterConfig + if b, err := json.Marshal(arg); err != nil { + return nil, err + } else { + if err = json.Unmarshal(b, &config); err != nil { + return nil, err + } + } + + setDiameterConfigDefaults(&config) + + log.Infof("Config %+v\n", config) + log.Infof("CE Config %+v\n", config.CapacityExchange) + + return &config, nil +} + +func setDiameterConfigDefaults(config *DiameterConfig) { + // Default values + var defaultMaxRetransmits uint = 1 + var defaultRetransmitInterval = 1 * time.Second + var defaultEnableWatchdog = true + var defaultWatchdogInterval = 5 * time.Second + var defaultWatchdogStream uint = 0 + + var defaultVendorID uint32 = 13 + var defaultProductName = "xk6-diameter" + var defaultOriginHost = "origin.host" + var defaultOriginRealm = "origin.realm" + + // Set defaults for DiameterConfig + if config.MaxRetransmits == nil { + config.MaxRetransmits = &defaultMaxRetransmits + } + if config.RetransmitInterval == nil { + config.RetransmitInterval = &defaultRetransmitInterval + } + if config.EnableWatchdog == nil { + config.EnableWatchdog = &defaultEnableWatchdog + } + if config.WatchdogInterval == nil { + config.WatchdogInterval = &defaultWatchdogInterval + } + if config.WatchdogStream == nil { + config.WatchdogStream = &defaultWatchdogStream + } + + // Set defaults for CapacityExchangeConfig + if config.CapacityExchange == nil { + config.CapacityExchange = &CapacityExchangeConfig{} + } + if config.CapacityExchange.VendorID == nil { + config.CapacityExchange.VendorID = &defaultVendorID + } + if config.CapacityExchange.ProductName == nil { + config.CapacityExchange.ProductName = &defaultProductName + } + if config.CapacityExchange.OriginHost == nil { + config.CapacityExchange.OriginHost = &defaultOriginHost + } + if config.CapacityExchange.OriginRealm == nil { + config.CapacityExchange.OriginRealm = &defaultOriginRealm + } +} diff --git a/diameter.go b/diameter.go index 9485466..00792a5 100644 --- a/diameter.go +++ b/diameter.go @@ -31,32 +31,18 @@ type AVP struct{} type Dict struct{} -type DiameterConfig struct { - // Settings - VendorID datatype.Unsigned32 `json:"vendorID"` - ProductName datatype.UTF8String `json:"productName"` - - // Client Config - MaxRetransmits uint `json:"maxRetransmits"` - RetransmitInterval time.Duration `json:"retransmitInterval"` - EnableWatchdog bool `json:"enableWatchdog"` - WatchdogInterval time.Duration `json:"watchdogInterval"` - WatchdogStream uint `json:"watchdogStream"` - - // SupportedVendorID []*diam.AVP // Supported vendor ID - // AcctApplicationID []*diam.AVP // Acct applications - // AuthApplicationID []*diam.AVP // Auth applications - // VendorSpecificApplicationID []*diam.AVP // Vendor specific applications -} +func (*Diameter) XClient(arg map[string]interface{}) (*DiameterClient, error) { -func (*Diameter) XClient() (*DiameterClient, error) { + config, err := processConfig(arg) + if err != nil { + return nil, err + } - // TODO make all this configurable later cfg := &sm.Settings{ - OriginHost: datatype.DiameterIdentity("diam.host"), - OriginRealm: datatype.DiameterIdentity("diam.realm"), - VendorID: 13, - ProductName: "xk6-diameter", + OriginHost: datatype.DiameterIdentity(*config.CapacityExchange.OriginHost), + OriginRealm: datatype.DiameterIdentity(*config.CapacityExchange.OriginRealm), + VendorID: datatype.Unsigned32(*config.CapacityExchange.VendorID), + ProductName: datatype.UTF8String(*config.CapacityExchange.ProductName), OriginStateID: datatype.Unsigned32(time.Now().Unix()), FirmwareRevision: 1, HostIPAddresses: []datatype.Address{ @@ -67,15 +53,15 @@ func (*Diameter) XClient() (*DiameterClient, error) { hopIds := make(map[uint32]chan *diam.Message) mux.Handle("CCA", handleCCA(hopIds)) - // TODO need to support other diameter CMD client := &sm.Client{ Dict: dict.Default, Handler: mux, - MaxRetransmits: 1, - RetransmitInterval: time.Second, - EnableWatchdog: true, - WatchdogInterval: 5 * time.Second, + MaxRetransmits: *config.MaxRetransmits, + RetransmitInterval: *config.RetransmitInterval, + EnableWatchdog: *config.EnableWatchdog, + WatchdogInterval: *config.WatchdogInterval, + WatchdogStream: *config.WatchdogStream, AuthApplicationID: []*diam.AVP{ diam.NewAVP(avp.AuthApplicationID, avp.Mbit, 0, datatype.Unsigned32(4)), }, diff --git a/example/example.js b/example/example.js index 4fe3ca2..2d08eda 100644 --- a/example/example.js +++ b/example/example.js @@ -14,7 +14,11 @@ export let options = { dict.load("dict/extra.xml") // Init Client -let client = diam.Client() +let client = diam.Client({ + MaxRetransmits: 1, + EnableWatchdog: false, +}) + let dataType = diam.DataType() export default function () { From d209ff84b04ad905f2094d9915f6ab45048125e0 Mon Sep 17 00:00:00 2001 From: Jason Lee Date: Sun, 19 Nov 2023 14:02:51 +0800 Subject: [PATCH 3/4] add custom Duration type to handle duration string --- config.go | 46 +++++++++++++++++++++++++++++++++++++--------- diameter.go | 27 ++++++++++++--------------- example/example.js | 4 +--- 3 files changed, 50 insertions(+), 27 deletions(-) diff --git a/config.go b/config.go index 05d1b74..a4263ff 100644 --- a/config.go +++ b/config.go @@ -2,9 +2,8 @@ package diameter import ( "encoding/json" + "errors" "time" - - log "github.com/sirupsen/logrus" ) type CapacityExchangeConfig struct { @@ -15,10 +14,11 @@ type CapacityExchangeConfig struct { } type DiameterConfig struct { + RequestTimeout *Duration `json:"requestTimeout,omitempty"` MaxRetransmits *uint `json:"maxRetransmits,omitempty"` - RetransmitInterval *time.Duration `json:"retransmitInterval,omitempty"` + RetransmitInterval *Duration `json:"retransmitInterval,omitempty"` EnableWatchdog *bool `json:"enableWatchdog,omitempty"` - WatchdogInterval *time.Duration `json:"watchdogInterval,omitempty"` + WatchdogInterval *Duration `json:"watchdogInterval,omitempty"` WatchdogStream *uint `json:"watchdogStream,omitempty"` CapacityExchange *CapacityExchangeConfig `json:"capacityExchange,omitempty"` } @@ -36,18 +36,16 @@ func processConfig(arg map[string]interface{}) (*DiameterConfig, error) { setDiameterConfigDefaults(&config) - log.Infof("Config %+v\n", config) - log.Infof("CE Config %+v\n", config.CapacityExchange) - return &config, nil } func setDiameterConfigDefaults(config *DiameterConfig) { // Default values + var defaultRequestTimeout = Duration{1 * time.Second} var defaultMaxRetransmits uint = 1 - var defaultRetransmitInterval = 1 * time.Second + var defaultRetransmitInterval = Duration{1 * time.Second} var defaultEnableWatchdog = true - var defaultWatchdogInterval = 5 * time.Second + var defaultWatchdogInterval = Duration{5 * time.Second} var defaultWatchdogStream uint = 0 var defaultVendorID uint32 = 13 @@ -56,6 +54,9 @@ func setDiameterConfigDefaults(config *DiameterConfig) { var defaultOriginRealm = "origin.realm" // Set defaults for DiameterConfig + if config.RequestTimeout == nil { + config.RequestTimeout = &defaultRequestTimeout + } if config.MaxRetransmits == nil { config.MaxRetransmits = &defaultMaxRetransmits } @@ -89,3 +90,30 @@ func setDiameterConfigDefaults(config *DiameterConfig) { config.CapacityExchange.OriginRealm = &defaultOriginRealm } } + +type Duration struct { + time.Duration +} + +func (d Duration) MarshalJSON() ([]byte, error) { + return json.Marshal(d.String()) +} + +func (d *Duration) UnmarshalJSON(b []byte) error { + var v interface{} + if err := json.Unmarshal(b, &v); err != nil { + return err + } + + switch value := v.(type) { + case string: + var err error + d.Duration, err = time.ParseDuration(value) + if err != nil { + return err + } + return nil + default: + return errors.New("invalid duration") + } +} diff --git a/diameter.go b/diameter.go index 00792a5..6d120ac 100644 --- a/diameter.go +++ b/diameter.go @@ -15,9 +15,10 @@ import ( ) type DiameterClient struct { - client *sm.Client - conn diam.Conn - hopIds map[uint32]chan *diam.Message + client *sm.Client + conn diam.Conn + hopIds map[uint32]chan *diam.Message + requestTimeout time.Duration } type DiameterMessage struct { @@ -58,9 +59,9 @@ func (*Diameter) XClient(arg map[string]interface{}) (*DiameterClient, error) { Dict: dict.Default, Handler: mux, MaxRetransmits: *config.MaxRetransmits, - RetransmitInterval: *config.RetransmitInterval, + RetransmitInterval: *&config.RetransmitInterval.Duration, EnableWatchdog: *config.EnableWatchdog, - WatchdogInterval: *config.WatchdogInterval, + WatchdogInterval: *&config.WatchdogInterval.Duration, WatchdogStream: *config.WatchdogStream, AuthApplicationID: []*diam.AVP{ diam.NewAVP(avp.AuthApplicationID, avp.Mbit, 0, datatype.Unsigned32(4)), @@ -68,9 +69,10 @@ func (*Diameter) XClient(arg map[string]interface{}) (*DiameterClient, error) { } return &DiameterClient{ - client: client, - conn: nil, - hopIds: hopIds, + client: client, + conn: nil, + hopIds: hopIds, + requestTimeout: config.RequestTimeout.Duration, }, nil } @@ -102,7 +104,7 @@ func (c *DiameterClient) Connect(address string) error { return nil } -func (c *DiameterClient) Send(msg *DiameterMessage, requestTimeoutMillis int) (uint32, error) { +func (c *DiameterClient) Send(msg *DiameterMessage) (uint32, error) { if c.conn == nil { return 0, errors.New("Not connected") @@ -115,12 +117,7 @@ func (c *DiameterClient) Send(msg *DiameterMessage, requestTimeoutMillis int) (u c.hopIds[hopByHopID] = make(chan *diam.Message) // Timeout settings - var timeout <-chan time.Time - if requestTimeoutMillis == 0 { - timeout = time.After(60 * time.Second) - } else { - timeout = time.After(time.Duration(requestTimeoutMillis) * time.Millisecond) - } + timeout := time.After(c.requestTimeout) // Send CCR _, err := req.WriteTo(c.conn) diff --git a/example/example.js b/example/example.js index 2d08eda..04d8c34 100644 --- a/example/example.js +++ b/example/example.js @@ -15,10 +15,8 @@ dict.load("dict/extra.xml") // Init Client let client = diam.Client({ - MaxRetransmits: 1, - EnableWatchdog: false, + requestTimeout: "50ms", }) - let dataType = diam.DataType() export default function () { From 39154c31021a50c5f68a820125e6dc968c0c0cdd Mon Sep 17 00:00:00 2001 From: Jason Lee Date: Sun, 19 Nov 2023 14:58:09 +0800 Subject: [PATCH 4/4] add more capability-exhcnage config --- config.go | 58 +++++++++++++++++++++++++++++++---------------------- diameter.go | 21 ++++++++++--------- 2 files changed, 46 insertions(+), 33 deletions(-) diff --git a/config.go b/config.go index a4263ff..3608236 100644 --- a/config.go +++ b/config.go @@ -6,21 +6,23 @@ import ( "time" ) -type CapacityExchangeConfig struct { - VendorID *uint32 `json:"vendorID"` - ProductName *string `json:"productName,omitempty"` - OriginHost *string `json:"originHost,omitempty"` - OriginRealm *string `json:"originRealm,omitempty"` +type DiameterConfig struct { + RequestTimeout *Duration `json:"requestTimeout,omitempty"` + MaxRetransmits *uint `json:"maxRetransmits,omitempty"` + RetransmitInterval *Duration `json:"retransmitInterval,omitempty"` + EnableWatchdog *bool `json:"enableWatchdog,omitempty"` + WatchdogInterval *Duration `json:"watchdogInterval,omitempty"` + WatchdogStream *uint `json:"watchdogStream,omitempty"` + CapabilityExchange *CapabilityExchangeConfig `json:"capabilityExchange,omitempty"` } -type DiameterConfig struct { - RequestTimeout *Duration `json:"requestTimeout,omitempty"` - MaxRetransmits *uint `json:"maxRetransmits,omitempty"` - RetransmitInterval *Duration `json:"retransmitInterval,omitempty"` - EnableWatchdog *bool `json:"enableWatchdog,omitempty"` - WatchdogInterval *Duration `json:"watchdogInterval,omitempty"` - WatchdogStream *uint `json:"watchdogStream,omitempty"` - CapacityExchange *CapacityExchangeConfig `json:"capacityExchange,omitempty"` +type CapabilityExchangeConfig struct { + VendorID *uint32 `json:"vendorID"` + ProductName *string `json:"productName,omitempty"` + OriginHost *string `json:"originHost,omitempty"` + OriginRealm *string `json:"originRealm,omitempty"` + FirmwareRevision *uint32 `json:"firmwareRevision,omitempty"` + HostIPAddresses *[]string `json:"hostIPAddresses,omitempty"` } func processConfig(arg map[string]interface{}) (*DiameterConfig, error) { @@ -52,6 +54,8 @@ func setDiameterConfigDefaults(config *DiameterConfig) { var defaultProductName = "xk6-diameter" var defaultOriginHost = "origin.host" var defaultOriginRealm = "origin.realm" + var defaultFirmwareRevision uint32 = 1 + var defaultHostIPAddresses = []string{"127.0.0.1"} // Set defaults for DiameterConfig if config.RequestTimeout == nil { @@ -73,21 +77,27 @@ func setDiameterConfigDefaults(config *DiameterConfig) { config.WatchdogStream = &defaultWatchdogStream } - // Set defaults for CapacityExchangeConfig - if config.CapacityExchange == nil { - config.CapacityExchange = &CapacityExchangeConfig{} + // Set defaults for CapabilityExchangeConfig + if config.CapabilityExchange == nil { + config.CapabilityExchange = &CapabilityExchangeConfig{} + } + if config.CapabilityExchange.VendorID == nil { + config.CapabilityExchange.VendorID = &defaultVendorID + } + if config.CapabilityExchange.ProductName == nil { + config.CapabilityExchange.ProductName = &defaultProductName } - if config.CapacityExchange.VendorID == nil { - config.CapacityExchange.VendorID = &defaultVendorID + if config.CapabilityExchange.OriginHost == nil { + config.CapabilityExchange.OriginHost = &defaultOriginHost } - if config.CapacityExchange.ProductName == nil { - config.CapacityExchange.ProductName = &defaultProductName + if config.CapabilityExchange.OriginRealm == nil { + config.CapabilityExchange.OriginRealm = &defaultOriginRealm } - if config.CapacityExchange.OriginHost == nil { - config.CapacityExchange.OriginHost = &defaultOriginHost + if config.CapabilityExchange.FirmwareRevision == nil { + config.CapabilityExchange.FirmwareRevision = &defaultFirmwareRevision } - if config.CapacityExchange.OriginRealm == nil { - config.CapacityExchange.OriginRealm = &defaultOriginRealm + if config.CapabilityExchange.HostIPAddresses == nil { + config.CapabilityExchange.HostIPAddresses = &defaultHostIPAddresses } } diff --git a/diameter.go b/diameter.go index 6d120ac..e95b175 100644 --- a/diameter.go +++ b/diameter.go @@ -39,16 +39,19 @@ func (*Diameter) XClient(arg map[string]interface{}) (*DiameterClient, error) { return nil, err } + hostIPAddresses := []datatype.Address{} + for _, ip := range *config.CapabilityExchange.HostIPAddresses { + hostIPAddresses = append(hostIPAddresses, datatype.Address(net.ParseIP(ip))) + } + cfg := &sm.Settings{ - OriginHost: datatype.DiameterIdentity(*config.CapacityExchange.OriginHost), - OriginRealm: datatype.DiameterIdentity(*config.CapacityExchange.OriginRealm), - VendorID: datatype.Unsigned32(*config.CapacityExchange.VendorID), - ProductName: datatype.UTF8String(*config.CapacityExchange.ProductName), + OriginHost: datatype.DiameterIdentity(*config.CapabilityExchange.OriginHost), + OriginRealm: datatype.DiameterIdentity(*config.CapabilityExchange.OriginRealm), + VendorID: datatype.Unsigned32(*config.CapabilityExchange.VendorID), + ProductName: datatype.UTF8String(*config.CapabilityExchange.ProductName), OriginStateID: datatype.Unsigned32(time.Now().Unix()), - FirmwareRevision: 1, - HostIPAddresses: []datatype.Address{ - datatype.Address(net.ParseIP("127.0.0.1")), - }, + FirmwareRevision: datatype.Unsigned32(*config.CapabilityExchange.FirmwareRevision), + HostIPAddresses: hostIPAddresses, } mux := sm.New(cfg) @@ -64,7 +67,7 @@ func (*Diameter) XClient(arg map[string]interface{}) (*DiameterClient, error) { WatchdogInterval: *&config.WatchdogInterval.Duration, WatchdogStream: *config.WatchdogStream, AuthApplicationID: []*diam.AVP{ - diam.NewAVP(avp.AuthApplicationID, avp.Mbit, 0, datatype.Unsigned32(4)), + diam.NewAVP(avp.AuthApplicationID, avp.Mbit, 0, datatype.Unsigned32(4)), // TODO make configurable }, }