/* This Source Code Form is subject to the terms of the Mozilla Public
 * License, v. 2.0. If a copy of the MPL was not distributed with this
 * file, You can obtain one at http://mozilla.org/MPL/2.0/. */

#ifndef ProtocolParser_h_
#define ProtocolParser_h_

#include "HashStore.h"
#include "chromium/safebrowsing.pb.h"
#include "chromium/safebrowsing_v5.pb.h"

namespace mozilla {
namespace safebrowsing {

/**
 * Abstract base class for parsing update data in multiple formats.
 */
class ProtocolParser {
 public:
  struct ForwardedUpdate {
    nsCString table;
    nsCString url;
  };

  ProtocolParser();
  virtual ~ProtocolParser();

  nsresult Status() const { return mUpdateStatus; }

  virtual void SetCurrentTable(const nsACString& aTable) = 0;

  void SetRequestedTables(const nsTArray<nsCString>& aRequestTables) {
    mRequestedTables = aRequestTables.Clone();
  }

  nsresult Begin(const nsACString& aTable,
                 const nsTArray<nsCString>& aUpdateTables);
  virtual nsresult AppendStream(const nsACString& aData) = 0;

  uint32_t UpdateWaitSec() { return mUpdateWaitSec; }

  // Notify that the inbound data is ready for parsing if progressive
  // parsing is not supported, for example in V4.
  virtual void End() = 0;

  RefPtr<TableUpdate> GetTableUpdate(const nsACString& aTable);
  void ForgetTableUpdates() { mTableUpdates.Clear(); }
  const TableUpdateArray& GetTableUpdates() { return mTableUpdates; }

  // These are only meaningful to V2. Since they were originally public,
  // moving them to ProtocolParserV2 requires a dymamic cast in the call
  // sites. As a result, we will leave them until we remove support
  // for V2 entirely..
  virtual const nsTArray<ForwardedUpdate>& Forwards() const {
    return mForwards;
  }
  bool ResetRequested() const { return !mTablesToReset.IsEmpty(); }
  const nsTArray<nsCString>& TablesToReset() const { return mTablesToReset; }

 protected:
  virtual RefPtr<TableUpdate> CreateTableUpdate(
      const nsACString& aTableName) const = 0;

  nsCString mPending;
  nsresult mUpdateStatus;

  // Keep track of updates to apply before passing them to the DBServiceWorkers.
  TableUpdateArray mTableUpdates;

  nsTArray<ForwardedUpdate> mForwards;

  // The table names that were requested from the client.
  nsTArray<nsCString> mRequestedTables;

  // The table names that failed to update and need to be reset.
  nsTArray<nsCString> mTablesToReset;

  // How long we should wait until the next update.
  uint32_t mUpdateWaitSec;
};

/**
 * Helpers to parse the "shavar", "digest256" and "simple" list formats.
 */
class ProtocolParserV2 final : public ProtocolParser {
 public:
  ProtocolParserV2();
  virtual ~ProtocolParserV2();

  virtual void SetCurrentTable(const nsACString& aTable) override;
  virtual nsresult AppendStream(const nsACString& aData) override;
  virtual void End() override;

  // Update information.
  virtual const nsTArray<ForwardedUpdate>& Forwards() const override {
    return mForwards;
  }

 private:
  virtual RefPtr<TableUpdate> CreateTableUpdate(
      const nsACString& aTableName) const override;

  nsresult ProcessControl(bool* aDone);
  nsresult ProcessExpirations(const nsCString& aLine);
  nsresult ProcessChunkControl(const nsCString& aLine);
  nsresult ProcessForward(const nsCString& aLine);
  nsresult AddForward(const nsACString& aUrl);
  nsresult ProcessChunk(bool* done);
  // Remove this, it's only used for testing
  nsresult ProcessPlaintextChunk(const nsACString& aChunk);
  nsresult ProcessShaChunk(const nsACString& aChunk);
  nsresult ProcessHostAdd(const Prefix& aDomain, uint8_t aNumEntries,
                          const nsACString& aChunk, uint32_t* aStart);
  nsresult ProcessHostSub(const Prefix& aDomain, uint8_t aNumEntries,
                          const nsACString& aChunk, uint32_t* aStart);
  nsresult ProcessHostAddComplete(uint8_t aNumEntries, const nsACString& aChunk,
                                  uint32_t* aStart);
  nsresult ProcessHostSubComplete(uint8_t numEntries, const nsACString& aChunk,
                                  uint32_t* start);
  // Digest chunks are very similar to shavar chunks, except digest chunks
  // always contain the full hash, so there is no need for chunk data to
  // contain prefix sizes.
  nsresult ProcessDigestChunk(const nsACString& aChunk);
  nsresult ProcessDigestAdd(const nsACString& aChunk);
  nsresult ProcessDigestSub(const nsACString& aChunk);
  bool NextLine(nsACString& aLine);

  enum ParserState { PROTOCOL_STATE_CONTROL, PROTOCOL_STATE_CHUNK };
  ParserState mState;

  enum ChunkType {
    // Types for shavar tables.
    CHUNK_ADD,
    CHUNK_SUB,
    // Types for digest256 tables. digest256 tables differ in format from
    // shavar tables since they only contain complete hashes.
    CHUNK_ADD_DIGEST,
    CHUNK_SUB_DIGEST
  };

  struct ChunkState {
    ChunkType type;
    uint32_t num;
    uint32_t hashSize;
    uint32_t length;
    void Clear() {
      num = 0;
      hashSize = 0;
      length = 0;
    }
  };
  ChunkState mChunkState;

  // Updates to apply to the current table being parsed.
  RefPtr<TableUpdateV2> mTableUpdate;
};

// Helpers to parse the "proto" list format.
class ProtocolParserProtobuf final : public ProtocolParser {
 public:
  typedef FetchThreatListUpdatesResponse_ListUpdateResponse ListUpdateResponse;
  typedef google::protobuf::RepeatedPtrField<ThreatEntrySet> ThreatEntrySetList;

 public:
  ProtocolParserProtobuf();

  virtual void SetCurrentTable(const nsACString& aTable) override;
  virtual nsresult AppendStream(const nsACString& aData) override;
  virtual void End() override;

 private:
  virtual ~ProtocolParserProtobuf();

  virtual RefPtr<TableUpdate> CreateTableUpdate(
      const nsACString& aTableName) const override;

  // For parsing update info.
  nsresult ProcessOneResponse(const ListUpdateResponse& aResponse,
                              nsACString& aListName);

  nsresult ProcessAdditionOrRemoval(TableUpdateV4& aTableUpdate,
                                    const ThreatEntrySetList& aUpdate,
                                    bool aIsAddition);

  nsresult ProcessRawAddition(TableUpdateV4& aTableUpdate,
                              const ThreatEntrySet& aAddition);

  nsresult ProcessRawRemoval(TableUpdateV4& aTableUpdate,
                             const ThreatEntrySet& aRemoval);

  nsresult ProcessEncodedAddition(TableUpdateV4& aTableUpdate,
                                  const ThreatEntrySet& aAddition);

  nsresult ProcessEncodedRemoval(TableUpdateV4& aTableUpdate,
                                 const ThreatEntrySet& aRemoval);
};

class ProtocolParserProtobufV5 final : public ProtocolParser {
 public:
  ProtocolParserProtobufV5();

  virtual void SetCurrentTable(const nsACString& aTable) override;
  virtual nsresult AppendStream(const nsACString& aData) override;
  virtual void End() override;

 private:
  virtual ~ProtocolParserProtobufV5();

  virtual RefPtr<TableUpdate> CreateTableUpdate(
      const nsACString& aTableName) const override;

  // Process a single hash list in the response.
  nsresult ProcessOneResponse(const v5::HashList& aHashList,
                              nsACString& aListName);

  // Process the additions for a 4-byte encoded prefixes.
  nsresult ProcessAddition4Bytes(TableUpdateV4& aTableUpdate,
                                 const v5::RiceDeltaEncoded32Bit& aAddition);

  // Process the additions for a 8-byte encoded prefixes.
  nsresult ProcessAddition8Bytes(TableUpdateV4& aTableUpdate,
                                 const v5::RiceDeltaEncoded64Bit& aAddition);

  // Process the additions for a 16-byte encoded prefixes.
  nsresult ProcessAddition16Bytes(TableUpdateV4& aTableUpdate,
                                  const v5::RiceDeltaEncoded128Bit& aAddition);

  // Process the additions for a 32-byte encoded prefixes.
  nsresult ProcessAddition32Bytes(TableUpdateV4& aTableUpdate,
                                  const v5::RiceDeltaEncoded256Bit& aAddition);

  // Process the removals for a encoded prefixes.
  nsresult ProcessRemoval(TableUpdateV4& aTableUpdate,
                          const v5::RiceDeltaEncoded32Bit& aRemoval);
};

}  // namespace safebrowsing
}  // namespace mozilla

#endif
