forkjo/vendor/github.com/denisenkom/go-mssqldb/buf.go

package mssql

import (
	"encoding/binary"
	"io"
	"errors"
)

type header struct {
	PacketType uint8
	Status     uint8
	Size       uint16
	Spid       uint16
	PacketNo   uint8
	Pad        uint8
}

type tdsBuffer struct {
	buf         []byte
	pos         uint16
	transport   io.ReadWriteCloser
	size        uint16
	final       bool
	packet_type uint8
	afterFirst  func()
}

func newTdsBuffer(bufsize int, transport io.ReadWriteCloser) *tdsBuffer {
	buf := make([]byte, bufsize)
	w := new(tdsBuffer)
	w.buf = buf
	w.pos = 8
	w.transport = transport
	w.size = 0
	return w
}

func (w *tdsBuffer) flush() (err error) {
	// writing packet size
	binary.BigEndian.PutUint16(w.buf[2:], w.pos)

	// writing packet into underlying transport
	if _, err = w.transport.Write(w.buf[:w.pos]); err != nil {
		return err
	}

	// execute afterFirst hook if it is set
	if w.afterFirst != nil {
		w.afterFirst()
		w.afterFirst = nil
	}

	w.pos = 8
	// packet number
	w.buf[6] += 1
	return nil
}

func (w *tdsBuffer) Write(p []byte) (total int, err error) {
	total = 0
	for {
		copied := copy(w.buf[w.pos:], p)
		w.pos += uint16(copied)
		total += copied
		if copied == len(p) {
			break
		}
		if err = w.flush(); err != nil {
			return
		}
		p = p[copied:]
	}
	return
}

func (w *tdsBuffer) WriteByte(b byte) error {
	if int(w.pos) == len(w.buf) {
		if err := w.flush(); err != nil {
			return err
		}
	}
	w.buf[w.pos] = b
	w.pos += 1
	return nil
}

func (w *tdsBuffer) BeginPacket(packet_type byte) {
	w.buf[0] = packet_type
	w.buf[1] = 0 // packet is incomplete
	w.buf[4] = 0 // spid
	w.buf[5] = 0
	w.buf[6] = 1 // packet id
	w.buf[7] = 0 // window
	w.pos = 8
}

func (w *tdsBuffer) FinishPacket() error {
	w.buf[1] = 1 // this is last packet
	return w.flush()
}

func (r *tdsBuffer) readNextPacket() error {
	header := header{}
	var err error
	err = binary.Read(r.transport, binary.BigEndian, &header)
	if err != nil {
		return err
	}
	offset := uint16(binary.Size(header))
	if int(header.Size) > len(r.buf) {
		return errors.New("Invalid packet size, it is longer than buffer size")
	}
	if int(offset) > int(header.Size) {
		return errors.New("Invalid packet size, it is shorter than header size")
	}
	_, err = io.ReadFull(r.transport, r.buf[offset:header.Size])
	if err != nil {
		return err
	}
	r.pos = offset
	r.size = header.Size
	r.final = header.Status != 0
	r.packet_type = header.PacketType
	return nil
}

func (r *tdsBuffer) BeginRead() (uint8, error) {
	err := r.readNextPacket()
	if err != nil {
		return 0, err
	}
	return r.packet_type, nil
}

func (r *tdsBuffer) ReadByte() (res byte, err error) {
	if r.pos == r.size {
		if r.final {
			return 0, io.EOF
		}
		err = r.readNextPacket()
		if err != nil {
			return 0, err
		}
	}
	res = r.buf[r.pos]
	r.pos++
	return res, nil
}

func (r *tdsBuffer) byte() byte {
	b, err := r.ReadByte()
	if err != nil {
		badStreamPanic(err)
	}
	return b
}

func (r *tdsBuffer) ReadFull(buf []byte) {
	_, err := io.ReadFull(r, buf[:])
	if err != nil {
		badStreamPanic(err)
	}
}

func (r *tdsBuffer) uint64() uint64 {
	var buf [8]byte
	r.ReadFull(buf[:])
	return binary.LittleEndian.Uint64(buf[:])
}

func (r *tdsBuffer) int32() int32 {
	return int32(r.uint32())
}

func (r *tdsBuffer) uint32() uint32 {
	var buf [4]byte
	r.ReadFull(buf[:])
	return binary.LittleEndian.Uint32(buf[:])
}

func (r *tdsBuffer) uint16() uint16 {
	var buf [2]byte
	r.ReadFull(buf[:])
	return binary.LittleEndian.Uint16(buf[:])
}

func (r *tdsBuffer) BVarChar() string {
	l := int(r.byte())
	return r.readUcs2(l)
}

func (r *tdsBuffer) UsVarChar() string {
	l := int(r.uint16())
	return r.readUcs2(l)
}

func (r *tdsBuffer) readUcs2(numchars int) string {
	b := make([]byte, numchars*2)
	r.ReadFull(b)
	res, err := ucs22str(b)
	if err != nil {
		badStreamPanic(err)
	}
	return res
}

func (r *tdsBuffer) Read(buf []byte) (copied int, err error) {
	copied = 0
	err = nil
	if r.pos == r.size {
		if r.final {
			return 0, io.EOF
		}
		err = r.readNextPacket()
		if err != nil {
			return
		}
	}
	copied = copy(buf, r.buf[r.pos:r.size])
	r.pos += uint16(copied)
	return
}