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 }