func TestTracingSpanEncoding(t *testing.T) {
s1 := Span{
traceID: 1,
parentID: 2,
spanID: 3,
flags: 4,
}
// Encoding is: spanid:8 parentid:8 traceid:8 traceflags:1
// http://tchannel.readthedocs.io/en/latest/protocol/#tracing
encoded := []byte{
0, 0, 0, 0, 0, 0, 0, 3, /* spanID */
0, 0, 0, 0, 0, 0, 0, 2, /* parentID */
0, 0, 0, 0, 0, 0, 0, 1, /* traceID */
4, /* flags */
}
buf := make([]byte, len(encoded))
writer := typed.NewWriteBuffer(buf)
require.NoError(t, s1.write(writer), "Failed to encode span")
assert.Equal(t, encoded, buf, "Encoded span mismatch")
var s2 Span
reader := typed.NewReadBuffer(buf)
require.NoError(t, s2.read(reader), "Failed to decode span")
assert.Equal(t, s1, s2, "Roundtrip of span failed")
}
// newFragment creates a new fragment for marshaling into
func (w *reqResWriter) newFragment(initial bool, checksum Checksum) (*writableFragment, error) {
if err := w.mex.ctx.Err(); err != nil {
return nil, w.failed(GetContextError(err))
}
message := w.messageForFragment(initial)
// Create the frame
frame := w.conn.framePool.Get()
frame.Header.ID = w.mex.msgID
frame.Header.messageType = message.messageType()
// Write the message into the fragment, reserving flags and checksum bytes
wbuf := typed.NewWriteBuffer(frame.Payload[:])
fragment := new(writableFragment)
fragment.frame = frame
fragment.flagsRef = wbuf.DeferByte()
if err := message.write(wbuf); err != nil {
return nil, err
}
wbuf.WriteSingleByte(byte(checksum.TypeCode()))
fragment.checksumRef = wbuf.DeferBytes(checksum.Size())
fragment.checksum = checksum
fragment.contents = wbuf
return fragment, wbuf.Err()
}
// TODO(prashant): Use a small buffer and then flush it when it's full.
func writeHeaders(w io.Writer, headers map[string]string) error {
// Calculate the size of the buffer that we need.
size := 2
for k, v := range headers {
size += 4 /* size of key/value lengths */
size += len(k) + len(v)
}
buf := make([]byte, size)
writeBuffer := typed.NewWriteBuffer(buf)
writeBuffer.WriteUint16(uint16(len(headers)))
for k, v := range headers {
writeBuffer.WriteLen16String(k)
writeBuffer.WriteLen16String(v)
}
if err := writeBuffer.Err(); err != nil {
return err
}
// Safety check to ensure the bytes written calculation is correct.
if writeBuffer.BytesWritten() != size {
return fmt.Errorf("writeHeaders size calculation wrong, expected to write %v bytes, only wrote %v bytes",
size, writeBuffer.BytesWritten())
}
_, err := writeBuffer.FlushTo(w)
return err
}
// WriteHeaders writes the given key-value pairs using the following encoding:
// len~2 (k~4 v~4)~len
func WriteHeaders(w io.Writer, headers map[string]string) error {
// TODO(prashant): Since we are not writing length-prefixed data here,
// we can write out to the buffer, and if it fills up, flush it.
// Right now, we calculate the size of the required buffer and write it out.
// Calculate the size of the buffer that we need.
size := 2
for k, v := range headers {
size += 4 /* size of key/value lengths */
size += len(k) + len(v)
}
buf := make([]byte, size)
writeBuffer := typed.NewWriteBuffer(buf)
writeBuffer.WriteUint16(uint16(len(headers)))
for k, v := range headers {
writeBuffer.WriteLen16String(k)
writeBuffer.WriteLen16String(v)
}
if err := writeBuffer.Err(); err != nil {
return err
}
// Safety check to ensure the bytes written calculation is correct.
if writeBuffer.BytesWritten() != size {
return fmt.Errorf(
"writeHeaders size calculation wrong, expected to write %v bytes, only wrote %v bytes",
size, writeBuffer.BytesWritten())
}
_, err := writeBuffer.FlushTo(w)
return err
}
func (ch fragmentChannel) newFragment(initial bool, checksum Checksum) (*writableFragment, error) {
wbuf := typed.NewWriteBuffer(make([]byte, testFragmentSize))
fragment := new(writableFragment)
fragment.flagsRef = wbuf.DeferByte()
wbuf.WriteSingleByte(byte(checksum.TypeCode()))
fragment.checksumRef = wbuf.DeferBytes(checksum.Size())
fragment.checksum = checksum
fragment.contents = wbuf
return fragment, wbuf.Err()
}
func (ec SystemErrCode) fakeErrFrame() lazyError {
f := NewFrame(100)
fh := FrameHeader{
size: uint16(0xFF34),
messageType: messageTypeError,
ID: invalidMessageID,
}
f.Header = fh
fh.write(typed.NewWriteBuffer(f.headerBuffer))
payload := typed.NewWriteBuffer(f.Payload)
payload.WriteSingleByte(byte(ec))
payload.WriteBytes(make([]byte, 25)) // tracing
msg := ec.String()
payload.WriteUint16(uint16(len(msg)))
payload.WriteBytes([]byte(msg))
return newLazyError(f)
}
func (cr testCallReq) req() 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()
f.Header = fh
fh.write(typed.NewWriteBuffer(f.headerBuffer))
payload := typed.NewWriteBuffer(f.Payload)
payload.WriteSingleByte(0) // flags
payload.WriteUint32(42) // TTL
payload.WriteBytes(make([]byte, 25)) // tracing
payload.WriteLen8String("bankmoji") // service
headers := make(map[string]string)
if cr&reqHasHeaders != 0 {
addRandomHeaders(headers)
}
if cr&reqHasCaller != 0 {
headers["cn"] = "fake-caller"
}
if cr&reqHasDelegate != 0 {
headers["rd"] = "fake-delegate"
}
if cr&reqHasRoutingKey != 0 {
headers["rk"] = "fake-routingkey"
}
writeHeaders(payload, headers)
if cr&reqHasChecksum == 0 {
payload.WriteSingleByte(byte(ChecksumTypeNone)) // checksum type
// no checksum contents for None
} else {
payload.WriteSingleByte(byte(ChecksumTypeCrc32C)) // checksum type
payload.WriteUint32(0) // checksum contents
}
payload.WriteLen16String("moneys") // method
return newLazyCallReq(f)
}
func (cr testCallRes) res() lazyCallRes {
f := NewFrame(100)
fh := FrameHeader{
size: uint16(0xFF34),
messageType: messageTypeCallRes,
ID: 0xDEADBEEF,
}
f.Header = fh
fh.write(typed.NewWriteBuffer(f.headerBuffer))
payload := typed.NewWriteBuffer(f.Payload)
if cr&resIsContinued == 0 {
payload.WriteSingleByte(0) // flags
} else {
payload.WriteSingleByte(hasMoreFragmentsFlag) // flags
}
if cr&resIsOK == 0 {
payload.WriteSingleByte(1) // code not ok
} else {
payload.WriteSingleByte(0) // code ok
}
headers := make(map[string]string)
if cr&resHasHeaders != 0 {
addRandomHeaders(headers)
}
writeHeaders(payload, headers)
if cr&resHasChecksum == 0 {
payload.WriteSingleByte(byte(ChecksumTypeNone)) // checksum type
// No contents for ChecksumTypeNone.
} else {
payload.WriteSingleByte(byte(ChecksumTypeCrc32C)) // checksum type
payload.WriteUint32(0) // checksum contents
}
payload.WriteUint16(0) // no arg1 for call res
return newLazyCallRes(f)
}
func TestFinishesCallResponses(t *testing.T) {
tests := []struct {
msgType messageType
flags byte
finishesCall bool
}{
{messageTypeCallRes, 0x00, true},
{messageTypeCallRes, 0x01, false},
{messageTypeCallRes, 0x02, true},
{messageTypeCallRes, 0x03, false},
{messageTypeCallRes, 0x04, true},
{messageTypeCallResContinue, 0x00, true},
{messageTypeCallResContinue, 0x01, false},
{messageTypeCallResContinue, 0x02, true},
{messageTypeCallResContinue, 0x03, false},
{messageTypeCallResContinue, 0x04, true},
// By definition, callreq should never terminate an RPC.
{messageTypeCallReq, 0x00, false},
{messageTypeCallReq, 0x01, false},
{messageTypeCallReq, 0x02, false},
{messageTypeCallReq, 0x03, false},
{messageTypeCallReq, 0x04, false},
}
for _, tt := range tests {
f := NewFrame(100)
fh := FrameHeader{
size: uint16(0xFF34),
messageType: tt.msgType,
ID: 0xDEADBEEF,
}
f.Header = fh
fh.write(typed.NewWriteBuffer(f.headerBuffer))
payload := typed.NewWriteBuffer(f.Payload)
payload.WriteSingleByte(tt.flags)
assert.Equal(t, tt.finishesCall, finishesCall(f), "Wrong isLast for flags %v and message type %v", tt.flags, tt.msgType)
}
}
func TestVarintString(t *testing.T) {
tests := []string{
"",
"short string",
testutils.RandString(1000),
}
for _, tt := range tests {
buf := make([]byte, 2000)
wb := typed.NewWriteBuffer(buf)
writeVarintString(wb, tt)
rb := typed.NewReadBuffer(buf)
got := readVarintString(rb)
assert.Equal(t, tt, got, "Varint string mismatch")
}
}
func TestHeaders(t *testing.T) {
tests := []http.Header{
http.Header{},
http.Header{
"K1": []string{"K1V1", "K1V2", "K1V3"},
"K2": []string{"K2V2", "K2V2"},
},
}
for _, tt := range tests {
buf := make([]byte, 1000)
wb := typed.NewWriteBuffer(buf)
writeHeaders(wb, tt)
newHeaders := make(http.Header)
rb := typed.NewReadBuffer(buf)
readHeaders(rb, newHeaders)
assert.Equal(t, tt, newHeaders, "Headers mismatch")
}
}