Skip to content

Commit

Permalink
Merge pull request #4 from MATRIXXSoftware/feature/configurable-options
Browse files Browse the repository at this point in the history
Feature/configurable options
  • Loading branch information
lwlee2608 authored Nov 19, 2023
2 parents 786faf3 + 39154c3 commit 5aeb6b7
Show file tree
Hide file tree
Showing 5 changed files with 188 additions and 57 deletions.
2 changes: 1 addition & 1 deletion Makefile
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
129 changes: 129 additions & 0 deletions config.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,129 @@
package diameter

import (
"encoding/json"
"errors"
"time"
)

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 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) {

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)

return &config, nil
}

func setDiameterConfigDefaults(config *DiameterConfig) {
// Default values
var defaultRequestTimeout = Duration{1 * time.Second}
var defaultMaxRetransmits uint = 1
var defaultRetransmitInterval = Duration{1 * time.Second}
var defaultEnableWatchdog = true
var defaultWatchdogInterval = Duration{5 * time.Second}
var defaultWatchdogStream uint = 0

var defaultVendorID uint32 = 13
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 {
config.RequestTimeout = &defaultRequestTimeout
}
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 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.CapabilityExchange.OriginHost == nil {
config.CapabilityExchange.OriginHost = &defaultOriginHost
}
if config.CapabilityExchange.OriginRealm == nil {
config.CapabilityExchange.OriginRealm = &defaultOriginRealm
}
if config.CapabilityExchange.FirmwareRevision == nil {
config.CapabilityExchange.FirmwareRevision = &defaultFirmwareRevision
}
if config.CapabilityExchange.HostIPAddresses == nil {
config.CapabilityExchange.HostIPAddresses = &defaultHostIPAddresses
}
}

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")
}
}
87 changes: 32 additions & 55 deletions diameter.go
Original file line number Diff line number Diff line change
Expand Up @@ -12,15 +12,13 @@ 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
hopIds map[uint32]chan *diam.Message
client *sm.Client
conn diam.Conn
hopIds map[uint32]chan *diam.Message
requestTimeout time.Duration
}

type DiameterMessage struct {
Expand All @@ -34,60 +32,50 @@ 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) {

config, err := processConfig(arg)
if err != nil {
return nil, err
}

func (*Diameter) XClient() (*DiameterClient, error) {
hostIPAddresses := []datatype.Address{}
for _, ip := range *config.CapabilityExchange.HostIPAddresses {
hostIPAddresses = append(hostIPAddresses, datatype.Address(net.ParseIP(ip)))
}

// 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.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)

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.Duration,
EnableWatchdog: *config.EnableWatchdog,
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
},
}

return &DiameterClient{
client: client,
conn: nil,
hopIds: hopIds,
client: client,
conn: nil,
hopIds: hopIds,
requestTimeout: config.RequestTimeout.Duration,
}, nil
}

Expand Down Expand Up @@ -119,7 +107,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")
Expand All @@ -132,12 +120,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)
Expand Down Expand Up @@ -274,9 +257,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{})
}
4 changes: 3 additions & 1 deletion example/example.js
Original file line number Diff line number Diff line change
Expand Up @@ -14,7 +14,9 @@ export let options = {
dict.load("dict/extra.xml")

// Init Client
let client = diam.Client()
let client = diam.Client({
requestTimeout: "50ms",
})
let dataType = diam.DataType()

export default function () {
Expand Down
23 changes: 23 additions & 0 deletions module.go
Original file line number Diff line number Diff line change
@@ -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{})
}

0 comments on commit 5aeb6b7

Please sign in to comment.