rest/rest_websocket/message.go

package rest_websocket

import (
	"bytes"
	"io"
	"time"

	"git.ali33.ru/fcg-xvii/go-tools/json"
	"git.ali33.ru/fcg-xvii/rest"
)

// Int64ToBytes упаковывает int64 в срез байтов заданной длины
func Int64ToBytes(num int64, byteCount int) []byte {
	bytes := make([]byte, byteCount)
	for i := 0; i < byteCount; i++ {
		shift := uint((byteCount - 1 - i) * 8)
		bytes[i] = byte(num >> shift)
	}
	return bytes
}

// BytesToInt64 конвертирует срез байтов в int64
func BytesToInt64(bytes []byte) int64 {
	var num int64
	for _, b := range bytes {
		num = (num << 8) | int64(b)
	}
	return num
}

func ioError(field string, err error) rest.IErrorArgs {
	return rest.NewError(
		"ErrIO",
		json.Map{
			"field": field,
			"error": err.Error(),
		},
	)
}

func ReadMessage(r io.Reader) (*Message, rest.IErrorArgs) {
	// todo
	// id
	sType := make([]byte, 1)
	if _, err := r.Read(sType); err != nil {
		return nil, ioError("type", err)
	}
	mes := Message{
		mType: rest.RequestType(sType[0]),
	}
	if mes.mType == rest.RequestTypeMessage || mes.mType == rest.RequestTypeAnswer {
		sID := make([]byte, 2)
		if _, err := r.Read(sID); err != nil {
			return nil, ioError("id", err)
		}
		mes.id = BytesToInt64(sID)
	}
	if mes.mType == rest.RequestTypeMessage {
		sTimeout := make([]byte, 8)
		if _, err := r.Read(sTimeout); err != nil {
			return nil, ioError("timeout", err)
		}
		mes.timeout = time.Unix(BytesToInt64(sTimeout), 0)
	}
	sCommandSize := make([]byte, 2)
	if _, err := r.Read(sCommandSize); err != nil {
		return nil, ioError("data_size", err)
	}
	if BytesToInt64(sCommandSize) > 0 {
		sCommand := make([]byte, BytesToInt64(sCommandSize))
		if _, err := r.Read(sCommand); err != nil {
			return nil, ioError("data", err)
		}
		mes.command = string(sCommand)
	}
	sDataSize := make([]byte, 8)
	if _, err := r.Read(sDataSize); err != nil {
		return nil, ioError("data_size", err)
	}
	sData := make([]byte, BytesToInt64(sDataSize))
	if _, err := r.Read(sData); err != nil {
		return nil, ioError("data", err)
	}
	if len(sData) > 0 {
		if err := json.Unmarshal(sData, &mes.data); err != nil {
			return nil, ioError("data", err)
		}
	}
	sFilesCount := make([]byte, 2)
	if _, err := r.Read(sFilesCount); err != nil {
		return nil, ioError("files_count", err)
	}
	filesCount := BytesToInt64(sFilesCount)
	files := make(map[string]rest.IReadCloserLen)
	for i := 0; i < int(filesCount); i++ {
		sFileNameLen := make([]byte, 2)
		if _, err := r.Read(sFileNameLen); err != nil {
			return nil, ioError("file_name_length", err)
		}
		sFileName := make([]byte, BytesToInt64(sFileNameLen))
		if _, err := r.Read(sFileName); err != nil {
			return nil, ioError("file_name", err)
		}
		sFileSize := make([]byte, 8)
		if _, err := r.Read(sFileSize); err != nil {
			return nil, ioError("file_size", err)
		}
		fileSize := BytesToInt64(sFileSize)
		if fileSize < 1024*1024 {
			// RAM buffer
			sFileData := make([]byte, fileSize)
			if _, err := r.Read(sFileData); err != nil {
				return nil, ioError("file_data", err)
			}
			buf := rest.NewReadCloserLen(
				io.NopCloser(bytes.NewBuffer(sFileData)),
				int64(len(sFileData)),
			)
			files[string(sFileName)] = buf
		} else {
			// temporary file
			tmpF, err := rest.NewTemporaryFile(fileSize, r)
			if err != nil {
				return nil, err
			}
			files[string(sFileName)] = tmpF
		}
	}
	return &mes, nil
}

func NewMessage(command string, data json.Map, files map[string]rest.IReadCloserLen, timeout time.Duration, mType rest.RequestType) *Message {
	return &Message{
		command: command,
		data:    data,
		files:   files,
		mType:   mType,
		timeout: time.Now().Add(timeout),
	}
}

type Message struct {
	id      int64
	command string
	mType   rest.RequestType
	timeout time.Time
	data    json.Map
	files   map[string]rest.IReadCloserLen
	owner   *WebSocket
}

func (s *Message) Data() json.Map {
	return s.data
}

func (s *Message) File(name string) (rest.IReadCloserLen, bool) {
	file, check := s.files[name]
	return file, check
}

func (s *Message) FileKeys() []string {
	keys := make([]string, 0, len(s.files))
	for k := range s.files {
		keys = append(keys, k)
	}
	return keys
}

func (s *Message) Command() string {
	return s.command
}

func (s *Message) IsBinagy() bool {
	return len(s.files) > 0
}

func (s *Message) Write(w io.Writer) rest.IErrorArgs {
	// id
	if _, err := w.Write(Int64ToBytes(s.id, 2)); err != nil {
		return ioError("id", err)
	}
	// type
	if _, err := w.Write(Int64ToBytes(int64(s.mType), 1)); err != nil {
		return ioError("type", err)
	}
	// timeout
	if _, err := w.Write(Int64ToBytes(s.timeout.Unix(), 8)); err != nil {
		return ioError("timeout", err)
	}
	// command length
	if _, err := w.Write(Int64ToBytes(int64(len(s.command)), 2)); err != nil {
		return ioError("command_length", err)
	}
	// command
	if len(s.command) > 0 {
		if _, err := w.Write([]byte(s.command)); err != nil {
			return ioError("command", err)
		}
	}
	// data
	data := s.data.JSON()
	// data size
	if _, err := w.Write(Int64ToBytes(int64(len(data)), 8)); err != nil {
		return ioError("data_size", err)
	}
	// data body
	if _, err := w.Write(data); err != nil {
		return ioError("data_body", err)
	}
	// files count
	filesCount := int64(len(s.files))
	if _, err := w.Write(Int64ToBytes(filesCount, 2)); err != nil {
		return ioError("files_count", err)
	}
	// files
	for name, file := range s.files {
		// file name size
		fileNameSize := int64(len(name))
		if _, err := w.Write(Int64ToBytes(fileNameSize, 2)); err != nil {
			return ioError("file_name_size", err)
		}
		// file name
		if _, err := w.Write([]byte(name)); err != nil {
			return ioError("file_name", err)
		}
		// file body size
		if _, err := w.Write(Int64ToBytes(file.Len(), 8)); err != nil {
			return ioError("file_body_size", err)
		}
		// file body
		if _, err := io.Copy(w, file); err != nil {
			return ioError("file_body", err)
		}
	}
	return nil
}

func (s *Message) Close() {
	for _, file := range s.files {
		file.Close()
	}
}