mysql-tunnel/internal/mysqlserver/session.go

package mysqlserver

import (
	"context"
	"fmt"
	"regexp"
	"strings"

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

var useDBRe = regexp.MustCompile(`(?is)^\s*USE\s+(?:` + "`" + `([^` + "`" + `]+)` + "`" + `|([A-Za-z0-9_]+))\s*;?\s*$`)

type Backend interface {
	Query(ctx context.Context, database string, sqls ...string) (*navicat.QueryResponse, error)
}

type session struct {
	pc       *packetConn
	backend  Backend
	log      *logger.Logger
	logSQL   bool
	timeout  ctxTimeout
	user     string
	db       string
	defaultDB string
	caps     uint32
	status   uint16
	scramble []byte
}

type ctxTimeout func() (context.Context, context.CancelFunc)

func (s *session) serve(ctx context.Context) {
	defer s.pc.close()
	for {
		select {
		case <-ctx.Done():
			return
		default:
		}
		s.pc.resetSeq()
		payload, err := s.pc.readPacket()
		if err != nil {
			return
		}
		if len(payload) == 0 {
			_ = s.pc.writeERR(errUnknownCommand, "empty command")
			continue
		}
		if err := s.dispatch(ctx, payload[0], payload[1:]); err != nil {
			return
		}
	}
}

func (s *session) dispatch(parent context.Context, cmd byte, body []byte) error {
	switch cmd {
	case ComQuit:
		return fmt.Errorf("quit")
	case ComPing:
		return s.pc.writeOK(0, 0, s.status, 0)
	case ComInitDB:
		db := strings.TrimSpace(string(body))
		return s.changeDB(parent, db)
	case ComQuery:
		return s.query(parent, string(body))
	case ComResetConnection:
		s.db = s.defaultDB
		return s.pc.writeOK(0, 0, s.status, 0)
	case ComFieldList, ComStmtPrepare, ComStmtExecute, ComStmtClose:
		return s.pc.writeERR(errNotSupported, "command not implemented yet")
	default:
		return s.pc.writeERR(errUnknownCommand, "unknown command")
	}
}

func (s *session) changeDB(parent context.Context, db string) error {
	if db == "" {
		return s.pc.writeERR(errNoDB, "no database selected")
	}
	s.db = db
	return s.pc.writeOK(0, 0, s.status, 0)
}

func (s *session) query(parent context.Context, sql string) error {
	sql = strings.TrimRight(sql, "\x00")
	s.logQuery(sql)
	pendingDB := ""
	if m := useDBRe.FindStringSubmatch(sql); m != nil {
		pendingDB = m[1]
		if pendingDB == "" {
			pendingDB = m[2]
		}
	}

	ctx, cancel := s.timeout()
	defer cancel()
	if ctx == nil {
		ctx = parent
	}

	resp, err := s.backend.Query(ctx, s.db, sql)
	if err != nil {
		return s.writeBackendErr(err)
	}
	for _, r := range resp.Results {
		if r.IsError() {
			return s.pc.writeERR(uint16(r.Errno), r.Message)
		}
	}
	if pendingDB != "" {
		s.db = pendingDB
	}
	res, ok := navicat.LastResultSet(resp)
	if !ok {
		return s.pc.writeOK(0, 0, s.status, 0)
	}
	if res.IsError() {
		return s.pc.writeERR(uint16(res.Errno), res.Message)
	}
	if res.IsResultSet() {
		if s.log != nil {
			s.log.Info("query completed", fmt.Sprintf("rows=%d", len(res.Rows)))
		}
		return s.pc.writeResultSet(res, s.caps, s.status)
	}
	if s.log != nil {
		s.log.Info("query completed", fmt.Sprintf("affected=%d", res.AffectedRows))
	}
	return s.pc.writeOK(uint64(res.AffectedRows), uint64(res.InsertID), s.status, 0)
}

func (s *session) writeBackendErr(err error) error {
	msg := logger.RedactError(err)
	if ne, ok := err.(*navicat.Error); ok {
		errno := uint16(ne.Errno)
		if errno == 0 {
			errno = errUnknown
		}
		if ne.Kind == "http" || ne.Kind == "tunnel" {
			errno = errUnknown
		}
		return s.pc.writeERR(errno, msg)
	}
	return s.pc.writeERR(errUnknown, msg)
}

func (s *session) logQuery(sql string) {
	if s.log == nil {
		return
	}
	if !s.logSQL {
		s.log.Info("query")
		return
	}
	s.log.Info("query", trimSQL(sql))
}

func trimSQL(sql string) string {
	s := logger.RedactText(sql)
	s = strings.ReplaceAll(s, "\n", " ")
	if len(s) > 240 {
		return s[:240] + "..."
	}
	return s
}