Avoid trying to access underlying_stream_ after it's been orphaned.

2 files changed, 15 insertions, 5 deletions

diff --git a/src/core/lib/transport/byte_stream.cc b/src/core/lib/transport/byte_stream.cc
index f3ad31ff79..1951492dbd 100644
--- a/src/core/lib/transport/byte_stream.cc
+++ b/src/core/lib/transport/byte_stream.cc
@@ -79,7 +79,9 @@ void SliceBufferByteStream::Shutdown(grpc_error* error) {
 //
 
 ByteStreamCache::ByteStreamCache(OrphanablePtr<ByteStream> underlying_stream)
-    : underlying_stream_(std::move(underlying_stream)) {
+    : underlying_stream_(std::move(underlying_stream)),
+      length_(underlying_stream_->length()),
+      flags_(underlying_stream_->flags()) {
   grpc_slice_buffer_init(&cache_buffer_);
 }
 
@@ -97,8 +99,7 @@ void ByteStreamCache::Destroy() {
 //
 
 ByteStreamCache::CachingByteStream::CachingByteStream(ByteStreamCache* cache)
-    : ByteStream(cache->underlying_stream_->length(),
-                 cache->underlying_stream_->flags()),
+    : ByteStream(cache->length_, cache->flags_),
       cache_(cache) {}
 
 ByteStreamCache::CachingByteStream::~CachingByteStream() {}
@@ -115,6 +116,7 @@ bool ByteStreamCache::CachingByteStream::Next(size_t max_size_hint,
                                               grpc_closure* on_complete) {
   if (shutdown_error_ != GRPC_ERROR_NONE) return true;
   if (cursor_ < cache_->cache_buffer_.count) return true;
+  GPR_ASSERT(cache_->underlying_stream_ != nullptr);
   return cache_->underlying_stream_->Next(max_size_hint, on_complete);
 }
 
@@ -128,6 +130,7 @@ grpc_error* ByteStreamCache::CachingByteStream::Pull(grpc_slice* slice) {
     offset_ += GRPC_SLICE_LENGTH(*slice);
     return GRPC_ERROR_NONE;
   }
+  GPR_ASSERT(cache_->underlying_stream_ != nullptr);
   grpc_error* error = cache_->underlying_stream_->Pull(slice);
   if (error == GRPC_ERROR_NONE) {
     grpc_slice_buffer_add(&cache_->cache_buffer_,
@@ -145,9 +148,14 @@ grpc_error* ByteStreamCache::CachingByteStream::Pull(grpc_slice* slice) {
 void ByteStreamCache::CachingByteStream::Shutdown(grpc_error* error) {
   GRPC_ERROR_UNREF(shutdown_error_);
   shutdown_error_ = GRPC_ERROR_REF(error);
-  cache_->underlying_stream_->Shutdown(error);
+  if (cache_->underlying_stream_ != nullptr) {
+    cache_->underlying_stream_->Shutdown(error);
+  }
 }
 
-void ByteStreamCache::CachingByteStream::Reset() { cursor_ = 0; }
+void ByteStreamCache::CachingByteStream::Reset() {
+  cursor_ = 0;
+  offset_ = 0;
+}
 
 }  // namespace grpc_core
diff --git a/src/core/lib/transport/byte_stream.h b/src/core/lib/transport/byte_stream.h
index f3247737f3..95a756a5e4 100644
--- a/src/core/lib/transport/byte_stream.h
+++ b/src/core/lib/transport/byte_stream.h
@@ -154,6 +154,8 @@ class ByteStreamCache {
 
  private:
   OrphanablePtr<ByteStream> underlying_stream_;
+  uint32_t length_;
+  uint32_t flags_;
   grpc_slice_buffer cache_buffer_;
 };