예제 #1
0
func TestConntracker(t *testing.T) {
	oldExecCmd, oldConntrackPresent := exec.Command, ConntrackModulePresent
	defer func() { exec.Command, ConntrackModulePresent = oldExecCmd, oldConntrackPresent }()

	ConntrackModulePresent = func() bool {
		return true
	}

	reader, writer := io.Pipe()
	exec.Command = func(name string, args ...string) exec.Cmd {
		return testExec.NewMockCmd(reader)
	}

	conntracker, err := NewConntracker(false)
	if err != nil {
		t.Fatal(err)
	}

	bw := bufio.NewWriter(writer)
	if _, err := bw.WriteString(XMLHeader); err != nil {
		t.Fatal(err)
	}
	if _, err := bw.WriteString(ConntrackOpenTag); err != nil {
		t.Fatal(err)
	}
	if err := bw.Flush(); err != nil {
		t.Fatal(err)
	}

	have := func() interface{} {
		result := []Flow{}
		conntracker.WalkFlows(func(f Flow) {
			f.Original = nil
			f.Reply = nil
			f.Independent = nil
			result = append(result, f)
		})
		return result
	}
	ts := 100 * time.Millisecond

	// First, assert we have no flows
	test.Poll(t, ts, []Flow{}, have)

	// Now add some flows
	xmlEncoder := xml.NewEncoder(bw)
	writeFlow := func(f Flow) {
		if err := xmlEncoder.Encode(f); err != nil {
			t.Fatal(err)
		}
		if _, err := bw.WriteString("\n"); err != nil {
			t.Fatal(err)
		}
		if err := bw.Flush(); err != nil {
			t.Fatal(err)
		}
	}

	flow1 := makeFlow(New)
	addMeta(&flow1, "original", "1.2.3.4", "2.3.4.5", 2, 3)
	addIndependant(&flow1, 1, "")
	writeFlow(flow1)
	test.Poll(t, ts, []Flow{flow1}, have)

	// Now check when we remove the flow, we still get it in the next Walk
	flow1.Type = Destroy
	writeFlow(flow1)
	test.Poll(t, ts, []Flow{flow1}, have)
	test.Poll(t, ts, []Flow{}, have)

	// This time we're not going to remove it, but put it in state TIME_WAIT
	flow1.Type = New
	writeFlow(flow1)
	test.Poll(t, ts, []Flow{flow1}, have)

	flow1.Metas[1].State = TimeWait
	writeFlow(flow1)
	test.Poll(t, ts, []Flow{flow1}, have)
	test.Poll(t, ts, []Flow{}, have)
}
예제 #2
0
func TestConntracker(t *testing.T) {
	oldExecCmd, oldConntrackPresent := exec.Command, ConntrackModulePresent
	defer func() { exec.Command, ConntrackModulePresent = oldExecCmd, oldConntrackPresent }()

	ConntrackModulePresent = func() bool {
		return true
	}

	first := true
	existingConnectionsReader, existingConnectionsWriter := io.Pipe()
	reader, writer := io.Pipe()
	exec.Command = func(name string, args ...string) exec.Cmd {
		if first {
			first = false
			return testexec.NewMockCmd(existingConnectionsReader)
		}
		return testexec.NewMockCmd(reader)
	}

	flowWalker := newConntrackFlowWalker(true)
	defer flowWalker.stop()

	// First write out some empty xml for the existing connections
	ecbw := bufio.NewWriter(existingConnectionsWriter)
	if _, err := ecbw.WriteString(xmlHeader); err != nil {
		t.Fatal(err)
	}
	if _, err := ecbw.WriteString(conntrackOpenTag); err != nil {
		t.Fatal(err)
	}
	if _, err := ecbw.WriteString(conntrackCloseTag); err != nil {
		t.Fatal(err)
	}
	if err := ecbw.Flush(); err != nil {
		t.Fatal(err)
	}

	// Then write out eventa
	bw := bufio.NewWriter(writer)
	if _, err := bw.WriteString(xmlHeader); err != nil {
		t.Fatal(err)
	}
	if _, err := bw.WriteString(conntrackOpenTag); err != nil {
		t.Fatal(err)
	}
	if err := bw.Flush(); err != nil {
		t.Fatal(err)
	}

	have := func() interface{} {
		result := []flow{}
		flowWalker.walkFlows(func(f flow) {
			f.Original = nil
			f.Reply = nil
			f.Independent = nil
			result = append(result, f)
		})
		return result
	}
	ts := 100 * time.Millisecond

	// First, assert we have no flows
	test.Poll(t, ts, []flow{}, have)

	// Now add some flows
	xmlEncoder := xml.NewEncoder(bw)
	writeFlow := func(f flow) {
		if err := xmlEncoder.Encode(f); err != nil {
			t.Fatal(err)
		}
		if _, err := bw.WriteString("\n"); err != nil {
			t.Fatal(err)
		}
		if err := bw.Flush(); err != nil {
			t.Fatal(err)
		}
	}

	flow1 := makeFlow(newType)
	addMeta(&flow1, "original", "1.2.3.4", "2.3.4.5", 2, 3)
	addIndependant(&flow1, 1, "")
	writeFlow(flow1)
	test.Poll(t, ts, []flow{flow1}, have)

	// Now check when we remove the flow, we still get it in the next Walk
	flow1.Type = destroyType
	writeFlow(flow1)
	test.Poll(t, ts, []flow{flow1}, have)
	test.Poll(t, ts, []flow{}, have)

	// This time we're not going to remove it, but put it in state TIME_WAIT
	flow1.Type = newType
	writeFlow(flow1)
	test.Poll(t, ts, []flow{flow1}, have)

	flow1.Metas[1].State = timeWait
	writeFlow(flow1)
	test.Poll(t, ts, []flow{flow1}, have)
	test.Poll(t, ts, []flow{}, have)
}