Add method to RelayFrame to reveal Arg2 offset

relay/relay.go

@@ -37,6 +37,13 @@ type CallFrame interface {
 	// RoutingKey may refer to an alternate traffic group instead of the
 	// traffic group identified by the service name.
 	RoutingKey() []byte
+	// Arg2StartOffset returns the offset from start of frame to the
+	// beginning of Arg2 in bytes.
+	Arg2StartOffset() int
+	// Arg2EndOffset returns the offset from start of frame to the end of
+	// Arg2 in bytes, and hasMore to indicate if there are mote data from
+	// other frames (i.e. arg2 is fragmented).
+	Arg2EndOffset() (_ int, hasMore bool)
 }
 
 // Conn contains information about the underlying connection.

relay_messages.go

@@ -88,6 +88,9 @@ type lazyCallReq struct {
 	*Frame
 
 	caller, method, delegate, key []byte
+
+	arg2StartOffset, arg2EndOffset int
+	isArg2Fragmented               bool
 }
 
 // TODO: Consider pooling lazyCallReq and using pointers to the struct.
@@ -132,6 +135,13 @@ func newLazyCallReq(f *Frame) lazyCallReq {
 	arg1Len := int(binary.BigEndian.Uint16(f.Payload[cur : cur+2]))
 	cur += 2
 	cr.method = f.Payload[cur : cur+arg1Len]
+
+	// arg2~2
+	cur += arg1Len
+	cr.arg2StartOffset = cur + 2
+	cr.arg2EndOffset = cur + 2 /*arg2 len*/ + int(binary.BigEndian.Uint16(f.Payload[cur:cur+2]))
+	// arg2 is fragmented if we don't see arg3 in this frame.
+	cr.isArg2Fragmented = int(cr.Header.FrameSize()) <= (FrameHeaderSize+cr.arg2EndOffset) && f.Payload[_flagsIndex]&hasMoreFragmentsFlag != 0
 	return cr
 }
 
@@ -183,6 +193,18 @@ func (f lazyCallReq) HasMoreFragments() bool {
 	return f.Payload[_flagsIndex]&hasMoreFragmentsFlag != 0
 }
 
+// Arg2EndOffset returns the offset from start of frame to the end of Arg2
+// in bytes, and whether there are more data from other frames.
+func (f lazyCallReq) Arg2EndOffset() (_ int, hasMore bool) {
+	return f.arg2EndOffset, f.isArg2Fragmented
+}
+
+// Arg2StartOffset returns the offset from start of frame to the beginning
+// of Arg2 in bytes.
+func (f lazyCallReq) Arg2StartOffset() int {
+	return f.arg2StartOffset
+}
+
 // finishesCall checks whether this frame is the last one we should expect for
 // this RPC req-res.
 func finishesCall(f *Frame) bool {

relay_messages_test.go

@@ -41,17 +41,27 @@ const (
 	reqHasAll testCallReq = reqTotalCombinations - 1
 )
 
+type testCallReqParams struct {
+	flags   byte
+	arg2Buf []byte
+}
+
 func (cr testCallReq) req() lazyCallReq {
+	return cr.reqWithParams(testCallReqParams{})
+}
+
+func (cr testCallReq) reqWithParams(p testCallReqParams) lazyCallReq {
 	// TODO: Constructing a frame is ugly because the initial flags byte is
 	// written in reqResWriter instead of callReq. We should instead handle that
 	// in callReq, which will allow our tests to be sane.
 	f := NewFrame(200)
 	fh := fakeHeader()
+	fh.size = 0xD8 // 200 + 16 bytes of header = 216 (0xD8)
 	f.Header = fh
 	fh.write(typed.NewWriteBuffer(f.headerBuffer))
 
 	payload := typed.NewWriteBuffer(f.Payload)
-	payload.WriteSingleByte(0)           // flags
+	payload.WriteSingleByte(p.flags)     // flags
 	payload.WriteUint32(42)              // TTL
 	payload.WriteBytes(make([]byte, 25)) // tracing
 	payload.WriteLen8String("bankmoji")  // service
@@ -79,6 +89,9 @@ func (cr testCallReq) req() lazyCallReq {
 		payload.WriteUint32(0)                            // checksum contents
 	}
 	payload.WriteLen16String("moneys") // method
+
+	payload.WriteUint16(uint16(len(p.arg2Buf)))
+	payload.WriteBytes(p.arg2Buf)
 	return newLazyCallReq(f)
 }
 
@@ -272,6 +285,64 @@ func TestLazyCallReqSetTTL(t *testing.T) {
 	})
 }
 
+func TestLazyCallArg2Offset(t *testing.T) {
+	wantArg2Buf := []byte("test arg2 buf")
+	tests := []struct {
+		msg     string
+		flags   byte
+		arg2Buf []byte
+	}{
+		{
+			msg:     "arg2 is fully contained in frame",
+			arg2Buf: wantArg2Buf,
+		},
+		{
+			msg: "has no arg2",
+		},
+		{
+			msg:     "frame fragmented but arg2 is fully contained",
+			flags:   hasMoreFragmentsFlag,
+			arg2Buf: wantArg2Buf,
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.msg, func(t *testing.T) {
+			withLazyCallReqCombinations(func(crt testCallReq) {
+				cr := crt.reqWithParams(testCallReqParams{
+					flags:   tt.flags,
+					arg2Buf: tt.arg2Buf,
+				})
+				arg2EndOffset, hasMore := cr.Arg2EndOffset()
+				assert.False(t, hasMore)
+				if len(tt.arg2Buf) == 0 {
+					assert.Zero(t, arg2EndOffset-cr.Arg2StartOffset())
+					return
+				}
+
+				arg2Payload := cr.Payload[cr.Arg2StartOffset():arg2EndOffset]
+				assert.Equal(t, tt.arg2Buf, arg2Payload)
+			})
+		})
+	}
+
+	t.Run("arg2 has been fragmented", func(t *testing.T) {
+		withLazyCallReqCombinations(func(crt testCallReq) {
+			// For each CallReq, we first get the remaining space left, and
+			// fill up the remaining space with arg2.
+			crNoArg2 := crt.req()
+			arg2Size := int(crNoArg2.Header.FrameSize()) - crNoArg2.Arg2StartOffset()
+			cr := crt.reqWithParams(testCallReqParams{
+				flags:   hasMoreFragmentsFlag,
+				arg2Buf: make([]byte, arg2Size),
+			})
+			endOffset, hasMore := cr.Arg2EndOffset()
+			assert.True(t, hasMore)
+			assert.EqualValues(t, crNoArg2.Header.FrameSize(), endOffset)
+		})
+	})
+}
+
 func TestLazyCallResRejectsOtherFrames(t *testing.T) {
 	assertWrappingPanics(
 		t,

testutils/call.go

@@ -122,3 +122,15 @@ func (f FakeCallFrame) RoutingKey() []byte {
 func (f FakeCallFrame) RoutingDelegate() []byte {
 	return []byte(f.RoutingDelegateF)
 }
+
+// Arg2StartOffset returns the offset from start of frame to
+// the beginning of Arg2.
+func (f FakeCallFrame) Arg2StartOffset() int {
+	return 0
+}
+
+// Arg2EndOffset returns the offset from start of frame to the end
+// of Arg2 and whether Arg2 is fragmented.
+func (f FakeCallFrame) Arg2EndOffset() (int, bool) {
+	return 0, false
+}