/*
 * Licensed to the Apache Software Foundation (ASF) under one or more
 * license agreements; and to You under the Apache License, version 2.0:
 *
 *   https://www.apache.org/licenses/LICENSE-2.0
 *
 * This file is part of the Apache Pekko project, which was derived from Akka.
 */

/*
 * Copyright (C) 2009-2022 Lightbend Inc. <https://www.lightbend.com>
 */

/*
 * Adapted from github.com/twitter/hpack with this license:
 *
 * Copyright 2014 Twitter, Inc.
 *
 * 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
 *
 *     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.pekko.http.shaded.com.twitter.hpack;

import java.io.IOException;
import java.io.InputStream;

import org.apache.pekko.http.impl.util.StringTools;
import org.apache.pekko.http.shaded.com.twitter.hpack.HpackUtil.IndexType;

import static org.apache.pekko.http.shaded.com.twitter.hpack.HeaderField.HEADER_ENTRY_OVERHEAD;

public final class Decoder {

  private static final IOException DECOMPRESSION_EXCEPTION =
      new IOException("decompression failure");
  private static final IOException ILLEGAL_INDEX_VALUE = new IOException("illegal index value");
  private static final IOException INVALID_MAX_DYNAMIC_TABLE_SIZE =
      new IOException("invalid max dynamic table size");
  private static final IOException MAX_DYNAMIC_TABLE_SIZE_CHANGE_REQUIRED =
      new IOException("max dynamic table size change required");

  private static final String EMPTY = "";

  private final DynamicTable dynamicTable;

  private int maxHeaderSize;
  private int maxDynamicTableSize;
  private int encoderMaxDynamicTableSize;
  private boolean maxDynamicTableSizeChangeRequired;

  private long headerSize;
  private State state;
  private IndexType indexType;
  private int index;
  private boolean huffmanEncoded;
  private int skipLength;
  private int nameLength;
  private int valueLength;
  private String name;

  private enum State {
    READ_HEADER_REPRESENTATION,
    READ_MAX_DYNAMIC_TABLE_SIZE,
    READ_INDEXED_HEADER,
    READ_INDEXED_HEADER_NAME,
    READ_LITERAL_HEADER_NAME_LENGTH_PREFIX,
    READ_LITERAL_HEADER_NAME_LENGTH,
    READ_LITERAL_HEADER_NAME,
    SKIP_LITERAL_HEADER_NAME,
    READ_LITERAL_HEADER_VALUE_LENGTH_PREFIX,
    READ_LITERAL_HEADER_VALUE_LENGTH,
    READ_LITERAL_HEADER_VALUE,
    SKIP_LITERAL_HEADER_VALUE
  }

  /** Creates a new decoder. */
  public Decoder(int maxHeaderSize, int maxHeaderTableSize) {
    dynamicTable = new DynamicTable(maxHeaderTableSize);
    this.maxHeaderSize = maxHeaderSize;
    maxDynamicTableSize = maxHeaderTableSize;
    encoderMaxDynamicTableSize = maxHeaderTableSize;
    maxDynamicTableSizeChangeRequired = false;
    reset();
  }

  private void reset() {
    headerSize = 0;
    state = State.READ_HEADER_REPRESENTATION;
    indexType = IndexType.NONE;
  }

  /** Decode the header block into header fields. */
  public void decode(InputStream in, HeaderListener headerListener) throws IOException {
    while (in.available() > 0) {
      switch (state) {
        case READ_HEADER_REPRESENTATION:
          byte b = (byte) in.read();
          if (maxDynamicTableSizeChangeRequired && (b & 0xE0) != 0x20) {
            // Encoder MUST signal maximum dynamic table size change
            throw MAX_DYNAMIC_TABLE_SIZE_CHANGE_REQUIRED;
          }
          if (b < 0) {
            // Indexed Header Field
            index = b & 0x7F;
            if (index == 0) {
              throw ILLEGAL_INDEX_VALUE;
            } else if (index == 0x7F) {
              state = State.READ_INDEXED_HEADER;
            } else {
              indexHeader(index, headerListener);
            }
          } else if ((b & 0x40) == 0x40) {
            // Literal Header Field with Incremental Indexing
            indexType = IndexType.INCREMENTAL;
            index = b & 0x3F;
            if (index == 0) {
              state = State.READ_LITERAL_HEADER_NAME_LENGTH_PREFIX;
            } else if (index == 0x3F) {
              state = State.READ_INDEXED_HEADER_NAME;
            } else {
              // Index was stored as the prefix
              readName(index);
              state = State.READ_LITERAL_HEADER_VALUE_LENGTH_PREFIX;
            }
          } else if ((b & 0x20) == 0x20) {
            // Dynamic Table Size Update
            index = b & 0x1F;
            if (index == 0x1F) {
              state = State.READ_MAX_DYNAMIC_TABLE_SIZE;
            } else {
              setDynamicTableSize(index);
              state = State.READ_HEADER_REPRESENTATION;
            }
          } else {
            // Literal Header Field without Indexing / never Indexed
            indexType = ((b & 0x10) == 0x10) ? IndexType.NEVER : IndexType.NONE;
            index = b & 0x0F;
            if (index == 0) {
              state = State.READ_LITERAL_HEADER_NAME_LENGTH_PREFIX;
            } else if (index == 0x0F) {
              state = State.READ_INDEXED_HEADER_NAME;
            } else {
              // Index was stored as the prefix
              readName(index);
              state = State.READ_LITERAL_HEADER_VALUE_LENGTH_PREFIX;
            }
          }
          break;

        case READ_MAX_DYNAMIC_TABLE_SIZE:
          int maxSize = decodeULE128(in);
          if (maxSize == -1) {
            return;
          }

          // Check for numerical overflow
          if (maxSize > Integer.MAX_VALUE - index) {
            throw DECOMPRESSION_EXCEPTION;
          }

          setDynamicTableSize(index + maxSize);
          state = State.READ_HEADER_REPRESENTATION;
          break;

        case READ_INDEXED_HEADER:
          int headerIndex = decodeULE128(in);
          if (headerIndex == -1) {
            return;
          }

          // Check for numerical overflow
          if (headerIndex > Integer.MAX_VALUE - index) {
            throw DECOMPRESSION_EXCEPTION;
          }

          indexHeader(index + headerIndex, headerListener);
          state = State.READ_HEADER_REPRESENTATION;
          break;

        case READ_INDEXED_HEADER_NAME:
          // Header Name matches an entry in the Header Table
          int nameIndex = decodeULE128(in);
          if (nameIndex == -1) {
            return;
          }

          // Check for numerical overflow
          if (nameIndex > Integer.MAX_VALUE - index) {
            throw DECOMPRESSION_EXCEPTION;
          }

          readName(index + nameIndex);
          state = State.READ_LITERAL_HEADER_VALUE_LENGTH_PREFIX;
          break;

        case READ_LITERAL_HEADER_NAME_LENGTH_PREFIX:
          b = (byte) in.read();
          huffmanEncoded = (b & 0x80) == 0x80;
          index = b & 0x7F;
          if (index == 0x7f) {
            state = State.READ_LITERAL_HEADER_NAME_LENGTH;
          } else {
            nameLength = index;

            // Disallow empty names -- they cannot be represented in HTTP/1.x
            if (nameLength == 0) {
              throw DECOMPRESSION_EXCEPTION;
            }

            // Check name length against max header size
            if (exceedsMaxHeaderSize(nameLength)) {

              if (indexType == IndexType.NONE) {
                // Name is unused so skip bytes
                name = EMPTY;
                skipLength = nameLength;
                state = State.SKIP_LITERAL_HEADER_NAME;
                break;
              }

              // Check name length against max dynamic table size
              if (nameLength + HEADER_ENTRY_OVERHEAD > dynamicTable.capacity()) {
                dynamicTable.clear();
                name = EMPTY;
                skipLength = nameLength;
                state = State.SKIP_LITERAL_HEADER_NAME;
                break;
              }
            }
            state = State.READ_LITERAL_HEADER_NAME;
          }
          break;

        case READ_LITERAL_HEADER_NAME_LENGTH:
          // Header Name is a Literal String
          nameLength = decodeULE128(in);
          if (nameLength == -1) {
            return;
          }

          // Check for numerical overflow
          if (nameLength > Integer.MAX_VALUE - index) {
            throw DECOMPRESSION_EXCEPTION;
          }
          nameLength += index;

          // Check name length against max header size
          if (exceedsMaxHeaderSize(nameLength)) {
            if (indexType == IndexType.NONE) {
              // Name is unused so skip bytes
              name = EMPTY;
              skipLength = nameLength;
              state = State.SKIP_LITERAL_HEADER_NAME;
              break;
            }

            // Check name length against max dynamic table size
            if (nameLength + HEADER_ENTRY_OVERHEAD > dynamicTable.capacity()) {
              dynamicTable.clear();
              name = EMPTY;
              skipLength = nameLength;
              state = State.SKIP_LITERAL_HEADER_NAME;
              break;
            }
          }
          state = State.READ_LITERAL_HEADER_NAME;
          break;

        case READ_LITERAL_HEADER_NAME:
          // Wait until entire name is readable
          if (in.available() < nameLength) {
            return;
          }

          name = readStringLiteral(in, nameLength);

          state = State.READ_LITERAL_HEADER_VALUE_LENGTH_PREFIX;
          break;

        case SKIP_LITERAL_HEADER_NAME:
          skipLength -= in.skip(skipLength);

          if (skipLength == 0) {
            state = State.READ_LITERAL_HEADER_VALUE_LENGTH_PREFIX;
          }
          break;

        case READ_LITERAL_HEADER_VALUE_LENGTH_PREFIX:
          b = (byte) in.read();
          huffmanEncoded = (b & 0x80) == 0x80;
          index = b & 0x7F;
          if (index == 0x7f) {
            state = State.READ_LITERAL_HEADER_VALUE_LENGTH;
          } else {
            valueLength = index;

            // Check new header size against max header size
            long newHeaderSize = (long) nameLength + (long) valueLength;
            if (exceedsMaxHeaderSize(newHeaderSize)) {
              // truncation will be reported during endHeaderBlock
              headerSize = maxHeaderSize + 1;

              if (indexType == IndexType.NONE) {
                // Value is unused so skip bytes
                state = State.SKIP_LITERAL_HEADER_VALUE;
                break;
              }

              // Check new header size against max dynamic table size
              if (newHeaderSize + HEADER_ENTRY_OVERHEAD > dynamicTable.capacity()) {
                dynamicTable.clear();
                state = State.SKIP_LITERAL_HEADER_VALUE;
                break;
              }
            }

            if (valueLength == 0) {
              insertHeader(headerListener, name, EMPTY, indexType);
              state = State.READ_HEADER_REPRESENTATION;
            } else {
              state = State.READ_LITERAL_HEADER_VALUE;
            }
          }

          break;

        case READ_LITERAL_HEADER_VALUE_LENGTH:
          // Header Value is a Literal String
          valueLength = decodeULE128(in);
          if (valueLength == -1) {
            return;
          }

          // Check for numerical overflow
          if (valueLength > Integer.MAX_VALUE - index) {
            throw DECOMPRESSION_EXCEPTION;
          }
          valueLength += index;

          // Check new header size against max header size
          long newHeaderSize = (long) nameLength + (long) valueLength;
          if (newHeaderSize + headerSize > maxHeaderSize) {
            // truncation will be reported during endHeaderBlock
            headerSize = maxHeaderSize + 1;

            if (indexType == IndexType.NONE) {
              // Value is unused so skip bytes
              state = State.SKIP_LITERAL_HEADER_VALUE;
              break;
            }

            // Check new header size against max dynamic table size
            if (newHeaderSize + HEADER_ENTRY_OVERHEAD > dynamicTable.capacity()) {
              dynamicTable.clear();
              state = State.SKIP_LITERAL_HEADER_VALUE;
              break;
            }
          }
          state = State.READ_LITERAL_HEADER_VALUE;
          break;

        case READ_LITERAL_HEADER_VALUE:
          // Wait until entire value is readable
          if (in.available() < valueLength) {
            return;
          }

          String value = readStringLiteral(in, valueLength);
          insertHeader(headerListener, name, value, indexType);
          state = State.READ_HEADER_REPRESENTATION;
          break;

        case SKIP_LITERAL_HEADER_VALUE:
          valueLength -= in.skip(valueLength);

          if (valueLength == 0) {
            state = State.READ_HEADER_REPRESENTATION;
          }
          break;

        default:
          throw new IllegalStateException("should not reach here");
      }
    }
  }

  /**
   * End the current header block. Returns if the header field has been truncated. This must be
   * called after the header block has been completely decoded.
   */
  public boolean endHeaderBlock() {
    boolean truncated = headerSize > maxHeaderSize;
    reset();
    return truncated;
  }

  /**
   * Set the maximum table size. If this is below the maximum size of the dynamic table used by the
   * encoder, the beginning of the next header block MUST signal this change.
   */
  public void setMaxHeaderTableSize(int maxHeaderTableSize) {
    maxDynamicTableSize = maxHeaderTableSize;
    if (maxDynamicTableSize < encoderMaxDynamicTableSize) {
      // decoder requires less space than encoder
      // encoder MUST signal this change
      maxDynamicTableSizeChangeRequired = true;
      dynamicTable.setCapacity(maxDynamicTableSize);
    }
  }

  /**
   * Return the maximum table size. This is the maximum size allowed by both the encoder and the
   * decoder.
   */
  public int getMaxHeaderTableSize() {
    return dynamicTable.capacity();
  }

  /** Return the number of header fields in the dynamic table. Exposed for testing. */
  int length() {
    return dynamicTable.length();
  }

  /** Return the size of the dynamic table. Exposed for testing. */
  int size() {
    return dynamicTable.size();
  }

  /** Return the header field at the given index. Exposed for testing. */
  HeaderField getHeaderField(int index) {
    return dynamicTable.getEntry(index + 1);
  }

  private void setDynamicTableSize(int dynamicTableSize) throws IOException {
    if (dynamicTableSize > maxDynamicTableSize) {
      throw INVALID_MAX_DYNAMIC_TABLE_SIZE;
    }
    encoderMaxDynamicTableSize = dynamicTableSize;
    maxDynamicTableSizeChangeRequired = false;
    dynamicTable.setCapacity(dynamicTableSize);
  }

  private void readName(int index) throws IOException {
    if (index <= StaticTable.length) {
      HeaderField headerField = StaticTable.getEntry(index);
      name = headerField.name;
    } else if (index - StaticTable.length <= dynamicTable.length()) {
      HeaderField headerField = dynamicTable.getEntry(index - StaticTable.length);
      name = headerField.name;
    } else {
      throw ILLEGAL_INDEX_VALUE;
    }
  }

  private void indexHeader(int index, HeaderListener headerListener) throws IOException {
    if (index <= StaticTable.length) {
      HeaderField headerField = StaticTable.getEntry(index);
      Object parsed =
          addHeader(
              headerListener, headerField.name, headerField.value, headerField.parsedValue, false);
      headerField.parsedValue = parsed;
    } else if (index - StaticTable.length <= dynamicTable.length()) {
      HeaderField headerField = dynamicTable.getEntry(index - StaticTable.length);
      Object parsed =
          addHeader(
              headerListener, headerField.name, headerField.value, headerField.parsedValue, false);
      headerField.parsedValue = parsed;
    } else {
      throw ILLEGAL_INDEX_VALUE;
    }
  }

  private void insertHeader(
      HeaderListener headerListener, String name, String value, IndexType indexType) {
    Object parsedValue = addHeader(headerListener, name, value, null, indexType == IndexType.NEVER);

    switch (indexType) {
      case NONE:
      case NEVER:
        break;

      case INCREMENTAL:
        dynamicTable.add(new HeaderField(name, value, parsedValue));
        break;

      default:
        throw new IllegalStateException("should not reach here");
    }
  }

  private Object addHeader(
      HeaderListener headerListener,
      String name,
      String value,
      Object parsedValue,
      boolean sensitive) {
    if (name.isEmpty()) {
      throw new AssertionError("name is empty");
    }
    long newSize = headerSize + name.length() + value.length();
    if (newSize <= maxHeaderSize) {
      parsedValue = headerListener.addHeader(name, value, parsedValue, sensitive);
      headerSize = (int) newSize;
    } else {
      // truncation will be reported during endHeaderBlock
      headerSize = maxHeaderSize + 1;
    }
    return parsedValue;
  }

  private boolean exceedsMaxHeaderSize(long size) {
    // Check new header size against max header size
    if (size + headerSize <= maxHeaderSize) {
      return false;
    }

    // truncation will be reported during endHeaderBlock
    headerSize = maxHeaderSize + 1;
    return true;
  }

  private String readStringLiteral(InputStream in, int length) throws IOException {
    byte[] buf = new byte[length];
    if (in.read(buf) != length) {
      throw DECOMPRESSION_EXCEPTION;
    }
    final byte[] result;

    if (huffmanEncoded) {
      result = Huffman.DECODER.decode(buf);
    } else {
      result = buf;
    }
    return StringTools.asciiStringFromBytes(result);
  }

  // Unsigned Little Endian Base 128 Variable-Length Integer Encoding
  private static int decodeULE128(InputStream in) throws IOException {
    in.mark(5);
    int result = 0;
    int shift = 0;
    while (shift < 32) {
      if (in.available() == 0) {
        // Buffer does not contain entire integer,
        // reset reader index and return -1.
        in.reset();
        return -1;
      }
      byte b = (byte) in.read();
      if (shift == 28 && (b & 0xF8) != 0) {
        break;
      }
      result |= (b & 0x7F) << shift;
      if ((b & 0x80) == 0) {
        return result;
      }
      shift += 7;
    }
    // Value exceeds Integer.MAX_VALUE
    in.reset();
    throw DECOMPRESSION_EXCEPTION;
  }
}
