mysql-tunnel/internal/mysqlserver/server.go

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)
}