52 lines
906 B
Go
52 lines
906 B
Go
package tcp
|
|
|
|
import (
|
|
"encoding/binary"
|
|
"io"
|
|
"net"
|
|
|
|
"go.uber.org/zap"
|
|
)
|
|
|
|
const (
|
|
LengthPrefixSize = 4
|
|
)
|
|
|
|
func readFrame(conn net.Conn) ([]byte, error) {
|
|
lengthBytes := make([]byte, LengthPrefixSize)
|
|
_, err := io.ReadFull(conn, lengthBytes)
|
|
if err != nil {
|
|
return []byte{}, err
|
|
}
|
|
|
|
messageLength := binary.BigEndian.Uint32(lengthBytes)
|
|
|
|
if messageLength > 65536 {
|
|
zap.L().Warn("message length exceed",
|
|
zap.Uint32("length", messageLength))
|
|
}
|
|
|
|
payload := make([]byte, messageLength)
|
|
_, err = io.ReadFull(conn, payload)
|
|
if err != nil {
|
|
return []byte{}, err
|
|
}
|
|
|
|
return payload, nil
|
|
}
|
|
|
|
func sendFrame(conn net.Conn, data []byte) error {
|
|
messageLength := uint32(len(data))
|
|
|
|
lengthBytes := make([]byte, LengthPrefixSize)
|
|
binary.BigEndian.PutUint32(lengthBytes, messageLength)
|
|
|
|
_, err := conn.Write(lengthBytes)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
|
|
_, err = conn.Write(data)
|
|
return err
|
|
}
|