Commit 2ed3f9b3 authored by Anders Jensen Løvig's avatar Anders Jensen Løvig
Browse files

Refactor network message

parent c5f97d4d
Pipeline #21751 passed with stages
in 1 minute
......@@ -134,7 +134,7 @@ func main() {
}
func closeCallback(reason election.Reason) {
msgClosed := network.NewMessage(election.MsgClosing, nil)
msgClosed := network.NewMessage(election.MsgClosing)
server.Server.Broadcast(msgClosed)
server.Done <- true
......
......@@ -30,7 +30,7 @@ func (e *Election) OnBallot(conn *network.Conn, data network.MessageData) {
voterID := UniqueID(conn.SerialNumber())
go e.handleBallot(voterID, &ballot)
err = conn.WriteMessage(network.NewMessage(MsgBallotAck, nil))
err = conn.WriteMessage(network.NewMessage(MsgBallotAck))
if err != nil {
log.Println("Election: error:", err)
}
......
package network
import (
"bufio"
"encoding/binary"
"encoding/json"
"io"
"log"
......@@ -30,32 +28,20 @@ type Message struct {
}
// NewMessage returns a message with given command and data v
func NewMessage(command string, v interface{}) *Message {
return &Message{
Command: command,
Data: marshal(v),
}
}
func readMessage(reader *bufio.Reader) (*Message, error) {
sizeBuffer := make([]byte, 4)
_, err := io.ReadFull(reader, sizeBuffer)
if err != nil {
return nil, err
func NewMessage(command string, v ...interface{}) *Message {
var data MessageData
if len(v) == 0 {
data = marshal(nil)
} else if len(v) == 1 {
data = marshal(v[0])
} else {
log.Panic("v must have size 1, was", len(v))
}
size := binary.BigEndian.Uint32(sizeBuffer)
data := make([]byte, size)
_, err = io.ReadFull(reader, data)
if err != nil {
return nil, err
return &Message{
Command: command,
Data: data,
}
message := new(Message)
err = unmarshal(data, message)
return message, err
}
func marshal(v interface{}) []byte {
......@@ -75,10 +61,9 @@ func unmarshal(data []byte, v interface{}) error {
}
func messageLoop(conn *Conn) {
reader := bufio.NewReader(conn)
for {
// First wait for a message and check the returned error.
message, err := readMessage(reader)
message, err := conn.ReadMessage()
if err == io.EOF {
log.Printf("%s: connection closed: %s\n", logName, conn.RemoteAddr())
break
......
Supports Markdown
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment