diff --git a/cmd/query.go b/cmd/query.go index 28cb50e..223c290 100644 --- a/cmd/query.go +++ b/cmd/query.go @@ -44,11 +44,12 @@ var queryCmd = &cobra.Command{ } ifacesConfigs = append(ifacesConfigs, ifaceConfig) } + var config respond.Config + config.Interfaces = ifacesConfigs nodes := runtime.NewNodes(&runtime.NodesConfig{}) - sitesDomains := make(map[string][]string) - collector := respond.NewCollector(nil, nodes, sitesDomains, ifacesConfigs) + collector := respond.NewCollector(nil, nodes, &config) defer collector.Close() collector.SendPacket(dstAddress) diff --git a/cmd/serve.go b/cmd/serve.go index 2b1f648..0fd283d 100644 --- a/cmd/serve.go +++ b/cmd/serve.go @@ -55,7 +55,7 @@ var serveCmd = &cobra.Command{ time.Sleep(delay) } - collector = respond.NewCollector(allDatabase.Conn, nodes, config.Respondd.SitesDomains(), config.Respondd.Interfaces) + collector = respond.NewCollector(allDatabase.Conn, nodes, &config.Respondd) collector.Start(config.Respondd.CollectInterval.Duration) defer collector.Close() } diff --git a/respond/collector.go b/respond/collector.go index 39d1d45..911b555 100644 --- a/respond/collector.go +++ b/respond/collector.go @@ -20,12 +20,12 @@ import ( type Collector struct { connections []multicastConn // UDP sockets - queue chan *Response // received responses - db database.Connection - nodes *runtime.Nodes - sitesDomains map[string][]string - interval time.Duration // Interval for multicast packets - stop chan interface{} + queue chan *Response // received responses + db database.Connection + nodes *runtime.Nodes + interval time.Duration // Interval for multicast packets + stop chan interface{} + config *Config } type multicastConn struct { @@ -35,17 +35,17 @@ type multicastConn struct { } // NewCollector creates a Collector struct -func NewCollector(db database.Connection, nodes *runtime.Nodes, sitesDomains map[string][]string, ifaces []InterfaceConfig) *Collector { +func NewCollector(db database.Connection, nodes *runtime.Nodes, config *Config) *Collector { coll := &Collector{ - db: db, - nodes: nodes, - sitesDomains: sitesDomains, - queue: make(chan *Response, 400), - stop: make(chan interface{}), + db: db, + nodes: nodes, + queue: make(chan *Response, 400), + stop: make(chan interface{}), + config: config, } - for _, iface := range ifaces { + for _, iface := range config.Interfaces { coll.listenUDP(iface) } @@ -349,7 +349,7 @@ func (coll *Collector) globalStatsWorker() { // saves global statistics func (coll *Collector) saveGlobalStats() { - stats := runtime.NewGlobalStats(coll.nodes, coll.sitesDomains) + stats := runtime.NewGlobalStats(coll.nodes, coll.config.SitesDomains()) for site, domains := range stats { for domain, stat := range domains { diff --git a/respond/collector_test.go b/respond/collector_test.go index b837259..6ee8bd4 100644 --- a/respond/collector_test.go +++ b/respond/collector_test.go @@ -16,8 +16,15 @@ const ( func TestCollector(t *testing.T) { nodes := runtime.NewNodes(&runtime.NodesConfig{}) + config := &Config{ + Sites: map[string]SiteConfig{ + SITE_TEST: { + Domains: []string{DOMAIN_TEST}, + }, + }, + } - collector := NewCollector(nil, nodes, map[string][]string{SITE_TEST: {DOMAIN_TEST}}, []InterfaceConfig{}) + collector := NewCollector(nil, nodes, config) collector.Start(time.Millisecond) time.Sleep(time.Millisecond * 10) collector.Close()