app-init

2026-09-21 16:46:55 +08:00
parent 01842c59c6
commit ad2b0a7d3e
58 changed files with 6906 additions and 3 deletions
@@ -0,0 +1,194 @@
+package mysqlserver
+
+import (
+	"context"
+	"fmt"
+	"net"
+	"strconv"
+	"sync"
+	"sync/atomic"
+	"time"
+
+	"mysql-tunnel/internal/logger"
+)
+
+type Options struct {
+	Host           string
+	Port           int
+	Username       string
+	Password       string
+	DefaultDB      string
+	MaxConnections int
+	Timeout        time.Duration
+	LogSQL         bool
+}
+
+type Server struct {
+	opt     Options
+	backend Backend
+	log     *logger.Logger
+
+	ln     net.Listener
+	connID atomic.Uint32
+
+	sem chan struct{}
+	wg  sync.WaitGroup
+}
+
+func New(opt Options, backend Backend, log *logger.Logger) *Server {
+	if opt.Host == "" {
+		opt.Host = "127.0.0.1"
+	}
+	if opt.Port == 0 {
+		opt.Port = 13306
+	}
+	if opt.MaxConnections < 1 {
+		opt.MaxConnections = 5
+	}
+	if opt.Timeout <= 0 {
+		opt.Timeout = 30 * time.Second
+	}
+	return &Server{
+		opt:     opt,
+		backend: backend,
+		log:     log,
+		sem:     make(chan struct{}, opt.MaxConnections),
+	}
+}
+
+func (s *Server) Addr() string {
+	if s.ln != nil {
+		return s.ln.Addr().String()
+	}
+	return net.JoinHostPort(s.opt.Host, strconv.Itoa(s.opt.Port))
+}
+
+func (s *Server) ListenAndServe(ctx context.Context) error {
+	ln, err := net.Listen("tcp", net.JoinHostPort(s.opt.Host, strconv.Itoa(s.opt.Port)))
+	if err != nil {
+		return err
+	}
+	return s.Serve(ctx, ln)
+}
+
+func (s *Server) Serve(ctx context.Context, ln net.Listener) error {
+	s.ln = ln
+	if s.log != nil {
+		s.log.Info("server started", ln.Addr().String())
+	}
+	go func() {
+		<-ctx.Done()
+		_ = ln.Close()
+	}()
+	for {
+		c, err := ln.Accept()
+		if err != nil {
+			select {
+			case <-ctx.Done():
+				s.wg.Wait()
+				return nil
+			default:
+				return err
+			}
+		}
+		select {
+		case s.sem <- struct{}{}:
+		default:
+			_ = c.Close()
+			if s.log != nil {
+				s.log.Warn("max connections reached, rejected", c.RemoteAddr().String())
+			}
+			continue
+		}
+		s.wg.Add(1)
+		go func(c net.Conn) {
+			defer s.wg.Done()
+			defer func() { <-s.sem }()
+			s.handle(ctx, c)
+		}(c)
+	}
+}
+
+func (s *Server) handle(ctx context.Context, c net.Conn) {
+	pc := newPacketConn(c, s.opt.Timeout)
+	remote := pc.remote()
+	if s.log != nil {
+		s.log.Info("client connected", remote)
+	}
+	defer func() {
+		_ = pc.close()
+		if s.log != nil {
+			s.log.Info("client disconnected", remote)
+		}
+	}()
+
+	scramble, err := randomScramble()
+	if err != nil {
+		return
+	}
+	id := s.connID.Add(1)
+	if err := pc.writePacket(buildHandshake(id, scramble)); err != nil {
+		return
+	}
+	payload, err := pc.readPacket()
+	if err != nil {
+		return
+	}
+	hr, err := parseHandshakeResponse(payload)
+	if err != nil {
+		_ = pc.writeERR(errUnknown, "bad handshake")
+		return
+	}
+	if hr.SSL {
+		_ = pc.writeERR(errNotSupported, "SSL is not supported on the local tunnel")
+		return
+	}
+	auth := hr.Auth
+	plugin := hr.Plugin
+	if plugin == "" {
+		plugin = AuthPluginNative
+	}
+	if plugin != AuthPluginNative {
+		if err := pc.writePacket(encodeAuthSwitch(AuthPluginNative, scramble)); err != nil {
+			return
+		}
+		auth, err = pc.readPacket()
+		if err != nil {
+			return
+		}
+		if len(auth) > 0 && auth[len(auth)-1] == 0 && len(auth) != 20 {
+			auth = auth[:len(auth)-1]
+		}
+	}
+	db := hr.Database
+	if db == "" {
+		db = s.opt.DefaultDB
+	}
+	if len(auth) == 21 && auth[20] == 0 {
+		auth = auth[:20]
+	}
+	if !usernamesEqual(hr.Username, s.opt.Username) || !verifyNativePassword(s.opt.Password, scramble, auth) {
+		_ = pc.writeERR(errAccessDenied, fmt.Sprintf("Access denied for user '%s'@'localhost'", hr.Username))
+		return
+	}
+	if err := pc.writeOK(0, 0, serverStatusAutocommit, 0); err != nil {
+		return
+	}
+
+	sess := &session{
+		pc:      pc,
+		backend: s.backend,
+		log:     s.log,
+		logSQL:  s.opt.LogSQL,
+		timeout: func() (context.Context, context.CancelFunc) {
+			return context.WithTimeout(ctx, s.opt.Timeout)
+		},
+		user:      hr.Username,
+		db:        db,
+		defaultDB: s.opt.DefaultDB,
+		caps:      hr.Capabilities,
+		status:   serverStatusAutocommit,
+		scramble: scramble,
+	}
+	sess.serve(ctx)
+}