mysql-tunnel/internal/navicat/mock.go

package navicat

import (
	"encoding/base64"
	"net/http"
	"strings"
)

// MockOption configures the in-process ntunnel_mysql.php stand-in.
type MockOption func(*mockState)

type mockState struct {
	token       string
	connectFail *Error
	queryFn     func(sql string) Result
	capture     *CapturedRequest
}

// CapturedRequest is the last POST seen by the mock.
type CapturedRequest struct {
	Action       string
	Host         string
	Port         string
	Login        string
	Database     string
	Queries      []string
	EncodeBase64 string
	HasPassword  bool
}

func WithToken(token string) MockOption {
	return func(m *mockState) { m.token = token }
}

func WithConnectError(errno uint32, msg string) MockOption {
	return func(m *mockState) {
		m.connectFail = &Error{Kind: "mysql", Errno: errno, Message: msg}
	}
}

func WithQueryHandler(fn func(sql string) Result) MockOption {
	return func(m *mockState) { m.queryFn = fn }
}

func WithCapture(dst *CapturedRequest) MockOption {
	return func(m *mockState) { m.capture = dst }
}

// NewMockHandler implements a subset of ntunnel_mysql.php for tests.
func NewMockHandler(opts ...MockOption) http.Handler {
	st := &mockState{queryFn: defaultQuery}
	for _, o := range opts {
		o(st)
	}
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		_ = r.ParseForm()
		if st.token != "" && r.URL.Query().Get("token") != st.token {
			if r.PostForm.Get("actn") != "" {
				writeBytes(w, encodeHeaderError(ErrnoInvalidToken, "invalid access token"))
				return
			}
			http.Error(w, "403 Forbidden: invalid or missing token", http.StatusForbidden)
			return
		}
		if r.Method != http.MethodPost {
			writeBytes(w, encodeHeaderError(ErrnoInvalidParams, "invalid parameters"))
			return
		}
		actn := r.PostForm.Get("actn")
		host := r.PostForm.Get("host")
		login := r.PostForm.Get("login")
		if actn == "" || host == "" || !r.PostForm.Has("port") || login == "" {
			writeBytes(w, encodeHeaderError(ErrnoInvalidParams, "invalid parameters"))
			return
		}
		if st.capture != nil {
			*st.capture = CapturedRequest{
				Action:       actn,
				Host:         host,
				Port:         r.PostForm.Get("port"),
				Login:        login,
				Database:     r.PostForm.Get("db"),
				Queries:      append([]string(nil), r.PostForm["q[]"]...),
				EncodeBase64: r.PostForm.Get("encodeBase64"),
				HasPassword:  r.PostForm.Has("password"),
			}
		}
		if st.connectFail != nil {
			writeBytes(w, encodeHeaderError(st.connectFail.Errno, st.connectFail.Message))
			return
		}
		switch actn {
		case ActionConnect:
			writeBytes(w, encodeConnectOK("127.0.0.1 via TCP/IP", "10", "8.0.36-mock"))
		case ActionQuery:
			queries := append([]string(nil), r.PostForm["q[]"]...)
			if r.PostForm.Get("encodeBase64") == "1" {
				decoded := make([]string, 0, len(queries))
				for _, q := range queries {
					b, err := base64.StdEncoding.DecodeString(q)
					if err != nil {
						writeBytes(w, encodeQueryResults([]Result{{Errno: 1064, Message: "bad base64 query"}}))
						return
					}
					decoded = append(decoded, string(b))
				}
				queries = decoded
			}
			var results []Result
			for _, q := range queries {
				if q == "" {
					continue
				}
				results = append(results, st.queryFn(q))
			}
			writeBytes(w, encodeQueryResults(results))
		default:
			writeBytes(w, AppendHeader(nil, Version206, 0))
		}
	})
}

func defaultQuery(sql string) Result {
	s := strings.TrimSpace(sql)
	u := strings.ToUpper(s)
	switch {
	case strings.HasPrefix(u, "SET NAMES"):
		return Result{}
	case u == "SELECT 1" || strings.HasPrefix(u, "SELECT 1 "):
		return selectLiteral("1", "1", 8, 20)
	case strings.Contains(u, "SELECT NULL"):
		return Result{
			NumFields: 1,
			NumRows:   1,
			Fields:    []Field{{Name: "n", Type: 6, Length: 0}},
			Rows:      [][]Cell{{{Null: true}}},
		}
	case strings.Contains(strings.ToLower(s), "bigint_val"):
		r := selectLiteral("bigint_val", "18446744073709551615", 8, 20)
		r.Fields[0].Flags = 32 // UNSIGNED
		return r
	case strings.Contains(strings.ToLower(s), "dec_val"):
		return selectLiteral("dec_val", "1234567890.123456789", 246, 32)
	case strings.Contains(s, "你好") || strings.Contains(strings.ToLower(s), "zh_val"):
		return selectLiteral("zh_val", "你好世界", 253, 12)
	case strings.HasPrefix(u, "INSERT") || strings.HasPrefix(u, "UPDATE") || strings.HasPrefix(u, "DELETE") || strings.HasPrefix(u, "REPLACE"):
		return Result{AffectedRows: 1, InsertID: 42, Info: "Rows matched: 1 Changed: 1 Warnings: 0"}
	case strings.Contains(u, "SYNTAX"):
		return Result{Errno: 1064, Message: "You have an error in your SQL syntax"}
	default:
		return Result{Errno: 1064, Message: "mock: unhandled SQL"}
	}
}

func selectLiteral(name, value string, typ, length uint32) Result {
	return Result{
		NumFields: 1,
		NumRows:   1,
		Fields:    []Field{{Name: name, Type: typ, Length: length}},
		Rows:      [][]Cell{{{Bytes: []byte(value)}}},
	}
}

func encodeHeaderError(errno uint32, msg string) []byte {
	b := AppendHeader(nil, Version206, errno)
	return AppendBlock(b, []byte(msg))
}

func encodeConnectOK(host, proto, server string) []byte {
	b := AppendHeader(nil, Version206, 0)
	b = AppendBlock(b, []byte(host))
	b = AppendBlock(b, []byte(proto))
	b = AppendBlock(b, []byte(server))
	return b
}

func encodeQueryResults(results []Result) []byte {
	b := AppendHeader(nil, Version206, 0)
	if len(results) == 0 {
		return b
	}
	for i, res := range results {
		b = appendQueryResult(b, res)
		if i < len(results)-1 {
			b = append(b, sepMore)
		} else {
			b = append(b, sepLast)
		}
	}
	return b
}

func appendQueryResult(b []byte, res Result) []byte {
	nfields := res.NumFields
	nrows := res.NumRows
	if nfields == 0 {
		nfields = uint32(len(res.Fields))
	}
	if nrows == 0 && len(res.Rows) > 0 {
		nrows = uint32(len(res.Rows))
	}
	b = AppendResultSetHeader(b, res.Errno, res.AffectedRows, res.InsertID, nfields, nrows)
	if res.Errno > 0 {
		return AppendBlock(b, []byte(res.Message))
	}
	if nfields == 0 {
		return AppendBlock(b, []byte(res.Info))
	}
	for _, f := range res.Fields {
		b = AppendField(b, f.Name, f.Table, f.Type, f.Flags, f.Length)
	}
	for _, row := range res.Rows {
		for _, cell := range row {
			if cell.Null {
				b = AppendNull(b)
			} else {
				b = AppendBlock(b, cell.Bytes)
			}
		}
	}
	return b
}

func writeBytes(w http.ResponseWriter, body []byte) {
	w.Header().Set("Content-Type", "text/plain; charset=x-user-defined")
	w.WriteHeader(http.StatusOK)
	_, _ = w.Write(body)
}