mysql-tunnel/internal/mysqlserver/server_test.go

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)
	}
}