mysql-tunnel/internal/navicat/client.go

package navicat

import (
	"context"
	"fmt"
	"net/http"
	"regexp"
	"strings"

	"mysql-tunnel/internal/logger"
	"mysql-tunnel/internal/tunnel"
)

var charsetIdent = regexp.MustCompile(`^[A-Za-z][A-Za-z0-9_]*$`)

// Target is the remote MySQL as seen by ntunnel_mysql.php.
type Target struct {
	Host     string
	Port     int
	User     string
	Password string
	Database string
	Charset  string
}

// Client talks Navicat HTTP Tunnel over the generic HTTP client.
type Client struct {
	http         *tunnel.HTTPClient
	url          string
	target       Target
	encodeBase64 bool
	retries      int
}

type ClientOptions struct {
	URL          string
	Target       Target
	EncodeBase64 bool
	// Retries is extra attempts on transport failure only (not SQL errors). 0 means a single try.
	Retries int
}

func NewClient(httpClient *tunnel.HTTPClient, opt ClientOptions) *Client {
	retries := opt.Retries
	if retries < 0 {
		retries = 0
	}
	return &Client{
		http:         httpClient,
		url:          opt.URL,
		target:       opt.Target,
		encodeBase64: opt.EncodeBase64,
		retries:      retries,
	}
}

func (c *Client) Connect(ctx context.Context) (*ConnectResponse, error) {
	body, status, err := c.post(ctx, Request{
		Action:       ActionConnect,
		Host:         c.target.Host,
		Port:         c.target.Port,
		Login:        c.target.User,
		Password:     c.target.Password,
		Database:     c.target.Database,
		EncodeBase64: c.encodeBase64,
	})
	if err != nil {
		return nil, err
	}
	resp, err := ParseConnect(body)
	if err != nil {
		return nil, err
	}
	if !resp.Header.OK() {
		return resp, headerError(resp.Header, status)
	}
	return resp, nil
}

func (c *Client) Query(ctx context.Context, database string, sqls ...string) (*QueryResponse, error) {
	db := database
	if db == "" {
		db = c.target.Database
	}
	queries := make([]string, 0, len(sqls)+1)
	if s := charsetSQL(c.target.Charset); s != "" {
		queries = append(queries, s)
	}
	queries = append(queries, nonEmpty(sqls)...)
	if len(queries) == 0 {
		return nil, fmt.Errorf("no SQL to execute")
	}
	body, status, err := c.post(ctx, Request{
		Action:       ActionQuery,
		Host:         c.target.Host,
		Port:         c.target.Port,
		Login:        c.target.User,
		Password:     c.target.Password,
		Database:     db,
		Queries:      queries,
		EncodeBase64: c.encodeBase64,
	})
	if err != nil {
		return nil, err
	}
	resp, err := ParseQuery(body)
	if err != nil {
		return nil, err
	}
	if !resp.Header.OK() {
		return resp, headerError(resp.Header, status)
	}
	return resp, nil
}

func (c *Client) post(ctx context.Context, req Request) (body []byte, status int, err error) {
	encoded := req.Encode()
	attempts := 1 + c.retries
	var last error
	for i := 0; i < attempts; i++ {
		if err := ctx.Err(); err != nil {
			return nil, 0, err
		}
		httpResp, err := c.http.Do(ctx, http.MethodPost, c.url, "application/x-www-form-urlencoded", strings.NewReader(encoded))
		if err != nil {
			last = err
			continue
		}
		status = httpResp.StatusCode
		body = httpResp.Body
		if status == http.StatusForbidden {
			return body, status, &Error{Kind: "http", HTTPStatus: status, Errno: ErrnoInvalidToken, Message: "invalid or missing token"}
		}
		if status < 200 || status >= 300 {
			return body, status, &Error{Kind: "http", HTTPStatus: status, Message: "unexpected status"}
		}
		if looksHTML(body) {
			return body, status, &Error{Kind: "http", HTTPStatus: status, Message: "tunnel returned HTML, not Navicat binary"}
		}
		return body, status, nil
	}
	return nil, 0, last
}

func charsetSQL(cs string) string {
	cs = strings.TrimSpace(cs)
	if cs == "" {
		return ""
	}
	if strings.EqualFold(cs, "utf8") || strings.EqualFold(cs, "utf8mb3") {
		return ""
	}
	if !charsetIdent.MatchString(cs) {
		return ""
	}
	return "SET NAMES " + cs
}

// LastResultSet returns the last result that has columns (skipping SET NAMES OK packets).
func LastResultSet(resp *QueryResponse) (Result, bool) {
	if resp == nil {
		return Result{}, false
	}
	for i := len(resp.Results) - 1; i >= 0; i-- {
		if resp.Results[i].IsResultSet() {
			return resp.Results[i], true
		}
	}
	if len(resp.Results) == 0 {
		return Result{}, false
	}
	return resp.Results[len(resp.Results)-1], true
}

func RedactedURL(raw string) string {
	return logger.RedactURL(raw)
}