rpcforward.go

package grpc

import (
	"context"
	"errors"
	"io"
	"time"

	"google.golang.org/grpc/codes"
	"google.golang.org/grpc/encoding"
	"google.golang.org/grpc/internal/binarylog"
	"google.golang.org/grpc/internal/channelz"
	"google.golang.org/grpc/internal/transport"
	"google.golang.org/grpc/peer"
	"google.golang.org/grpc/stats"
	"google.golang.org/grpc/status"
)

func recvForward(p *parser, c baseCodec, s *transport.Stream, dc Decompressor, maxReceiveMessageSize int, payInfo *payloadInfo, compressor encoding.Compressor) (msgBytes []byte, err error) {
	d, err := recvAndDecompress(p, s, dc, maxReceiveMessageSize, payInfo, compressor)
	if err != nil {
		return nil, err
	}
	//if err := c.Unmarshal(d, m); err != nil {
	//	return status.Errorf(codes.Internal, "grpc: failed to unmarshal the received message %v", err)
	//}
	msgBytes = d
	if payInfo != nil {
		payInfo.uncompressedBytes = d
	}
	return msgBytes, nil
}

func (cs *clientStream) SendMsgForward(mBytes []byte) (err error) {
	defer func() {
		if err != nil && err != io.EOF {
			// Call finish on the client stream for errors generated by this SendMsg
			// call, as these indicate problems created by this client.  (Transport
			// errors are converted to an io.EOF error in csAttempt.sendMsg; the real
			// error will be returned from RecvMsg eventually in that case, or be
			// retried.)
			cs.finish(err)
		}
	}()
	if cs.sentLast {
		return status.Errorf(codes.Internal, "SendMsg called after CloseSend")
	}
	if !cs.desc.ClientStreams {
		cs.sentLast = true
	}

	// load hdr, payload, data
	hdr, payload, data, err := prepareMsgForward(mBytes, cs.codec, cs.cp, cs.comp)
	if err != nil {
		return err
	}

	// TODO(dfawley): should we be checking len(data) instead?
	if len(payload) > *cs.callInfo.maxSendMessageSize {
		return status.Errorf(codes.ResourceExhausted, "trying to send message larger than max (%d vs. %d)", len(payload), *cs.callInfo.maxSendMessageSize)
	}
	op := func(a *csAttempt) error {
		return a.sendMsgForward(mBytes, hdr, payload, data)
	}
	err = cs.withRetry(op, func() { cs.bufferForRetryLocked(len(hdr)+len(payload), op) })
	if len(cs.binlogs) != 0 && err == nil {
		cm := &binarylog.ClientMessage{
			OnClientSide: true,
			Message:      data,
		}
		for _, binlog := range cs.binlogs {
			binlog.Log(cm)
		}
	}
	return err
}
func (cs *clientStream) RecvMsgForward() (mBytes []byte, err error) {
	if len(cs.binlogs) != 0 && !cs.serverHeaderBinlogged {
		// Call Header() to binary log header if it's not already logged.
		cs.Header()
	}
	var recvInfo *payloadInfo
	if len(cs.binlogs) != 0 {
		recvInfo = &payloadInfo{}
	}
	mBytes, err = cs.withRetryForward(func(a *csAttempt) ([]byte, error) {
		return a.recvMsgForward(recvInfo)
	}, cs.commitAttemptLocked)
	if len(cs.binlogs) != 0 && err == nil {
		sm := &binarylog.ServerMessage{
			OnClientSide: true,
			Message:      recvInfo.uncompressedBytes,
		}
		for _, binlog := range cs.binlogs {
			binlog.Log(sm)
		}
	}
	if err != nil || !cs.desc.ServerStreams {
		// err != nil or non-server-streaming indicates end of stream.
		cs.finish(err)

		if len(cs.binlogs) != 0 {
			// finish will not log Trailer. Log Trailer here.
			logEntry := &binarylog.ServerTrailer{
				OnClientSide: true,
				Trailer:      cs.Trailer(),
				Err:          err,
			}
			if logEntry.Err == io.EOF {
				logEntry.Err = nil
			}
			if peer, ok := peer.FromContext(cs.Context()); ok {
				logEntry.PeerAddr = peer.Addr
			}
			for _, binlog := range cs.binlogs {
				binlog.Log(logEntry)
			}
		}
	}
	return mBytes, err
}
func (a *csAttempt) sendMsgForward(mBytes []byte, hdr, payld, data []byte) error {
	cs := a.cs
	if a.trInfo != nil {
		a.mu.Lock()
		if a.trInfo.tr != nil {
			a.trInfo.tr.LazyLog(&payload{sent: true, msg: nil}, true)
		}
		a.mu.Unlock()
	}
	if err := a.t.Write(a.s, hdr, payld, &transport.Options{Last: !cs.desc.ClientStreams}); err != nil {
		if !cs.desc.ClientStreams {
			// For non-client-streaming RPCs, we return nil instead of EOF on error
			// because the generated code requires it.  finish is not called; RecvMsg()
			// will call it with the stream's status independently.
			return nil
		}
		return io.EOF
	}
	for _, sh := range a.statsHandlers {
		sh.HandleRPC(a.ctx, outPayload(true, nil, data, payld, time.Now()))
	}
	if channelz.IsOn() {
		a.t.IncrMsgSent()
	}
	return nil
}
func (a *csAttempt) recvMsgForward(payInfo *payloadInfo) (mBytes []byte, err error) {
	cs := a.cs
	if len(a.statsHandlers) != 0 && payInfo == nil {
		payInfo = &payloadInfo{}
	}

	if !a.decompSet {
		// Block until we receive headers containing received message encoding.
		if ct := a.s.RecvCompress(); ct != "" && ct != encoding.Identity {
			if a.dc == nil || a.dc.Type() != ct {
				// No configured decompressor, or it does not match the incoming
				// message encoding; attempt to find a registered compressor that does.
				a.dc = nil
				a.decomp = encoding.GetCompressor(ct)
			}
		} else {
			// No compression is used; disable our decompressor.
			a.dc = nil
		}
		// Only initialize this state once per stream.
		a.decompSet = true
	}
	mBytes, err = recvForward(a.p, cs.codec, a.s, a.dc, *cs.callInfo.maxReceiveMessageSize, payInfo, a.decomp)
	if err != nil {
		if err == io.EOF {
			if statusErr := a.s.Status().Err(); statusErr != nil {
				return mBytes, statusErr
			}
			return mBytes, io.EOF // indicates successful end of stream.
		}
		return mBytes, toRPCErr(err)
	}
	if a.trInfo != nil {
		a.mu.Lock()
		if a.trInfo.tr != nil {
			a.trInfo.tr.LazyLog(&payload{sent: false, msg: nil}, true)
		}
		a.mu.Unlock()
	}
	for _, sh := range a.statsHandlers {
		sh.HandleRPC(a.ctx, &stats.InPayload{
			Client:   true,
			RecvTime: time.Now(),
			Payload:  nil,
			// TODO truncate large payload.
			Data:       payInfo.uncompressedBytes,
			WireLength: payInfo.wireLength + headerLen,
			Length:     len(payInfo.uncompressedBytes),
		})
	}
	if channelz.IsOn() {
		a.t.IncrMsgRecv()
	}
	if cs.desc.ServerStreams {
		// Subsequent messages should be received by subsequent RecvMsg calls.
		return mBytes, nil
	}
	// Special handling for non-server-stream rpcs.
	// This recv expects EOF or errors, so we don't collect inPayload.
	_, err = recvForward(a.p, cs.codec, a.s, a.dc, *cs.callInfo.maxReceiveMessageSize, nil, a.decomp)
	if err == nil {
		return nil, toRPCErr(errors.New("grpc: client streaming protocol violation: get <nil>, want <EOF>"))
	}
	if err == io.EOF {
		return mBytes, a.s.Status().Err() // non-server streaming Recv returns nil on success
	}
	return mBytes, toRPCErr(err)
}

func (cc *ClientConn) InvokeForward(ctx context.Context, method string, args []byte, opts ...CallOption) (reply []byte, err error) {
	// allow interceptor to see all applicable call options, which means those
	// configured as defaults from dial option as well as per-call options
	opts = combine(cc.dopts.callOptions, opts)

	//if cc.dopts.unaryInt != nil {
	//	return cc.dopts.unaryInt(ctx, method, args, reply, cc, invoke, opts...)
	//}
	return invokeForward(ctx, method, args, cc, opts...)
}

func invokeForward(ctx context.Context, method string, req []byte, cc *ClientConn, opts ...CallOption) (reply []byte, err error) {
	cs, err := newClientStream(ctx, unaryStreamDesc, cc, method, opts...)
	if err != nil {
		return nil, err
	}
	if err := cs.SendMsgForward(req); err != nil {
		return nil, err
	}
	return cs.RecvMsgForward()
}

func (as *addrConnStream) SendMsgForward(mBytes []byte) (err error) {
	defer func() {
		if err != nil && err != io.EOF {
			// Call finish on the client stream for errors generated by this SendMsg
			// call, as these indicate problems created by this client.  (Transport
			// errors are converted to an io.EOF error in csAttempt.sendMsg; the real
			// error will be returned from RecvMsg eventually in that case, or be
			// retried.)
			as.finish(err)
		}
	}()
	if as.sentLast {
		return status.Errorf(codes.Internal, "SendMsg called after CloseSend")
	}
	if !as.desc.ClientStreams {
		as.sentLast = true
	}

	// load hdr, payload, data
	hdr, payld, _, err := prepareMsgForward(mBytes, as.codec, as.cp, as.comp)
	if err != nil {
		return err
	}

	// TODO(dfawley): should we be checking len(data) instead?
	if len(payld) > *as.callInfo.maxSendMessageSize {
		return status.Errorf(codes.ResourceExhausted, "trying to send message larger than max (%d vs. %d)", len(payld), *as.callInfo.maxSendMessageSize)
	}

	if err := as.t.Write(as.s, hdr, payld, &transport.Options{Last: !as.desc.ClientStreams}); err != nil {
		if !as.desc.ClientStreams {
			// For non-client-streaming RPCs, we return nil instead of EOF on error
			// because the generated code requires it.  finish is not called; RecvMsg()
			// will call it with the stream's status independently.
			return nil
		}
		return io.EOF
	}

	if channelz.IsOn() {
		as.t.IncrMsgSent()
	}
	return nil
}

func (as *addrConnStream) RecvMsgForward() (mBytes []byte, err error) {
	defer func() {
		if err != nil || !as.desc.ServerStreams {
			// err != nil or non-server-streaming indicates end of stream.
			as.finish(err)
		}
	}()

	if !as.decompSet {
		// Block until we receive headers containing received message encoding.
		if ct := as.s.RecvCompress(); ct != "" && ct != encoding.Identity {
			if as.dc == nil || as.dc.Type() != ct {
				// No configured decompressor, or it does not match the incoming
				// message encoding; attempt to find a registered compressor that does.
				as.dc = nil
				as.decomp = encoding.GetCompressor(ct)
			}
		} else {
			// No compression is used; disable our decompressor.
			as.dc = nil
		}
		// Only initialize this state once per stream.
		as.decompSet = true
	}
	mBytes, err = recvForward(as.p, as.codec, as.s, as.dc, *as.callInfo.maxReceiveMessageSize, nil, as.decomp)
	if err != nil {
		if err == io.EOF {
			if statusErr := as.s.Status().Err(); statusErr != nil {
				return mBytes, statusErr
			}
			return mBytes, io.EOF // indicates successful end of stream.
		}
		return mBytes, toRPCErr(err)
	}

	if channelz.IsOn() {
		as.t.IncrMsgRecv()
	}
	if as.desc.ServerStreams {
		// Subsequent messages should be received by subsequent RecvMsg calls.
		return mBytes, nil
	}

	// Special handling for non-server-stream rpcs.
	// This recv expects EOF or errors, so we don't collect inPayload.
	mBytes, err = recvForward(as.p, as.codec, as.s, as.dc, *as.callInfo.maxReceiveMessageSize, nil, as.decomp)
	if err == nil {
		return mBytes, toRPCErr(errors.New("grpc: client streaming protocol violation: get <nil>, want <EOF>"))
	}
	if err == io.EOF {
		return mBytes, as.s.Status().Err() // non-server streaming Recv returns nil on success
	}
	return mBytes, toRPCErr(err)
}

func (cs *clientStream) withRetryForward(op func(a *csAttempt) ([]byte, error), onSuccess func()) ([]byte, error) {
	cs.mu.Lock()
	for {
		if cs.committed {
			cs.mu.Unlock()
			// toRPCErr is used in case the error from the attempt comes from
			// NewClientStream, which intentionally doesn't return a status
			// error to allow for further inspection; all other errors should
			// already be status errors.
			data, err := op(cs.attempt)
			return data, toRPCErr(err)
		}
		if len(cs.buffer) == 0 {
			// For the first op, which controls creation of the stream and
			// assigns cs.attempt, we need to create a new attempt inline
			// before executing the first op.  On subsequent ops, the attempt
			// is created immediately before replaying the ops.
			var err error
			if cs.attempt, err = cs.newAttemptLocked(false /* isTransparent */); err != nil {
				cs.mu.Unlock()
				cs.finish(err)
				return nil, err
			}
		}
		a := cs.attempt
		cs.mu.Unlock()
		data, err := op(a)
		cs.mu.Lock()
		if a != cs.attempt {
			// We started another attempt already.
			continue
		}
		if err == io.EOF {
			<-a.s.Done()
		}
		if err == nil || (err == io.EOF && a.s.Status().Code() == codes.OK) {
			onSuccess()
			cs.mu.Unlock()
			return data, err
		}
		if err := cs.retryLocked(a, err); err != nil {
			cs.mu.Unlock()
			return data, err
		}
	}
}

func (ss *serverStream) SendMsgForward(req []byte) (err error) {
	defer func() {
		if ss.trInfo != nil {
			ss.mu.Lock()
			if ss.trInfo.tr != nil {
				if err == nil {
					ss.trInfo.tr.LazyLog(&payload{sent: true, msg: nil, msgBytes: req}, true)
				} else {
					ss.trInfo.tr.LazyLog(&fmtStringer{"%v", []interface{}{err}}, true)
					ss.trInfo.tr.SetError()
				}
			}
			ss.mu.Unlock()
		}
		if err != nil && err != io.EOF {
			st, _ := status.FromError(toRPCErr(err))
			ss.t.WriteStatus(ss.s, st)
			// Non-user specified status was sent out. This should be an error
			// case (as a server side Cancel maybe).
			//
			// This is not handled specifically now. User will return a final
			// status from the service handler, we will log that error instead.
			// This behavior is similar to an interceptor.
		}
		if channelz.IsOn() && err == nil {
			ss.t.IncrMsgSent()
		}
	}()

	// load hdr, payload, data
	hdr, payload, data, err := prepareMsgForward(req, ss.codec, ss.cp, ss.comp)
	if err != nil {
		return err
	}

	// TODO(dfawley): should we be checking len(data) instead?
	if len(payload) > ss.maxSendMessageSize {
		return status.Errorf(codes.ResourceExhausted, "trying to send message larger than max (%d vs. %d)", len(payload), ss.maxSendMessageSize)
	}
	if err := ss.t.Write(ss.s, hdr, payload, &transport.Options{Last: false}); err != nil {
		return toRPCErr(err)
	}
	if len(ss.binlogs) != 0 {
		if !ss.serverHeaderBinlogged {
			h, _ := ss.s.Header()
			sh := &binarylog.ServerHeader{
				Header: h,
			}
			ss.serverHeaderBinlogged = true
			for _, binlog := range ss.binlogs {
				binlog.Log(sh)
			}
		}
		sm := &binarylog.ServerMessage{
			Message: data,
		}
		for _, binlog := range ss.binlogs {
			binlog.Log(sm)
		}
	}
	if len(ss.statsHandler) != 0 {
		for _, sh := range ss.statsHandler {
			sh.HandleRPC(ss.s.Context(), outPayload(false, nil, data, payload, time.Now()))
		}
	}
	return nil
}

func (ss *serverStream) RecvMsgForward() (msgBytes []byte, err error) {
	defer func() {
		if ss.trInfo != nil {
			ss.mu.Lock()
			if ss.trInfo.tr != nil {
				if err == nil {
					ss.trInfo.tr.LazyLog(&payload{sent: false, msg: nil, msgBytes: msgBytes}, true)
				} else if err != io.EOF {
					ss.trInfo.tr.LazyLog(&fmtStringer{"%v", []interface{}{err}}, true)
					ss.trInfo.tr.SetError()
				}
			}
			ss.mu.Unlock()
		}
		if err != nil && err != io.EOF {
			st, _ := status.FromError(toRPCErr(err))
			ss.t.WriteStatus(ss.s, st)
			// Non-user specified status was sent out. This should be an error
			// case (as a server side Cancel maybe).
			//
			// This is not handled specifically now. User will return a final
			// status from the service handler, we will log that error instead.
			// This behavior is similar to an interceptor.
		}
		if channelz.IsOn() && err == nil {
			ss.t.IncrMsgRecv()
		}
	}()
	var payInfo *payloadInfo
	if len(ss.statsHandler) != 0 || len(ss.binlogs) != 0 {
		payInfo = &payloadInfo{}
	}
	if msgBytes, err = recvForward(ss.p, ss.codec, ss.s, ss.dc, ss.maxReceiveMessageSize, payInfo, ss.decomp); err != nil {
		if err == io.EOF {
			if len(ss.binlogs) != 0 {
				chc := &binarylog.ClientHalfClose{}
				for _, binlog := range ss.binlogs {
					binlog.Log(chc)
				}
			}
			return nil, err
		}
		if err == io.ErrUnexpectedEOF {
			err = status.Errorf(codes.Internal, io.ErrUnexpectedEOF.Error())
		}
		return nil, toRPCErr(err)
	}
	if len(ss.statsHandler) != 0 {
		for _, sh := range ss.statsHandler {
			sh.HandleRPC(ss.s.Context(), &stats.InPayload{
				RecvTime: time.Now(),
				Payload:  msgBytes,
				// TODO truncate large payload.
				Data:       payInfo.uncompressedBytes,
				WireLength: payInfo.wireLength + headerLen,
				Length:     len(payInfo.uncompressedBytes),
			})
		}
	}
	if len(ss.binlogs) != 0 {
		cm := &binarylog.ClientMessage{
			Message: payInfo.uncompressedBytes,
		}
		for _, binlog := range ss.binlogs {
			binlog.Log(cm)
		}
	}
	return msgBytes, nil
}

// prepareMsg returns the hdr, payload and data
// using the compressors passed or using the
// passed preparedmsg
func prepareMsgForward(mBytes []byte, codec baseCodec, cp Compressor, comp encoding.Compressor) (hdr, payload, data []byte, err error) {
	//if preparedMsg, ok := m.(*PreparedMsg); ok {
	//	return preparedMsg.hdr, preparedMsg.payload, preparedMsg.encodedData, nil
	//}
	// The input interface is not a prepared msg.
	// Marshal and Compress the data at this point
	data, err = mBytes, nil //encode(codec, mBytes)
	//if err != nil {
	//	return nil, nil, nil, err
	//}
	compData, err := compress(data, cp, comp)
	if err != nil {
		return nil, nil, nil, err
	}
	hdr, payload = msgHeader(data, compData)
	return hdr, payload, data, nil
}