app-init

internal/mysqlserver/session.go

@@ -0,0 +1,168 @@
+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
+}