158 lines
3.7 KiB
Go
158 lines
3.7 KiB
Go
// Copyright 2013 The Go-MySQL-Driver Authors. All rights reserved.
|
|
//
|
|
// This Source Code Form is subject to the terms of the Mozilla Public
|
|
// License, v. 2.0. If a copy of the MPL was not distributed with this file,
|
|
// You can obtain one at http://mozilla.org/MPL/2.0/.
|
|
|
|
// The MIT License (MIT)
|
|
//
|
|
// Copyright (c) 2014 wandoulabs
|
|
// Copyright (c) 2014 siddontang
|
|
//
|
|
// Permission is hereby granted, free of charge, to any person obtaining a copy of
|
|
// this software and associated documentation files (the "Software"), to deal in
|
|
// the Software without restriction, including without limitation the rights to
|
|
// use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of
|
|
// the Software, and to permit persons to whom the Software is furnished to do so,
|
|
// subject to the following conditions:
|
|
//
|
|
// The above copyright notice and this permission notice shall be included in all
|
|
// copies or substantial portions of the Software.
|
|
|
|
// Copyright 2015 PingCAP, Inc.
|
|
//
|
|
// Licensed under the Apache License, Version 2.0 (the "License");
|
|
// you may not use this file except in compliance with the License.
|
|
// You may obtain a copy of the License at
|
|
//
|
|
// http://www.apache.org/licenses/LICENSE-2.0
|
|
//
|
|
// Unless required by applicable law or agreed to in writing, software
|
|
// distributed under the License is distributed on an "AS IS" BASIS,
|
|
// See the License for the specific language governing permissions and
|
|
// limitations under the License.
|
|
|
|
package server
|
|
|
|
import (
|
|
"bufio"
|
|
"io"
|
|
"net"
|
|
|
|
"github.com/juju/errors"
|
|
"github.com/pingcap/tidb/mysql"
|
|
)
|
|
|
|
const (
|
|
defaultReaderSize = 16 * 1024
|
|
defaultWriterSize = 16 * 1024
|
|
)
|
|
|
|
// packetIO is a helper to read and write data in packet format.
|
|
type packetIO struct {
|
|
rb *bufio.Reader
|
|
wb *bufio.Writer
|
|
|
|
sequence uint8
|
|
}
|
|
|
|
func newPacketIO(conn net.Conn) *packetIO {
|
|
p := &packetIO{
|
|
rb: bufio.NewReaderSize(conn, defaultReaderSize),
|
|
wb: bufio.NewWriterSize(conn, defaultWriterSize),
|
|
}
|
|
|
|
return p
|
|
}
|
|
|
|
func (p *packetIO) readOnePacket() ([]byte, error) {
|
|
var header [4]byte
|
|
|
|
if _, err := io.ReadFull(p.rb, header[:]); err != nil {
|
|
return nil, errors.Trace(err)
|
|
}
|
|
|
|
sequence := uint8(header[3])
|
|
if sequence != p.sequence {
|
|
return nil, errInvalidSequence.Gen("invalid sequence %d != %d", sequence, p.sequence)
|
|
}
|
|
|
|
p.sequence++
|
|
|
|
length := int(uint32(header[0]) | uint32(header[1])<<8 | uint32(header[2])<<16)
|
|
|
|
data := make([]byte, length)
|
|
if _, err := io.ReadFull(p.rb, data); err != nil {
|
|
return nil, errors.Trace(err)
|
|
}
|
|
return data, nil
|
|
}
|
|
|
|
func (p *packetIO) readPacket() ([]byte, error) {
|
|
data, err := p.readOnePacket()
|
|
if err != nil {
|
|
return nil, errors.Trace(err)
|
|
}
|
|
|
|
if len(data) < mysql.MaxPayloadLen {
|
|
return data, nil
|
|
}
|
|
|
|
// handle muliti-packet
|
|
for {
|
|
buf, err := p.readOnePacket()
|
|
if err != nil {
|
|
return nil, errors.Trace(err)
|
|
}
|
|
|
|
data = append(data, buf...)
|
|
|
|
if len(buf) < mysql.MaxPayloadLen {
|
|
break
|
|
}
|
|
}
|
|
|
|
return data, nil
|
|
}
|
|
|
|
// writePacket writes data that already have header
|
|
func (p *packetIO) writePacket(data []byte) error {
|
|
length := len(data) - 4
|
|
|
|
for length >= mysql.MaxPayloadLen {
|
|
data[0] = 0xff
|
|
data[1] = 0xff
|
|
data[2] = 0xff
|
|
|
|
data[3] = p.sequence
|
|
|
|
if n, err := p.wb.Write(data[:4+mysql.MaxPayloadLen]); err != nil {
|
|
return mysql.ErrBadConn
|
|
} else if n != (4 + mysql.MaxPayloadLen) {
|
|
return mysql.ErrBadConn
|
|
} else {
|
|
p.sequence++
|
|
length -= mysql.MaxPayloadLen
|
|
data = data[mysql.MaxPayloadLen:]
|
|
}
|
|
}
|
|
|
|
data[0] = byte(length)
|
|
data[1] = byte(length >> 8)
|
|
data[2] = byte(length >> 16)
|
|
data[3] = p.sequence
|
|
|
|
if n, err := p.wb.Write(data); err != nil {
|
|
return errors.Trace(mysql.ErrBadConn)
|
|
} else if n != len(data) {
|
|
return errors.Trace(mysql.ErrBadConn)
|
|
} else {
|
|
p.sequence++
|
|
return nil
|
|
}
|
|
}
|
|
|
|
func (p *packetIO) flush() error {
|
|
return p.wb.Flush()
|
|
}
|