<commit_before>#pragma once

#include "boost_defs.hpp"

#include "cf_utility.hpp"
#include "filesystem.hpp"
#include "logger.hpp"
#include <CoreServices/CoreServices.h>
#include <boost/signals2.hpp>
#include <utility>
#include <vector>

namespace krbn {
class file_monitor final {
public:
  // Signals

  boost::signals2::signal<void(const std::string& file_path)> file_changed;

  // Methods

  file_monitor(const std::vector<std::string>& files) : directories_(cf_utility::create_cfmutablearray()),
                                                        stream_(nullptr) {
    run_loop_thread_ = std::make_unique<cf_utility::run_loop_thread>();

    std::vector<std::string> directories;
    for (const auto& f : files) {
      auto directory = filesystem::dirname(f);
      if (std::none_of(std::begin(directories),
                       std::end(directories),
                       [&](auto&& d) {
                         return directory == d;
                       })) {
        directories.push_back(directory);
      }

      // Do not apply `realpath` here since `realpath(f)` will be failed if `f` does not exist.
      files_.push_back(f);
    }

    for (const auto& d : directories) {
      if (auto directory = cf_utility::create_cfstring(d)) {
        if (directories_) {
          CFArrayAppendValue(directories_, directory);
        }
        CFRelease(directory);
      }
    }
  }

  ~file_monitor(void) {
    unregister_stream();

    if (directories_) {
      CFRelease(directories_);
      directories_ = nullptr;
    }

    run_loop_thread_ = nullptr;
  }

  void start(void) {
    register_stream();
  }

  void enqueue_file_changed(const std::string& file_path) {
    run_loop_thread_->enqueue(^{
      call_file_changed_slots(file_path);
    });
  }

private:
  void register_stream(void) {
    std::lock_guard<std::mutex> lock(stream_mutex_);

    // Skip if already started.

    if (stream_) {
      return;
    }

    // ----------------------------------------
    // File System Events API does not call the callback if the root directory and files are moved at the same time.
    //
    // Example:
    //
    //   FSEventStreamCreate(... ,{"target/file1", "target/file2"})
    //
    //   $ mkdir target.new
    //   $ echo  target.new/file1
    //   $ mv    target.new target
    //
    //   In this case, the callback will not be called.
    //
    // Thus, we should signal manually once.

    for (const auto& file : files_) {
      run_loop_thread_->enqueue(^{
        call_file_changed_slots(file);
      });
    }

    // ----------------------------------------

    FSEventStreamContext context{0};
    context.info = this;

    // kFSEventStreamCreateFlagWatchRoot and kFSEventStreamCreateFlagFileEvents are required in the following case.
    // (When directory == ~/.karabiner.d/configuration, file == ~/.karabiner.d/configuration/xxx.json)
    //
    // $ mkdir ~/.karabiner.d/configuration
    // $ touch ~/.karabiner.d/configuration/xxx.json
    // $ mv ~/.karabiner.d/configuration ~/.karabiner.d/configuration.back
    // $ ln -s ~/file-synchronisation-service/karabiner.d/configuration ~/.karabiner.d/
    // $ touch ~/.karabiner.d/configuration/xxx.json

    auto flags = FSEventStreamCreateFlags(0);
    flags |= kFSEventStreamCreateFlagWatchRoot;
    flags |= kFSEventStreamCreateFlagIgnoreSelf;
    flags |= kFSEventStreamCreateFlagFileEvents;

    stream_ = FSEventStreamCreate(kCFAllocatorDefault,
                                  static_stream_callback,
                                  &context,
                                  directories_,
                                  kFSEventStreamEventIdSinceNow,
                                  0.1, // 100 ms
                                  flags);
    if (!stream_) {
      logger::get_logger().error("FSEventStreamCreate error @ {0}", __PRETTY_FUNCTION__);
    } else {
      FSEventStreamScheduleWithRunLoop(stream_,
                                       run_loop_thread_->get_run_loop(),
                                       kCFRunLoopDefaultMode);
      if (!FSEventStreamStart(stream_)) {
        logger::get_logger().error("FSEventStreamStart error @ {0}", __PRETTY_FUNCTION__);
      }
    }
  }

  void unregister_stream(void) {
    std::lock_guard<std::mutex> lock(stream_mutex_);

    if (stream_) {
      FSEventStreamStop(stream_);
      FSEventStreamInvalidate(stream_);
      FSEventStreamRelease(stream_);
      stream_ = nullptr;
    }
  }

  static void static_stream_callback(ConstFSEventStreamRef stream,
                                     void* client_callback_info,
                                     size_t num_events,
                                     void* event_paths,
                                     const FSEventStreamEventFlags event_flags[],
                                     const FSEventStreamEventId event_ids[]) {
    auto self = reinterpret_cast<file_monitor*>(client_callback_info);
    if (self) {
      self->stream_callback(num_events,
                            static_cast<const char**>(event_paths),
                            event_flags,
                            event_ids);
    }
  }

  void stream_callback(size_t num_events,
                       const char* event_paths[],
                       const FSEventStreamEventFlags event_flags[],
                       const FSEventStreamEventId event_ids[]) {
    for (size_t i = 0; i < num_events; ++i) {
      if (event_flags[i] & kFSEventStreamEventFlagRootChanged) {
        logger::get_logger().info("The configuration directory is updated.");
        // re-register stream
        unregister_stream();
        register_stream();

      } else {
        call_file_changed_slots(event_paths[i]);
      }
    }
  }

  void call_file_changed_slots(const std::string& file_path) const {
    // FSEvents passes realpathed file path to callback.
    // Thus, we have to compare realpathed file paths.

    if (auto path = filesystem::realpath(file_path)) {
      if (filesystem::exists(*path)) {
        if (std::any_of(std::begin(files_),
                        std::end(files_),
                        [&](auto&& p) {
                          return *path == filesystem::realpath(p);
                        })) {
          file_changed(*path);
        }
      }
    }
  }

  CFMutableArrayRef directories_;
  std::vector<std::string> files_;
  std::unique_ptr<cf_utility::run_loop_thread> run_loop_thread_;

  FSEventStreamRef stream_;
  std::mutex stream_mutex_;
};
} // namespace krbn
<commit_msg>add file_monitor::get_run_loop_thread<commit_after>#pragma once

#include "boost_defs.hpp"

#include "cf_utility.hpp"
#include "filesystem.hpp"
#include "logger.hpp"
#include <CoreServices/CoreServices.h>
#include <boost/signals2.hpp>
#include <utility>
#include <vector>

namespace krbn {
class file_monitor final {
public:
  // Signals

  boost::signals2::signal<void(const std::string& file_path)> file_changed;

  // Methods

  file_monitor(const std::vector<std::string>& files) : directories_(cf_utility::create_cfmutablearray()),
                                                        stream_(nullptr) {
    run_loop_thread_ = std::make_shared<cf_utility::run_loop_thread>();

    std::vector<std::string> directories;
    for (const auto& f : files) {
      auto directory = filesystem::dirname(f);
      if (std::none_of(std::begin(directories),
                       std::end(directories),
                       [&](auto&& d) {
                         return directory == d;
                       })) {
        directories.push_back(directory);
      }

      // Do not apply `realpath` here since `realpath(f)` will be failed if `f` does not exist.
      files_.push_back(f);
    }

    for (const auto& d : directories) {
      if (auto directory = cf_utility::create_cfstring(d)) {
        if (directories_) {
          CFArrayAppendValue(directories_, directory);
        }
        CFRelease(directory);
      }
    }
  }

  ~file_monitor(void) {
    unregister_stream();

    if (directories_) {
      CFRelease(directories_);
      directories_ = nullptr;
    }

    run_loop_thread_ = nullptr;
  }

  std::shared_ptr<cf_utility::run_loop_thread> get_run_loop_thread(void) const {
    return run_loop_thread_;
  }

  void start(void) {
    register_stream();
  }

  void enqueue_file_changed(const std::string& file_path) {
    run_loop_thread_->enqueue(^{
      call_file_changed_slots(file_path);
    });
  }

private:
  void register_stream(void) {
    std::lock_guard<std::mutex> lock(stream_mutex_);

    // Skip if already started.

    if (stream_) {
      return;
    }

    // ----------------------------------------
    // File System Events API does not call the callback if the root directory and files are moved at the same time.
    //
    // Example:
    //
    //   FSEventStreamCreate(... ,{"target/file1", "target/file2"})
    //
    //   $ mkdir target.new
    //   $ echo  target.new/file1
    //   $ mv    target.new target
    //
    //   In this case, the callback will not be called.
    //
    // Thus, we should signal manually once.

    for (const auto& file : files_) {
      run_loop_thread_->enqueue(^{
        call_file_changed_slots(file);
      });
    }

    // ----------------------------------------

    FSEventStreamContext context{0};
    context.info = this;

    // kFSEventStreamCreateFlagWatchRoot and kFSEventStreamCreateFlagFileEvents are required in the following case.
    // (When directory == ~/.karabiner.d/configuration, file == ~/.karabiner.d/configuration/xxx.json)
    //
    // $ mkdir ~/.karabiner.d/configuration
    // $ touch ~/.karabiner.d/configuration/xxx.json
    // $ mv ~/.karabiner.d/configuration ~/.karabiner.d/configuration.back
    // $ ln -s ~/file-synchronisation-service/karabiner.d/configuration ~/.karabiner.d/
    // $ touch ~/.karabiner.d/configuration/xxx.json

    auto flags = FSEventStreamCreateFlags(0);
    flags |= kFSEventStreamCreateFlagWatchRoot;
    flags |= kFSEventStreamCreateFlagIgnoreSelf;
    flags |= kFSEventStreamCreateFlagFileEvents;

    stream_ = FSEventStreamCreate(kCFAllocatorDefault,
                                  static_stream_callback,
                                  &context,
                                  directories_,
                                  kFSEventStreamEventIdSinceNow,
                                  0.1, // 100 ms
                                  flags);
    if (!stream_) {
      logger::get_logger().error("FSEventStreamCreate error @ {0}", __PRETTY_FUNCTION__);
    } else {
      FSEventStreamScheduleWithRunLoop(stream_,
                                       run_loop_thread_->get_run_loop(),
                                       kCFRunLoopDefaultMode);
      if (!FSEventStreamStart(stream_)) {
        logger::get_logger().error("FSEventStreamStart error @ {0}", __PRETTY_FUNCTION__);
      }
    }
  }

  void unregister_stream(void) {
    std::lock_guard<std::mutex> lock(stream_mutex_);

    if (stream_) {
      FSEventStreamStop(stream_);
      FSEventStreamInvalidate(stream_);
      FSEventStreamRelease(stream_);
      stream_ = nullptr;
    }
  }

  static void static_stream_callback(ConstFSEventStreamRef stream,
                                     void* client_callback_info,
                                     size_t num_events,
                                     void* event_paths,
                                     const FSEventStreamEventFlags event_flags[],
                                     const FSEventStreamEventId event_ids[]) {
    auto self = reinterpret_cast<file_monitor*>(client_callback_info);
    if (self) {
      self->stream_callback(num_events,
                            static_cast<const char**>(event_paths),
                            event_flags,
                            event_ids);
    }
  }

  void stream_callback(size_t num_events,
                       const char* event_paths[],
                       const FSEventStreamEventFlags event_flags[],
                       const FSEventStreamEventId event_ids[]) {
    for (size_t i = 0; i < num_events; ++i) {
      if (event_flags[i] & kFSEventStreamEventFlagRootChanged) {
        logger::get_logger().info("The configuration directory is updated.");
        // re-register stream
        unregister_stream();
        register_stream();

      } else {
        call_file_changed_slots(event_paths[i]);
      }
    }
  }

  void call_file_changed_slots(const std::string& file_path) const {
    // FSEvents passes realpathed file path to callback.
    // Thus, we have to compare realpathed file paths.

    if (auto path = filesystem::realpath(file_path)) {
      if (filesystem::exists(*path)) {
        if (std::any_of(std::begin(files_),
                        std::end(files_),
                        [&](auto&& p) {
                          return *path == filesystem::realpath(p);
                        })) {
          file_changed(*path);
        }
      }
    }
  }

  CFMutableArrayRef directories_;
  std::vector<std::string> files_;
  std::shared_ptr<cf_utility::run_loop_thread> run_loop_thread_;

  FSEventStreamRef stream_;
  std::mutex stream_mutex_;
};
} // namespace krbn
<|endoftext|>


<commit_before>// 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.

#include "net/base/keygen_handler.h"

#include <windows.h>
#include <wincrypt.h>
#pragma comment(lib, "crypt32.lib")
#include <rpc.h>
#pragma comment(lib, "rpcrt4.lib")

#include <list>
#include <string>
#include <vector>

#include "base/base64.h"
#include "base/basictypes.h"
#include "base/logging.h"
#include "base/string_piece.h"
#include "base/utf_string_conversions.h"

namespace net {

// TODO(rsleevi): The following encoding functions are adapted from
// base/crypto/rsa_private_key.h and can/should probably be refactored.
static const uint8 kSequenceTag = 0x30;

void PrependLength(size_t size, std::list<BYTE>* data) {
  // The high bit is used to indicate whether additional octets are needed to
  // represent the length.
  if (size < 0x80) {
    data->push_front(static_cast<BYTE>(size));
  } else {
    uint8 num_bytes = 0;
    while (size > 0) {
      data->push_front(static_cast<BYTE>(size & 0xFF));
      size >>= 8;
      num_bytes++;
    }
    CHECK_LE(num_bytes, 4);
    data->push_front(0x80 | num_bytes);
  }
}

void PrependTypeHeaderAndLength(uint8 type, uint32 length,
                                std::vector<BYTE>* output) {
  std::list<BYTE> type_and_length;

  PrependLength(length, &type_and_length);
  type_and_length.push_front(type);

  output->insert(output->begin(), type_and_length.begin(),
                 type_and_length.end());
}

bool EncodeAndAppendType(LPCSTR type, const void* to_encode,
                         std::vector<BYTE>* output) {
  BOOL ok;
  DWORD size = 0;
  ok = CryptEncodeObject(X509_ASN_ENCODING, type, to_encode, NULL, &size);
  DCHECK(ok);
  if (!ok)
    return false;

  std::vector<BYTE>::size_type old_size = output->size();
  output->resize(old_size + size);

  ok = CryptEncodeObject(X509_ASN_ENCODING, type, to_encode,
                         &(*output)[old_size], &size);
  DCHECK(ok);
  if (!ok)
    return false;

  // Sometimes the initial call to CryptEncodeObject gave a generous estimate
  // of the size, so shrink back to what was actually used.
  output->resize(old_size + size);

  return true;
}

// Appends a DER IA5String containing |challenge| to |output|.
// Returns true if encoding was successful.
bool EncodeChallenge(const std::string& challenge, std::vector<BYTE>* output) {
  CERT_NAME_VALUE challenge_nv;
  challenge_nv.dwValueType = CERT_RDN_IA5_STRING;
  challenge_nv.Value.pbData = const_cast<BYTE*>(
      reinterpret_cast<const BYTE*>(challenge.data()));
  challenge_nv.Value.cbData = challenge.size();

  return EncodeAndAppendType(X509_ANY_STRING, &challenge_nv, output);
}

// Appends a DER SubjectPublicKeyInfo structure for the signing key in |prov|
// to |output|.
// Returns true if encoding was successful.
bool EncodeSubjectPublicKeyInfo(HCRYPTPROV prov, std::vector<BYTE>* output) {
  BOOL ok;
  DWORD size = 0;

  // From the private key stored in HCRYPTPROV, obtain the public key, stored
  // as a CERT_PUBLIC_KEY_INFO structure. Currently, only RSA public keys are
  // supported.
  ok = CryptExportPublicKeyInfoEx(prov, AT_KEYEXCHANGE, X509_ASN_ENCODING,
                                  szOID_RSA_RSA, 0, NULL, NULL, &size);
  DCHECK(ok);
  if (!ok)
    return false;

  std::vector<BYTE> public_key_info(size);
  PCERT_PUBLIC_KEY_INFO public_key_casted =
      reinterpret_cast<PCERT_PUBLIC_KEY_INFO>(&public_key_info[0]);
  ok = CryptExportPublicKeyInfoEx(prov, AT_KEYEXCHANGE, X509_ASN_ENCODING,
                                  szOID_RSA_RSA, 0, NULL, public_key_casted,
                                  &size);
  DCHECK(ok);
  if (!ok)
    return false;

  public_key_info.resize(size);

  return EncodeAndAppendType(X509_PUBLIC_KEY_INFO, &public_key_info[0],
                             output);
}

// Generates an ASN.1 DER representation of the PublicKeyAndChallenge structure
// from the signing key of |prov| and the specified |challenge| and appends it
// to |output|.
// True if the the encoding was successfully generated.
bool GetPublicKeyAndChallenge(HCRYPTPROV prov, const std::string& challenge,
                              std::vector<BYTE>* output) {
  if (!EncodeSubjectPublicKeyInfo(prov, output) ||
      !EncodeChallenge(challenge, output)) {
    return false;
  }

  PrependTypeHeaderAndLength(kSequenceTag, output->size(), output);
  return true;
}

// Generates a DER encoded SignedPublicKeyAndChallenge structure from the
// signing key of |prov| and the specified |challenge| string and appends it
// to |output|.
// True if the encoding was successfully generated.
bool GetSignedPublicKeyAndChallenge(HCRYPTPROV prov,
                                    const std::string& challenge,
                                    std::string* output) {
  std::vector<BYTE> pkac;
  if (!GetPublicKeyAndChallenge(prov, challenge, &pkac))
    return false;

  std::vector<BYTE> signature;
  std::vector<BYTE> signed_pkac;
  DWORD size = 0;
  BOOL ok;

  // While the MSDN documentation states that CERT_SIGNED_CONTENT_INFO should
  // be an X.509 certificate type, for encoding this is not necessary. The
  // result of encoding this structure will be a DER-encoded structure with
  // the ASN.1 equivalent of
  //  ::= SEQUENCE {
  //    ToBeSigned IMPLICIT OCTET STRING,
  //    SignatureAlgorithm AlgorithmIdentifier,
  //    Signature BIT STRING
  //  }
  //
  // This happens to be the same naive type as an SPKAC, so this works.
  CERT_SIGNED_CONTENT_INFO info;
  info.ToBeSigned.cbData = pkac.size();
  info.ToBeSigned.pbData = &pkac[0];
  info.SignatureAlgorithm.pszObjId = szOID_RSA_MD5RSA;
  info.SignatureAlgorithm.Parameters.cbData = 0;
  info.SignatureAlgorithm.Parameters.pbData = NULL;

  ok = CryptSignCertificate(prov, AT_KEYEXCHANGE, X509_ASN_ENCODING,
                            info.ToBeSigned.pbData, info.ToBeSigned.cbData,
                            &info.SignatureAlgorithm, NULL, NULL, &size);
  DCHECK(ok);
  if (!ok)
    return false;

  signature.resize(size);
  info.Signature.cbData = signature.size();
  info.Signature.pbData = &signature[0];
  info.Signature.cUnusedBits = 0;

  ok = CryptSignCertificate(prov, AT_KEYEXCHANGE, X509_ASN_ENCODING,
                            info.ToBeSigned.pbData, info.ToBeSigned.cbData,
                            &info.SignatureAlgorithm, NULL,
                            info.Signature.pbData, &info.Signature.cbData);
  DCHECK(ok);
  if (!ok || !EncodeAndAppendType(X509_CERT, &info, &signed_pkac))
    return false;

  output->assign(reinterpret_cast<char*>(&signed_pkac[0]),
                 signed_pkac.size());

  return true;
}

// Generates a unique name for the container which will store the key that is
// generated. The traditional Windows approach is to use a GUID here.
std::wstring GetNewKeyContainerId() {
  RPC_STATUS status = RPC_S_OK;
  std::wstring result;

  UUID id = { 0 };
  status = UuidCreateSequential(&id);
  if (status != RPC_S_OK && status != RPC_S_UUID_LOCAL_ONLY)
    return result;

  RPC_WSTR rpc_string = NULL;
  status = UuidToString(&id, &rpc_string);
  if (status != RPC_S_OK)
    return result;

  // RPC_WSTR is unsigned short*.  wchar_t is a built-in type of Visual C++,
  // so the type cast is necessary.
  result.assign(reinterpret_cast<wchar_t*>(rpc_string));
  RpcStringFree(&rpc_string);

  return result;
}

void StoreKeyLocationInCache(HCRYPTPROV prov) {
  BOOL ok;
  DWORD size = 0;

  // Though it is known the container and provider name, as they are supplied
  // during GenKeyAndSignChallenge, explicitly resolving them via
  // CryptGetProvParam ensures that any defaults (such as provider name being
  // NULL) or any CSP modifications to the container name are properly
  // reflected.

  // Find the container name. Though the MSDN documentation states it will
  // return the exact same value as supplied when the provider was aquired, it
  // also notes the return type will be CHAR, /not/ WCHAR.
  ok = CryptGetProvParam(prov, PP_CONTAINER, NULL, &size, 0);
  if (!ok)
    return;

  std::vector<BYTE> buffer(size);
  ok = CryptGetProvParam(prov, PP_CONTAINER, &buffer[0], &size, 0);
  if (!ok)
    return;

  KeygenHandler::KeyLocation key_location;
  UTF8ToWide(reinterpret_cast<char*>(&buffer[0]), size,
             &key_location.container_name);

  // Get the provider name. This will always resolve, even if NULL (indicating
  // the default provider) was supplied to the CryptAcquireContext.
  size = 0;
  ok = CryptGetProvParam(prov, PP_NAME, NULL, &size, 0);
  if (!ok)
    return;

  buffer.resize(size);
  ok = CryptGetProvParam(prov, PP_NAME, &buffer[0], &size, 0);
  if (!ok)
    return;

  UTF8ToWide(reinterpret_cast<char*>(&buffer[0]), size,
             &key_location.provider_name);

  std::vector<BYTE> public_key_info;
  if (!EncodeSubjectPublicKeyInfo(prov, &public_key_info))
    return;

  KeygenHandler::Cache* cache = KeygenHandler::Cache::GetInstance();
  cache->Insert(std::string(public_key_info.begin(), public_key_info.end()),
                key_location);
}

bool KeygenHandler::KeyLocation::Equals(
    const KeygenHandler::KeyLocation& location) const {
  return container_name == location.container_name &&
         provider_name == location.provider_name;
}

std::string KeygenHandler::GenKeyAndSignChallenge() {
  std::string result;

  bool is_success = true;
  HCRYPTPROV prov = NULL;
  HCRYPTKEY key = NULL;
  DWORD flags = (key_size_in_bits_ << 16) | CRYPT_EXPORTABLE;
  std::string spkac;

  std::wstring new_key_id;

  // TODO(rsleevi): Have the user choose which provider they should use, which
  // needs to be filtered by those providers which can provide the key type
  // requested or the key size requested. This is especially important for
  // generating certificates that will be stored on smart cards.
  const int kMaxAttempts = 5;
  BOOL ok = FALSE;
  for (int attempt = 0; attempt < kMaxAttempts; ++attempt) {
    // Per MSDN documentation for CryptAcquireContext, if applications will be
    // creating their own keys, they should ensure unique naming schemes to
    // prevent overlap with any other applications or consumers of CSPs, and
    // *should not* store new keys within the default, NULL key container.
    new_key_id = GetNewKeyContainerId();
    if (new_key_id.empty())
      return result;

    // Only create new key containers, so that existing key containers are not
    // overwritten.
    ok = CryptAcquireContext(&prov, new_key_id.c_str(), NULL, PROV_RSA_FULL,
                             CRYPT_SILENT | CRYPT_NEWKEYSET);

    if (ok || GetLastError() != NTE_BAD_KEYSET)
      break;
  }
  if (!ok) {
    LOG(ERROR) << "Couldn't acquire a CryptoAPI provider context: "
               << GetLastError();
    is_success = false;
    goto failure;
  }

  if (!CryptGenKey(prov, CALG_RSA_KEYX, flags, &key)) {
    LOG(ERROR) << "Couldn't generate an RSA key";
    is_success = false;
    goto failure;
  }

  if (!GetSignedPublicKeyAndChallenge(prov, challenge_, &spkac)) {
    LOG(ERROR) << "Couldn't generate the signed public key and challenge";
    is_success = false;
    goto failure;
  }

  if (!base::Base64Encode(spkac, &result)) {
    LOG(ERROR) << "Couldn't convert signed key into base64";
    is_success = false;
    goto failure;
  }

  StoreKeyLocationInCache(prov);

 failure:
  if (!is_success) {
    LOG(ERROR) << "SSL Keygen failed";
  } else {
    LOG(INFO) << "SSL Key succeeded";
  }
  if (key) {
    // Securely destroys the handle, but leaves the underlying key alone. The
    // key can be obtained again by resolving the key location. If
    // |stores_key_| is false, the underlying key will be destroyed below.
    CryptDestroyKey(key);
  }

  if (prov) {
    CryptReleaseContext(prov, 0);
    prov = NULL;
    if (!stores_key_) {
      // Fully destroys any of the keys that were created and releases prov.
      CryptAcquireContext(&prov, new_key_id.c_str(), NULL, PROV_RSA_FULL,
                          CRYPT_SILENT | CRYPT_DELETEKEYSET);
    }
  }

  return result;
}

}  // namespace net
<commit_msg>Simplify the Windows <keygen> implementation by making better use of CryptoAPI.<commit_after>// 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.

#include "net/base/keygen_handler.h"

#include <windows.h>
#include <wincrypt.h>
#pragma comment(lib, "crypt32.lib")
#include <rpc.h>
#pragma comment(lib, "rpcrt4.lib")

#include <list>
#include <string>
#include <vector>

#include "base/base64.h"
#include "base/basictypes.h"
#include "base/logging.h"
#include "base/string_piece.h"
#include "base/string_util.h"
#include "base/utf_string_conversions.h"

namespace net {

bool EncodeAndAppendType(LPCSTR type, const void* to_encode,
                         std::vector<BYTE>* output) {
  BOOL ok;
  DWORD size = 0;
  ok = CryptEncodeObject(X509_ASN_ENCODING, type, to_encode, NULL, &size);
  DCHECK(ok);
  if (!ok)
    return false;

  std::vector<BYTE>::size_type old_size = output->size();
  output->resize(old_size + size);

  ok = CryptEncodeObject(X509_ASN_ENCODING, type, to_encode,
                         &(*output)[old_size], &size);
  DCHECK(ok);
  if (!ok)
    return false;

  // Sometimes the initial call to CryptEncodeObject gave a generous estimate
  // of the size, so shrink back to what was actually used.
  output->resize(old_size + size);

  return true;
}

// Assigns the contents of a CERT_PUBLIC_KEY_INFO structure for the signing
// key in |prov| to |output|. Returns true if encoding was successful.
bool GetSubjectPublicKeyInfo(HCRYPTPROV prov, std::vector<BYTE>* output) {
  BOOL ok;
  DWORD size = 0;

  // From the private key stored in HCRYPTPROV, obtain the public key, stored
  // as a CERT_PUBLIC_KEY_INFO structure. Currently, only RSA public keys are
  // supported.
  ok = CryptExportPublicKeyInfoEx(prov, AT_KEYEXCHANGE, X509_ASN_ENCODING,
                                  szOID_RSA_RSA, 0, NULL, NULL, &size);
  DCHECK(ok);
  if (!ok)
    return false;

  output->resize(size);

  PCERT_PUBLIC_KEY_INFO public_key_casted =
      reinterpret_cast<PCERT_PUBLIC_KEY_INFO>(&(*output)[0]);
  ok = CryptExportPublicKeyInfoEx(prov, AT_KEYEXCHANGE, X509_ASN_ENCODING,
                                  szOID_RSA_RSA, 0, NULL, public_key_casted,
                                  &size);
  DCHECK(ok);
  if (!ok)
    return false;

  output->resize(size);

  return true;
}

// Appends a DER SubjectPublicKeyInfo structure for the signing key in |prov|
// to |output|.
// Returns true if encoding was successful.
bool EncodeSubjectPublicKeyInfo(HCRYPTPROV prov, std::vector<BYTE>* output) {
  std::vector<BYTE> public_key_info;
  if (!GetSubjectPublicKeyInfo(prov, &public_key_info))
    return false;

  return EncodeAndAppendType(X509_PUBLIC_KEY_INFO, &public_key_info[0],
                             output);
}

// Generates a DER encoded SignedPublicKeyAndChallenge structure from the
// signing key of |prov| and the specified ASCII |challenge| string and
// appends it to |output|.
// True if the encoding was successfully generated.
bool GetSignedPublicKeyAndChallenge(HCRYPTPROV prov,
                                    const std::string& challenge,
                                    std::string* output) {
  std::wstring wide_challenge = ASCIIToWide(challenge);
  std::vector<BYTE> spki;

  if (!GetSubjectPublicKeyInfo(prov, &spki))
    return false;

  // PublicKeyAndChallenge ::= SEQUENCE {
  //     spki SubjectPublicKeyInfo,
  //     challenge IA5STRING
  // }
  CERT_KEYGEN_REQUEST_INFO pkac;
  pkac.dwVersion = CERT_KEYGEN_REQUEST_V1;
  pkac.SubjectPublicKeyInfo =
      *reinterpret_cast<PCERT_PUBLIC_KEY_INFO>(&spki[0]);
  pkac.pwszChallengeString = const_cast<wchar_t*>(wide_challenge.c_str());

  CRYPT_ALGORITHM_IDENTIFIER sig_alg;
  memset(&sig_alg, 0, sizeof(sig_alg));
  sig_alg.pszObjId = szOID_RSA_MD5RSA;

  BOOL ok;
  DWORD size = 0;
  std::vector<BYTE> signed_pkac;
  ok = CryptSignAndEncodeCertificate(prov, AT_KEYEXCHANGE, X509_ASN_ENCODING,
                                     X509_KEYGEN_REQUEST_TO_BE_SIGNED,
                                     &pkac, &sig_alg, NULL,
                                     NULL, &size);
  DCHECK(ok);
  if (!ok)
    return false;

  signed_pkac.resize(size);
  ok = CryptSignAndEncodeCertificate(prov, AT_KEYEXCHANGE, X509_ASN_ENCODING,
                                     X509_KEYGEN_REQUEST_TO_BE_SIGNED,
                                     &pkac, &sig_alg, NULL,
                                     &signed_pkac[0], &size);
  DCHECK(ok);
  if (!ok)
    return false;

  output->assign(reinterpret_cast<char*>(&signed_pkac[0]), size);
  return true;
}

// Generates a unique name for the container which will store the key that is
// generated. The traditional Windows approach is to use a GUID here.
std::wstring GetNewKeyContainerId() {
  RPC_STATUS status = RPC_S_OK;
  std::wstring result;

  UUID id = { 0 };
  status = UuidCreateSequential(&id);
  if (status != RPC_S_OK && status != RPC_S_UUID_LOCAL_ONLY)
    return result;

  RPC_WSTR rpc_string = NULL;
  status = UuidToString(&id, &rpc_string);
  if (status != RPC_S_OK)
    return result;

  // RPC_WSTR is unsigned short*.  wchar_t is a built-in type of Visual C++,
  // so the type cast is necessary.
  result.assign(reinterpret_cast<wchar_t*>(rpc_string));
  RpcStringFree(&rpc_string);

  return result;
}

void StoreKeyLocationInCache(HCRYPTPROV prov) {
  BOOL ok;
  DWORD size = 0;

  // Though it is known the container and provider name, as they are supplied
  // during GenKeyAndSignChallenge, explicitly resolving them via
  // CryptGetProvParam ensures that any defaults (such as provider name being
  // NULL) or any CSP modifications to the container name are properly
  // reflected.

  // Find the container name. Though the MSDN documentation states it will
  // return the exact same value as supplied when the provider was aquired, it
  // also notes the return type will be CHAR, /not/ WCHAR.
  ok = CryptGetProvParam(prov, PP_CONTAINER, NULL, &size, 0);
  if (!ok)
    return;

  std::vector<BYTE> buffer(size);
  ok = CryptGetProvParam(prov, PP_CONTAINER, &buffer[0], &size, 0);
  if (!ok)
    return;

  KeygenHandler::KeyLocation key_location;
  UTF8ToWide(reinterpret_cast<char*>(&buffer[0]), size,
             &key_location.container_name);

  // Get the provider name. This will always resolve, even if NULL (indicating
  // the default provider) was supplied to the CryptAcquireContext.
  size = 0;
  ok = CryptGetProvParam(prov, PP_NAME, NULL, &size, 0);
  if (!ok)
    return;

  buffer.resize(size);
  ok = CryptGetProvParam(prov, PP_NAME, &buffer[0], &size, 0);
  if (!ok)
    return;

  UTF8ToWide(reinterpret_cast<char*>(&buffer[0]), size,
             &key_location.provider_name);

  std::vector<BYTE> public_key_info;
  if (!EncodeSubjectPublicKeyInfo(prov, &public_key_info))
    return;

  KeygenHandler::Cache* cache = KeygenHandler::Cache::GetInstance();
  cache->Insert(std::string(public_key_info.begin(), public_key_info.end()),
                key_location);
}

bool KeygenHandler::KeyLocation::Equals(
    const KeygenHandler::KeyLocation& location) const {
  return container_name == location.container_name &&
         provider_name == location.provider_name;
}

std::string KeygenHandler::GenKeyAndSignChallenge() {
  std::string result;

  bool is_success = true;
  HCRYPTPROV prov = NULL;
  HCRYPTKEY key = NULL;
  DWORD flags = (key_size_in_bits_ << 16) | CRYPT_EXPORTABLE;
  std::string spkac;

  std::wstring new_key_id;

  // TODO(rsleevi): Have the user choose which provider they should use, which
  // needs to be filtered by those providers which can provide the key type
  // requested or the key size requested. This is especially important for
  // generating certificates that will be stored on smart cards.
  const int kMaxAttempts = 5;
  BOOL ok = FALSE;
  for (int attempt = 0; attempt < kMaxAttempts; ++attempt) {
    // Per MSDN documentation for CryptAcquireContext, if applications will be
    // creating their own keys, they should ensure unique naming schemes to
    // prevent overlap with any other applications or consumers of CSPs, and
    // *should not* store new keys within the default, NULL key container.
    new_key_id = GetNewKeyContainerId();
    if (new_key_id.empty())
      return result;

    // Only create new key containers, so that existing key containers are not
    // overwritten.
    ok = CryptAcquireContext(&prov, new_key_id.c_str(), NULL, PROV_RSA_FULL,
                             CRYPT_SILENT | CRYPT_NEWKEYSET);

    if (ok || GetLastError() != NTE_BAD_KEYSET)
      break;
  }
  if (!ok) {
    LOG(ERROR) << "Couldn't acquire a CryptoAPI provider context: "
               << GetLastError();
    is_success = false;
    goto failure;
  }

  if (!CryptGenKey(prov, CALG_RSA_KEYX, flags, &key)) {
    LOG(ERROR) << "Couldn't generate an RSA key";
    is_success = false;
    goto failure;
  }

  if (!GetSignedPublicKeyAndChallenge(prov, challenge_, &spkac)) {
    LOG(ERROR) << "Couldn't generate the signed public key and challenge";
    is_success = false;
    goto failure;
  }

  if (!base::Base64Encode(spkac, &result)) {
    LOG(ERROR) << "Couldn't convert signed key into base64";
    is_success = false;
    goto failure;
  }

  StoreKeyLocationInCache(prov);

 failure:
  if (!is_success) {
    LOG(ERROR) << "SSL Keygen failed";
  } else {
    LOG(INFO) << "SSL Key succeeded";
  }
  if (key) {
    // Securely destroys the handle, but leaves the underlying key alone. The
    // key can be obtained again by resolving the key location. If
    // |stores_key_| is false, the underlying key will be destroyed below.
    CryptDestroyKey(key);
  }

  if (prov) {
    CryptReleaseContext(prov, 0);
    prov = NULL;
    if (!stores_key_) {
      // Fully destroys any of the keys that were created and releases prov.
      CryptAcquireContext(&prov, new_key_id.c_str(), NULL, PROV_RSA_FULL,
                          CRYPT_SILENT | CRYPT_DELETEKEYSET);
    }
  }

  return result;
}

}  // namespace net
<|endoftext|>


<commit_before>/*
    Copyright (C) 2012  Aleix Pol Gonzalez <aleixpol@blue-systems.com>

    This library is free software; you can redistribute it and/or
    modify it under the terms of the GNU Lesser General Public
    License as published by the Free Software Foundation; either
    version 2.1 of the License, or (at your option) any later version.

    This library is distributed in the hope that it will be useful,
    but WITHOUT ANY WARRANTY; without even the implied warranty of
    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
    Lesser General Public License for more details.

    You should have received a copy of the GNU Lesser General Public
    License along with this library; if not, write to the Free Software
    Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA  02110-1301  USA
*/

#include "matchessolver.h"
#include "match.h"
#include "personsmodel.h"
#include <nepomuk2/datamanagement.h>
#include <QUrl>

using namespace KPeople;

MatchesSolver::MatchesSolver(const QList<Match> &matches, PersonsModel *model, QObject *parent)
    : KJob(parent)
    , m_matches(matches)
    , m_model(model)
{
}

void MatchesSolver::start()
{
    QMetaObject::invokeMethod(this, "startMatching", Qt::QueuedConnection);
}

void MatchesSolver::startMatching()
{
    Q_FOREACH(const Match &m, m_matches) {
        QPersistentModelIndex idxDestination = m.indexA;
        QPersistentModelIndex idxOrigin = m.indexB;

        QList<QUrl> persons;
        persons << idxDestination.data(PersonsModel::UriRole).toUrl();
        persons << idxOrigin.data(PersonsModel::UriRole).toUrl();

        KJob *job = Nepomuk2::mergeResources(persons);
        m_pending.insert(job);
        connect(job, SIGNAL(finished(KJob*)), SLOT(jobDone(KJob*)));
    }
    if(m_matches.isEmpty())
        jobDone(0);
}

QList<QUrl> MatchesSolver::contactUris(const QModelIndex &idxOrigin)
{
    QList<QUrl> ret;
    QModelIndex idx=idxOrigin.child(0,0);
    for (; idx.isValid(); idx = idx.sibling(idx.row() + 1, 0)) {
        ret += idx.data(PersonsModel::UriRole).toUrl();
    }
    return ret;
}

void MatchesSolver::jobDone(KJob *job)
{
    m_pending.remove(job);
    if (m_pending.isEmpty()) {
        emitResult();
    }
}
<commit_msg>Separate the data fetching and the merge job spawning<commit_after>/*
    Copyright (C) 2012  Aleix Pol Gonzalez <aleixpol@blue-systems.com>

    This library is free software; you can redistribute it and/or
    modify it under the terms of the GNU Lesser General Public
    License as published by the Free Software Foundation; either
    version 2.1 of the License, or (at your option) any later version.

    This library is distributed in the hope that it will be useful,
    but WITHOUT ANY WARRANTY; without even the implied warranty of
    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
    Lesser General Public License for more details.

    You should have received a copy of the GNU Lesser General Public
    License along with this library; if not, write to the Free Software
    Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA  02110-1301  USA
*/

#include "matchessolver.h"
#include "match.h"
#include "personsmodel.h"
#include <nepomuk2/datamanagement.h>
#include <QUrl>
#include <QDebug>

using namespace KPeople;

MatchesSolver::MatchesSolver(const QList<Match> &matches, PersonsModel *model, QObject *parent)
    : KJob(parent)
    , m_matches(matches)
    , m_model(model)
{
}

void MatchesSolver::start()
{
    QMetaObject::invokeMethod(this, "startMatching", Qt::QueuedConnection);
}

void MatchesSolver::startMatching()
{
    QMap<int, QSet<QUrl> > jobsData;
    Q_FOREACH(const Match &m, m_matches) {
        QPersistentModelIndex idxDestination = m.indexA;
        QPersistentModelIndex idxOrigin = m.indexB;

        QSet<QUrl>& jobs = jobsData[m.indexA.row()];
        jobs.insert(idxDestination.data(PersonsModel::UriRole).toUrl());
        jobs.insert(idxOrigin.data(PersonsModel::UriRole).toUrl());
        Q_ASSERT(jobs.size()>=2);
    }

    foreach(const QSet<QUrl>& uris, jobsData) {
        KJob* job = m_model->createPersonFromUris(uris.toList());
        m_pending.insert(job);
        connect(job, SIGNAL(finished(KJob*)), SLOT(jobDone(KJob*)));
    }
    bool wrongJobsCount = m_pending.remove(0);
    if(wrongJobsCount >= 0) {
        qDebug() << "error: some of the jobs couldn't be performed" << wrongJobsCount;
    }

    if(jobsData.isEmpty())
        jobDone(0);
}

QList<QUrl> MatchesSolver::contactUris(const QModelIndex &idxOrigin)
{
    QList<QUrl> ret;
    QModelIndex idx=idxOrigin.child(0,0);
    for (; idx.isValid(); idx = idx.sibling(idx.row() + 1, 0)) {
        ret += idx.data(PersonsModel::UriRole).toUrl();
    }
    return ret;
}

void MatchesSolver::jobDone(KJob *job)
{
    m_pending.remove(job);
    if (m_pending.isEmpty()) {
        emitResult();
    }
}
<|endoftext|>


<commit_before>/*
 * Copyright (c) 2012 ARM Limited
 * All rights reserved.
 *
 * The license below extends only to copyright in the software and shall
 * not be construed as granting a license to any other intellectual
 * property including but not limited to intellectual property relating
 * to a hardware implementation of the functionality of the software
 * licensed hereunder.  You may use the software subject to the license
 * terms below provided that you ensure that this notice is replicated
 * unmodified and in its entirety in all distributions of the software,
 * modified or unmodified, in source code or in binary form.
 *
 * Copyright (c) 2002-2005 The Regents of The University of Michigan
 * Copyright (c) 2010 Advanced Micro Devices, Inc.
 * All rights reserved.
 *
 * Redistribution and use in source and binary forms, with or without
 * modification, are permitted provided that the following conditions are
 * met: redistributions of source code must retain the above copyright
 * notice, this list of conditions and the following disclaimer;
 * redistributions in binary form must reproduce the above copyright
 * notice, this list of conditions and the following disclaimer in the
 * documentation and/or other materials provided with the distribution;
 * neither the name of the copyright holders nor the names of its
 * contributors may be used to endorse or promote products derived from
 * this software without specific prior written permission.
 *
 * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
 * "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
 * LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR
 * A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT
 * OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,
 * SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT
 * LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
 * DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY
 * THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
 * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
 * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
 *
 * Authors: Erik Hallnor
 *          Dave Greene
 */

/**
 * @file
 * Miss Status and Handling Register (MSHR) definitions.
 */

#include <algorithm>
#include <cassert>
#include <string>
#include <vector>

#include "base/misc.hh"
#include "base/types.hh"
#include "debug/Cache.hh"
#include "mem/cache/cache.hh"
#include "mem/cache/mshr.hh"
#include "sim/core.hh"

using namespace std;

MSHR::MSHR()
{
    inService = false;
    ntargets = 0;
    threadNum = InvalidThreadID;
    targets = new TargetList();
    deferredTargets = new TargetList();
}

MSHR::TargetList::TargetList()
    : needsExclusive(false), hasUpgrade(false)
{}

inline void
MSHR::TargetList::add(PacketPtr pkt, Tick readyTime,
                      Counter order, Target::Source source, bool markPending)
{
    if (source != Target::FromSnoop) {
        if (pkt->needsExclusive()) {
            needsExclusive = true;
        }

        // StoreCondReq is effectively an upgrade if it's in an MSHR
        // since it would have been failed already if we didn't have a
        // read-only copy
        if (pkt->isUpgrade() || pkt->cmd == MemCmd::StoreCondReq) {
            hasUpgrade = true;
        }
    }

    if (markPending) {
        // Iterate over the SenderState stack and see if we find
        // an MSHR entry. If we do, set the downstreamPending
        // flag. Otherwise, do nothing.
        MSHR *mshr = pkt->findNextSenderState<MSHR>();
        if (mshr != NULL) {
            assert(!mshr->downstreamPending);
            mshr->downstreamPending = true;
        }
    }

    push_back(Target(pkt, readyTime, order, source, markPending));
}

static void
replaceUpgrade(PacketPtr pkt)
{
    if (pkt->cmd == MemCmd::UpgradeReq) {
        pkt->cmd = MemCmd::ReadExReq;
        DPRINTF(Cache, "Replacing UpgradeReq with ReadExReq\n");
    } else if (pkt->cmd == MemCmd::SCUpgradeReq) {
        pkt->cmd = MemCmd::SCUpgradeFailReq;
        DPRINTF(Cache, "Replacing SCUpgradeReq with SCUpgradeFailReq\n");
    } else if (pkt->cmd == MemCmd::StoreCondReq) {
        pkt->cmd = MemCmd::StoreCondFailReq;
        DPRINTF(Cache, "Replacing StoreCondReq with StoreCondFailReq\n");
    }
}

void
MSHR::TargetList::replaceUpgrades()
{
    if (!hasUpgrade)
        return;

    Iterator end_i = end();
    for (Iterator i = begin(); i != end_i; ++i) {
        replaceUpgrade(i->pkt);
    }

    hasUpgrade = false;
}

void
MSHR::TargetList::clearDownstreamPending()
{
    Iterator end_i = end();
    for (Iterator i = begin(); i != end_i; ++i) {
        if (i->markedPending) {
            // Iterate over the SenderState stack and see if we find
            // an MSHR entry. If we find one, clear the
            // downstreamPending flag by calling
            // clearDownstreamPending(). This recursively clears the
            // downstreamPending flag in all caches this packet has
            // passed through.
            MSHR *mshr = i->pkt->findNextSenderState<MSHR>();
            if (mshr != NULL) {
                mshr->clearDownstreamPending();
            }
        }
    }
}

bool
MSHR::TargetList::checkFunctional(PacketPtr pkt)
{
    Iterator end_i = end();
    for (Iterator i = begin(); i != end_i; ++i) {
        if (pkt->checkFunctional(i->pkt)) {
            return true;
        }
    }

    return false;
}

void
MSHR::TargetList::
print(std::ostream &os, int verbosity, const std::string &prefix) const
{
    ConstIterator end_i = end();
    for (ConstIterator i = begin(); i != end_i; ++i) {
        const char *s;
        switch (i->source) {
          case Target::FromCPU:
            s = "FromCPU";
            break;
          case Target::FromSnoop:
            s = "FromSnoop";
            break;
          case Target::FromPrefetcher:
            s = "FromPrefetcher";
            break;
          default:
            s = "";
            break;
        }
        ccprintf(os, "%s%s: ", prefix, s);
        i->pkt->print(os, verbosity, "");
    }
}

void
MSHR::allocate(Addr _addr, int _size, PacketPtr target,
               Tick whenReady, Counter _order)
{
    addr = _addr;
    size = _size;
    readyTime = whenReady;
    order = _order;
    assert(target);
    isForward = false;
    _isUncacheable = target->req->isUncacheable();
    inService = false;
    downstreamPending = false;
    threadNum = 0;
    ntargets = 1;
    assert(targets->isReset());
    // Don't know of a case where we would allocate a new MSHR for a
    // snoop (mem-side request), so set source according to request here
    Target::Source source = (target->cmd == MemCmd::HardPFReq) ?
        Target::FromPrefetcher : Target::FromCPU;
    targets->add(target, whenReady, _order, source, true);
    assert(deferredTargets->isReset());
    data = NULL;
}

void
MSHR::clearDownstreamPending()
{
    assert(downstreamPending);
    downstreamPending = false;
    // recursively clear flag on any MSHRs we will be forwarding
    // responses to
    targets->clearDownstreamPending();
}

bool
MSHR::markInService(PacketPtr pkt)
{
    assert(!inService);
    if (isForwardNoResponse()) {
        // we just forwarded the request packet & don't expect a
        // response, so get rid of it
        assert(getNumTargets() == 1);
        popTarget();
        return true;
    }
    inService = true;
    pendingDirty = (targets->needsExclusive ||
                    (!pkt->sharedAsserted() && pkt->memInhibitAsserted()));
    postInvalidate = postDowngrade = false;

    if (!downstreamPending) {
        // let upstream caches know that the request has made it to a
        // level where it's going to get a response
        targets->clearDownstreamPending();
    }
    return false;
}

void
MSHR::deallocate()
{
    assert(targets->empty());
    targets->resetFlags();
    assert(deferredTargets->isReset());
    assert(ntargets == 0);
    inService = false;
}

/*
 * Adds a target to an MSHR
 */
void
MSHR::allocateTarget(PacketPtr pkt, Tick whenReady, Counter _order)
{
    // if there's a request already in service for this MSHR, we will
    // have to defer the new target until after the response if any of
    // the following are true:
    // - there are other targets already deferred
    // - there's a pending invalidate to be applied after the response
    //   comes back (but before this target is processed)
    // - this target requires an exclusive block and either we're not
    //   getting an exclusive block back or we have already snooped
    //   another read request that will downgrade our exclusive block
    //   to shared

    // assume we'd never issue a prefetch when we've got an
    // outstanding miss
    assert(pkt->cmd != MemCmd::HardPFReq);

    if (inService &&
        (!deferredTargets->empty() || hasPostInvalidate() ||
         (pkt->needsExclusive() &&
          (!isPendingDirty() || hasPostDowngrade() || isForward)))) {
        // need to put on deferred list
        if (hasPostInvalidate())
            replaceUpgrade(pkt);
        deferredTargets->add(pkt, whenReady, _order, Target::FromCPU, true);
    } else {
        // No request outstanding, or still OK to append to
        // outstanding request: append to regular target list.  Only
        // mark pending if current request hasn't been issued yet
        // (isn't in service).
        targets->add(pkt, whenReady, _order, Target::FromCPU, !inService);
    }

    ++ntargets;
}

bool
MSHR::handleSnoop(PacketPtr pkt, Counter _order)
{
    DPRINTF(Cache, "%s for %s address %x size %d\n", __func__,
            pkt->cmdString(), pkt->getAddr(), pkt->getSize());
    if (!inService || (pkt->isExpressSnoop() && downstreamPending)) {
        // Request has not been issued yet, or it's been issued
        // locally but is buffered unissued at some downstream cache
        // which is forwarding us this snoop.  Either way, the packet
        // we're snooping logically precedes this MSHR's request, so
        // the snoop has no impact on the MSHR, but must be processed
        // in the standard way by the cache.  The only exception is
        // that if we're an L2+ cache buffering an UpgradeReq from a
        // higher-level cache, and the snoop is invalidating, then our
        // buffered upgrades must be converted to read exclusives,
        // since the upper-level cache no longer has a valid copy.
        // That is, even though the upper-level cache got out on its
        // local bus first, some other invalidating transaction
        // reached the global bus before the upgrade did.
        if (pkt->needsExclusive()) {
            targets->replaceUpgrades();
            deferredTargets->replaceUpgrades();
        }

        return false;
    }

    // From here on down, the request issued by this MSHR logically
    // precedes the request we're snooping.
    if (pkt->needsExclusive()) {
        // snooped request still precedes the re-request we'll have to
        // issue for deferred targets, if any...
        deferredTargets->replaceUpgrades();
    }

    if (hasPostInvalidate()) {
        // a prior snoop has already appended an invalidation, so
        // logically we don't have the block anymore; no need for
        // further snooping.
        return true;
    }

    if (isPendingDirty() || pkt->isInvalidate()) {
        // We need to save and replay the packet in two cases:
        // 1. We're awaiting an exclusive copy, so ownership is pending,
        //    and we need to respond after we receive data.
        // 2. It's an invalidation (e.g., UpgradeReq), and we need
        //    to forward the snoop up the hierarchy after the current
        //    transaction completes.

        // Actual target device (typ. a memory) will delete the
        // packet on reception, so we need to save a copy here.
        PacketPtr cp_pkt = new Packet(pkt, true);
        targets->add(cp_pkt, curTick(), _order, Target::FromSnoop,
                     downstreamPending && targets->needsExclusive);
        ++ntargets;

        if (isPendingDirty()) {
            pkt->assertMemInhibit();
            pkt->setSupplyExclusive();
        }

        if (pkt->needsExclusive()) {
            // This transaction will take away our pending copy
            postInvalidate = true;
        }
    }

    if (!pkt->needsExclusive()) {
        // This transaction will get a read-shared copy, downgrading
        // our copy if we had an exclusive one
        postDowngrade = true;
        pkt->assertShared();
    }

    return true;
}

bool
MSHR::promoteDeferredTargets()
{
    assert(targets->empty());
    if (deferredTargets->empty()) {
        return false;
    }

    // swap targets & deferredTargets lists
    TargetList *tmp = targets;
    targets = deferredTargets;
    deferredTargets = tmp;

    assert(targets->size() == ntargets);

    // clear deferredTargets flags
    deferredTargets->resetFlags();

    order = targets->front().order;
    readyTime = std::max(curTick(), targets->front().readyTime);

    return true;
}

void
MSHR::handleFill(Packet *pkt, CacheBlk *blk)
{
    if (!pkt->sharedAsserted()
        && !(hasPostInvalidate() || hasPostDowngrade())
        && deferredTargets->needsExclusive) {
        // We got an exclusive response, but we have deferred targets
        // which are waiting to request an exclusive copy (not because
        // of a pending invalidate).  This can happen if the original
        // request was for a read-only (non-exclusive) block, but we
        // got an exclusive copy anyway because of the E part of the
        // MOESI/MESI protocol.  Since we got the exclusive copy
        // there's no need to defer the targets, so move them up to
        // the regular target list.
        assert(!targets->needsExclusive);
        targets->needsExclusive = true;
        // if any of the deferred targets were upper-level cache
        // requests marked downstreamPending, need to clear that
        assert(!downstreamPending);  // not pending here anymore
        deferredTargets->clearDownstreamPending();
        // this clears out deferredTargets too
        targets->splice(targets->end(), *deferredTargets);
        deferredTargets->resetFlags();
    }
}

bool
MSHR::checkFunctional(PacketPtr pkt)
{
    // For printing, we treat the MSHR as a whole as single entity.
    // For other requests, we iterate over the individual targets
    // since that's where the actual data lies.
    if (pkt->isPrint()) {
        pkt->checkFunctional(this, addr, size, NULL);
        return false;
    } else {
        return (targets->checkFunctional(pkt) ||
                deferredTargets->checkFunctional(pkt));
    }
}

void
MSHR::print(std::ostream &os, int verbosity, const std::string &prefix) const
{
    ccprintf(os, "%s[%x:%x] %s %s %s state: %s %s %s %s\n",
             prefix, addr, addr+size-1,
             isForward ? "Forward" : "",
             isForwardNoResponse() ? "ForwNoResp" : "",
             needsExclusive() ? "Excl" : "",
             _isUncacheable ? "Unc" : "",
             inService ? "InSvc" : "",
             downstreamPending ? "DwnPend" : "",
             hasPostInvalidate() ? "PostInv" : "",
             hasPostDowngrade() ? "PostDowngr" : "");

    ccprintf(os, "%s  Targets:\n", prefix);
    targets->print(os, verbosity, prefix + "    ");
    if (!deferredTargets->empty()) {
        ccprintf(os, "%s  Deferred Targets:\n", prefix);
        deferredTargets->print(os, verbosity, prefix + "      ");
    }
}

std::string
MSHR::print() const
{
    ostringstream str;
    print(str);
    return str.str();
}

MSHR::~MSHR()
{
    delete[] targets;
    delete[] deferredTargets;
}
<commit_msg>mem: Fix MSHR print format<commit_after>/*
 * Copyright (c) 2012 ARM Limited
 * All rights reserved.
 *
 * The license below extends only to copyright in the software and shall
 * not be construed as granting a license to any other intellectual
 * property including but not limited to intellectual property relating
 * to a hardware implementation of the functionality of the software
 * licensed hereunder.  You may use the software subject to the license
 * terms below provided that you ensure that this notice is replicated
 * unmodified and in its entirety in all distributions of the software,
 * modified or unmodified, in source code or in binary form.
 *
 * Copyright (c) 2002-2005 The Regents of The University of Michigan
 * Copyright (c) 2010 Advanced Micro Devices, Inc.
 * All rights reserved.
 *
 * Redistribution and use in source and binary forms, with or without
 * modification, are permitted provided that the following conditions are
 * met: redistributions of source code must retain the above copyright
 * notice, this list of conditions and the following disclaimer;
 * redistributions in binary form must reproduce the above copyright
 * notice, this list of conditions and the following disclaimer in the
 * documentation and/or other materials provided with the distribution;
 * neither the name of the copyright holders nor the names of its
 * contributors may be used to endorse or promote products derived from
 * this software without specific prior written permission.
 *
 * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
 * "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
 * LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR
 * A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT
 * OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,
 * SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT
 * LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
 * DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY
 * THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
 * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
 * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
 *
 * Authors: Erik Hallnor
 *          Dave Greene
 */

/**
 * @file
 * Miss Status and Handling Register (MSHR) definitions.
 */

#include <algorithm>
#include <cassert>
#include <string>
#include <vector>

#include "base/misc.hh"
#include "base/types.hh"
#include "debug/Cache.hh"
#include "mem/cache/cache.hh"
#include "mem/cache/mshr.hh"
#include "sim/core.hh"

using namespace std;

MSHR::MSHR()
{
    inService = false;
    ntargets = 0;
    threadNum = InvalidThreadID;
    targets = new TargetList();
    deferredTargets = new TargetList();
}

MSHR::TargetList::TargetList()
    : needsExclusive(false), hasUpgrade(false)
{}

inline void
MSHR::TargetList::add(PacketPtr pkt, Tick readyTime,
                      Counter order, Target::Source source, bool markPending)
{
    if (source != Target::FromSnoop) {
        if (pkt->needsExclusive()) {
            needsExclusive = true;
        }

        // StoreCondReq is effectively an upgrade if it's in an MSHR
        // since it would have been failed already if we didn't have a
        // read-only copy
        if (pkt->isUpgrade() || pkt->cmd == MemCmd::StoreCondReq) {
            hasUpgrade = true;
        }
    }

    if (markPending) {
        // Iterate over the SenderState stack and see if we find
        // an MSHR entry. If we do, set the downstreamPending
        // flag. Otherwise, do nothing.
        MSHR *mshr = pkt->findNextSenderState<MSHR>();
        if (mshr != NULL) {
            assert(!mshr->downstreamPending);
            mshr->downstreamPending = true;
        }
    }

    push_back(Target(pkt, readyTime, order, source, markPending));
}

static void
replaceUpgrade(PacketPtr pkt)
{
    if (pkt->cmd == MemCmd::UpgradeReq) {
        pkt->cmd = MemCmd::ReadExReq;
        DPRINTF(Cache, "Replacing UpgradeReq with ReadExReq\n");
    } else if (pkt->cmd == MemCmd::SCUpgradeReq) {
        pkt->cmd = MemCmd::SCUpgradeFailReq;
        DPRINTF(Cache, "Replacing SCUpgradeReq with SCUpgradeFailReq\n");
    } else if (pkt->cmd == MemCmd::StoreCondReq) {
        pkt->cmd = MemCmd::StoreCondFailReq;
        DPRINTF(Cache, "Replacing StoreCondReq with StoreCondFailReq\n");
    }
}

void
MSHR::TargetList::replaceUpgrades()
{
    if (!hasUpgrade)
        return;

    Iterator end_i = end();
    for (Iterator i = begin(); i != end_i; ++i) {
        replaceUpgrade(i->pkt);
    }

    hasUpgrade = false;
}

void
MSHR::TargetList::clearDownstreamPending()
{
    Iterator end_i = end();
    for (Iterator i = begin(); i != end_i; ++i) {
        if (i->markedPending) {
            // Iterate over the SenderState stack and see if we find
            // an MSHR entry. If we find one, clear the
            // downstreamPending flag by calling
            // clearDownstreamPending(). This recursively clears the
            // downstreamPending flag in all caches this packet has
            // passed through.
            MSHR *mshr = i->pkt->findNextSenderState<MSHR>();
            if (mshr != NULL) {
                mshr->clearDownstreamPending();
            }
        }
    }
}

bool
MSHR::TargetList::checkFunctional(PacketPtr pkt)
{
    Iterator end_i = end();
    for (Iterator i = begin(); i != end_i; ++i) {
        if (pkt->checkFunctional(i->pkt)) {
            return true;
        }
    }

    return false;
}

void
MSHR::TargetList::
print(std::ostream &os, int verbosity, const std::string &prefix) const
{
    ConstIterator end_i = end();
    for (ConstIterator i = begin(); i != end_i; ++i) {
        const char *s;
        switch (i->source) {
          case Target::FromCPU:
            s = "FromCPU";
            break;
          case Target::FromSnoop:
            s = "FromSnoop";
            break;
          case Target::FromPrefetcher:
            s = "FromPrefetcher";
            break;
          default:
            s = "";
            break;
        }
        ccprintf(os, "%s%s: ", prefix, s);
        i->pkt->print(os, verbosity, "");
    }
}

void
MSHR::allocate(Addr _addr, int _size, PacketPtr target,
               Tick whenReady, Counter _order)
{
    addr = _addr;
    size = _size;
    readyTime = whenReady;
    order = _order;
    assert(target);
    isForward = false;
    _isUncacheable = target->req->isUncacheable();
    inService = false;
    downstreamPending = false;
    threadNum = 0;
    ntargets = 1;
    assert(targets->isReset());
    // Don't know of a case where we would allocate a new MSHR for a
    // snoop (mem-side request), so set source according to request here
    Target::Source source = (target->cmd == MemCmd::HardPFReq) ?
        Target::FromPrefetcher : Target::FromCPU;
    targets->add(target, whenReady, _order, source, true);
    assert(deferredTargets->isReset());
    data = NULL;
}

void
MSHR::clearDownstreamPending()
{
    assert(downstreamPending);
    downstreamPending = false;
    // recursively clear flag on any MSHRs we will be forwarding
    // responses to
    targets->clearDownstreamPending();
}

bool
MSHR::markInService(PacketPtr pkt)
{
    assert(!inService);
    if (isForwardNoResponse()) {
        // we just forwarded the request packet & don't expect a
        // response, so get rid of it
        assert(getNumTargets() == 1);
        popTarget();
        return true;
    }
    inService = true;
    pendingDirty = (targets->needsExclusive ||
                    (!pkt->sharedAsserted() && pkt->memInhibitAsserted()));
    postInvalidate = postDowngrade = false;

    if (!downstreamPending) {
        // let upstream caches know that the request has made it to a
        // level where it's going to get a response
        targets->clearDownstreamPending();
    }
    return false;
}

void
MSHR::deallocate()
{
    assert(targets->empty());
    targets->resetFlags();
    assert(deferredTargets->isReset());
    assert(ntargets == 0);
    inService = false;
}

/*
 * Adds a target to an MSHR
 */
void
MSHR::allocateTarget(PacketPtr pkt, Tick whenReady, Counter _order)
{
    // if there's a request already in service for this MSHR, we will
    // have to defer the new target until after the response if any of
    // the following are true:
    // - there are other targets already deferred
    // - there's a pending invalidate to be applied after the response
    //   comes back (but before this target is processed)
    // - this target requires an exclusive block and either we're not
    //   getting an exclusive block back or we have already snooped
    //   another read request that will downgrade our exclusive block
    //   to shared

    // assume we'd never issue a prefetch when we've got an
    // outstanding miss
    assert(pkt->cmd != MemCmd::HardPFReq);

    if (inService &&
        (!deferredTargets->empty() || hasPostInvalidate() ||
         (pkt->needsExclusive() &&
          (!isPendingDirty() || hasPostDowngrade() || isForward)))) {
        // need to put on deferred list
        if (hasPostInvalidate())
            replaceUpgrade(pkt);
        deferredTargets->add(pkt, whenReady, _order, Target::FromCPU, true);
    } else {
        // No request outstanding, or still OK to append to
        // outstanding request: append to regular target list.  Only
        // mark pending if current request hasn't been issued yet
        // (isn't in service).
        targets->add(pkt, whenReady, _order, Target::FromCPU, !inService);
    }

    ++ntargets;
}

bool
MSHR::handleSnoop(PacketPtr pkt, Counter _order)
{
    DPRINTF(Cache, "%s for %s address %x size %d\n", __func__,
            pkt->cmdString(), pkt->getAddr(), pkt->getSize());
    if (!inService || (pkt->isExpressSnoop() && downstreamPending)) {
        // Request has not been issued yet, or it's been issued
        // locally but is buffered unissued at some downstream cache
        // which is forwarding us this snoop.  Either way, the packet
        // we're snooping logically precedes this MSHR's request, so
        // the snoop has no impact on the MSHR, but must be processed
        // in the standard way by the cache.  The only exception is
        // that if we're an L2+ cache buffering an UpgradeReq from a
        // higher-level cache, and the snoop is invalidating, then our
        // buffered upgrades must be converted to read exclusives,
        // since the upper-level cache no longer has a valid copy.
        // That is, even though the upper-level cache got out on its
        // local bus first, some other invalidating transaction
        // reached the global bus before the upgrade did.
        if (pkt->needsExclusive()) {
            targets->replaceUpgrades();
            deferredTargets->replaceUpgrades();
        }

        return false;
    }

    // From here on down, the request issued by this MSHR logically
    // precedes the request we're snooping.
    if (pkt->needsExclusive()) {
        // snooped request still precedes the re-request we'll have to
        // issue for deferred targets, if any...
        deferredTargets->replaceUpgrades();
    }

    if (hasPostInvalidate()) {
        // a prior snoop has already appended an invalidation, so
        // logically we don't have the block anymore; no need for
        // further snooping.
        return true;
    }

    if (isPendingDirty() || pkt->isInvalidate()) {
        // We need to save and replay the packet in two cases:
        // 1. We're awaiting an exclusive copy, so ownership is pending,
        //    and we need to respond after we receive data.
        // 2. It's an invalidation (e.g., UpgradeReq), and we need
        //    to forward the snoop up the hierarchy after the current
        //    transaction completes.

        // Actual target device (typ. a memory) will delete the
        // packet on reception, so we need to save a copy here.
        PacketPtr cp_pkt = new Packet(pkt, true);
        targets->add(cp_pkt, curTick(), _order, Target::FromSnoop,
                     downstreamPending && targets->needsExclusive);
        ++ntargets;

        if (isPendingDirty()) {
            pkt->assertMemInhibit();
            pkt->setSupplyExclusive();
        }

        if (pkt->needsExclusive()) {
            // This transaction will take away our pending copy
            postInvalidate = true;
        }
    }

    if (!pkt->needsExclusive()) {
        // This transaction will get a read-shared copy, downgrading
        // our copy if we had an exclusive one
        postDowngrade = true;
        pkt->assertShared();
    }

    return true;
}

bool
MSHR::promoteDeferredTargets()
{
    assert(targets->empty());
    if (deferredTargets->empty()) {
        return false;
    }

    // swap targets & deferredTargets lists
    TargetList *tmp = targets;
    targets = deferredTargets;
    deferredTargets = tmp;

    assert(targets->size() == ntargets);

    // clear deferredTargets flags
    deferredTargets->resetFlags();

    order = targets->front().order;
    readyTime = std::max(curTick(), targets->front().readyTime);

    return true;
}

void
MSHR::handleFill(Packet *pkt, CacheBlk *blk)
{
    if (!pkt->sharedAsserted()
        && !(hasPostInvalidate() || hasPostDowngrade())
        && deferredTargets->needsExclusive) {
        // We got an exclusive response, but we have deferred targets
        // which are waiting to request an exclusive copy (not because
        // of a pending invalidate).  This can happen if the original
        // request was for a read-only (non-exclusive) block, but we
        // got an exclusive copy anyway because of the E part of the
        // MOESI/MESI protocol.  Since we got the exclusive copy
        // there's no need to defer the targets, so move them up to
        // the regular target list.
        assert(!targets->needsExclusive);
        targets->needsExclusive = true;
        // if any of the deferred targets were upper-level cache
        // requests marked downstreamPending, need to clear that
        assert(!downstreamPending);  // not pending here anymore
        deferredTargets->clearDownstreamPending();
        // this clears out deferredTargets too
        targets->splice(targets->end(), *deferredTargets);
        deferredTargets->resetFlags();
    }
}

bool
MSHR::checkFunctional(PacketPtr pkt)
{
    // For printing, we treat the MSHR as a whole as single entity.
    // For other requests, we iterate over the individual targets
    // since that's where the actual data lies.
    if (pkt->isPrint()) {
        pkt->checkFunctional(this, addr, size, NULL);
        return false;
    } else {
        return (targets->checkFunctional(pkt) ||
                deferredTargets->checkFunctional(pkt));
    }
}

void
MSHR::print(std::ostream &os, int verbosity, const std::string &prefix) const
{
    ccprintf(os, "%s[%x:%x] %s %s %s state: %s %s %s %s %s\n",
             prefix, addr, addr+size-1,
             isForward ? "Forward" : "",
             isForwardNoResponse() ? "ForwNoResp" : "",
             needsExclusive() ? "Excl" : "",
             _isUncacheable ? "Unc" : "",
             inService ? "InSvc" : "",
             downstreamPending ? "DwnPend" : "",
             hasPostInvalidate() ? "PostInv" : "",
             hasPostDowngrade() ? "PostDowngr" : "");

    ccprintf(os, "%s  Targets:\n", prefix);
    targets->print(os, verbosity, prefix + "    ");
    if (!deferredTargets->empty()) {
        ccprintf(os, "%s  Deferred Targets:\n", prefix);
        deferredTargets->print(os, verbosity, prefix + "      ");
    }
}

std::string
MSHR::print() const
{
    ostringstream str;
    print(str);
    return str.str();
}

MSHR::~MSHR()
{
    delete[] targets;
    delete[] deferredTargets;
}
<|endoftext|>


<commit_before>#include "mode_calc_BAR.hpp"

#include <mill/util/logger.hpp>
#include <toml/toml.hpp>
#include <numeric>
#include <cmath>

namespace mill
{

const char* mode_calc_BAR_usage() noexcept
{
    return "usage: mill calc BAR [input.toml]\n"
           "       It calculates BAR from a column of a space-separated file.\n"
           "       Lines starting with `#` will be skipped as a comment.\n"
           "       Each file must contain 2 columns, energy values of self potential and other's potential, respectively.\n"
           "       ```toml\n"
           "       [input]\n"
           "       traj1 = \"traj1.dat\"\n"
           "       traj2 = \"traj2.dat\"\n"
           "       kBT   = 0.6 # ~300K in kcal/mol\n"
           "\n"
           "       [output]\n"
           "       filename = \"BAR.dat\"\n"
           "       C.start = -6.0\n"
           "       C.stop  =  6.0\n"
           "       C.delta =  0.01\n"
           "       ```\n"
           ;
}

std::vector<std::pair<double, double>> read_energies(const std::string& fname)
{
    std::ifstream ifs(fname);
    if(not ifs.good())
    {
        log::fatal("file open error: ", fname);
    }

    std::vector<std::pair<double, double>> es;
    std::size_t n_line = 0;

    log::info("reading ", fname, "...");
    while(not ifs.eof())
    {
        n_line += 1;

        std::string line;
        std::getline(ifs, line);
        if(line.empty() || line.at(0) == '#')
        {
            continue;
        }

        std::istringstream iss(line);
        double e1 = std::numeric_limits<double>::quiet_NaN();
        double e2 = std::numeric_limits<double>::quiet_NaN();
        iss >> e1 >> e2;
        if(iss.fail())
        {
            log::fatal("line ", n_line, " does not have 2 column.");
        }
        es.emplace_back(e1, e2);
    }
    log::info("done.");

    return es;
}

// <1 / (1 + exp(βU_other - βU_self - C))>
double calc_weights(const std::vector<std::pair<double, double>>& energies, const double C, const double beta)
{
    double W = 0.0;
    for(const auto [Uself, Uother] : energies)
    {
        W += 1.0 / (1.0 + std::exp(beta * (Uother - Uself) - C));
    }
    return W;
}

int mode_calc_BAR(std::deque<std::string_view> args)
{
    if(args.empty())
    {
        log::error("mill calc BAR: too few arguments.");
        log::error(mode_calc_BAR_usage());
        return 1;
    }

    const std::string fname(args.front());
    args.pop_front();
    if(fname == "help")
    {
        log::info(mode_calc_BAR_usage());
        return 0;
    }

    const auto input = toml::parse(fname);

    const auto kBT  = toml::find<double>(input, "input", "kBT");
    const auto beta = 1.0 / kBT;

    const auto es1 = read_energies(toml::find<std::string>(input, "input", "traj1"));
    const auto es2 = read_energies(toml::find<std::string>(input, "input", "traj2"));

    const double logn1n2 = std::log(double(es2.size()) / double(es1.size()));

    const auto out = toml::find<std::string>(input, "output", "filename");
    std::ofstream ofs(out);
    if(not ofs.good())
    {
        log::fatal("file open error: ", out);
    }

    const auto C_beg = toml::find<double>(input, "output", "C", "start");
    const auto C_end = toml::find<double>(input, "output", "C", "stop");
    const auto dC    = toml::find<double>(input, "output", "C", "delta");
    const std::size_t nC = (C_end - C_beg) / dC;
    const double rnC = 1.0 / nC;

    ofs << "# C   weight1   weight2\n";
    for(std::size_t i=0; i<=nC; ++i)
    {
        std::ostringstream oss;
        for(std::size_t j=0; j<       (i*100/nC); ++j) {oss << '#';}
        for(std::size_t j=0; j< 100 - (i*100/nC); ++j) {oss << ' ';}
        std::cerr << "\r[" << oss.str() << ']';

        const auto C = C_beg * (1.0 - i * rnC) + C_end * (i * rnC);
        ofs << C << ' ' << calc_weights(es1,  C, beta)
                 << ' ' << calc_weights(es2, -C, beta)
                 << ' ' << logn1n2 << '\n';
    }

    return 0;
}

} // mill
<commit_msg>:children_crossing: add checking<commit_after>#include "mode_calc_BAR.hpp"

#include <mill/util/logger.hpp>
#include <toml/toml.hpp>
#include <numeric>
#include <cmath>

namespace mill
{

const char* mode_calc_BAR_usage() noexcept
{
    return "usage: mill calc BAR [input.toml]\n"
           "       It calculates BAR from a column of a space-separated file.\n"
           "       Lines starting with `#` will be skipped as a comment.\n"
           "       Each file must contain 2 columns, energy values of self potential and other's potential, respectively.\n"
           "       ```toml\n"
           "       [input]\n"
           "       traj1 = \"traj1.dat\"\n"
           "       traj2 = \"traj2.dat\"\n"
           "       kBT   = 0.6 # ~300K in kcal/mol\n"
           "\n"
           "       [output]\n"
           "       filename = \"BAR.dat\"\n"
           "       C.start = -6.0\n"
           "       C.stop  =  6.0\n"
           "       C.delta =  0.01\n"
           "       ```\n"
           ;
}

std::vector<std::pair<double, double>> read_energies(const std::string& fname)
{
    std::ifstream ifs(fname);
    if(not ifs.good())
    {
        log::fatal("file open error: ", fname);
    }

    std::vector<std::pair<double, double>> es;
    std::size_t n_line = 0;

    log::info("reading ", fname, "...");
    while(not ifs.eof())
    {
        n_line += 1;

        std::string line;
        std::getline(ifs, line);
        if(line.empty() || line.at(0) == '#')
        {
            continue;
        }

        std::istringstream iss(line);
        double e1 = std::numeric_limits<double>::quiet_NaN();
        double e2 = std::numeric_limits<double>::quiet_NaN();
        iss >> e1 >> e2;
        if(iss.fail())
        {
            log::fatal("line ", n_line, " does not have 2 column.");
        }
        es.emplace_back(e1, e2);
    }
    log::info("done.");

    return es;
}

// <1 / (1 + exp(βU_other - βU_self - C))>
double calc_weights(const std::vector<std::pair<double, double>>& energies, const double C, const double beta)
{
    double W = 0.0;
    for(const auto [Uself, Uother] : energies)
    {
        W += 1.0 / (1.0 + std::exp(beta * (Uother - Uself) - C));
    }
    return W;
}

int mode_calc_BAR(std::deque<std::string_view> args)
{
    if(args.empty())
    {
        log::error("mill calc BAR: too few arguments.");
        log::error(mode_calc_BAR_usage());
        return 1;
    }

    const std::string fname(args.front());
    args.pop_front();
    if(fname == "help")
    {
        log::info(mode_calc_BAR_usage());
        return 0;
    }

    const auto input = toml::parse(fname);

    const auto kBT  = toml::find<double>(input, "input", "kBT");
    const auto beta = 1.0 / kBT;

    const auto es1 = read_energies(toml::find<std::string>(input, "input", "traj1"));
    const auto es2 = read_energies(toml::find<std::string>(input, "input", "traj2"));

    const double logn1n2 = std::log(double(es2.size()) / double(es1.size()));

    const auto out = toml::find<std::string>(input, "output", "filename");
    std::ofstream ofs(out);
    if(not ofs.good())
    {
        log::fatal("file open error: ", out);
    }

    const auto C_beg = toml::find<double>(input, "output", "C", "start");
    const auto C_end = toml::find<double>(input, "output", "C", "stop");

    if(C_end < C_beg)
    {
        std::cerr << toml::format_error("C_beg < C_end is not satisfied",
                toml::find(input, "output", "C", "start"), "C_begin",
                toml::find(input, "output", "C", "stop"), "C_end") << std::endl;
        return 1;
    }

    const auto dC    = toml::find<double>(input, "output", "C", "delta");
    const std::size_t nC = (C_end - C_beg) / dC;
    const double rnC = 1.0 / nC;

    ofs << "# C   weight1   weight2\n";
    for(std::size_t i=0; i<=nC; ++i)
    {
        std::ostringstream oss;
        for(std::size_t j=0; j<       (i*100/nC); ++j) {oss << '#';}
        for(std::size_t j=0; j< 100 - (i*100/nC); ++j) {oss << ' ';}
        std::cerr << "\r[" << oss.str() << ']' << i * 100.0 / nC << '%' << std::flush;

        const auto C = C_beg * (1.0 - i * rnC) + C_end * (i * rnC);
        ofs << C << ' ' << calc_weights(es1,  C, beta)
                 << ' ' << calc_weights(es2, -C, beta)
                 << ' ' << logn1n2 << std::endl;
    }

    return 0;
}

} // mill
<|endoftext|>