// Copyright 2016 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.

#include "stdafx.h"

#include <algorithm>

#include "system_information_sampler.h"

// From ntdef.h
typedef struct _UNICODE_STRING {
  USHORT Length;
  USHORT MaximumLength;
  PWCH Buffer;
} UNICODE_STRING;

// From <wdm.h>
typedef LONG KPRIORITY;
typedef LONG KWAIT_REASON;  // Full definition is in wdm.h

// From ntddk.h
typedef struct _VM_COUNTERS {
  SIZE_T PeakVirtualSize;
  SIZE_T VirtualSize;
  ULONG PageFaultCount;
  // Padding here in 64-bit
  SIZE_T PeakWorkingSetSize;
  SIZE_T WorkingSetSize;
  SIZE_T QuotaPeakPagedPoolUsage;
  SIZE_T QuotaPagedPoolUsage;
  SIZE_T QuotaPeakNonPagedPoolUsage;
  SIZE_T QuotaNonPagedPoolUsage;
  SIZE_T PagefileUsage;
  SIZE_T PeakPagefileUsage;
} VM_COUNTERS;

// Two possibilities available from here:
// http://stackoverflow.com/questions/28858849/where-is-system-information-class-defined

typedef enum _SYSTEM_INFORMATION_CLASS {
  SystemBasicInformation = 0,
  SystemPerformanceInformation = 2,
  SystemTimeOfDayInformation = 3,
  SystemProcessInformation = 5,  // This is the number that we need
  SystemProcessorPerformanceInformation = 8,
  SystemInterruptInformation = 23,
  SystemExceptionInformation = 33,
  SystemRegistryQuotaInformation = 37,
  SystemLookasideInformation = 45
} SYSTEM_INFORMATION_CLASS;

// https://msdn.microsoft.com/en-us/library/gg750647.aspx?f=255&MSPPError=-2147217396
typedef struct {
  HANDLE UniqueProcess;  // Actually process ID
  HANDLE UniqueThread;   // Actually thread ID
} CLIENT_ID;

// From http://alax.info/blog/1182, with corrections and modifications
// Originally from
// http://undocumented.ntinternals.net/index.html?page=UserMode%2FUndocumented%20Functions%2FSystem%20Information%2FStructures%2FSYSTEM_THREAD.html
struct SYSTEM_THREAD_INFORMATION {
  ULONGLONG KernelTime;
  ULONGLONG UserTime;
  ULONGLONG CreateTime;
  ULONG WaitTime;
  // Padding here in 64-bit
  PVOID StartAddress;
  CLIENT_ID ClientId;
  KPRIORITY Priority;
  LONG BasePriority;
  ULONG ContextSwitchCount;
  ULONG State;
  KWAIT_REASON WaitReason;
};
#if _M_X64
static_assert(sizeof(SYSTEM_THREAD_INFORMATION) == 80,
              "Structure size mismatch");
#else
static_assert(sizeof(SYSTEM_THREAD_INFORMATION) == 64,
              "Structure size mismatch");
#endif

// From http://alax.info/blog/1182, with corrections and modifications
// Originally from
// http://undocumented.ntinternals.net/index.html?page=UserMode%2FUndocumented%20Functions%2FSystem%20Information%2FStructures%2FSYSTEM_THREAD.html
struct SYSTEM_PROCESS_INFORMATION {
  ULONG NextEntryOffset;
  ULONG NumberOfThreads;
  // http://processhacker.sourceforge.net/doc/struct___s_y_s_t_e_m___p_r_o_c_e_s_s___i_n_f_o_r_m_a_t_i_o_n.html
  ULONGLONG WorkingSetPrivateSize;
  ULONG HardFaultCount;
  ULONG Reserved1;
  ULONGLONG CycleTime;
  ULONGLONG CreateTime;
  ULONGLONG UserTime;
  ULONGLONG KernelTime;
  UNICODE_STRING ImageName;
  KPRIORITY BasePriority;
  HANDLE ProcessId;
  HANDLE ParentProcessId;
  ULONG HandleCount;
  ULONG Reserved2[2];
  // Padding here in 64-bit
  VM_COUNTERS VirtualMemoryCounters;
  size_t Reserved3;
  IO_COUNTERS IoCounters;
  SYSTEM_THREAD_INFORMATION Threads[1];
};
#if _M_X64
static_assert(sizeof(SYSTEM_PROCESS_INFORMATION) == 336,
              "Structure size mismatch");
#else
static_assert(sizeof(SYSTEM_PROCESS_INFORMATION) == 248,
              "Structure size mismatch");
#endif

// ntstatus.h conflicts with windows.h so define this locally.
#define STATUS_SUCCESS ((NTSTATUS)0x00000000L)
#define STATUS_BUFFER_TOO_SMALL ((NTSTATUS)0xC0000023L)
#define STATUS_INFO_LENGTH_MISMATCH ((NTSTATUS)0xC0000004L)

typedef NTSTATUS(WINAPI* NTQUERYSYSTEMINFORMATION)(
    SYSTEM_INFORMATION_CLASS SystemInformationClass,
    PVOID SystemInformation,
    ULONG SystemInformationLength,
    PULONG ReturnLength);

__declspec(noreturn) void oops(const char* pMessage) {
  printf("%s\n", pMessage);
  exit(0);
}

// Simple memory buffer wrapper for passing the data out of
// QuerySystemProcessInformation.
class ByteBuffer {
 public:
  explicit ByteBuffer(size_t capacity) : size_(0), capacity_(0) {
    if (capacity > 0)
      grow(capacity);
  }

  ~ByteBuffer() {}

  BYTE* data() { return data_.get(); }

  size_t size() { return size_; }

  void set_size(size_t new_size) { size_ = new_size; }

  size_t capacity() { return capacity_; }

  void grow(size_t new_capacity) {
    capacity_ = new_capacity;
    data_.reset(new BYTE[new_capacity]);
  }

 private:
  std::unique_ptr<BYTE[]> data_;
  size_t size_;
  size_t capacity_;

  ByteBuffer& operator=(const ByteBuffer&) = delete;
  ByteBuffer(const ByteBuffer&) = delete;
};

// Wrapper for NtQuerySystemProcessInformation with buffer reallocation logic.
bool QuerySystemProcessInformation(ByteBuffer* buffer) {
  typedef NTSTATUS(WINAPI * NTQUERYSYSTEMINFORMATION)(
      SYSTEM_INFORMATION_CLASS SystemInformationClass, PVOID SystemInformation,
      ULONG SystemInformationLength, PULONG ReturnLength);

  HMODULE ntdll = GetModuleHandle(L"ntdll.dll");
  if (!ntdll) {
    oops("Couldn't load ntdll.dll");
  }

  NTQUERYSYSTEMINFORMATION nt_query_system_information_ptr =
      reinterpret_cast<NTQUERYSYSTEMINFORMATION>(
          GetProcAddress(ntdll, "NtQuerySystemInformation"));
  if (!nt_query_system_information_ptr)
    oops("Couldn't find NtQuerySystemInformation");

  NTSTATUS result;

  // There is a potential race condition between growing the buffer and new
  // processes being created. Try a few times before giving up.
  for (int i = 0; i < 10; i++) {
    ULONG data_size = 0;
    ULONG buffer_size = static_cast<ULONG>(buffer->capacity());
    result = nt_query_system_information_ptr(
        SystemProcessInformation, buffer->data(), buffer_size, &data_size);

    if (result == STATUS_SUCCESS) {
      buffer->set_size(data_size);
      break;
    }

    if (result == STATUS_INFO_LENGTH_MISMATCH ||
        result == STATUS_BUFFER_TOO_SMALL) {
      // Insufficient buffer. Grow to the returned |data_size| plus 10% extra
      // to avoid frequent reallocations and try again.
      buffer->grow(static_cast<ULONG>(data_size * 1.1));
    } else {
      // An error other than the two above.
      break;
    }
  }

  return result == STATUS_SUCCESS;
}

SystemInformationSampler::SystemInformationSampler(
    const wchar_t* process_name) {
  lstrcpyn(target_process_name_, process_name,
           sizeof(target_process_name_) / sizeof(wchar_t));

  QueryPerformanceFrequency(&perf_frequency_);
  QueryPerformanceCounter(&initial_counter_);
}

SystemInformationSampler::~SystemInformationSampler() {}

std::unique_ptr<ProcessDataSnapshot> SystemInformationSampler::TakeSnapshot() {
  // Preallocate the buffer with the size determined on the previous call to
  // QuerySystemProcessInformation. This should be sufficient most of the time.
  // QuerySystemProcessInformation will grow the buffer if necessary.
  ByteBuffer data_buffer(previous_buffer_size_);

  if (!QuerySystemProcessInformation(&data_buffer))
    return std::unique_ptr<ProcessDataSnapshot>();

  previous_buffer_size_ = data_buffer.capacity();

  std::unique_ptr<ProcessDataSnapshot> snapshot(new ProcessDataSnapshot);

  LARGE_INTEGER perf_counter_value;
  QueryPerformanceCounter(&perf_counter_value);
  snapshot->timestamp = static_cast<double>(
      (perf_counter_value.QuadPart - initial_counter_.QuadPart) /
      perf_frequency_.QuadPart);

  for (size_t offset = 0; offset < data_buffer.size();) {
    auto pi = reinterpret_cast<const SYSTEM_PROCESS_INFORMATION*>(
        data_buffer.data() + offset);

    // Validate that the offset is valid and all needed data is within
    // the buffer boundary.
    if (offset + sizeof(SYSTEM_PROCESS_INFORMATION) > data_buffer.size())
      break;
    if (offset + sizeof(SYSTEM_PROCESS_INFORMATION) +
            (pi->NumberOfThreads - 1) * sizeof(SYSTEM_THREAD_INFORMATION) >
        data_buffer.size())
      break;

    if (pi->ImageName.Buffer) {
      // Validate that the image name is within the buffer boundary.
      // ImageName.Length seems to be in bytes rather than characters.
      size_t image_name_offset =
          reinterpret_cast<BYTE*>(pi->ImageName.Buffer) - data_buffer.data();
      if (image_name_offset + pi->ImageName.Length > data_buffer.size())
        break;

      // Check if this is a chrome process. Ignore all other processes.
      if (wcsncmp(target_process_name_filter(), pi->ImageName.Buffer,
                  lstrlen(target_process_name_filter())) == 0) {
        // Collect enough data to be able to do a diff between two snapshots.
        // Some threads might stop or new threads might be created between two
        // snapshots. If a thread with a large number of context switches gets
        // terminated the total number of context switches for the process might
        // go down and the delta would be negative.
        // To avoid that we need to compare thread IDs between two snapshots and
        // not count context switches for threads that are missing in the most
        // recent snapshot.
        ProcessData process_data;

        process_data.cpu_time = pi->KernelTime + pi->UserTime;
        process_data.working_set = pi->WorkingSetPrivateSize;

        // Iterate over threads and store each thread's ID and number of context
        // switches.
        for (ULONG thread_index = 0; thread_index < pi->NumberOfThreads;
             ++thread_index) {
          const SYSTEM_THREAD_INFORMATION* ti = &pi->Threads[thread_index];
          if (ti->ClientId.UniqueProcess != pi->ProcessId)
            continue;

          ThreadData thread_data;
          thread_data.thread_id = ti->ClientId.UniqueThread;
          thread_data.context_switches = ti->ContextSwitchCount;
          process_data.threads.push_back(thread_data);
        }

        // Order thread data by thread ID to help diff two snapshots.
        std::sort(process_data.threads.begin(), process_data.threads.end(),
                  [](const ThreadData& l, const ThreadData r) {
                    return l.thread_id < r.thread_id;
                  });

        snapshot->processes.insert(
            std::make_pair(pi->ProcessId, std::move(process_data)));
      }
    }

    // Check for end of the list.
    if (!pi->NextEntryOffset)
      break;

    // Jump to the next entry.
    offset += pi->NextEntryOffset;
  }

  return snapshot;
}