more features

message.go

@@ -3,12 +3,8 @@ package main
 import (
 	"bytes"
 	"encoding/binary"
-	"fmt"
 	"io"
 	"time"
-
-	messages "git.tornberg.me/go-cart-actor/proto"
-	"google.golang.org/protobuf/proto"
 )
 
 type StorableMessage interface {
@@ -43,21 +39,12 @@ func GetData(fn func(w io.Writer) error) ([]byte, error) {
 }
 
 func (m Message) Write(w io.Writer) error {
-	data, err := GetData(func(wr io.Writer) error {
-		if m.Type == AddRequestType {
-			messageBytes, err := proto.Marshal(m.Content.(*messages.AddRequest))
-			if err != nil {
-				return err
-			}
-			wr.Write(messageBytes)
-		} else if m.Type == AddItemType {
-			messageBytes, err := proto.Marshal(m.Content.(*messages.AddItem))
-			if err != nil {
-				return err
-			}
-			wr.Write(messageBytes)
-		}
-		return nil
+	h, err := GetMessageHandler(m.Type)
+	if err != nil {
+		return err
+	}
+	data, err := GetData(func(w io.Writer) error {
+		return h.Write(&m, w)
 	})
 	if err != nil {
 		return err
@@ -79,8 +66,12 @@ func (m Message) Write(w io.Writer) error {
 }
 
 func ReadMessage(reader io.Reader, m *Message) error {
+	h, err := GetMessageHandler(m.Type)
+	if err != nil {
+		return err
+	}
 	header := StorableMessageHeader{}
-	err := binary.Read(reader, binary.LittleEndian, &header)
+	err = binary.Read(reader, binary.LittleEndian, &header)
 	if err != nil {
 		return err
 	}
@@ -89,21 +80,13 @@ func ReadMessage(reader io.Reader, m *Message) error {
 	if err != nil {
 		return err
 	}
-	switch header.Type {
-	case AddRequestType:
-		msg := &messages.AddRequest{}
-		err = proto.Unmarshal(messageBytes, msg)
-		m.Content = msg
-	case AddItemType:
-		msg := &messages.AddItem{}
-		err = proto.Unmarshal(messageBytes, msg)
-		m.Content = msg
-	default:
-		return fmt.Errorf("unknown message type")
-	}
+
+	content, err := h.Read(messageBytes)
 	if err != nil {
 		return err
 	}
+	m.Content = content
+
 	m.Type = header.Type
 	m.TimeStamp = &header.TimeStamp