content: /*
   * Licensed to the Apache Software Foundation (ASF) under one or more
   * contributor license agreements.  See the NOTICE file distributed with
   * this work for additional information regarding copyright ownership.
   * The ASF licenses this file to You 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
   *
   *    http://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.
   */
  
  package org.apache.celeborn.common.network.protocol;
  
  import java.io.EOFException;
  import java.io.InputStream;
  
  import javax.annotation.Nullable;
  
  import com.google.common.base.Preconditions;
  import io.netty.buffer.ByteBuf;
  import io.netty.buffer.ByteBufAllocator;
  import io.netty.channel.ChannelHandlerContext;
  import io.netty.handler.stream.ChunkedInput;
  import io.netty.handler.stream.ChunkedStream;
  
  import org.apache.celeborn.common.network.buffer.ManagedBuffer;
  
  /**
   * A wrapper message that holds two separate pieces (a header and a body).
   *
   * <p>The header must be a ByteBuf, while the body can be any InputStream or ChunkedStream Based on
   * common/network-common/org.apache.spark.network.protocol.EncryptedMessageWithHeader
   */
  public class EncryptedMessageWithHeader implements ChunkedInput<ByteBuf> {
  
    @Nullable private final ManagedBuffer managedBuffer;
    private final ByteBuf header;
    private final int headerLength;
    private final Object body;
    private final long bodyLength;
    private long totalBytesTransferred;
  
    /**
     * Construct a new EncryptedMessageWithHeader.
     *
     * @param managedBuffer the {@link ManagedBuffer} that the message body came from. This needs to
     *     be passed in so that the buffer can be freed when this message is deallocated. Ownership of
     *     the caller's reference to this buffer is transferred to this class, so if the caller wants
     *     to continue to use the ManagedBuffer in other messages then they will need to call retain()
     *     on it before passing it to this constructor.
     * @param header the message header.
     * @param body the message body.
     * @param bodyLength the length of the message body, in bytes.
     */
    public EncryptedMessageWithHeader(
        @Nullable ManagedBuffer managedBuffer, ByteBuf header, Object body, long bodyLength) {
      Preconditions.checkArgument(
          body instanceof InputStream || body instanceof ChunkedStream,
          "Body must be an InputStream or a ChunkedStream.");
      this.managedBuffer = managedBuffer;
      this.header = header;
      this.headerLength = header.readableBytes();
      this.body = body;
      this.bodyLength = bodyLength;
      this.totalBytesTransferred = 0;
    }
  
    @Override
    public ByteBuf readChunk(ChannelHandlerContext ctx) throws Exception {
      return readChunk(ctx.alloc());
    }
  
    @Override
    public ByteBuf readChunk(ByteBufAllocator allocator) throws Exception {
      if (isEndOfInput()) {
        return null;
      }
  
      if (totalBytesTransferred < headerLength) {
        totalBytesTransferred += headerLength;
        return header.retain();
      } else if (body instanceof InputStream) {
        InputStream stream = (InputStream) body;
        int available = stream.available();
        if (available <= 0) {
          available = (int) (length() - totalBytesTransferred);
        } else {
          available = (int) Math.min(available, length() - totalBytesTransferred);
        }
        ByteBuf buffer = allocator.buffer(available);
        int toRead = Math.min(available, buffer.writableBytes());
        int read = buffer.writeBytes(stream, toRead);
        if (read >= 0) {
          totalBytesTransferred += read;
          return buffer;
        } else {
          throw new EOFException("Unable to read bytes from InputStream");
        }
      } else if (body instanceof ChunkedStream) {
        ChunkedStream stream = (ChunkedStream) body;
        long old = stream.transferredBytes();
        ByteBuf buffer = stream.readChunk(allocator);
        long read = stream.transferredBytes() - old;
        if (read >= 0) {
          totalBytesTransferred += read;
          assert (totalBytesTransferred <= length());
          return buffer;
        } else {
          throw new EOFException("Unable to read bytes from ChunkedStream");
        }
      } else {
        return null;
      }
    }
  
    @Override
    public long length() {
      return headerLength + bodyLength;
    }
  
    @Override
    public long progress() {
      return totalBytesTransferred;
    }
  
    @Override
    public boolean isEndOfInput() throws Exception {
      return (headerLength + bodyLength) == totalBytesTransferred;
    }
  
    @Override
    public void close() throws Exception {
      header.release();
      if (managedBuffer != null) {
        managedBuffer.release();
      }
      if (body instanceof InputStream) {
        ((InputStream) body).close();
      } else if (body instanceof ChunkedStream) {
        ((ChunkedStream) body).close();
      }
    }
  }
  