Fix decoding messages of identical types in stream

Change-Id: Id6f6bc3de5b9ff7230108a4171810207c94f9a38
Signed-off-by: Ondrej Fabry <ofabry@cisco.com>

1 files changed, 24 insertions, 17 deletions

diff --git a/core/stream.go b/core/stream.go
index 171b201..61a9965 100644
--- a/core/stream.go
+++ b/core/stream.go
@@ -69,15 +69,14 @@ func (c *Connection) Invoke(ctx context.Context, req api.Message, reply api.Mess
 	if err := stream.SendMsg(req); err != nil {
 		return err
 	}
-	msg, err := stream.RecvMsg()
+	s := stream.(*Stream)
+	rep, err := s.recvReply()
 	if err != nil {
 		return err
 	}
-	if msg.GetMessageName() != reply.GetMessageName() ||
-		msg.GetCrcString() != reply.GetCrcString() {
-		return fmt.Errorf("unexpected reply: %T %+v", msg, msg)
+	if err := s.channel.msgCodec.DecodeMsg(rep.data, reply); err != nil {
+		return err
 	}
-	reflect.ValueOf(reply).Elem().Set(reflect.ValueOf(msg).Elem())
 	return nil
 }
 
@@ -106,6 +105,25 @@ func (s *Stream) SendMsg(msg api.Message) error {
 }
 
 func (s *Stream) RecvMsg() (api.Message, error) {
+	reply, err := s.recvReply()
+	if err != nil {
+		return nil, err
+	}
+	// resolve message type
+	msg, err := s.channel.msgIdentifier.LookupByID(reply.msgID)
+	if err != nil {
+		return nil, err
+	}
+	// allocate message instance
+	msg = reflect.New(reflect.TypeOf(msg).Elem()).Interface().(api.Message)
+	// decode message data
+	if err := s.channel.msgCodec.DecodeMsg(reply.data, msg); err != nil {
+		return nil, err
+	}
+	return msg, nil
+}
+
+func (s *Stream) recvReply() (*vppReply, error) {
 	if s.conn == nil {
 		return nil, errors.New("stream closed")
 	}
@@ -120,18 +138,7 @@ func (s *Stream) RecvMsg() (api.Message, error) {
 			// and stream does not use it
 			return nil, reply.err
 		}
-		// resolve message type
-		msg, err := s.channel.msgIdentifier.LookupByID(reply.msgID)
-		if err != nil {
-			return nil, err
-		}
-		// allocate message instance
-		msg = reflect.New(reflect.TypeOf(msg).Elem()).Interface().(api.Message)
-		// decode message data
-		if err := s.channel.msgCodec.DecodeMsg(reply.data, msg); err != nil {
-			return nil, err
-		}
-		return msg, nil
+		return reply, nil
 
 	case <-s.ctx.Done():
 		return nil, s.ctx.Err()