1.2.0

server/tcp_server_test.go

@@ -0,0 +1,120 @@
+package main
+
+import (
+	"bufio"
+	"encoding/json"
+	"fmt"
+	"net"
+	"strings"
+	"testing"
+	"time"
+)
+
+func TestAgentProtocolAndTrafficState(t *testing.T) {
+	app := newTestApp(t, minimalTestConfig())
+	listener, err := net.Listen("tcp", "127.0.0.1:0")
+	if err != nil {
+		t.Fatal(err)
+	}
+	serveDone := make(chan error, 1)
+	go func() { serveDone <- NewAgentServer(app).Serve(listener) }()
+
+	connection, err := net.DialTimeout("tcp", listener.Addr().String(), time.Second)
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer connection.Close()
+	reader := bufio.NewReader(connection)
+	readContains(t, reader, "Authentication required")
+	if _, err := fmt.Fprintln(connection, "s01:secret"); err != nil {
+		t.Fatal(err)
+	}
+	readContains(t, reader, "Authentication successful")
+	readContains(t, reader, "You are connecting via: IPv4")
+	readContains(t, reader, `"monitor":0`)
+
+	if _, err := fmt.Fprintln(connection, "pong on"); err != nil {
+		t.Fatal(err)
+	}
+	online6 := true
+	update := AgentStats{
+		Uptime: 90061, Load1: 1.25, Load5: 1, Load15: 0.75, CPU: 33.5, CPUCores: 4, CPUModel: "Test CPU",
+		MemoryTotal: 1024, MemoryUsed: 512, SwapTotal: 128, SwapUsed: 2, HDDTotal: 10000, HDDUsed: 4000,
+		NetworkRX: 123, NetworkTX: 456, NetworkIn: 1_000_000, NetworkOut: 2_000_000,
+		Ping10010: 1, Ping189: 2, Ping10086: 3, Time10010: 10, Time189: 20, Time10086: 30,
+		TCPCount: 10, UDPCount: 2, ProcessCount: 30, ThreadCount: 60, IORead: 1, IOWrite: 2,
+		OS: "linux", Custom: "example=12", Online6: &online6,
+	}
+	payload, _ := json.Marshal(update)
+	if _, err := fmt.Fprintf(connection, "update %s\n", payload); err != nil {
+		t.Fatal(err)
+	}
+	readContains(t, reader, "0")
+
+	eventually(t, time.Second, func() bool {
+		servers := app.SnapshotStats()["servers"].([]any)
+		server := servers[0].(map[string]any)
+		return server["online4"] == true && server["online6"] == true && server["cpu_model"] == "Test CPU" && server["uptime"] == "1 天"
+	})
+
+	result, apiErr := app.ResetTraffic("s01")
+	if apiErr != nil {
+		t.Fatal(apiErr)
+	}
+	stats := result["stats"].(map[string]any)
+	if stats["last_network_in"] != int64(1_000_000) {
+		t.Fatalf("traffic reset result: %#v", stats)
+	}
+
+	duplicate, err := net.DialTimeout("tcp", listener.Addr().String(), time.Second)
+	if err != nil {
+		t.Fatal(err)
+	}
+	duplicateReader := bufio.NewReader(duplicate)
+	readContains(t, duplicateReader, "Authentication required")
+	_, _ = fmt.Fprintln(duplicate, "s01:secret")
+	readContains(t, duplicateReader, "Only one connection per user")
+	_ = duplicate.Close()
+
+	wrong, err := net.DialTimeout("tcp", listener.Addr().String(), time.Second)
+	if err != nil {
+		t.Fatal(err)
+	}
+	wrongReader := bufio.NewReader(wrong)
+	readContains(t, wrongReader, "Authentication required")
+	_, _ = fmt.Fprintln(wrong, "s01:wrong")
+	readContains(t, wrongReader, "Wrong username")
+	_ = wrong.Close()
+
+	if apiErr := app.ReloadConfig(); apiErr != nil {
+		t.Fatal(apiErr)
+	}
+	readContains(t, reader, "Server reloading")
+	eventually(t, time.Second, func() bool {
+		server := app.SnapshotStats()["servers"].([]any)[0].(map[string]any)
+		return server["online4"] == false && server["online6"] == false
+	})
+
+	app.cancel()
+	_ = listener.Close()
+	select {
+	case err := <-serveDone:
+		if err != nil {
+			t.Fatal(err)
+		}
+	case <-time.After(time.Second):
+		t.Fatal("agent server did not stop")
+	}
+}
+
+func readContains(t *testing.T, reader *bufio.Reader, expected string) string {
+	t.Helper()
+	line, err := reader.ReadString('\n')
+	if err != nil {
+		t.Fatalf("read %q: %v", expected, err)
+	}
+	if !strings.Contains(line, expected) {
+		t.Fatalf("expected %q in %q", expected, line)
+	}
+	return line
+}