limited_stream.h

// Copyright 2025 The Pigweed Authors
//
// Licensed under the Apache License, Version 2.0 (the "License"); you may not
// use this file except in compliance with the License. You may obtain a copy of
// the License at
//
//     https://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
// WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
// License for the specific language governing permissions and limitations under
// the License.
#pragma once
 
#include "pw_status/try.h"
#include "pw_stream/stream.h"
 
namespace pw::stream {
 
 
class LimitedStreamWriter : public NonSeekableWriter {
 public:
  LimitedStreamWriter(Writer& writer, size_t limit = kUnlimited)
      : writer_(writer), limit_(limit) {}
 
  size_t limit() const { return limit_; }
 
  void set_limit(size_t limit) { limit_ = limit; }
 
  size_t bytes_written() const { return written_; }
 
 private:
  Writer& writer_;
  size_t limit_;
  size_t written_ = 0;
 
  size_t remain() const {
    if (limit() == kUnlimited) {
      return kUnlimited;
    }
    if (written_ > limit_) {
      return 0;
    }
    return limit_ - written_;
  }
 
  Status DoWrite(ConstByteSpan data) override {
    // Write all or nothing.
    if (data.size() > remain()) {
      return Status::OutOfRange();
    }
    PW_TRY(writer_.Write(data));
    written_ += data.size();
    return OkStatus();
  }
 
  size_t ConservativeLimit(LimitType limit_type) const override {
    switch (limit_type) {
      case LimitType::kRead:
        return 0;
      case LimitType::kWrite:
        return std::min(remain(), writer_.ConservativeWriteLimit());
    }
  }
};
 
 
}  // namespace pw::stream