1 files changed, 117 insertions, 75 deletions

diff --git a/core/data.go b/core/data.go
index ddf1db8..16775db 100644
--- a/core/data.go
+++ b/core/data.go
@@ -1,122 +1,164 @@
 package core
 
 import (
+	"bytes"
 	"encoding/binary"
 	"fmt"
 	"io"
-	"strings"
 	"time"
 )
 
-type Marshaler interface {
-	Marshal() []any
+type Wrapper interface {
+	Wrap() (PayloadType, []any)
 }
 
-type DataType uint8
+type PayloadType uint8
 
 const (
-	TypeUnknown   DataType = 0x00
-	TypeHandshake          = 0x01
-	TypePing               = 0x02
-	TypePong               = 0x03
-	TypeTest               = 0xFF
+	PayloadUnknown   PayloadType = 0x00
+	PayloadHandshake             = 0x01
+	PayloadPing                  = 0x02
+	PayloadPong                  = 0x03
+	PayloadMessage               = 0x04
+	PayloadTest                  = 0xFF
 )
 
-func ReadDataType(r io.Reader) (DataType, error) {
-	var data uint8
-	if err := read(r, &data); err != nil {
-		return TypeUnknown, err
-	}
-
-	dType := DataType(data)
-
-	return dType, nil
+type Frame struct {
+	Length  uint16
+	Type    PayloadType
+	Payload []any
 }
 
-type Handshake struct {
-	Version uint8
+type Binary struct {
+	Length  uint16
+	Payload []byte
 }
 
-func (t Handshake) Marshal() []any {
-	return []any{t.Version}
-}
+func Encode(w Wrapper) ([]byte, error) {
+	payloadType, payload := w.Wrap()
+	buf := bytes.NewBuffer(make([]byte, 3))
+
+	for _, p := range payload {
+		switch v := p.(type) {
+		case string:
+			err := binary.Write(buf, binary.BigEndian, uint16(len(v)))
+			if err != nil {
+				return []byte{}, fmt.Errorf("cannot encode string length: %v", err)
+			}
+
+			_, err = buf.WriteString(v)
+			if err != nil {
+				return []byte{}, fmt.Errorf("cannot encode string: %v", err)
+			}
+
+		case []byte:
+			err := binary.Write(buf, binary.BigEndian, uint16(len(v)))
+			if err != nil {
+				return []byte{}, fmt.Errorf("cannot encode byte slice length: %v", err)
+			}
+
+			_, err = buf.Write(v)
+			if err != nil {
+				return []byte{}, fmt.Errorf("cannot encode byte slice: %v", err)
+			}
+
+		case time.Time:
+			err := binary.Write(buf, binary.BigEndian, uint64(v.Unix()))
+			if err != nil {
+				return []byte{}, fmt.Errorf("cannot encode time: %v", err)
+			}
+
+		default:
+			err := binary.Write(buf, binary.BigEndian, v)
+			if err != nil {
+				return []byte{}, fmt.Errorf("cannot encode: %v", err)
+			}
+		}
+	}
 
-func ReadHandshake(r io.Reader) (Handshake, error) {
-	var t Handshake
-	err := read(r, &t.Version)
-	return t, err
-}
+	frame := buf.Bytes()
+	frame[0] = uint8(payloadType)
+	binary.BigEndian.PutUint16(frame[1:], uint16(len(frame)-3))
 
-type Test struct {
-	Message string
+	return frame, nil
 }
 
-func (t Test) Marshal() []any {
-	return []any{append([]byte(t.Message), 0x0)}
-}
+func Decode(buf io.Reader) (any, error) {
+	var pTypeByte uint8
 
-func ReadTest(r io.Reader) (Test, error) {
-	var t Test
-	err := readString(r, &t.Message)
-	return t, err
-}
+	err := binary.Read(buf, binary.BigEndian, &pTypeByte)
+	if err != nil {
+		return nil, fmt.Errorf("cannot read payload type: %v", err)
+	}
 
-type Ping struct{}
+	pType := PayloadType(pTypeByte)
 
-func (t Ping) Marshal() []any {
-	return []any{}
-}
+	var pLen uint16
+	err = binary.Read(buf, binary.BigEndian, &pLen)
+	if err != nil {
+		return nil, fmt.Errorf("cannot read payload length: %v", err)
+	}
 
-type Pong struct {
-	Timestamp time.Time
+	switch pType {
+	case PayloadHandshake:
+		return DecodeHandshake(buf)
+	case PayloadPing:
+	case PayloadPong:
+	case PayloadMessage:
+	case PayloadTest:
+		return DecodeTest(buf)
+	default:
+		return nil, fmt.Errorf("invalid payload type: %v", pType)
+	}
+
+	return pType, nil
 }
 
-func (t Pong) Marshal() []any {
-	return []any{uint64(t.Timestamp.Unix())}
+func decodeNumeric(buf io.Reader, ptr any) error {
+	err := binary.Read(buf, binary.BigEndian, ptr)
+	if err != nil {
+		return fmt.Errorf("cannot decode: %v", err)
+	}
+	return nil
 }
 
-func ReadPong(r io.Reader) (Pong, error) {
-	var (
-		t         Pong
-		timestamp uint64
-	)
+func decodeTime(buf io.Reader, ptr *time.Time) error {
+	var timeBuf uint64
+	err := binary.Read(buf, binary.BigEndian, &timeBuf)
+	if err != nil {
+		return fmt.Errorf("cannot decode time: %v", err)
+	}
 
-	err := read(r, &timestamp)
-	t.Timestamp = time.Unix(int64(timestamp), 0)
-	return t, err
+	*ptr = time.Unix(int64(timeBuf), 0)
+	return nil
 }
 
-//
+func decodeBin(buf io.Reader, ptr *[]byte) error {
+	var payloadLen uint16
 
-func readString(r io.Reader, ptr *string) error {
-	var (
-		sb  strings.Builder
-		buf byte
-	)
+	err := binary.Read(buf, binary.BigEndian, &payloadLen)
+	if err != nil {
+		return fmt.Errorf("cannot decode payload length: %v", err)
+	}
 
-	for {
-		if err := read(r, &buf); err != nil {
-			return err
-		}
+	payload := make([]byte, payloadLen)
 
-		if buf == 0x0 {
-			break
-		}
-
-		if err := sb.WriteByte(buf); err != nil {
-			fmt.Errorf("cannot write byte to string buffer: %v", err)
-		}
+	_, err = buf.Read(payload)
+	if err != nil {
+		return fmt.Errorf("cannot decode payload: %v", err)
 	}
 
-	*ptr = sb.String()
+	*ptr = payload
 	return nil
 }
 
-func read(r io.Reader, ptr any) error {
-	err := binary.Read(r, binary.BigEndian, ptr)
+func decodeString(buf io.Reader, ptr *string) error {
+	var strBytes []byte
+	err := decodeBin(buf, &strBytes)
 	if err != nil {
-		return fmt.Errorf("cannot read: %v", err)
+		return err
 	}
 
+	*ptr = string(strBytes)
 	return nil
 }