// Copyright (c) 2010 The Chromium Authors. All rights reserved.
// Use of this source code is governed by a BSD-style license that can be
// found in the LICENSE file.

#ifndef CHROME_BROWSER_SAFE_BROWSING_SAFE_BROWSING_STORE_SQLITE_H_
#define CHROME_BROWSER_SAFE_BROWSING_SAFE_BROWSING_STORE_SQLITE_H_

#include <set>
#include <vector>

#include "base/callback.h"
#include "chrome/browser/safe_browsing/safe_browsing_store.h"
#include "testing/gtest/include/gtest/gtest_prod.h"

struct sqlite3;
class SqliteCompiledStatement;
class SqliteStatementCache;
class SQLTransaction;

class SafeBrowsingStoreSqlite : public SafeBrowsingStore {
 public:
  SafeBrowsingStoreSqlite();
  virtual ~SafeBrowsingStoreSqlite();

  virtual bool Delete();

  virtual void Init(const FilePath& filename,
                    Callback0::Type* corruption_callback);

  virtual bool BeginChunk() {
    return true;
  }
  virtual bool WriteAddPrefix(int32 chunk_id, SBPrefix prefix) {
    const std::vector<SBAddPrefix> prefixes(1, SBAddPrefix(chunk_id, prefix));
    return WriteAddPrefixes(prefixes);
  }
  virtual bool WriteAddHash(int32 chunk_id,
                            base::Time receive_time, SBFullHash full_hash) {
    const std::vector<SBAddFullHash>
        hashes(1, SBAddFullHash(chunk_id, receive_time, full_hash));
    return WriteAddHashes(hashes);
  }
  virtual bool WriteSubPrefix(int32 chunk_id,
                              int32 add_chunk_id, SBPrefix prefix) {
    const std::vector<SBSubPrefix>
        prefixes(1, SBSubPrefix(chunk_id, add_chunk_id, prefix));
    return WriteSubPrefixes(prefixes);
  }
  virtual bool WriteSubHash(int32 chunk_id, int32 add_chunk_id,
                            SBFullHash full_hash) {
    const std::vector<SBSubFullHash>
        hashes(1, SBSubFullHash(chunk_id, add_chunk_id, full_hash));
    return WriteSubHashes(hashes);
  }
  virtual bool FinishChunk() {
    return true;
  }

  virtual bool BeginUpdate();
  // TODO(shess): Should not be public.
  virtual bool DoUpdate(const std::vector<SBAddFullHash>& pending_adds,
                        std::vector<SBAddPrefix>* add_prefixes_result,
                        std::vector<SBAddFullHash>* add_full_hashes_result);
  virtual bool FinishUpdate(const std::vector<SBAddFullHash>& pending_adds,
                            std::vector<SBAddPrefix>* add_prefixes_result,
                            std::vector<SBAddFullHash>* add_full_hashes_result);
  virtual bool CancelUpdate();

  virtual void SetAddChunk(int32 chunk_id) {
    add_chunks_cache_.insert(chunk_id);
  }
  virtual bool CheckAddChunk(int32 chunk_id) {
    return add_chunks_cache_.count(chunk_id) > 0;
  }
  virtual void GetAddChunks(std::vector<int32>* out) {
    out->clear();
    out->insert(out->end(), add_chunks_cache_.begin(), add_chunks_cache_.end());
  }

  virtual void SetSubChunk(int32 chunk_id) {
    sub_chunks_cache_.insert(chunk_id);
  }
  virtual bool CheckSubChunk(int32 chunk_id) {
    return sub_chunks_cache_.count(chunk_id) > 0;
  }
  virtual void GetSubChunks(std::vector<int32>* out) {
    out->clear();
    out->insert(out->end(), sub_chunks_cache_.begin(), sub_chunks_cache_.end());
  }

  virtual void DeleteAddChunk(int32 chunk_id) {
    add_del_cache_.insert(chunk_id);
  }
  virtual void DeleteSubChunk(int32 chunk_id) {
    sub_del_cache_.insert(chunk_id);
  }

  // Returns the name of the SQLite journal file for |filename|.
  // Exported for unit tests.
  static const FilePath JournalFileForFilename(const FilePath& filename) {
    return FilePath(filename.value() + FILE_PATH_LITERAL("-journal"));
  }

 private:
  // For on-the-fly migration.
  // TODO(shess): Remove (entire class) after migration.
  friend class SafeBrowsingStoreFile;

  // The following routines return true on success, or false on
  // failure.  Failure is presumed to be persistent, so the caller
  // should stop trying and unwind the transaction.
  // OnCorruptDatabase() is called if SQLite returns SQLITE_CORRUPT.

  // Open |db_| from |filename_|, creating if necessary.
  bool Open();

  // Close |db_|, rolling back any in-progress transaction.
  bool Close();

  // Execute all statements in sql, returning true if every one of
  // them returns SQLITE_OK.
  bool ExecSql(const char* sql);

  bool SetupDatabase();
  bool CheckCompatibleVersion();

  bool CreateTables();
  bool RenameTables();
  bool DeleteOldTables();

  // Read and write the chunks-seen data from |*_chunks_cache_|.
  // Chunk deletions are not accounted for.
  bool ReadAddChunks();
  bool ReadSubChunks();
  bool WriteAddChunks();
  bool WriteSubChunks();

  // Read the various types of data, skipping items which belong to
  // deleted chunks.  New data is appended to the vectors.
  bool ReadAddPrefixes(std::vector<SBAddPrefix>* add_prefixes);
  bool ReadSubPrefixes(std::vector<SBSubPrefix>* sub_prefixes);
  bool ReadAddHashes(std::vector<SBAddFullHash>* add_hashes);
  bool ReadSubHashes(std::vector<SBSubFullHash>* sub_hashes);

  // Write the various types of data.  The existing data is not
  // cleared.
  bool WriteAddPrefixes(const std::vector<SBAddPrefix>& add_prefixes);
  bool WriteSubPrefixes(const std::vector<SBSubPrefix>& sub_prefixes);
  bool WriteAddHashes(const std::vector<SBAddFullHash>& add_hashes);
  bool WriteSubHashes(const std::vector<SBSubFullHash>& sub_hashes);

  // Calls |corruption_callback_| if non-NULL, always returns false as
  // a convenience to the caller.
  bool OnCorruptDatabase();

  // The database path from Init().
  FilePath filename_;

  // Between BeginUpdate() and FinishUpdate(), this will be the SQLite
  // database connection.  Otherwise NULL.
  sqlite3 *db_;

  // Cache of compiled statements for |db_|.
  // TODO(shess): Probably doesn't gain us much.
  scoped_ptr<SqliteStatementCache> statement_cache_;

  // Transaction for protecting database integrity between
  // BeginUpdate() and FinishUpdate().
  scoped_ptr<SQLTransaction> insert_transaction_;

  // The set of chunks which the store has seen.  Elements are added
  // by SetAddChunk() and SetSubChunk(), and deleted on write for
  // chunks that have been deleted.
  std::set<int32> add_chunks_cache_;
  std::set<int32> sub_chunks_cache_;

  // Cache the DeletedAddChunk() and DeleteSubChunk() chunks for later
  // use in FinishUpdate().
  base::hash_set<int32> add_del_cache_;
  base::hash_set<int32> sub_del_cache_;

  // Called when SQLite returns SQLITE_CORRUPT.
  scoped_ptr<Callback0::Type> corruption_callback_;

  DISALLOW_COPY_AND_ASSIGN(SafeBrowsingStoreSqlite);
};

#endif  // CHROME_BROWSER_SAFE_BROWSING_SAFE_BROWSING_STORE_SQLITE_H_