160 lines
3.9 KiB
Go
160 lines
3.9 KiB
Go
// Licensed to Elasticsearch B.V. under one or more contributor
|
|
// license agreements. See the NOTICE file distributed with
|
|
// this work for additional information regarding copyright
|
|
// ownership. Elasticsearch B.V. licenses this file to you under
|
|
// the Apache License, Version 2.0 (the "License"); you may
|
|
// not use this file except in compliance with the License.
|
|
// You may obtain a copy of the License at
|
|
//
|
|
// http://www.apache.org/licenses/LICENSE-2.0
|
|
//
|
|
// Unless required by applicable law or agreed to in writing,
|
|
// software distributed under the License is distributed on an
|
|
// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
|
|
// KIND, either express or implied. See the License for the
|
|
// specific language governing permissions and limitations
|
|
// under the License.
|
|
|
|
// +build !integration
|
|
|
|
package flows
|
|
|
|
import (
|
|
"net"
|
|
"testing"
|
|
"time"
|
|
|
|
"github.com/stretchr/testify/assert"
|
|
|
|
"github.com/elastic/beats/libbeat/beat"
|
|
"github.com/elastic/beats/libbeat/common"
|
|
"github.com/elastic/beats/libbeat/logp"
|
|
"github.com/elastic/beats/packetbeat/config"
|
|
)
|
|
|
|
type flowsChan struct {
|
|
ch chan []beat.Event
|
|
}
|
|
|
|
func (f *flowsChan) PublishFlows(events []beat.Event) {
|
|
f.ch <- events
|
|
}
|
|
|
|
func TestFlowsCounting(t *testing.T) {
|
|
logp.TestingSetup()
|
|
|
|
mac1 := []byte{1, 2, 3, 4, 5, 6}
|
|
mac2 := []byte{6, 5, 4, 3, 2, 1}
|
|
ip1 := []byte{127, 0, 0, 1}
|
|
ip2 := []byte{128, 0, 1, 2}
|
|
port1 := []byte{0, 1}
|
|
port2 := []byte{0, 2}
|
|
|
|
module, err := NewFlows(nil, &config.Flows{})
|
|
assert.NoError(t, err)
|
|
|
|
uint1, err := module.NewUint("uint1")
|
|
uint2, err := module.NewUint("uint2")
|
|
int1, err := module.NewInt("int1")
|
|
int2, err := module.NewInt("int2")
|
|
float1, err := module.NewFloat("float1")
|
|
float2, err := module.NewFloat("float2")
|
|
|
|
assert.NoError(t, err)
|
|
|
|
pub := &flowsChan{make(chan []beat.Event, 1)}
|
|
|
|
processor := &flowsProcessor{
|
|
table: module.table,
|
|
counters: module.counterReg,
|
|
timeout: 20 * time.Millisecond,
|
|
}
|
|
processor.spool.init(pub.PublishFlows, 1)
|
|
|
|
worker, err := makeWorker(
|
|
processor,
|
|
10*time.Millisecond,
|
|
1,
|
|
-1,
|
|
0)
|
|
if err != nil {
|
|
t.Fatalf("Failed to create flow worker: %v", err)
|
|
}
|
|
|
|
worker.Start()
|
|
defer worker.Stop()
|
|
|
|
idForward := newFlowID()
|
|
addrForward := addAll(
|
|
addEther(mac1, mac2),
|
|
addIP(ip1, ip2),
|
|
addTCP(port1, port2),
|
|
)
|
|
addrForward(idForward)
|
|
|
|
idRev := newFlowID()
|
|
addrRev := addAll(
|
|
addEther(mac2, mac1),
|
|
addIP(ip2, ip1),
|
|
addTCP(port2, port1),
|
|
)
|
|
addrRev(idRev)
|
|
assert.True(t, FlowIDsEqual(idForward, idRev))
|
|
|
|
{
|
|
module.Lock()
|
|
|
|
flow := module.Get(idForward)
|
|
flowRev := module.Get(idRev)
|
|
|
|
int1.Add(flow, -1)
|
|
uint1.Add(flow, 1)
|
|
float1.Add(flow, 3.14)
|
|
|
|
int2.Set(flowRev, -1)
|
|
uint2.Set(flowRev, 5)
|
|
float2.Set(flowRev, 1.4142)
|
|
|
|
module.Unlock()
|
|
}
|
|
|
|
var events []beat.Event
|
|
select {
|
|
case events = <-pub.ch:
|
|
case <-time.After(5 * time.Second):
|
|
}
|
|
|
|
if events == nil {
|
|
t.Fatalf("no event received in time")
|
|
}
|
|
event := events[0].Fields
|
|
t.Logf("event: %v", event)
|
|
|
|
source := event["source"].(common.MapStr)
|
|
dest := event["dest"].(common.MapStr)
|
|
|
|
// validate generated event
|
|
assert.Equal(t, net.HardwareAddr(mac1).String(), source["mac"])
|
|
assert.Equal(t, net.HardwareAddr(mac2).String(), dest["mac"])
|
|
assert.Equal(t, net.IP(ip1).String(), source["ip"])
|
|
assert.Equal(t, net.IP(ip2).String(), dest["ip"])
|
|
assert.Equal(t, uint16(256), source["port"])
|
|
assert.Equal(t, uint16(512), dest["port"])
|
|
assert.Equal(t, "tcp", event["transport"])
|
|
|
|
stat := source["stats"].(map[string]interface{})
|
|
assert.Equal(t, int64(-1), stat["int1"])
|
|
assert.Equal(t, nil, stat["int2"])
|
|
assert.Equal(t, uint64(1), stat["uint1"])
|
|
assert.Equal(t, nil, stat["uint2"])
|
|
assert.Equal(t, 3.14, stat["float1"])
|
|
assert.Equal(t, nil, stat["float2"])
|
|
|
|
stat = dest["stats"].(map[string]interface{})
|
|
assert.Equal(t, nil, stat["int1"])
|
|
assert.Equal(t, int64(-1), stat["int2"])
|
|
assert.Equal(t, nil, stat["uint1"])
|
|
assert.Equal(t, uint64(5), stat["uint2"])
|
|
assert.Equal(t, nil, stat["float1"])
|
|
assert.Equal(t, 1.4142, stat["float2"])
|
|
}
|