network.go

package network

import (
    "bytes"
	"encoding/binary"
	"fmt"
	"io"
)

type Serializer interface {
	Serialize(out io.Writer) error
}

type Deserializer interface {
	Deserialize(in io.Reader) error
}

type MessageContainer interface {
	Serializer
	Deserializer
	MessageId() uint32
}

type MessageBuilder func() MessageContainer

type TypeContainer interface {
	TypeId() uint16
}

type TypeBuilder func() TypeContainer

var messages map[uint32]MessageBuilder

func init() {
	messages[1] = func() MessageBuilder {
		return new(ProtocolRequired)
	}
}

type ProtocolRequired struct {
	RequiredVersion int16
	CurrentVersion  int16
}

func (msg *ProtocolRequired) MessageId() int {
	return 1
}

func (msg *ProtocolRequired) Serialize(out io.Writer) error {
	binary.Write(out, binary.BigEndian, msg.RequiredVersion)
	binary.Write(out, binary.BigEndian, msg.CurrentVersion)
	return nil
}

func (msg *ProtocolRequired) Deserialize(in io.Reader) error {
	binary.Read(in, binary.BigEndian, &msg.RequiredVersion)
	binary.Read(in, binary.BigEndian, &msg.CurrentVersion)
	return nil
}

const (
	bitRightShiftLenPacketId = 2
	bitMask                  = 3
)

func writeStaticHeader(out io.Writer, typeLen int, id uint32) {
	header := (id << bitRightShiftLenPacketId) | typeLen
	binary.Write(out, binary.BigEndian, uint16(header))
}

func writeHeader(out io.Writer, len int, id uint32) {
	switch {
	case len > 65535:
		writeStaticHeader(out, 3, id)
		binary.Write(out, binary.BigEndian, uint8(len>>16&0xFF))
		binary.Write(out, binary.BigEndian, uint16(len&0xFFFF))
	case len > 255:
		writeStaticHeader(out, 2, id)
		binary.Write(out, binary.BigEndian, uint16(len))
	case len > 0:
		writeStaticHeader(out, 1, id)
		binary.Write(out, binary.BigEndian, uint8(len))
	default:
		writeStaticHeader(out, 0, id)
	}
}

func readHeader(in io.Reader) (id uint32, len int) {
	var header uint16
	binary.Read(in, binary.BigEndian, &header)

	id = uint32(header >> bitRightShiftLenPacketId)
	typeLen = header & bitMask

	switch typeLen {
	case 3:
		var tmp uint8

		binary.Read(in, binary.BigEndian, &tmp)
		len += tmp << 16
		binary.Read(in, binary.BigEndian, &tmp)
		len += tmp << 8
		binary.Read(in, binary.BigEndian, &tmp)
		len += tmp

	case 2:
		var tmp uint16
		binary.Read(in, binary.BigEndian, &tmp)
		len = tmp

	case 1:
		var tmp uint8
		binary.Read(in, binary.BigEndian, &tmp)
		len = tmp
	}

	return
}

func Pack(msg MessageContainer) *bytes.Buffer {
	var msgbuf bytes.Buffer
	msg.Serialize(&msgbuf)

	var buf bytes.Buffer
	writeHeader(buf, msgbuf.Len(), msg.MessageId())
	msgbuf.WriteTo(buf)

	return &buf
}

func Unpack(in io.Reader) MessageContainer {
	id, _ := readHeader(in)

	builder, ok := messages[id]
	if !ok {
		panic(fmt.Sprintf("unknown message %d", id))
	}

	msg := builder()
	msg.Deserialize(in)

	return msg
}