major refactor

tcp-connection.go

@@ -15,12 +15,23 @@ type Connection struct {
 }
 
 type FrameType uint32
+type StatusCode uint32
+type CheckSum uint32
 
 type Frame struct {
-	Id         uint64
 	Type       FrameType
-	StatusCode uint32
+	StatusCode StatusCode
 	Length     uint32
+	Checksum   CheckSum
+}
+
+func (f *Frame) IsValid() bool {
+	return f.Checksum == MakeChecksum(f.Type, f.StatusCode, f.Length)
+}
+
+func MakeChecksum(msg FrameType, statusCode StatusCode, length uint32) CheckSum {
+	sum := CheckSum((uint32(msg) + uint32(statusCode) + length) / 8)
+	return sum
 }
 
 type FrameWithPayload struct {
@@ -28,6 +39,19 @@ type FrameWithPayload struct {
 	Payload []byte
 }
 
+func MakeFrameWithPayload(msg FrameType, statusCode StatusCode, payload []byte) FrameWithPayload {
+	len := uint32(len(payload))
+	return FrameWithPayload{
+		Frame: Frame{
+			Type:       msg,
+			StatusCode: 0,
+			Length:     len,
+			Checksum:   MakeChecksum(msg, 0, len),
+		},
+		Payload: payload,
+	}
+}
+
 type FrameData interface {
 	ToBytes() []byte
 	FromBytes([]byte) error
@@ -41,11 +65,7 @@ func NewConnection(address string) *Connection {
 }
 
 func SendFrame(conn net.Conn, data *FrameWithPayload) error {
-	_, err := conn.Write(header[:])
-	if err != nil {
-		return err
-	}
-	err = binary.Write(conn, binary.LittleEndian, data.Frame)
+	err := binary.Write(conn, binary.LittleEndian, data.Frame)
 	if err != nil {
 		return err
 	}
@@ -53,68 +73,67 @@ func SendFrame(conn net.Conn, data *FrameWithPayload) error {
 	return err
 }
 
-func (c *Connection) CallAsync(msg FrameType, data FrameData, ch chan<- *FrameWithPayload) error {
+func (c *Connection) CallAsync(msg FrameType, payload []byte, ch chan<- FrameWithPayload) (net.Conn, error) {
 	conn, err := net.Dial("tcp", c.address)
 	go WaitForFrame(conn, ch)
 	if err != nil {
-		return err
-	}
-	payload := data.ToBytes()
-	toSend := &FrameWithPayload{
-		Frame: Frame{
-			Id:         c.count,
-			Type:       msg,
-			StatusCode: 0,
-			Length:     uint32(len(payload)),
-		},
-		Payload: payload,
+		return conn, err
 	}
+	toSend := MakeFrameWithPayload(msg, 1, payload)
 
-	err = SendFrame(conn, toSend)
+	err = SendFrame(conn, &toSend)
 	if err != nil {
+		conn.Close()
 		close(ch)
-		return err
+		return nil, err
 	}
 
 	c.count++
-	return nil
+	return conn, nil
 }
 
-func (c *Connection) Call(msg FrameType, data FrameData) (*FrameWithPayload, error) {
-	ch := make(chan *FrameWithPayload, 1)
-	c.CallAsync(msg, data, ch)
+func (c *Connection) Call(msg FrameType, data []byte) (*FrameWithPayload, error) {
+	ch := make(chan FrameWithPayload, 1)
+	conn, err := c.CallAsync(msg, data, ch)
+	if err != nil {
+		return nil, err
+	}
+	defer conn.Close()
 	select {
 	case ret := <-ch:
-		return ret, nil
-	case <-time.After(5 * time.Second):
+		return &ret, nil
+	case <-time.After(MaxCallDuration):
 		return nil, fmt.Errorf("timeout")
 	}
 }
 
-func WaitForFrame(conn net.Conn, resultChan chan<- *FrameWithPayload) error {
-	defer conn.Close()
+func WaitForFrame(conn net.Conn, resultChan chan<- FrameWithPayload) error {
 	var err error
+	var frame Frame
 	r := bufio.NewReader(conn)
-	h := make([]byte, 4)
-	r.Read(h)
-	if h[0] == header[0] && h[1] == header[1] && h[2] == header[2] && h[3] == header[3] {
-		frame := Frame{}
-		err = binary.Read(r, binary.LittleEndian, &frame)
+
+	err = binary.Read(r, binary.LittleEndian, &frame)
+	if err != nil {
+		return err
+	}
+	if frame.IsValid() {
 		payload := make([]byte, frame.Length)
 		_, err = r.Read(payload)
-		resultChan <- &FrameWithPayload{
+		if err != nil {
+			return err
+		}
+		resultChan <- FrameWithPayload{
 			Frame:   frame,
 			Payload: payload,
 		}
 		return err
 	}
-	resultChan <- nil
-	return err
+	return fmt.Errorf("checksum mismatch")
 }
 
 type GenericListener struct {
 	Closed   bool
-	handlers map[FrameType]func(*FrameWithPayload, chan<- *FrameWithPayload) error
+	handlers map[FrameType]func(*FrameWithPayload, chan<- FrameWithPayload) error
 }
 
 func (c *Connection) Listen() (*GenericListener, error) {
@@ -123,7 +142,7 @@ func (c *Connection) Listen() (*GenericListener, error) {
 		return nil, err
 	}
 	ret := &GenericListener{
-		handlers: make(map[FrameType]func(*FrameWithPayload, chan<- *FrameWithPayload) error),
+		handlers: make(map[FrameType]func(*FrameWithPayload, chan<- FrameWithPayload) error),
 	}
 	go func() {
 		for !ret.Closed {
@@ -137,36 +156,44 @@ func (c *Connection) Listen() (*GenericListener, error) {
 	return ret, nil
 }
 
+const (
+	MaxCallDuration = 500 * time.Millisecond
+)
+
 func (l *GenericListener) HandleConnection(conn net.Conn) {
-	ch := make(chan *FrameWithPayload, 1)
+	ch := make(chan FrameWithPayload, 1)
 	go WaitForFrame(conn, ch)
 	select {
 	case frame := <-ch:
-		go l.HandleFrame(conn, frame)
-	case <-time.After(1 * time.Second):
+		go l.HandleFrame(conn, &frame)
+	case <-time.After(MaxCallDuration):
 		close(ch)
 		log.Printf("Timeout waiting for frame\n")
 	}
 }
 
-func (l *GenericListener) AddHandler(msg FrameType, handler func(*FrameWithPayload, chan<- *FrameWithPayload) error) {
+func (l *GenericListener) AddHandler(msg FrameType, handler func(*FrameWithPayload, chan<- FrameWithPayload) error) {
 	l.handlers[msg] = handler
 }
 
 func (l *GenericListener) HandleFrame(conn net.Conn, frame *FrameWithPayload) {
 	handler, ok := l.handlers[frame.Type]
-	defer conn.Close()
 	if ok {
 		go func() {
-			resultChan := make(chan *FrameWithPayload, 1)
+			resultChan := make(chan FrameWithPayload, 1)
 			defer close(resultChan)
 			err := handler(frame, resultChan)
 			if err != nil {
 				log.Fatalf("Error handling frame: %v\n", err)
 			}
-			SendFrame(conn, <-resultChan)
+			result := <-resultChan
+			err = SendFrame(conn, &result)
+			if err != nil {
+				log.Fatalf("Error sending frame: %v\n", err)
+			}
 		}()
 	} else {
+		conn.Close()
 		log.Fatalf("No handler for frame type %d\n", frame.Type)
 	}
 }