package util
import (
"bufio"
"bytes"
"encoding/binary"
"errors"
"io"
"net"
"time"
)
type Packet interface {
Serialize() []byte
IsNull() bool
IsPing() bool
}
type MessagePacket struct {
data string
isPing bool
}
var errClosed = errors.New("conn is closed")
func (m *MessagePacket) Serialize() []byte {
return []byte(m.data)
}
func (m *MessagePacket) IsNull() bool {
return len(m.data) == 0 && !m.isPing
}
func (m *MessagePacket) IsPing() bool {
return m.isPing
}
type Protocol interface {
SetConn(conn *net.TCPConn)
ReadPacket() (Packet, error)
}
type MessageProtocol struct {
conn *net.TCPConn
reader *bufio.Reader
cache *bytes.Buffer
cacheSize int64
}
func (m *MessageProtocol) SetConn(conn *net.TCPConn) {
m.conn = conn
m.reader = bufio.NewReader(conn)
m.cache = bytes.NewBuffer(nil)
}
func (m *MessageProtocol) ReadPacket() (Packet, error) {
if m.reader != nil {
message, err := m.Decode()
if err != nil {
return nil, err
}
if m.isPing(message) {
return &MessagePacket{isPing: true}, nil
}
return &MessagePacket{data: message}, nil
}
return nil, errClosed
}
func (m *MessageProtocol) isPing(s string) bool {
return s == "0x00ping"
}
const maxConsecutiveEmptyReads = 100
func (m *MessageProtocol) Decode() (string, error) {
lengthByte, err := m.reader.Peek(4)
if err != nil {
return "", err
}
lengthBuff := bytes.NewBuffer(lengthByte)
var length int32
err = binary.Read(lengthBuff, binary.LittleEndian, &length)
if err != nil {
return "", err
}
if length == 0 {
return "", errClosed
}
if int32(m.reader.Buffered()) < length+4 {
var retry = 0
for m.cacheSize < int64(length+4) {
readSize := int64(length+4) - m.cacheSize
if readSize > int64(m.reader.Buffered()) {
readSize = int64(m.reader.Buffered())
}
buffer := make([]byte, readSize)
size, err := m.reader.Read(buffer)
if err != nil {
return "", err
}
if size <= 0 {
retry++
if retry > maxConsecutiveEmptyReads {
return "", io.ErrNoProgress
}
time.Sleep(time.Millisecond * 10)
} else {
m.cacheSize += int64(size)
m.cache.Write(buffer)
}
}
result := m.cache.Bytes()[4:]
m.cache.Reset()
m.cacheSize = 0
return string(result), nil
}
pack := make([]byte, int(4+length))
size, err := m.reader.Read(pack)
if err != nil {
return "", err
}
if size == 0 {
return "", io.ErrNoProgress
}
return string(pack[4:]), nil
}