Files
2026-09-21 16:46:55 +08:00

106 lines
2.1 KiB
Go

package mysqlserver
import (
"fmt"
"io"
"net"
"time"
)
type packetConn struct {
c net.Conn
seq uint8
max int
readTO time.Duration
writeTO time.Duration
}
func newPacketConn(c net.Conn, timeout time.Duration) *packetConn {
return &packetConn{c: c, max: MaxPacketSize, readTO: timeout, writeTO: timeout}
}
func (p *packetConn) resetSeq() { p.seq = 0 }
func (p *packetConn) readPacket() ([]byte, error) {
var payload []byte
for {
if p.readTO > 0 {
_ = p.c.SetReadDeadline(time.Now().Add(p.readTO))
}
var hdr [4]byte
if _, err := io.ReadFull(p.c, hdr[:]); err != nil {
return nil, err
}
n := int(readU24(hdr[:3]))
seq := hdr[3]
if seq != p.seq {
return nil, fmt.Errorf("sequence mismatch: got %d want %d", seq, p.seq)
}
p.seq++
if n == 0 {
return payload, nil
}
if n > p.max {
return nil, fmt.Errorf("packet too large: %d", n)
}
chunk := make([]byte, n)
if _, err := io.ReadFull(p.c, chunk); err != nil {
return nil, err
}
payload = append(payload, chunk...)
if n < 0xffffff {
return payload, nil
}
}
}
func (p *packetConn) writePacket(payload []byte) error {
off := 0
for {
n := len(payload) - off
if n > 0xffffff {
n = 0xffffff
}
var hdr [4]byte
hdr[0] = byte(n)
hdr[1] = byte(n >> 8)
hdr[2] = byte(n >> 16)
hdr[3] = p.seq
p.seq++
if p.writeTO > 0 {
_ = p.c.SetWriteDeadline(time.Now().Add(p.writeTO))
}
if _, err := p.c.Write(hdr[:]); err != nil {
return err
}
if n > 0 {
if _, err := p.c.Write(payload[off : off+n]); err != nil {
return err
}
}
off += n
if n < 0xffffff {
return nil
}
}
}
func (p *packetConn) writeOK(affected, lastInsertID uint64, status, warnings uint16) error {
return p.writePacket(encodeOK(affected, lastInsertID, status, warnings))
}
func (p *packetConn) writeERR(errno uint16, msg string) error {
return p.writePacket(encodeERR(errno, msg))
}
func (p *packetConn) remote() string {
if p.c.RemoteAddr() == nil {
return ""
}
return p.c.RemoteAddr().String()
}
func (p *packetConn) close() error {
return p.c.Close()
}