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")
}
// ReadResponse reads a http.Response from the given readers.
func ReadResponse(call argReader) (*http.Response, error) {
var arg2 []byte
if err := tchannel.NewArgReader(call.Arg2Reader()).Read(&arg2); err != nil {
return nil, err
}
rb := typed.NewReadBuffer(arg2)
statusCode := rb.ReadUint16()
message := readVarintString(rb)
response := &http.Response{
StatusCode: int(statusCode),
Status: fmt.Sprintf("%v %v", statusCode, message),
Proto: "HTTP/1.1",
ProtoMajor: 1,
ProtoMinor: 1,
Header: make(http.Header),
}
readHeaders(rb, response.Header)
if err := rb.Err(); err != nil {
return nil, err
}
arg3Reader, err := call.Arg3Reader()
if err != nil {
return nil, err
}
response.Body = arg3Reader
return response, nil
}
// handleError handles an error coming back from the peer. If the error is a
// protocol level error, the entire connection will be closed. If the error is
// a request specific error, it will be written to the request's response
// channel and converted into a SystemError returned from the next reader or
// access call.
// The return value is whether the frame should be released immediately.
func (c *Connection) handleError(frame *Frame) bool {
errMsg := errorMessage{
id: frame.Header.ID,
}
rbuf := typed.NewReadBuffer(frame.SizedPayload())
if err := errMsg.read(rbuf); err != nil {
c.log.WithFields(
LogField{"remotePeer", c.remotePeerInfo},
ErrField(err),
).Warn("Unable to read error frame.")
c.connectionError("parsing error frame", err)
return true
}
if errMsg.errCode == ErrCodeProtocol {
c.log.WithFields(
LogField{"remotePeer", c.remotePeerInfo},
LogField{"error", errMsg.message},
).Warn("Peer reported protocol error.")
c.connectionError("received protocol error", errMsg.AsSystemError())
return true
}
if err := c.outbound.forwardPeerFrame(frame); err != nil {
c.log.Infof("Failed to forward error frame %v to mex, error: %v", frame.Header, errMsg)
return true
}
// If the frame was forwarded, then the other side is responsible for releasing the frame.
return false
}
func (ch fragmentChannel) recvNextFragment(initial bool) (*readableFragment, error) {
rbuf := typed.NewReadBuffer(<-ch)
fragment := new(readableFragment)
fragment.done = func() {}
fragment.flags = rbuf.ReadSingleByte()
fragment.checksumType = ChecksumType(rbuf.ReadSingleByte())
fragment.checksum = rbuf.ReadBytes(fragment.checksumType.ChecksumSize())
fragment.contents = rbuf
return fragment, rbuf.Err()
}
// Handles an incoming InitReq. If we are waiting for the peer to send us an
// InitReq, and the InitReq is valid, send a corresponding InitRes and mark
// ourselves as active
func (c *Connection) handleInitReq(frame *Frame) {
id := frame.Header.ID
var req initReq
rbuf := typed.NewReadBuffer(frame.SizedPayload())
if err := req.read(rbuf); err != nil {
// TODO(mmihic): Technically probably a protocol error
c.connectionError(err)
return
}
if req.Version != CurrentProtocolVersion {
c.protocolError(id, fmt.Errorf("Unsupported protocol version %d from peer", req.Version))
return
}
var ok bool
if c.remotePeerInfo.HostPort, ok = req.initParams[InitParamHostPort]; !ok {
c.protocolError(id, fmt.Errorf("Header %v is required", InitParamHostPort))
return
}
if c.remotePeerInfo.ProcessName, ok = req.initParams[InitParamProcessName]; !ok {
c.protocolError(id, fmt.Errorf("Header %v is required", InitParamProcessName))
return
}
if c.remotePeerInfo.IsEphemeral() {
// TODO(prashant): Add an IsEphemeral bool to the peer info.
c.remotePeerInfo.HostPort = c.conn.RemoteAddr().String()
}
res := initRes{initMessage{id: frame.Header.ID}}
res.initParams = initParams{
InitParamHostPort: c.localPeerInfo.HostPort,
InitParamProcessName: c.localPeerInfo.ProcessName,
}
res.Version = CurrentProtocolVersion
if err := c.sendMessage(&res); err != nil {
c.connectionError(err)
return
}
c.withStateLock(func() error {
switch c.state {
case connectionWaitingToRecvInitReq:
c.state = connectionActive
}
return nil
})
c.callOnActive()
}
// parseInboundFragment parses an incoming fragment based on the given message
func parseInboundFragment(framePool FramePool, frame *Frame, message message) (*readableFragment, error) {
rbuf := typed.NewReadBuffer(frame.SizedPayload())
fragment := new(readableFragment)
fragment.flags = rbuf.ReadSingleByte()
if err := message.read(rbuf); err != nil {
return nil, err
}
fragment.checksumType = ChecksumType(rbuf.ReadSingleByte())
fragment.checksum = rbuf.ReadBytes(fragment.checksumType.ChecksumSize())
fragment.contents = rbuf
fragment.done = func() {
framePool.Release(frame)
}
return fragment, rbuf.Err()
}
func TestFraming(t *testing.T) {
fh := fakeHeader()
wbuf := typed.NewWriteBufferWithSize(1024)
require.Nil(t, fh.write(wbuf))
var b bytes.Buffer
if _, err := wbuf.FlushTo(&b); err != nil {
require.Nil(t, err)
}
rbuf := typed.NewReadBuffer(b.Bytes())
var fh2 FrameHeader
require.Nil(t, fh2.read(rbuf))
assert.Equal(t, fh, fh2)
}
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")
}
}
// TODO(prashant): Allow typed.ReadBuffer to read directly from the reader.
func readHeaders(r io.Reader) (map[string]string, error) {
bs, err := ioutil.ReadAll(r)
if err != nil {
return nil, err
}
buffer := typed.NewReadBuffer(bs)
numHeaders := buffer.ReadUint16()
if numHeaders == 0 {
return nil, buffer.Err()
}
headers := make(map[string]string, numHeaders)
for i := 0; i < int(numHeaders) && buffer.Err() == nil; i++ {
k := buffer.ReadLen16String()
v := buffer.ReadLen16String()
headers[k] = v
}
return headers, buffer.Err()
}
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")
}
}
// handleError andles an error coming back from the peer. If the error is a
// protocol level error, the entire connection will be closed. If the error is
// a request specific error, it will be written to the request's response
// channel and converted into a SystemError returned from the next reader or
// access call.
func (c *Connection) handleError(frame *Frame) {
errMsg := errorMessage{
id: frame.Header.ID,
}
rbuf := typed.NewReadBuffer(frame.SizedPayload())
if err := errMsg.read(rbuf); err != nil {
c.log.Warnf("Unable to read Error frame from %s: %v", c.remotePeerInfo, err)
c.connectionError(err)
return
}
if errMsg.errCode == ErrCodeProtocol {
c.log.Warnf("Peer %s reported protocol error: %s", c.remotePeerInfo, errMsg.message)
c.connectionError(errMsg.AsSystemError())
return
}
if err := c.outbound.forwardPeerFrame(frame); err != nil {
c.outbound.removeExchange(frame.Header.ID)
}
}
func TestFraming(t *testing.T) {
fh := FrameHeader{
size: uint16(0xFF34),
messageType: messageTypeCallReq,
ID: 0xDEADBEEF,
}
wbuf := typed.NewWriteBufferWithSize(1024)
require.Nil(t, fh.write(wbuf))
var b bytes.Buffer
if _, err := wbuf.FlushTo(&b); err != nil {
require.Nil(t, err)
}
rbuf := typed.NewReadBuffer(b.Bytes())
var fh2 FrameHeader
require.Nil(t, fh2.read(rbuf))
assert.Equal(t, fh, fh2)
}
// ReadRequest reads a http.Request from the given readers.
func ReadRequest(call argReader) (*http.Request, error) {
var arg2 []byte
if err := tchannel.NewArgReader(call.Arg2Reader()).Read(&arg2); err != nil {
return nil, err
}
rb := typed.NewReadBuffer(arg2)
method := rb.ReadLen8String()
url := readVarintString(rb)
r, err := http.NewRequest(method, url, nil)
if err != nil {
return nil, err
}
readHeaders(rb, r.Header)
if err := rb.Err(); err != nil {
return nil, err
}
r.Body, err = call.Arg3Reader()
return r, err
}
func TestReservedBytes(t *testing.T) {
// Set up a frame with non-zero values
f := NewFrame(MaxFramePayloadSize)
reader := testreader.Looper([]byte{^byte(0)})
io.ReadFull(reader, f.Payload)
f.Header.read(typed.NewReadBuffer(f.Payload))
m := &pingRes{id: 1}
f.write(m)
buf := &bytes.Buffer{}
f.WriteOut(buf)
assert.Equal(t,
[]byte{
0x0, 0x10, // size
0xd1, // type
0x0, // reserved should always be 0
0x0, 0x0, 0x0, 0x1, // id
0x0, 0x0, 0x0, 0x0, // reserved should always be 0
0x0, 0x0, 0x0, 0x0, // reserved should always be 0
},
buf.Bytes(), "Unexpected bytes")
}
func callReqSpan(f *Frame) Span {
rdr := typed.NewReadBuffer(f.Payload[_spanIndex : _spanIndex+_spanLength])
var s Span
s.read(rdr)
return s
}