forwardClientToServer method may lose the response header when got rpc error?

[LugiaChang]

func (s *handler) forwardClientToServer(src grpc.ClientStream, dst grpc.ServerStream) chan error {
    ret := make(chan error, 1)
    go func() {
        f := &emptypb.Empty{}
        for i := 0; ; i++ {
            if err := src.RecvMsg(f); err != nil {
                ret <- err // this can be io.EOF which is happy case
                break
            }
            if i == 0 {
                // This is a bit of a hack, but client to server headers are only readable after first client msg is
                // received but must be written to server stream before the first msg is flushed.
                // This is the only place to do it nicely.
                md, err := src.Header()
                if err != nil {
                    ret <- err
                    break
                }
                if err := dst.SendHeader(md); err != nil {
                    ret <- err
                    break
                }
            }
            if err := dst.SendMsg(f); err != nil {
                ret <- err
                break
            }
        }
    }()
    return ret
}

Is the response header information lost when a remote RPC returns an error here? I found that this is indeed the case during my testing. Can the following logic be added? But I am not sure if src.Header is blocked under any circumstances.

func (s *handler) forwardClientToServer(src grpc.ClientStream, dst grpc.ServerStream) chan error {
    ret := make(chan error, 1)
    go func() {
        f := &emptypb.Empty{}
        for i := 0; ; i++ {
            if err := src.RecvMsg(f); err != nil {
                // set response header
                if i == 0 {
                    md, err := src.Header()
                    if err != nil {
                        ret <- err
                        break
                    }
                    if err := dst.SetHeader(md); err != nil {
                        ret <- err
                        break
                    }
                }
                ret <- err // this can be io.EOF which is happy case
                break
            }
            if i == 0 {
                // This is a bit of a hack, but client to server headers are only readable after first client msg is
                // received but must be written to server stream before the first msg is flushed.
                // This is the only place to do it nicely.
                md, err := src.Header()
                if err != nil {
                    ret <- err
                    break
                }
                if err := dst.SendHeader(md); err != nil {
                    ret <- err
                    break
                }
            }
            if err := dst.SendMsg(f); err != nil {
                ret <- err
                break
            }
        }
    }()
    return ret
}