opsxdb/apps/server/internal/db/agent_test.go

package db

// agent_test.go —— JDBC agent 桥的协议层测试。
// 用 helper process 扮演 JVM agent（TestMain 拦截 -agent-stub），
// 无需本机安装 Java 即可验证：握手、请求/响应路由、超时、进程退出、并发串行写。

import (
	"archive/zip"
	"bufio"
	"context"
	"crypto/sha1"
	"crypto/sha256"
	"encoding/hex"
	"encoding/json"
	"errors"
	"fmt"
	"net/http"
	"net/http/httptest"
	"os"
	"os/exec"
	"path/filepath"
	"strings"
	"sync"
	"testing"
	"time"
)

// TestMain 拦截 stub 模式：测试二进制被当成 agent 子进程启动时直接进入协议循环。
func TestMain(m *testing.M) {
	for _, a := range os.Args {
		if a == "-agent-stub" {
			runAgentStub()
			return
		}
	}
	os.Exit(m.Run())
}

// runAgentStub 扮演 OpsxJdbcAgent 的最小协议面（与 agents/jdbc/OpsxJdbcAgent.java 对应）。
func runAgentStub() {
	out := bufio.NewWriter(os.Stdout)
	announce := agentResponse{ID: 0, OK: true}
	announce.Result = json.RawMessage(`{"agent":"opsx-agent-jdbc","version":"0.1.0","protocol":1}`)
	write := func(r agentResponse) {
		b, _ := json.Marshal(r)
		_, _ = out.WriteString(string(b) + "\n")
		out.Flush()
	}
	write(announce)
	in := bufio.NewScanner(os.Stdin)
	in.Buffer(make([]byte, 1024*1024), 1024*1024)
	for in.Scan() {
		var req agentRequest
		if json.Unmarshal(in.Bytes(), &req) != nil {
			continue
		}
		args := req.Args
		fail := func(msg string) agentResponse { return agentResponse{ID: req.ID, OK: false, Error: msg} }
		switch req.Op {
		case "hello":
			r := agentResponse{ID: req.ID, OK: true}
			r.Result = json.RawMessage(`{"agent":"opsx-agent-jdbc","version":"0.1.0","protocol":1}`)
			write(r)
		case "open":
			if args["url"] == "" || args["url"] == nil {
				write(fail("missing url"))
				break
			}
			r := agentResponse{ID: req.ID, OK: true}
			r.Result = json.RawMessage(`{"handle":"h1"}`)
			write(r)
		case "ping":
			write(agentResponse{ID: req.ID, OK: true})
		case "query":
			if strings.Contains(fmt.Sprintf("%v", args["sql"]), "SLOW") {
				time.Sleep(2 * time.Second)
			}
			r := agentResponse{ID: req.ID, OK: true}
			r.Result = json.RawMessage(`{"columns":["id","name"],"rows":[{"id":"1","name":"demo"}],"truncated":false}`)
			write(r)
		case "exec":
			r := agentResponse{ID: req.ID, OK: true}
			r.Result = json.RawMessage(`{"affectedRows":3}`)
			write(r)
		case "metaDatabases":
			r := agentResponse{ID: req.ID, OK: true}
			r.Result = json.RawMessage(`{"names":["MAIN","SYS"]}`)
			write(r)
		case "close":
			write(agentResponse{ID: req.ID, OK: true})
		case "shutdown":
			write(agentResponse{ID: req.ID, OK: true})
			os.Exit(0)
		default:
			write(fail("unknown op: " + req.Op))
		}
	}
	os.Exit(0)
}

// TestAgentStubHelper 占位：stub 子进程以 -test.run=TestAgentStubHelper 启动时
// 不会走到这里（TestMain 已拦截），但名字必须存在以匹配 -test.run。
func TestAgentStubHelper(t *testing.T) {}

func startStubAgent(t *testing.T) *AgentClient {
	t.Helper()
	cmd := exec.Command(os.Args[0], "-test.run=TestAgentStubHelper", "--", "-agent-stub")
	cli, err := StartAgentFromCmd(cmd, agentDriverSpecs["dameng"])
	if err != nil {
		t.Fatalf("StartAgentFromCmd: %v", err)
	}
	t.Cleanup(func() { cli.Shutdown(context.Background()) })
	return cli
}

func TestAgentClientRoundTrip(t *testing.T) {
	cli := startStubAgent(t)
	ctx := context.Background()

	// open → query → exec → metaDatabases → ping
	handle := ""
	{
		var res struct {
			Handle string `json:"handle"`
		}
		if err := cli.call(ctx, "open", map[string]any{"url": "jdbc:dm://127.0.0.1:5236", "user": "SYSDBA"}, &res); err != nil {
			t.Fatalf("open: %v", err)
		}
		if res.Handle != "h1" {
			t.Fatalf("handle = %q, want h1", res.Handle)
		}
		handle = res.Handle
	}
	{
		var res struct {
			Columns []string                 `json:"columns"`
			Rows    []map[string]interface{} `json:"rows"`
		}
		if err := cli.call(ctx, "query", map[string]any{"handle": handle, "sql": "SELECT 1", "maxRows": 100}, &res); err != nil {
			t.Fatalf("query: %v", err)
		}
		if len(res.Columns) != 2 || res.Rows[0]["name"] != "demo" {
			t.Fatalf("query result mismatch: %+v", res)
		}
	}
	{
		var res struct {
			AffectedRows int64 `json:"affectedRows"`
		}
		if err := cli.call(ctx, "exec", map[string]any{"handle": handle, "sql": "UPDATE t SET a=1"}, &res); err != nil {
			t.Fatalf("exec: %v", err)
		}
		if res.AffectedRows != 3 {
			t.Fatalf("affectedRows = %d, want 3", res.AffectedRows)
		}
	}
	{
		var res struct {
			Names []string `json:"names"`
		}
		if err := cli.call(ctx, "metaDatabases", map[string]any{"handle": handle}, &res); err != nil {
			t.Fatalf("metaDatabases: %v", err)
		}
		if len(res.Names) != 2 || res.Names[0] != "MAIN" {
			t.Fatalf("names = %v", res.Names)
		}
	}
	if err := cli.call(ctx, "ping", map[string]any{"handle": handle}, nil); err != nil {
		t.Fatalf("ping: %v", err)
	}
}

func TestAgentCallTimeout(t *testing.T) {
	cli := startStubAgent(t)
	ctx, cancel := context.WithTimeout(context.Background(), 100*time.Millisecond)
	defer cancel()
	var res map[string]any
	err := cli.call(ctx, "query", map[string]any{"handle": "h1", "sql": "SELECT SLOW", "maxRows": 10}, &res)
	if !errors.Is(err, context.DeadlineExceeded) {
		t.Fatalf("err = %v, want DeadlineExceeded", err)
	}
}

func TestAgentProcessExit(t *testing.T) {
	cli := startStubAgent(t)
	cli.Kill() // 杀进程后再调用，必须报进程退出而非挂死
	var res map[string]any
	err := cli.call(context.Background(), "ping", map[string]any{"handle": "h1"}, &res)
	if err == nil {
		t.Fatal("期望进程退出错误，got nil")
	}
}

func TestAgentConcurrentCalls(t *testing.T) {
	cli := startStubAgent(t)
	ctx := context.Background()
	if err := cli.call(ctx, "open", map[string]any{"url": "jdbc:dm://h:5236"}, nil); err != nil {
		t.Fatalf("open: %v", err)
	}
	const n = 24
	var wg sync.WaitGroup
	errs := make([]error, n)
	for i := 0; i < n; i++ {
		wg.Add(1)
		go func(i int) {
			defer wg.Done()
			var res struct {
				Columns []string `json:"columns"`
			}
			errs[i] = cli.call(ctx, "query", map[string]any{"handle": "h1", "sql": "SELECT 1", "maxRows": 1}, &res)
		}(i)
	}
	wg.Wait()
	for i, err := range errs {
		if err != nil {
			t.Fatalf("call %d: %v", i, err)
		}
	}
}

func TestRenderAgentURL(t *testing.T) {
	spec := agentDriverSpecs["dameng"]
	url, err := renderAgentURL(spec, ConnReq{Host: "127.0.0.1", Port: 5236})
	if err != nil {
		t.Fatalf("render: %v", err)
	}
	if url != "jdbc:dm://127.0.0.1:5236" {
		t.Fatalf("url = %q", url)
	}
	if _, err := renderAgentURL(spec, ConnReq{Host: "", Port: 5236}); err == nil {
		t.Fatal("空主机必须报错")
	}
	if _, err := renderAgentURL(spec, ConnReq{Host: "127.0.0.1", Port: 0}); err == nil {
		t.Fatal("零端口必须报错")
	}
}

func TestDownloadPinnedChecksum(t *testing.T) {
	payload := []byte("fake-gson-bytes")
	good := sha1.Sum(payload)
	cases := []struct {
		name    string
		body    []byte
		want    string
		wantErr bool
	}{
		{"校验通过", payload, hex.EncodeToString(good[:]), false},
		{"校验失败拒绝落盘", []byte("tampered"), hex.EncodeToString(good[:]), true},
	}
	for _, tc := range cases {
		t.Run(tc.name, func(t *testing.T) {
			srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
				_, _ = w.Write(tc.body)
			}))
			defer srv.Close()
			dest := t.TempDir() + "/gson.jar"
			err := downloadPinned(srv.URL, dest, tc.want)
			if tc.wantErr && err == nil {
				t.Fatal("期望校验失败")
			}
			if !tc.wantErr {
				if err != nil {
					t.Fatalf("download: %v", err)
				}
				if _, statErr := os.Stat(dest); statErr != nil {
					t.Fatalf("文件未落盘: %v", statErr)
				}
			} else {
				if _, statErr := os.Stat(dest); statErr == nil {
					t.Fatal("校验失败时文件不应存在")
				}
			}
		})
	}
}

func TestAgentConfigured(t *testing.T) {
	dir := t.TempDir()
	t.Setenv("OPSXDB_DAMENG_JARS", dir)
	if ok, _ := AgentConfigured("dameng"); ok {
		t.Fatal("空目录不应判为已配置")
	}
	if err := os.WriteFile(dir+"/DmJdbcDriver-8.1.3.140.jar", []byte("x"), 0o644); err != nil {
		t.Fatal(err)
	}
	if ok, _ := AgentConfigured("dameng"); !ok {
		t.Fatal("放置 jar 后应判为已配置")
	}
	if IsAgentDriver("mysql") {
		t.Fatal("mysql 不应是 agent 驱动")
	}
	if !IsAgentDriver("dameng") {
		t.Fatal("dameng 应是 agent 驱动")
	}
}

// TestEnsureAgentRuntimeOfflinePack —— 网络不可用（镜像指向无效 + sha256 不匹配）时，
// ensureAgentRuntime 必须从离线驱动包导入 jar（pack.sh 产物）。
func TestEnsureAgentRuntimeOfflinePack(t *testing.T) {
	home := t.TempDir()
	t.Setenv("OPSXDB_HOME", home)
	t.Setenv("OPSXDB_AGENT_JAR", filepath.Join(home, "agent", "opsx-agent-jdbc.jar"))
	// 网络必失败：镜像指向本地空端口 + 错误的 sha256
	t.Setenv("MAVEN_MIRROR", "http://127.0.0.1:1/none")
	badSHA := sha1.Sum([]byte("bad"))
	spec := AgentSpec{
		DriverKey:   "e2e-offline",
		DriverClass: "x",
		URLTemplate: "jdbc:x://{{.Host}}:{{.Port}}",
		JarsDirEnv:  "OPSXDB_OFFLINE_JARS",
	}
	agentDriverSpecs[spec.DriverKey] = spec
	t.Cleanup(func() { delete(agentDriverSpecs, spec.DriverKey) })
	agentDriverArtifacts[spec.DriverKey] = []agentArtifact{
		{"demo.jar", "com.nope:demo:1.0", hex.EncodeToString(badSHA[:])},
	}

	// 构造离线包：manifest + 正确的 jar 内容（sha256 以包内 manifest 为准）
	jarBytes := []byte("fake-jar-content")
	files := map[string]string{"drivers/e2e-offline/demo.jar": string(jarBytes)}
	manifestFiles := []map[string]string{{"path": "drivers/e2e-offline/demo.jar", "sha256": fmt.Sprintf("%x", sha256.Sum256(jarBytes))}}
	zipPath := filepath.Join(home, "offline.zip")
	if err := writeTestZip(zipPath, manifestFiles, files); err != nil {
		t.Fatal(err)
	}
	t.Setenv("OPSXDB_OFFLINE_ZIP", zipPath)

	// agent jar 也在包里（确保 ensureAgentJar 的离线路径不拦路）
	if err := importFromOfflinePack(map[string]string{"pack/agent/opsx-agent-jdbc.jar": filepath.Join(home, "agent", "opsx-agent-jdbc.jar")}); err != nil {
		// agent jar 不在此包属预期——单独构造一份最小 agent 条目
		agentBytes := []byte("fake-agent")
		mf2 := []map[string]string{{"path": "agent/opsx-agent-jdbc.jar", "sha256": fmt.Sprintf("%x", sha256.Sum256(agentBytes))}}
		f2 := map[string]string{"agent/opsx-agent-jdbc.jar": string(agentBytes)}
		if werr := writeTestZip(zipPath, append(manifestFiles, mf2...), mergeMaps(files, f2)); werr != nil {
			t.Fatal(werr)
		}
	}

	if err := ensureAgentRuntime(spec); err != nil {
		t.Fatalf("ensureAgentRuntime 应走离线包成功: %v", err)
	}
	got, err := os.ReadFile(filepath.Join(home, "drivers", "e2e-offline", "demo.jar"))
	if err != nil || string(got) != string(jarBytes) {
		t.Fatalf("离线导入内容不符: %v", err)
	}
	// 负例：manifest hash 不符必须拒
	jarBytes2 := []byte("tampered")
	mf3 := []map[string]string{{"path": "drivers/e2e-offline/demo.jar", "sha256": fmt.Sprintf("%x", sha256.Sum256(jarBytes2))}}
	if err := writeTestZip(zipPath, mf3, map[string]string{"drivers/e2e-offline/demo.jar": "not-the-same-content"}); err != nil {
		t.Fatal(err)
	}
	os.Remove(filepath.Join(home, "drivers", "e2e-offline", "demo.jar"))
	if err := ensureAgentRuntime(spec); err == nil {
		t.Fatal("manifest 与实际内容不符时必须失败")
	}
}

func writeTestZip(path string, manifestFiles []map[string]string, files map[string]string) error {
	zf, err := os.Create(path)
	if err != nil {
		return err
	}
	defer zf.Close()
	zw := zip.NewWriter(zf)
	mf, err := zw.Create("pack/manifest.json")
	if err != nil {
		return err
	}
	b, _ := json.Marshal(map[string]any{"bundleVersion": 1, "files": manifestFiles})
	if _, err := mf.Write(b); err != nil {
		return err
	}
	for rel, content := range files {
		w, err := zw.Create("pack/" + rel)
		if err != nil {
			return err
		}
		if _, err := w.Write([]byte(content)); err != nil {
			return err
		}
	}
	return zw.Close()
}

func mergeMaps(a, b map[string]string) map[string]string {
	out := map[string]string{}
	for k, v := range a {
		out[k] = v
	}
	for k, v := range b {
		out[k] = v
	}
	return out
}