app-init

2026-09-21 16:46:55 +08:00
parent 01842c59c6
commit ad2b0a7d3e
58 changed files with 6906 additions and 3 deletions
@@ -0,0 +1,423 @@
+package mysqlserver
+
+import (
+	"context"
+	"fmt"
+	"net"
+	"net/http/httptest"
+	"strings"
+	"testing"
+	"time"
+
+	"mysql-tunnel/internal/navicat"
+	"mysql-tunnel/internal/tunnel"
+)
+
+type stubBackend struct{}
+
+func (stubBackend) Query(_ context.Context, _ string, sqls ...string) (*navicat.QueryResponse, error) {
+	var results []navicat.Result
+	for _, sql := range sqls {
+		s := strings.TrimSpace(sql)
+		u := strings.ToUpper(s)
+		switch {
+		case strings.HasPrefix(u, "SET NAMES"), strings.HasPrefix(u, "USE "):
+			results = append(results, navicat.Result{})
+		case u == "SELECT 1" || strings.HasPrefix(u, "SELECT 1 "):
+			results = append(results, navicat.Result{
+				NumFields: 1, NumRows: 1,
+				Fields: []navicat.Field{{Name: "1", Type: 8, Length: 20}},
+				Rows:   [][]navicat.Cell{{{Bytes: []byte("1")}}},
+			})
+		case strings.Contains(u, "SELECT NULL"):
+			results = append(results, navicat.Result{
+				NumFields: 1, NumRows: 1,
+				Fields: []navicat.Field{{Name: "n", Type: 6}},
+				Rows:   [][]navicat.Cell{{{Null: true}}},
+			})
+		case strings.Contains(strings.ToLower(s), "bigint_val"):
+			results = append(results, navicat.Result{
+				NumFields: 1, NumRows: 1,
+				Fields: []navicat.Field{{Name: "bigint_val", Type: 8, Flags: 32, Length: 20}},
+				Rows:   [][]navicat.Cell{{{Bytes: []byte("18446744073709551615")}}},
+			})
+		case strings.Contains(strings.ToLower(s), "dec_val"):
+			results = append(results, navicat.Result{
+				NumFields: 1, NumRows: 1,
+				Fields: []navicat.Field{{Name: "dec_val", Type: 246, Length: 32}},
+				Rows:   [][]navicat.Cell{{{Bytes: []byte("1234567890.123456789")}}},
+			})
+		case strings.Contains(s, "你好") || strings.Contains(strings.ToLower(s), "zh_val"):
+			results = append(results, navicat.Result{
+				NumFields: 1, NumRows: 1,
+				Fields: []navicat.Field{{Name: "zh_val", Type: 253, Length: 12}},
+				Rows:   [][]navicat.Cell{{{Bytes: []byte("你好世界")}}},
+			})
+		case strings.HasPrefix(u, "UPDATE"), strings.HasPrefix(u, "INSERT"):
+			results = append(results, navicat.Result{AffectedRows: 2, InsertID: 7})
+		case strings.Contains(u, "SYNTAX"):
+			results = append(results, navicat.Result{Errno: 1064, Message: "You have an error in your SQL syntax"})
+		default:
+			results = append(results, navicat.Result{Errno: 1064, Message: "stub: unhandled SQL"})
+		}
+	}
+	return &navicat.QueryResponse{Results: results}, nil
+}
+
+func startServer(t *testing.T, backend Backend, user, pass string) (addr string, cancel context.CancelFunc) {
+	t.Helper()
+	ln, err := net.Listen("tcp", "127.0.0.1:0")
+	if err != nil {
+		t.Fatal(err)
+	}
+	ctx, cancel := context.WithCancel(context.Background())
+	s := New(Options{
+		Host: "127.0.0.1", Port: 1, Username: user, Password: pass,
+		DefaultDB: "app", MaxConnections: 5, Timeout: 5 * time.Second,
+	}, backend, nil)
+	errCh := make(chan error, 1)
+	go func() { errCh <- s.Serve(ctx, ln) }()
+	t.Cleanup(func() {
+		cancel()
+		_ = ln.Close()
+	})
+	return ln.Addr().String(), cancel
+}
+
+type testClient struct {
+	pc   *packetConn
+	caps uint32
+}
+
+func dialTest(t *testing.T, addr, user, pass, db string) *testClient {
+	t.Helper()
+	c, err := net.DialTimeout("tcp", addr, 3*time.Second)
+	if err != nil {
+		t.Fatal(err)
+	}
+	pc := newPacketConn(c, 5*time.Second)
+	hs, err := pc.readPacket()
+	if err != nil {
+		t.Fatal(err)
+	}
+	scramble, err := scrambleFromHandshake(hs)
+	if err != nil {
+		t.Fatal(err)
+	}
+	caps := clientProtocol41 | clientSecureConnection | clientPluginAuth |
+		clientConnectWithDB | clientPluginAuthLenencClientData | clientDeprecateEOF | clientMultiResults
+	token := nativePasswordToken(pass, scramble)
+	if err := pc.writePacket(buildClientAuth(user, token, db, caps)); err != nil {
+		t.Fatal(err)
+	}
+	ok, err := pc.readPacket()
+	if err != nil {
+		t.Fatal(err)
+	}
+	if len(ok) == 0 || ok[0] != markerOK {
+		t.Fatalf("auth failed: %x %q", ok, ok)
+	}
+	return &testClient{pc: pc, caps: caps}
+}
+
+func scrambleFromHandshake(payload []byte) ([]byte, error) {
+	if len(payload) < 20 || payload[0] != ProtocolVersion {
+		return nil, fmt.Errorf("bad handshake")
+	}
+	off := 1
+	_, off, err := readNUL(payload, off)
+	if err != nil {
+		return nil, err
+	}
+	off += 4 // conn id
+	part1 := payload[off : off+8]
+	off += 8 + 1 + 2 + 1 + 2 + 2 + 1 + 10
+	if off+13 > len(payload) {
+		return nil, fmt.Errorf("short scramble")
+	}
+	part2 := payload[off : off+12]
+	s := append(append([]byte{}, part1...), part2...)
+	return s, nil
+}
+
+func (c *testClient) query(t *testing.T, sql string) (cols []string, rows [][]*string, affected, lastID uint64, errno int, errMsg string) {
+	t.Helper()
+	c.pc.resetSeq()
+	payload := append([]byte{ComQuery}, []byte(sql)...)
+	if err := c.pc.writePacket(payload); err != nil {
+		t.Fatal(err)
+	}
+	pkt, err := c.pc.readPacket()
+	if err != nil {
+		t.Fatal(err)
+	}
+	if len(pkt) == 0 {
+		t.Fatal("empty")
+	}
+	if pkt[0] == markerERR {
+		errno = int(pkt[1]) | int(pkt[2])<<8
+		if len(pkt) > 9 {
+			errMsg = string(pkt[9:])
+		}
+		return
+	}
+	if pkt[0] == markerOK {
+		n, off, _ := readLenencInt(pkt, 1)
+		id, _, _ := readLenencInt(pkt, off)
+		return nil, nil, n, id, 0, ""
+	}
+	nCols, _, err := readLenencInt(pkt, 0)
+	if err != nil {
+		t.Fatal(err)
+	}
+	for i := uint64(0); i < nCols; i++ {
+		colPkt, err := c.pc.readPacket()
+		if err != nil {
+			t.Fatal(err)
+		}
+		name := columnName(colPkt)
+		cols = append(cols, name)
+	}
+	// metadata end OK/EOF
+	if _, err := c.pc.readPacket(); err != nil {
+		t.Fatal(err)
+	}
+	for {
+		rowPkt, err := c.pc.readPacket()
+		if err != nil {
+			t.Fatal(err)
+		}
+		if len(rowPkt) == 0 {
+			break
+		}
+		if rowPkt[0] == markerERR {
+			errno = int(rowPkt[1]) | int(rowPkt[2])<<8
+			return
+		}
+		if rowPkt[0] == markerEOF && len(rowPkt) < 9 {
+			break
+		}
+		if rowPkt[0] == markerOK && len(rowPkt) >= 7 {
+			break
+		}
+		off := 0
+		row := make([]*string, nCols)
+		for i := uint64(0); i < nCols; i++ {
+			if off >= len(rowPkt) {
+				t.Fatal("short row")
+			}
+			if rowPkt[off] == markerNull {
+				row[i] = nil
+				off++
+				continue
+			}
+			b, nOff, err := readLenencString(rowPkt, off)
+			if err != nil {
+				t.Fatal(err)
+			}
+			s := string(b)
+			row[i] = &s
+			off = nOff
+		}
+		rows = append(rows, row)
+	}
+	return
+}
+
+func (c *testClient) ping(t *testing.T) {
+	t.Helper()
+	c.pc.resetSeq()
+	if err := c.pc.writePacket([]byte{ComPing}); err != nil {
+		t.Fatal(err)
+	}
+	pkt, err := c.pc.readPacket()
+	if err != nil {
+		t.Fatal(err)
+	}
+	if pkt[0] != markerOK {
+		t.Fatalf("ping %x", pkt)
+	}
+}
+
+func (c *testClient) quit() {
+	c.pc.resetSeq()
+	_ = c.pc.writePacket([]byte{ComQuit})
+	_ = c.pc.close()
+}
+
+func columnName(pkt []byte) string {
+	off := 0
+	skip := func() {
+		_, n, err := readLenencString(pkt, off)
+		if err != nil {
+			return
+		}
+		off = n
+	}
+	skip() // catalog
+	skip() // schema
+	skip() // table
+	skip() // org table
+	b, _, err := readLenencString(pkt, off)
+	if err != nil {
+		return ""
+	}
+	return string(b)
+}
+
+func TestServerAuthPingQueryQuit(t *testing.T) {
+	addr, _ := startServer(t, stubBackend{}, "root", "123456")
+	cl := dialTest(t, addr, "root", "123456", "app")
+	cl.ping(t)
+	_, rows, _, _, errno, msg := cl.query(t, "SELECT 1")
+	if errno != 0 {
+		t.Fatalf("errno %d %s", errno, msg)
+	}
+	if len(rows) != 1 || rows[0][0] == nil || *rows[0][0] != "1" {
+		t.Fatalf("rows=%v", rows)
+	}
+	cl.quit()
+}
+
+func TestServerBadPassword(t *testing.T) {
+	addr, _ := startServer(t, stubBackend{}, "root", "123456")
+	c, err := net.DialTimeout("tcp", addr, 3*time.Second)
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer c.Close()
+	pc := newPacketConn(c, 5*time.Second)
+	hs, err := pc.readPacket()
+	if err != nil {
+		t.Fatal(err)
+	}
+	scramble, err := scrambleFromHandshake(hs)
+	if err != nil {
+		t.Fatal(err)
+	}
+	caps := clientProtocol41 | clientSecureConnection | clientPluginAuth | clientPluginAuthLenencClientData
+	tok := nativePasswordToken("wrong", scramble)
+	if err := pc.writePacket(buildClientAuth("root", tok, "", caps)); err != nil {
+		t.Fatal(err)
+	}
+	pkt, err := pc.readPacket()
+	if err != nil {
+		t.Fatal(err)
+	}
+	if pkt[0] != markerERR {
+		t.Fatalf("want ERR got %x", pkt)
+	}
+	errno := int(pkt[1]) | int(pkt[2])<<8
+	if errno != int(errAccessDenied) {
+		t.Fatalf("errno %d", errno)
+	}
+}
+
+func TestServerResultTypes(t *testing.T) {
+	addr, _ := startServer(t, stubBackend{}, "root", "p")
+	cl := dialTest(t, addr, "root", "p", "app")
+	defer cl.quit()
+
+	_, rows, _, _, errno, _ := cl.query(t, "SELECT NULL AS n")
+	if errno != 0 || rows[0][0] != nil {
+		t.Fatalf("null errno=%d row=%v", errno, rows)
+	}
+
+	_, rows, _, _, errno, _ = cl.query(t, "SELECT bigint_val")
+	if errno != 0 || *rows[0][0] != "18446744073709551615" {
+		t.Fatalf("bigint %v errno=%d", rows, errno)
+	}
+
+	_, rows, _, _, errno, _ = cl.query(t, "SELECT dec_val")
+	if errno != 0 || *rows[0][0] != "1234567890.123456789" {
+		t.Fatalf("dec %v", rows)
+	}
+
+	_, rows, _, _, errno, _ = cl.query(t, "SELECT zh_val")
+	if errno != 0 || *rows[0][0] != "你好世界" {
+		t.Fatalf("zh %v", rows)
+	}
+
+	_, _, affected, lastID, errno, msg := cl.query(t, "UPDATE t SET x=1")
+	if errno != 0 {
+		t.Fatal(msg)
+	}
+	if affected != 2 || lastID != 7 {
+		t.Fatalf("affected=%d id=%d", affected, lastID)
+	}
+
+	_, _, _, _, errno, msg = cl.query(t, "SELECT SYNTAX")
+	if errno != 1064 || !strings.Contains(msg, "syntax") {
+		t.Fatalf("errno=%d msg=%s", errno, msg)
+	}
+}
+
+func TestServerAuthSwitch(t *testing.T) {
+	addr, _ := startServer(t, stubBackend{}, "root", "p")
+	c, err := net.DialTimeout("tcp", addr, 3*time.Second)
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer c.Close()
+	pc := newPacketConn(c, 5*time.Second)
+	hs, err := pc.readPacket()
+	if err != nil {
+		t.Fatal(err)
+	}
+	scramble, err := scrambleFromHandshake(hs)
+	if err != nil {
+		t.Fatal(err)
+	}
+	caps := clientProtocol41 | clientSecureConnection | clientPluginAuth | clientPluginAuthLenencClientData
+	b := appendU32(nil, caps)
+	b = appendU32(b, MaxPacketSize)
+	b = append(b, charsetUTF8MB4)
+	b = append(b, make([]byte, 23)...)
+	b = appendNUL(b, "root")
+	b = appendLenencInt(b, 0) // empty, will switch
+	b = appendNUL(b, AuthPluginSHA2)
+	if err := pc.writePacket(b); err != nil {
+		t.Fatal(err)
+	}
+	sw, err := pc.readPacket()
+	if err != nil {
+		t.Fatal(err)
+	}
+	if sw[0] != markerEOF {
+		t.Fatalf("want auth switch %x", sw)
+	}
+	tok := nativePasswordToken("p", scramble)
+	if err := pc.writePacket(tok); err != nil {
+		t.Fatal(err)
+	}
+	ok, err := pc.readPacket()
+	if err != nil {
+		t.Fatal(err)
+	}
+	if ok[0] != markerOK {
+		t.Fatalf("auth after switch %x", ok)
+	}
+}
+
+func TestPhase4NavicatStack(t *testing.T) {
+	mock := httptest.NewServer(navicat.NewMockHandler(navicat.WithToken("t")))
+	t.Cleanup(mock.Close)
+	httpClient := tunnel.NewHTTPClient(tunnel.Options{Timeout: 5 * time.Second})
+	n := navicat.NewClient(httpClient, navicat.ClientOptions{
+		URL: mock.URL + "?token=t",
+		Target: navicat.Target{
+			Host: "mysql", Port: 3306, User: "remote", Password: "rp",
+			Database: "app", Charset: "utf8mb4",
+		},
+	})
+	addr, _ := startServer(t, n, "root", "local")
+	cl := dialTest(t, addr, "root", "local", "app")
+	defer cl.quit()
+	_, rows, _, _, errno, msg := cl.query(t, "SELECT 1")
+	if errno != 0 {
+		t.Fatal(msg)
+	}
+	if *rows[0][0] != "1" {
+		t.Fatalf("%v", rows)
+	}
+}