<commit_before>// WorkspaceMenu.cc for Fluxbox
// Copyright (c) 2004 Henrik Kinnunen (fluxgen at users.sourceforge.net)
//
// Permission is hereby granted, free of charge, to any person obtaining a
// copy of this software and associated documentation files (the "Software"),
// to deal in the Software without restriction, including without limitation
// the rights to use, copy, modify, merge, publish, distribute, sublicense,
// and/or sell copies of the Software, and to permit persons to whom the
// Software is furnished to do so, subject to the following conditions:
//
// The above copyright notice and this permission notice shall be included in
// all copies or substantial portions of the Software.
//
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT.  IN NO EVENT SHALL
// THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
// FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER
// DEALINGS IN THE SOFTWARE.

// $Id: WorkspaceMenu.cc,v 1.4 2004/06/08 13:15:30 rathnor Exp $

#include "WorkspaceMenu.hh"

#include "Screen.hh"
#include "fluxbox.hh"
#include "Workspace.hh"
#include "MenuCreator.hh"

#include "FbTk/I18n.hh"
#include "FbTk/SimpleCommand.hh"
#include "FbTk/RefCount.hh"
#include "FbTk/MenuItem.hh"

#include <typeinfo>

WorkspaceMenu::WorkspaceMenu(BScreen &screen):
   FbMenu(screen.menuTheme(),
           screen.imageControl(),
           *screen.layerManager().
           getLayer(Fluxbox::instance()->getMenuLayer())) {

    init(screen);
}

void WorkspaceMenu::update(FbTk::Subject *subj) {
    if (subj != 0 && typeid(*subj) == typeid(BScreen::ScreenSubject)) {
        BScreen::ScreenSubject &screen_subj = *static_cast<BScreen::ScreenSubject *>(subj);
        BScreen &screen = screen_subj.screen();
        if (subj == &screen.currentWorkspaceSig()) {
            FbTk::MenuItem *item = 0;
            for (unsigned int i = 2; i < numberOfItems(); ++i) {
                item = find(i);
                if (item && item->isSelected()) {
                    setItemSelected(i, false);
                    FbTk::Menu::update(i);
                    break;
                }
            }
            setItemSelected(screen.currentWorkspace()->workspaceID() + 2, true);
            FbTk::Menu::update(screen.currentWorkspace()->workspaceID() + 2);
        } else if (subj == &screen.workspaceCountSig() ||
                   subj == &screen.workspaceNamesSig()) {
            while (numberOfItems() != 3) {
                remove(2);
            }
            // for each workspace add workspace name and it's menu
            // to our workspace menu
            for (size_t workspace = 0; workspace < screen.getCount();
                 ++workspace) {
                Workspace *wkspc = screen.getWorkspace(workspace);
                wkspc->menu().setInternalMenu();
                insert(wkspc->name().c_str(), &wkspc->menu(), numberOfItems() - 1);
            }

            FbTk::Menu::update(-1);
        }
    } else {
        FbTk::Menu::update(subj);
    }
}

void WorkspaceMenu::init(BScreen &screen) {
    screen.currentWorkspaceSig().attach(this);
    screen.workspaceCountSig().attach(this);
    screen.workspaceNamesSig().attach(this);
    using namespace FbTk;
    _FB_USES_NLS;

    removeAll();

    setLabel(_FBTEXT(Workspace, MenuTitle, "Workspaces", "Title of main workspace menu"));
    RefCount<Command> new_workspace(new SimpleCommand<BScreen, int>(screen, &BScreen::addWorkspace));
    RefCount<Command> remove_last(new SimpleCommand<BScreen, int>(screen, &BScreen::removeLastWorkspace));
    insert(_FBTEXT(Workspace, NewWorkspace, "New Workspace", "Add a new workspace"),
           new_workspace);
    insert(_FBTEXT(Workspace, RemoveLast, "Remove Last", "Remove the last workspace"),
           remove_last);
    // for each workspace add workspace name and it's menu to our workspace menu
    for (size_t workspace = 0; workspace < screen.getCount(); ++workspace) {
        Workspace *wkspc = screen.getWorkspace(workspace);
        wkspc->menu().setInternalMenu();
        insert(wkspc->name().c_str(), &wkspc->menu());
    }
    setItemSelected(screen.currentWorkspace()->workspaceID() + 2, true);

    insert(_FBTEXT(Menu, Icons, "Icons", "Iconic windows menu title"),
           MenuCreator::createMenuType("iconmenu", screen.screenNumber()));
    FbMenu::update();
}
<commit_msg>added button2-on-workspacename in workspacemenu jumps to workspacadded button2-on-workspacename in workspacemenu jumps to workspacee<commit_after>// WorkspaceMenu.cc for Fluxbox
// Copyright (c) 2004 Henrik Kinnunen (fluxgen at users.sourceforge.net)
//
// Permission is hereby granted, free of charge, to any person obtaining a
// copy of this software and associated documentation files (the "Software"),
// to deal in the Software without restriction, including without limitation
// the rights to use, copy, modify, merge, publish, distribute, sublicense,
// and/or sell copies of the Software, and to permit persons to whom the
// Software is furnished to do so, subject to the following conditions:
//
// The above copyright notice and this permission notice shall be included in
// all copies or substantial portions of the Software.
//
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT.  IN NO EVENT SHALL
// THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
// FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER
// DEALINGS IN THE SOFTWARE.

// $Id: WorkspaceMenu.cc,v 1.5 2004/09/06 21:12:10 akir Exp $

#include "WorkspaceMenu.hh"

#include "Screen.hh"
#include "fluxbox.hh"
#include "Workspace.hh"
#include "WorkspaceCmd.hh"
#include "MenuCreator.hh"

#include "FbTk/I18n.hh"
#include "FbTk/SimpleCommand.hh"
#include "FbTk/RefCount.hh"
#include "FbTk/MenuItem.hh"
#include "FbTk/MultiButtonMenuItem.hh"

#include <typeinfo>

WorkspaceMenu::WorkspaceMenu(BScreen &screen):
   FbMenu(screen.menuTheme(),
           screen.imageControl(),
           *screen.layerManager().
           getLayer(Fluxbox::instance()->getMenuLayer())) {

    init(screen);
}

void WorkspaceMenu::update(FbTk::Subject *subj) {
    if (subj != 0 && typeid(*subj) == typeid(BScreen::ScreenSubject)) {
        BScreen::ScreenSubject &screen_subj = *static_cast<BScreen::ScreenSubject *>(subj);
        BScreen &screen = screen_subj.screen();
        if (subj == &screen.currentWorkspaceSig()) {
            FbTk::MenuItem *item = 0;
            for (unsigned int i = 2; i < numberOfItems(); ++i) {
                item = find(i);
                if (item && item->isSelected()) {
                    setItemSelected(i, false);
                    FbTk::Menu::update(i);
                    break;
                }
            }
            setItemSelected(screen.currentWorkspace()->workspaceID() + 2, true);
            FbTk::Menu::update(screen.currentWorkspace()->workspaceID() + 2);
        } else if (subj == &screen.workspaceCountSig() ||
                   subj == &screen.workspaceNamesSig()) {
            while (numberOfItems() != 3) {
                remove(2);
            }
            // for each workspace add workspace name and it's menu
            // to our workspace menu
            for (size_t workspace = 0; workspace < screen.getCount();
                 ++workspace) {
                Workspace *wkspc = screen.getWorkspace(workspace);
                wkspc->menu().setInternalMenu();
                FbTk::MultiButtonMenuItem* mb_menu = new FbTk::MultiButtonMenuItem(5,
                                                                                wkspc->name().c_str(),
                                                                                &wkspc->menu());
                FbTk::RefCount<FbTk::Command> jump_cmd(new JumpToWorkspaceCmd(wkspc->workspaceID()));
                mb_menu->setCommand(2, jump_cmd);
                insert(mb_menu);
            }

            FbTk::Menu::update(-1);
        }
    } else {
        FbTk::Menu::update(subj);
    }
}

void WorkspaceMenu::init(BScreen &screen) {
    screen.currentWorkspaceSig().attach(this);
    screen.workspaceCountSig().attach(this);
    screen.workspaceNamesSig().attach(this);
    using namespace FbTk;
    _FB_USES_NLS;

    removeAll();

    setLabel(_FBTEXT(Workspace, MenuTitle, "Workspaces", "Title of main workspace menu"));
    RefCount<Command> new_workspace(new SimpleCommand<BScreen, int>(screen, &BScreen::addWorkspace));
    RefCount<Command> remove_last(new SimpleCommand<BScreen, int>(screen, &BScreen::removeLastWorkspace));
    insert(_FBTEXT(Workspace, NewWorkspace, "New Workspace", "Add a new workspace"),
           new_workspace);
    insert(_FBTEXT(Workspace, RemoveLast, "Remove Last", "Remove the last workspace"),
           remove_last);
    // for each workspace add workspace name and it's menu to our workspace menu
    for (size_t workspace = 0; workspace < screen.getCount(); ++workspace) {
        Workspace *wkspc = screen.getWorkspace(workspace);
        wkspc->menu().setInternalMenu();
        FbTk::MultiButtonMenuItem* mb_menu = new FbTk::MultiButtonMenuItem(5,
                                                                           wkspc->name().c_str(),
                                                                           &wkspc->menu());
        FbTk::RefCount<FbTk::Command> jump_cmd(new JumpToWorkspaceCmd(wkspc->workspaceID()));
        mb_menu->setCommand(2, jump_cmd);
        insert(mb_menu);
    }
    setItemSelected(screen.currentWorkspace()->workspaceID() + 2, true);

    insert(_FBTEXT(Menu, Icons, "Icons", "Iconic windows menu title"),
           MenuCreator::createMenuType("iconmenu", screen.screenNumber()));
    FbMenu::update();
}
<|endoftext|>


<commit_before>//  File: tpm_seal.cc
//  Author: Tom Roeder <tmroeder@google.com>
//
//  Description: A test of the TrouSerS API
//
//  Copyright (c) 2013, Google Inc.  All rights reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

#include <gflags/gflags.h>
#include <glog/logging.h>
#include <keyczar/keyczar.h>

#include <netinet/in.h>

#include <openssl/pem.h>
#include <openssl/x509.h>
#include <openssl/sha.h>

#include <stdio.h>
#include <string.h>
#include <stdlib.h>
#include <sys/stat.h>
#include <sys/types.h>

#include <tss/tss_error.h>
#include <tss/platform.h>
#include <tss/tss_defines.h>
#include <tss/tss_typedef.h>
#include <tss/tss_structs.h>
#include <tss/tspi.h>
#include <trousers/trousers.h>

#include <fstream>
#include <list>
#include <sstream>
#include <string>

using std::ifstream;
using std::list;
using std::string;
using std::stringstream;

#define PCR_LEN 20

DEFINE_string(
    aikblobfile, "aikblob",
    "A file containing an AIK blob that has been loaded into the TPM");

int main(int argc, char **argv) {
  google::ParseCommandLineFlags(&argc, &argv, true);
  FLAGS_alsologtostderr = true;
  google::InitGoogleLogging(argv[0]);

  TSS_HCONTEXT tss_ctx;
  TSS_HTPM tpm;
  TSS_RESULT result;
  TSS_HKEY srk = 0;
  TSS_HPOLICY srk_policy = 0;
  TSS_UUID srk_uuid = {0x00000000, 0x0000, 0x0000, 0x00, 0x00,
                       {0x00, 0x00, 0x00, 0x00, 0x00, 0x01}};
  BYTE secret[20];

  // Use the well-known secret of 20 zeroes.
  memset(secret, 0, 20);

  // Set up the TSS context and the SRK + policy (with the right secret).
  result = Tspi_Context_Create(&tss_ctx);
  CHECK_EQ(result, TSS_SUCCESS) << "Could not create a TSS context.";

  result = Tspi_Context_Connect(tss_ctx, NULL /* Default TPM */);
  CHECK_EQ(result, TSS_SUCCESS) << "Could not connect to the default TPM";

  result = Tspi_Context_GetTpmObject(tss_ctx, &tpm);
  CHECK_EQ(result, TSS_SUCCESS) << "Could not get a handle to the TPM";

  result =
      Tspi_Context_LoadKeyByUUID(tss_ctx, TSS_PS_TYPE_SYSTEM, srk_uuid, &srk);
  CHECK_EQ(result, TSS_SUCCESS) << "Could not load the SRK handle";

  result = Tspi_GetPolicyObject(srk, TSS_POLICY_USAGE, &srk_policy);
  CHECK_EQ(result, TSS_SUCCESS) << "Could not get the SRK policy handle";

  result = Tspi_Policy_SetSecret(srk_policy, TSS_SECRET_MODE_SHA1, 20, secret);
  CHECK_EQ(result, TSS_SUCCESS) << "Could not set the well-known secret";

  // Create and fill the PCR information.
  TSS_HPCRS pcrs;
  result = Tspi_Context_CreateObject(tss_ctx, TSS_OBJECT_TYPE_PCRS, 0, &pcrs);
  CHECK_EQ(result, TSS_SUCCESS) << "Could not create a PCRs object";

  // This seal operation is meant to be used with DRTM, so the only PCRs that it
  // reads are 17 and 18. This is where you can set other PCRs to use.
  list<UINT32> pcrs_to_seal{17, 18};
  BYTE *pcr_value = NULL;
  UINT32 pcr_value_len = 0;
  for (UINT32 ui : pcrs_to_seal) {
    result = Tspi_TPM_PcrRead(tpm, ui, &pcr_value_len, &pcr_value);
    CHECK_EQ(result, TSS_SUCCESS) << "Could not read the value of PCR " << ui;

    result = Tspi_PcrComposite_SetPcrValue(pcrs, ui, pcr_value_len, pcr_value);
    CHECK_EQ(result, TSS_SUCCESS) << "Could not set the PCR value" << ui
                                  << " for sealing";
  }

  TSS_HENCDATA enc_data;
  result = Tspi_Context_CreateObject(tss_ctx, TSS_OBJECT_TYPE_ENCDATA,
                                     TSS_ENCDATA_SEAL, &enc_data);
  CHECK_EQ(result, TSS_SUCCESS) << "Could not create the data for sealing";

  BYTE data[16] = {0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15};
  result = Tspi_Data_Seal(enc_data, srk, 16, data, pcrs);
  CHECK_EQ(result, TSS_SUCCESS) << "Could not seal the test data";

  // Extract the sealed data, then try to unseal it.
  BYTE *sealed_data;
  UINT32 sealed_data_len;
  result = Tspi_GetAttribData(enc_data, TSS_TSPATTRIB_ENCDATA_BLOB,
                              TSS_TSPATTRIB_ENCDATABLOB_BLOB, &sealed_data_len,
                              &sealed_data);
  CHECK_EQ(result, TSS_SUCCESS) << "Could not get the sealed bits";

  BYTE *unsealed_data;
  UINT32 unsealed_data_len;
  result = Tspi_Data_Unseal(enc_data, srk, &unsealed_data_len, &unsealed_data);
  CHECK_EQ(result, TSS_SUCCESS) << "Could not unseal the data";

  // Check that the data was unsealed correctly.
  CHECK_EQ(unsealed_data_len, 16U) << "The unsealed data was the wrong length";
  CHECK_EQ(memcmp(unsealed_data, data, 16), 0)
      << "The unsealed data did not match the original data";

  // Get the public key blob from the AIK.
  // Load the blob and try to load the AIK
  ifstream blob_stream(FLAGS_aikblobfile, ifstream::in);
  stringstream blob_buf;
  blob_buf << blob_stream.rdbuf();
  string blob = blob_buf.str();
  UINT32 blob_len = (UINT32)blob.size();
  TSS_HKEY aik;
  result = Tspi_Context_LoadKeyByBlob(
      tss_ctx, srk, blob_len,
      reinterpret_cast<BYTE *>(const_cast<char *>(blob.data())), &aik);
  CHECK_EQ(result, TSS_SUCCESS) << "Could not load the AIK";

  // Generate a Quote using this AIK.
  BYTE hash_to_sign[20];
  memset(hash_to_sign, 0, sizeof(hash_to_sign));

  // First get the max number of PCRs in the TPM
  UINT32 tpm_property = TSS_TPMCAP_PROP_PCR;
  UINT32 npcrs_len;
  BYTE *npcrs;
  result =
      Tspi_TPM_GetCapability(tpm, TSS_TPMCAP_PROPERTY, sizeof(tpm_property),
                             (BYTE *)&tpm_property, &npcrs_len, &npcrs);
  CHECK_EQ(result, TSS_SUCCESS) << "Could not get the number of PCRs";

  UINT32 pcr_max = *(UINT32 *)npcrs;
  Tspi_Context_FreeMemory(tss_ctx, npcrs);

  // The total number of bytes needed to store the PCR mask
  UINT32 pcr_mask_len = (pcr_max + 7) / 8;

  TSS_HPCRS quote_pcrs;
  result = Tspi_Context_CreateObject(tss_ctx, TSS_OBJECT_TYPE_PCRS,
                                     TSS_PCRS_STRUCT_INFO, &quote_pcrs);
  CHECK_EQ(result, TSS_SUCCESS)
      << "Could not create a PCRs object for the Quote";

  // The following code for setting up the composite hash is based on
  // aikquote.c, the sample AIK quote code.
  scoped_array<BYTE> serialized_pcrs(new BYTE[
      sizeof(UINT16) + pcr_mask_len + sizeof(UINT32) + PCR_LEN * pcr_max]);

  // So, the format is:
  // UINT16: size of pcr mask (in network byte order)
  // pcr mask
  // UINT32: size of serialized pcrs
  // serialized pcrs

  BYTE *pcr_buf = serialized_pcrs.get();
  UINT32 index = 0;

  // size of pcr mask
  *(UINT16 *)(pcr_buf + index) = htons(pcr_mask_len);
  index += sizeof(UINT16);
  memset(pcr_buf + index, 0, pcr_mask_len);

  // Set the appropriate bits for the PCRs
  for (UINT32 pcr_val : pcrs_to_seal) {
    result = Tspi_PcrComposite_SelectPcrIndex(quote_pcrs, pcr_val);
    CHECK_EQ(result, TSS_SUCCESS) << "Could not select PCR " << (int)pcr_val;
    pcr_buf[index + (pcr_val / 8)] |= 1 << (pcr_val % 8);
  }
  index += pcr_mask_len;

  // Write the length of the set of PCRs
  *(UINT32 *)(pcr_buf + index) = htonl(PCR_LEN * pcrs_to_seal.size());
  index += sizeof(UINT32);

  TSS_VALIDATION valid;
  valid.ulExternalDataLength = sizeof(hash_to_sign);
  valid.rgbExternalData = hash_to_sign;

  result = Tspi_TPM_Quote(tpm, aik, quote_pcrs, &valid);
  CHECK_EQ(result, TSS_SUCCESS) << "Could not quote data with the AIK";

  // Check the hash from the quote
  TPM_QUOTE_INFO *quote_info = (TPM_QUOTE_INFO *)valid.rgbData;

  BYTE *temp_pcr;
  UINT32 temp_pcr_len;
  for (UINT32 pcr_val : pcrs_to_seal) {
    result = Tspi_PcrComposite_GetPcrValue(quote_pcrs, pcr_val, &temp_pcr_len,
                                           &temp_pcr);
    CHECK_EQ(result, TSS_SUCCESS) << "Could not get PCR " << (int)pcr_val;

    memcpy(pcr_buf + index, temp_pcr, temp_pcr_len);
    index += temp_pcr_len;

    Tspi_Context_FreeMemory(tss_ctx, temp_pcr);
  }

  // Compute the composite hash to check against the hash in the quote info.
  BYTE pcr_digest[20];
  SHA1(pcr_buf, index, pcr_digest);

  if (memcmp(pcr_digest, quote_info->compositeHash.digest,
             sizeof(pcr_digest)) !=
      0) {
    // aikquote.c here says "Try with a smaller digest length". I don't know
    // why. This code removes one of the bytes in the pcr mask and shifts
    // everything over to account for the difference, then hashes and tries
    // again.
    *(UINT16 *)pcr_buf = htons(pcr_mask_len - 1);
    memmove(pcr_buf + sizeof(UINT16) + pcr_mask_len - 1,
            pcr_buf + sizeof(UINT16) + pcr_mask_len,
            index - (sizeof(UINT16) + pcr_mask_len));
    index -= 1;
    SHA1(pcr_buf, index, pcr_digest);
    if (memcmp(pcr_digest, quote_info->compositeHash.digest,
               sizeof(pcr_digest)) !=
        0) {
      LOG(FATAL) << "Neither size of hash input worked";
      return 1;
    }
  }

  // At this point, the quote is in pcr_buf with length index.

  // Extract the modulus from the AIK
  UINT32 aik_mod_len;
  BYTE *aik_mod;
  result = Tspi_GetAttribData(aik, TSS_TSPATTRIB_RSAKEY_INFO,
                              TSS_TSPATTRIB_KEYINFO_RSA_MODULUS, &aik_mod_len,
                              &aik_mod);
  CHECK_EQ(result, TSS_SUCCESS) << "Could not extract the RSA modulus";

  // Set up an OpenSSL RSA public key to use to verify the Quote
  RSA *aik_rsa = RSA_new();
  aik_rsa->n = BN_bin2bn(aik_mod, aik_mod_len, NULL);
  aik_rsa->e = BN_new();
  BN_set_word(aik_rsa->e, 0x10001);

  // Try hashing and verifying the TPM_QUOTE_INFO itself
  BYTE quote_hash[20];

  // The quote can be verified in a qinfo, which has a header of 8 bytes, and
  // two hashes.
  // The first hash is the hash of the external data, and the second is the hash
  // of the
  // quote itself (pcr_buf above with length in index). This can be hashed and
  // verified
  // directly by OpenSSL.

  BYTE qinfo[8 + 2 * 20];
  qinfo[0] = 1;
  qinfo[1] = 1;
  qinfo[2] = 0;
  qinfo[3] = 0;
  qinfo[4] = 'Q';
  qinfo[5] = 'U';
  qinfo[6] = 'O';
  qinfo[7] = 'T';
  SHA1(pcr_buf, index, qinfo + 8);
  memcpy(qinfo + 8 + 20, hash_to_sign, 20);

  SHA1(qinfo, sizeof(qinfo), quote_hash);
  BYTE *sig = valid.rgbValidationData;
  UINT32 sig_len = valid.ulValidationDataLength;
  CHECK_EQ(RSA_verify(NID_sha1, quote_hash, sizeof(quote_hash), sig, sig_len,
                      aik_rsa),
           1) << "The RSA signature did not pass verification";

  // Clean-up code.
  result = Tspi_Context_FreeMemory(tss_ctx, NULL);
  CHECK_EQ(result, TSS_SUCCESS) << "Could not free the context";

  result = Tspi_Context_Close(tss_ctx);
  CHECK_EQ(result, TSS_SUCCESS) << "Could not clean up the context";

  return 0;
}
<commit_msg>Fixed the file name in the header to match the actual file name<commit_after>//  File: tpm_test.cc
//  Author: Tom Roeder <tmroeder@google.com>
//
//  Description: A test of the TrouSerS API
//
//  Copyright (c) 2013, Google Inc.  All rights reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

#include <gflags/gflags.h>
#include <glog/logging.h>
#include <keyczar/keyczar.h>

#include <netinet/in.h>

#include <openssl/pem.h>
#include <openssl/x509.h>
#include <openssl/sha.h>

#include <stdio.h>
#include <string.h>
#include <stdlib.h>
#include <sys/stat.h>
#include <sys/types.h>

#include <tss/tss_error.h>
#include <tss/platform.h>
#include <tss/tss_defines.h>
#include <tss/tss_typedef.h>
#include <tss/tss_structs.h>
#include <tss/tspi.h>
#include <trousers/trousers.h>

#include <fstream>
#include <list>
#include <sstream>
#include <string>

using std::ifstream;
using std::list;
using std::string;
using std::stringstream;

#define PCR_LEN 20

DEFINE_string(
    aikblobfile, "aikblob",
    "A file containing an AIK blob that has been loaded into the TPM");

int main(int argc, char **argv) {
  google::ParseCommandLineFlags(&argc, &argv, true);
  FLAGS_alsologtostderr = true;
  google::InitGoogleLogging(argv[0]);

  TSS_HCONTEXT tss_ctx;
  TSS_HTPM tpm;
  TSS_RESULT result;
  TSS_HKEY srk = 0;
  TSS_HPOLICY srk_policy = 0;
  TSS_UUID srk_uuid = {0x00000000, 0x0000, 0x0000, 0x00, 0x00,
                       {0x00, 0x00, 0x00, 0x00, 0x00, 0x01}};
  BYTE secret[20];

  // Use the well-known secret of 20 zeroes.
  memset(secret, 0, 20);

  // Set up the TSS context and the SRK + policy (with the right secret).
  result = Tspi_Context_Create(&tss_ctx);
  CHECK_EQ(result, TSS_SUCCESS) << "Could not create a TSS context.";

  result = Tspi_Context_Connect(tss_ctx, NULL /* Default TPM */);
  CHECK_EQ(result, TSS_SUCCESS) << "Could not connect to the default TPM";

  result = Tspi_Context_GetTpmObject(tss_ctx, &tpm);
  CHECK_EQ(result, TSS_SUCCESS) << "Could not get a handle to the TPM";

  result =
      Tspi_Context_LoadKeyByUUID(tss_ctx, TSS_PS_TYPE_SYSTEM, srk_uuid, &srk);
  CHECK_EQ(result, TSS_SUCCESS) << "Could not load the SRK handle";

  result = Tspi_GetPolicyObject(srk, TSS_POLICY_USAGE, &srk_policy);
  CHECK_EQ(result, TSS_SUCCESS) << "Could not get the SRK policy handle";

  result = Tspi_Policy_SetSecret(srk_policy, TSS_SECRET_MODE_SHA1, 20, secret);
  CHECK_EQ(result, TSS_SUCCESS) << "Could not set the well-known secret";

  // Create and fill the PCR information.
  TSS_HPCRS pcrs;
  result = Tspi_Context_CreateObject(tss_ctx, TSS_OBJECT_TYPE_PCRS, 0, &pcrs);
  CHECK_EQ(result, TSS_SUCCESS) << "Could not create a PCRs object";

  // This seal operation is meant to be used with DRTM, so the only PCRs that it
  // reads are 17 and 18. This is where you can set other PCRs to use.
  list<UINT32> pcrs_to_seal{17, 18};
  BYTE *pcr_value = NULL;
  UINT32 pcr_value_len = 0;
  for (UINT32 ui : pcrs_to_seal) {
    result = Tspi_TPM_PcrRead(tpm, ui, &pcr_value_len, &pcr_value);
    CHECK_EQ(result, TSS_SUCCESS) << "Could not read the value of PCR " << ui;

    result = Tspi_PcrComposite_SetPcrValue(pcrs, ui, pcr_value_len, pcr_value);
    CHECK_EQ(result, TSS_SUCCESS) << "Could not set the PCR value" << ui
                                  << " for sealing";
  }

  TSS_HENCDATA enc_data;
  result = Tspi_Context_CreateObject(tss_ctx, TSS_OBJECT_TYPE_ENCDATA,
                                     TSS_ENCDATA_SEAL, &enc_data);
  CHECK_EQ(result, TSS_SUCCESS) << "Could not create the data for sealing";

  BYTE data[16] = {0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15};
  result = Tspi_Data_Seal(enc_data, srk, 16, data, pcrs);
  CHECK_EQ(result, TSS_SUCCESS) << "Could not seal the test data";

  // Extract the sealed data, then try to unseal it.
  BYTE *sealed_data;
  UINT32 sealed_data_len;
  result = Tspi_GetAttribData(enc_data, TSS_TSPATTRIB_ENCDATA_BLOB,
                              TSS_TSPATTRIB_ENCDATABLOB_BLOB, &sealed_data_len,
                              &sealed_data);
  CHECK_EQ(result, TSS_SUCCESS) << "Could not get the sealed bits";

  BYTE *unsealed_data;
  UINT32 unsealed_data_len;
  result = Tspi_Data_Unseal(enc_data, srk, &unsealed_data_len, &unsealed_data);
  CHECK_EQ(result, TSS_SUCCESS) << "Could not unseal the data";

  // Check that the data was unsealed correctly.
  CHECK_EQ(unsealed_data_len, 16U) << "The unsealed data was the wrong length";
  CHECK_EQ(memcmp(unsealed_data, data, 16), 0)
      << "The unsealed data did not match the original data";

  // Get the public key blob from the AIK.
  // Load the blob and try to load the AIK
  ifstream blob_stream(FLAGS_aikblobfile, ifstream::in);
  stringstream blob_buf;
  blob_buf << blob_stream.rdbuf();
  string blob = blob_buf.str();
  UINT32 blob_len = (UINT32)blob.size();
  TSS_HKEY aik;
  result = Tspi_Context_LoadKeyByBlob(
      tss_ctx, srk, blob_len,
      reinterpret_cast<BYTE *>(const_cast<char *>(blob.data())), &aik);
  CHECK_EQ(result, TSS_SUCCESS) << "Could not load the AIK";

  // Generate a Quote using this AIK.
  BYTE hash_to_sign[20];
  memset(hash_to_sign, 0, sizeof(hash_to_sign));

  // First get the max number of PCRs in the TPM
  UINT32 tpm_property = TSS_TPMCAP_PROP_PCR;
  UINT32 npcrs_len;
  BYTE *npcrs;
  result =
      Tspi_TPM_GetCapability(tpm, TSS_TPMCAP_PROPERTY, sizeof(tpm_property),
                             (BYTE *)&tpm_property, &npcrs_len, &npcrs);
  CHECK_EQ(result, TSS_SUCCESS) << "Could not get the number of PCRs";

  UINT32 pcr_max = *(UINT32 *)npcrs;
  Tspi_Context_FreeMemory(tss_ctx, npcrs);

  // The total number of bytes needed to store the PCR mask
  UINT32 pcr_mask_len = (pcr_max + 7) / 8;

  TSS_HPCRS quote_pcrs;
  result = Tspi_Context_CreateObject(tss_ctx, TSS_OBJECT_TYPE_PCRS,
                                     TSS_PCRS_STRUCT_INFO, &quote_pcrs);
  CHECK_EQ(result, TSS_SUCCESS)
      << "Could not create a PCRs object for the Quote";

  // The following code for setting up the composite hash is based on
  // aikquote.c, the sample AIK quote code.
  scoped_array<BYTE> serialized_pcrs(new BYTE[
      sizeof(UINT16) + pcr_mask_len + sizeof(UINT32) + PCR_LEN * pcr_max]);

  // So, the format is:
  // UINT16: size of pcr mask (in network byte order)
  // pcr mask
  // UINT32: size of serialized pcrs
  // serialized pcrs

  BYTE *pcr_buf = serialized_pcrs.get();
  UINT32 index = 0;

  // size of pcr mask
  *(UINT16 *)(pcr_buf + index) = htons(pcr_mask_len);
  index += sizeof(UINT16);
  memset(pcr_buf + index, 0, pcr_mask_len);

  // Set the appropriate bits for the PCRs
  for (UINT32 pcr_val : pcrs_to_seal) {
    result = Tspi_PcrComposite_SelectPcrIndex(quote_pcrs, pcr_val);
    CHECK_EQ(result, TSS_SUCCESS) << "Could not select PCR " << (int)pcr_val;
    pcr_buf[index + (pcr_val / 8)] |= 1 << (pcr_val % 8);
  }
  index += pcr_mask_len;

  // Write the length of the set of PCRs
  *(UINT32 *)(pcr_buf + index) = htonl(PCR_LEN * pcrs_to_seal.size());
  index += sizeof(UINT32);

  TSS_VALIDATION valid;
  valid.ulExternalDataLength = sizeof(hash_to_sign);
  valid.rgbExternalData = hash_to_sign;

  result = Tspi_TPM_Quote(tpm, aik, quote_pcrs, &valid);
  CHECK_EQ(result, TSS_SUCCESS) << "Could not quote data with the AIK";

  // Check the hash from the quote
  TPM_QUOTE_INFO *quote_info = (TPM_QUOTE_INFO *)valid.rgbData;

  BYTE *temp_pcr;
  UINT32 temp_pcr_len;
  for (UINT32 pcr_val : pcrs_to_seal) {
    result = Tspi_PcrComposite_GetPcrValue(quote_pcrs, pcr_val, &temp_pcr_len,
                                           &temp_pcr);
    CHECK_EQ(result, TSS_SUCCESS) << "Could not get PCR " << (int)pcr_val;

    memcpy(pcr_buf + index, temp_pcr, temp_pcr_len);
    index += temp_pcr_len;

    Tspi_Context_FreeMemory(tss_ctx, temp_pcr);
  }

  // Compute the composite hash to check against the hash in the quote info.
  BYTE pcr_digest[20];
  SHA1(pcr_buf, index, pcr_digest);

  if (memcmp(pcr_digest, quote_info->compositeHash.digest,
             sizeof(pcr_digest)) !=
      0) {
    // aikquote.c here says "Try with a smaller digest length". I don't know
    // why. This code removes one of the bytes in the pcr mask and shifts
    // everything over to account for the difference, then hashes and tries
    // again.
    *(UINT16 *)pcr_buf = htons(pcr_mask_len - 1);
    memmove(pcr_buf + sizeof(UINT16) + pcr_mask_len - 1,
            pcr_buf + sizeof(UINT16) + pcr_mask_len,
            index - (sizeof(UINT16) + pcr_mask_len));
    index -= 1;
    SHA1(pcr_buf, index, pcr_digest);
    if (memcmp(pcr_digest, quote_info->compositeHash.digest,
               sizeof(pcr_digest)) !=
        0) {
      LOG(FATAL) << "Neither size of hash input worked";
      return 1;
    }
  }

  // At this point, the quote is in pcr_buf with length index.

  // Extract the modulus from the AIK
  UINT32 aik_mod_len;
  BYTE *aik_mod;
  result = Tspi_GetAttribData(aik, TSS_TSPATTRIB_RSAKEY_INFO,
                              TSS_TSPATTRIB_KEYINFO_RSA_MODULUS, &aik_mod_len,
                              &aik_mod);
  CHECK_EQ(result, TSS_SUCCESS) << "Could not extract the RSA modulus";

  // Set up an OpenSSL RSA public key to use to verify the Quote
  RSA *aik_rsa = RSA_new();
  aik_rsa->n = BN_bin2bn(aik_mod, aik_mod_len, NULL);
  aik_rsa->e = BN_new();
  BN_set_word(aik_rsa->e, 0x10001);

  // Try hashing and verifying the TPM_QUOTE_INFO itself
  BYTE quote_hash[20];

  // The quote can be verified in a qinfo, which has a header of 8 bytes, and
  // two hashes.
  // The first hash is the hash of the external data, and the second is the hash
  // of the
  // quote itself (pcr_buf above with length in index). This can be hashed and
  // verified
  // directly by OpenSSL.

  BYTE qinfo[8 + 2 * 20];
  qinfo[0] = 1;
  qinfo[1] = 1;
  qinfo[2] = 0;
  qinfo[3] = 0;
  qinfo[4] = 'Q';
  qinfo[5] = 'U';
  qinfo[6] = 'O';
  qinfo[7] = 'T';
  SHA1(pcr_buf, index, qinfo + 8);
  memcpy(qinfo + 8 + 20, hash_to_sign, 20);

  SHA1(qinfo, sizeof(qinfo), quote_hash);
  BYTE *sig = valid.rgbValidationData;
  UINT32 sig_len = valid.ulValidationDataLength;
  CHECK_EQ(RSA_verify(NID_sha1, quote_hash, sizeof(quote_hash), sig, sig_len,
                      aik_rsa),
           1) << "The RSA signature did not pass verification";

  // Clean-up code.
  result = Tspi_Context_FreeMemory(tss_ctx, NULL);
  CHECK_EQ(result, TSS_SUCCESS) << "Could not free the context";

  result = Tspi_Context_Close(tss_ctx);
  CHECK_EQ(result, TSS_SUCCESS) << "Could not clean up the context";

  return 0;
}
<|endoftext|>


<commit_before>/*
 * substitutionMap.cpp
 *
 */

#include "SubstitutionMap.h"
#include "simplifier.h"
#include "../AST/ArrayTransformer.h"

namespace BEEV
{

SubstitutionMap::~SubstitutionMap()
{
	delete SolverMap;
}

//if a is READ(Arr,const) or SYMBOL, and b is BVCONST then return 1
//if b is READ(Arr,const) or SYMBOL, and a is BVCONST then return -1
//
//else return 0 by default
int TermOrder(const ASTNode& a, const ASTNode& b)
{
  Kind k1 = a.GetKind();
  Kind k2 = b.GetKind();

  //a is of the form READ(Arr,const), and b is const, or
  //a is of the form var, and b is const
  if ((k1 == READ
       && a[0].GetKind() == SYMBOL
       && a[1].GetKind() == BVCONST
       && (k2 == BVCONST)))
    // || k2 == READ && b[0].GetKind() == SYMBOL && b[1].GetKind()
    // == BVCONST)))
    return 1;

  if (SYMBOL == k1 && (BVCONST == k2 || TRUE == k2 || FALSE == k2))
    return 1;

  //b is of the form READ(Arr,const), and a is const, or
  //b is of the form var, and a is const
  if ((k1 == BVCONST)
      && ((k2 == READ
           && b[0].GetKind() == SYMBOL
           && b[1].GetKind() == BVCONST)))
    return -1;

  if (SYMBOL == k2
      && (BVCONST == k1
          || TRUE == k1
          || FALSE == k1))
    return -1;

  return 0;
} //End of TermOrder()

//This finds everything which is (= SYMBOL BVCONST) and everything that is (READ SYMBOL BVCONST).
//i.e. anything that has a termorder of 1 or -1.
// The bvsolver puts expressions of the form (= SYMBOL TERM) into the solverMap.

ASTNode SubstitutionMap::CreateSubstitutionMap(const ASTNode& a,  ArrayTransformer*at
)
{
  if (!bm->UserFlags.wordlevel_solve_flag)
    return a;

  ASTNode output = a;
  //if the variable has been solved for, then simply return it
  if (CheckSubstitutionMap(a, output))
    return output;

  //traverse a and populate the SubstitutionMap
  const Kind k = a.GetKind();
  if (SYMBOL == k && BOOLEAN_TYPE == a.GetType())
    {
      bool updated = UpdateSubstitutionMap(a, ASTTrue);
      output = updated ? ASTTrue : a;
      return output;
    }
  if (NOT == k && SYMBOL == a[0].GetKind())
    {
      bool updated = UpdateSubstitutionMap(a[0], ASTFalse);
      output = updated ? ASTTrue : a;
      return output;
    }

  if (IFF == k || EQ == k)
    {
      ASTVec c = a.GetChildren();
      SortByArith(c);

      ASTNode c1;
      if (EQ == k)
    	  c1 = simp->SimplifyTerm(c[1]);
      else// (IFF == k )
    	  c1= simp->SimplifyFormula_NoRemoveWrites(c[1], false);

      //fill the arrayname readindices vector if e0 is a
      //READ(Arr,index) and index is a BVCONST
      int to;
      if ((to =TermOrder(c[0],c1))==1 && c[0].GetKind() == READ)
    	  at->FillUp_ArrReadIndex_Vec(c[0], c1);
      else if (to ==-1 && c1.GetKind() == READ)
		at->FillUp_ArrReadIndex_Vec(c1,c[0]);

      bool updated = UpdateSubstitutionMap(c[0], c1);
      output = updated ? ASTTrue : a;
      return output;
    }

  if (AND == k)
    {
      ASTVec o;
      ASTVec c = a.GetChildren();
      for (ASTVec::iterator
             it = c.begin(), itend = c.end();
           it != itend; it++)
        {
          simp->UpdateAlwaysTrueFormMap(*it);
          ASTNode aaa = CreateSubstitutionMap(*it,at);

          if (ASTTrue != aaa)
            {
              if (ASTFalse == aaa)
                return ASTFalse;
              else
                o.push_back(aaa);
            }
        }
      if (o.size() == 0)
        return ASTTrue;

      if (o.size() == 1)
        return o[0];

      return bm->CreateNode(AND, o);
    }

  //printf("I gave up on kind: %d node: %d\n", k, a.GetNodeNum());
  return output;
} //end of CreateSubstitutionMap()

};
<commit_msg>Don't perform simplifications when creating the substitution map. I'm not sure if this is justified, but it makes STP easier to understand. I'll check if it's useful.<commit_after>/*
 * substitutionMap.cpp
 *
 */

#include "SubstitutionMap.h"
#include "simplifier.h"
#include "../AST/ArrayTransformer.h"

namespace BEEV
{

SubstitutionMap::~SubstitutionMap()
{
	delete SolverMap;
}

// if false. Don't simplify while creating the substitution map.
// This makes it easier to spot how long is spent in the simplifier.
const bool simplify_during_create_subM = false;

//if a is READ(Arr,const) or SYMBOL, and b is BVCONST then return 1
//if b is READ(Arr,const) or SYMBOL, and a is BVCONST then return -1
//
//else return 0 by default
int TermOrder(const ASTNode& a, const ASTNode& b)
{
  Kind k1 = a.GetKind();
  Kind k2 = b.GetKind();

  //a is of the form READ(Arr,const), and b is const, or
  //a is of the form var, and b is const
  if ((k1 == READ
       && a[0].GetKind() == SYMBOL
       && a[1].GetKind() == BVCONST
       && (k2 == BVCONST)))
    // || k2 == READ && b[0].GetKind() == SYMBOL && b[1].GetKind()
    // == BVCONST)))
    return 1;

  if (SYMBOL == k1 && (BVCONST == k2 || TRUE == k2 || FALSE == k2))
    return 1;

  //b is of the form READ(Arr,const), and a is const, or
  //b is of the form var, and a is const
  if ((k1 == BVCONST)
      && ((k2 == READ
           && b[0].GetKind() == SYMBOL
           && b[1].GetKind() == BVCONST)))
    return -1;

  if (SYMBOL == k2
      && (BVCONST == k1
          || TRUE == k1
          || FALSE == k1))
    return -1;

  return 0;
} //End of TermOrder()

//This finds everything which is (= SYMBOL BVCONST) and everything that is (READ SYMBOL BVCONST).
//i.e. anything that has a termorder of 1 or -1.
// The bvsolver puts expressions of the form (= SYMBOL TERM) into the solverMap.

ASTNode SubstitutionMap::CreateSubstitutionMap(const ASTNode& a,  ArrayTransformer*at
)
{
  if (!bm->UserFlags.wordlevel_solve_flag)
    return a;

  ASTNode output = a;
  //if the variable has been solved for, then simply return it
  if (CheckSubstitutionMap(a, output))
    return output;

  //traverse a and populate the SubstitutionMap
  const Kind k = a.GetKind();
  if (SYMBOL == k && BOOLEAN_TYPE == a.GetType())
    {
      bool updated = UpdateSubstitutionMap(a, ASTTrue);
      output = updated ? ASTTrue : a;
      return output;
    }
  if (NOT == k && SYMBOL == a[0].GetKind())
    {
      bool updated = UpdateSubstitutionMap(a[0], ASTFalse);
      output = updated ? ASTTrue : a;
      return output;
    }

  if (IFF == k || EQ == k)
    {
      ASTVec c = a.GetChildren();
      SortByArith(c);

      ASTNode c1;
      if (EQ == k)
    	  c1 = simplify_during_create_subM ? simp->SimplifyTerm(c[1]) : c[1];
      else// (IFF == k )
    	  c1= simplify_during_create_subM ?  simp->SimplifyFormula_NoRemoveWrites(c[1], false) : c[1];

      //fill the arrayname readindices vector if e0 is a
      //READ(Arr,index) and index is a BVCONST
      int to;
      if ((to =TermOrder(c[0],c1))==1 && c[0].GetKind() == READ)
    	  at->FillUp_ArrReadIndex_Vec(c[0], c1);
      else if (to ==-1 && c1.GetKind() == READ)
		at->FillUp_ArrReadIndex_Vec(c1,c[0]);

      bool updated = UpdateSubstitutionMap(c[0], c1);
      output = updated ? ASTTrue : a;
      return output;
    }

  if (AND == k)
    {
      ASTVec o;
      ASTVec c = a.GetChildren();
      for (ASTVec::iterator
             it = c.begin(), itend = c.end();
           it != itend; it++)
        {
          simp->UpdateAlwaysTrueFormMap(*it);
          ASTNode aaa = CreateSubstitutionMap(*it,at);

          if (ASTTrue != aaa)
            {
              if (ASTFalse == aaa)
                return ASTFalse;
              else
                o.push_back(aaa);
            }
        }
      if (o.size() == 0)
        return ASTTrue;

      if (o.size() == 1)
        return o[0];

      return bm->CreateNode(AND, o);
    }

  //printf("I gave up on kind: %d node: %d\n", k, a.GetNodeNum());
  return output;
} //end of CreateSubstitutionMap()

};
<|endoftext|>


<commit_before>#include "ostree_http_repo.h"

#include <fcntl.h>

#include <iomanip>
#include <iostream>
#include <sstream>
#include <string>

#include <boost/filesystem.hpp>
#include <boost/property_tree/ini_parser.hpp>

#include "logging/logging.h"

using std::string;
namespace pt = boost::property_tree;

bool OSTreeHttpRepo::LooksValid() const {
  if (Get("config")) {
    pt::ptree config;
    try {
      pt::read_ini((root_ / "config").string(), config);
      if (config.get<string>("core.mode") != "archive-z2") {
        LOG_WARNING << "OSTree repo is not in archive-z2 format";
        return false;
      }
      return true;

    } catch (const pt::ini_parser_error &error) {
      LOG_WARNING << "Couldn't parse OSTree config file: " << (root_ / "config").string();
      return false;
    } catch (const pt::ptree_error &error) {
      LOG_WARNING << "Could not find core.mode in OSTree config file";
      return false;
    }
  } else {
    return false;
  }
}

OSTreeRef OSTreeHttpRepo::GetRef(const std::string &refname) const { return OSTreeRef(*server_, refname); }

OSTreeObject::ptr OSTreeHttpRepo::GetObject(const uint8_t sha256[32]) const { return GetObject(OSTreeHash(sha256)); }

bool OSTreeHttpRepo::Get(const boost::filesystem::path &path) const {
  CURLcode err = CURLE_OK;
  CurlEasyWrapper easy_handle;
  curl_easy_setopt(easy_handle.get(), CURLOPT_VERBOSE, get_curlopt_verbose());
  server_->InjectIntoCurl(path.string(), easy_handle.get());
  curl_easy_setopt(easy_handle.get(), CURLOPT_WRITEFUNCTION, &OSTreeHttpRepo::curl_handle_write);
  boost::filesystem::create_directories((root_ / path).parent_path());
  std::string filename = (root_ / path).string();
  int fp = open(filename.c_str(), O_WRONLY | O_CREAT, S_IWUSR | S_IRUSR | S_IRGRP | S_IROTH);
  if (fp == -1) {
    LOG_ERROR << "Failed to open file: " << filename;
    return false;
  }
  curl_easy_setopt(easy_handle.get(), CURLOPT_WRITEDATA, &fp);
  curl_easy_setopt(easy_handle.get(), CURLOPT_FAILONERROR, true);
  err = curl_easy_perform(easy_handle.get());
  close(fp);

  if (err == CURLE_HTTP_RETURNED_ERROR) {
    // http error (error code >= 400)
    // verbose mode will display the details
    return true;
  } else if (err != CURLE_OK) {
    // other unexpected error
    char *last_url = nullptr;
    curl_easy_getinfo(easy_handle.get(), CURLINFO_EFFECTIVE_URL, &last_url);
    LOG_ERROR << "Failed to get object:" << curl_easy_strerror(err);
    if (last_url != nullptr) {
      LOG_ERROR << "Url: " << last_url;
    }
    remove((root_ / path).c_str());
    return false;
  }

  return true;
}

OSTreeObject::ptr OSTreeHttpRepo::GetObject(const OSTreeHash hash) const {
  otable::const_iterator it;
  it = ObjectTable.find(hash);
  if (it != ObjectTable.end()) {
    return it->second;
  }

  const std::string exts[] = {".filez", ".dirtree", ".dirmeta", ".commit"};
  const std::string objpath = hash.string().insert(2, 1, '/');

  for (int i = 0; i < 3; ++i) {
    if (i > 0) {
      LOG_WARNING << "OSTree hash " << hash << " not found. Retrying (attempt " << i << " of 3)";
    }
    for (const std::string &ext : exts) {
      // NOLINTNEXTLINE(performance-inefficient-string-concatenation)
      if (Get(std::string("objects/") + objpath + ext)) {
        OSTreeObject::ptr obj(new OSTreeObject(*this, objpath + ext));
        ObjectTable[hash] = obj;
        return obj;
      }
    }
  }
  throw OSTreeObjectMissing(hash);
}

size_t OSTreeHttpRepo::curl_handle_write(void *buffer, size_t size, size_t nmemb, void *userp) {
  return static_cast<size_t>(write(*static_cast<int *>(userp), buffer, nmemb * size));
}

// vim: set tabstop=2 shiftwidth=2 expandtab:
<commit_msg>Fix return value from HTTP 400s.<commit_after>#include "ostree_http_repo.h"

#include <fcntl.h>

#include <iomanip>
#include <iostream>
#include <sstream>
#include <string>

#include <boost/filesystem.hpp>
#include <boost/property_tree/ini_parser.hpp>

#include "logging/logging.h"

using std::string;
namespace pt = boost::property_tree;

bool OSTreeHttpRepo::LooksValid() const {
  if (Get("config")) {
    pt::ptree config;
    try {
      pt::read_ini((root_ / "config").string(), config);
      if (config.get<string>("core.mode") != "archive-z2") {
        LOG_WARNING << "OSTree repo is not in archive-z2 format";
        return false;
      }
      return true;

    } catch (const pt::ini_parser_error &error) {
      LOG_WARNING << "Couldn't parse OSTree config file: " << (root_ / "config").string();
      return false;
    } catch (const pt::ptree_error &error) {
      LOG_WARNING << "Could not find core.mode in OSTree config file";
      return false;
    }
  } else {
    return false;
  }
}

OSTreeRef OSTreeHttpRepo::GetRef(const std::string &refname) const { return OSTreeRef(*server_, refname); }

OSTreeObject::ptr OSTreeHttpRepo::GetObject(const uint8_t sha256[32]) const { return GetObject(OSTreeHash(sha256)); }

bool OSTreeHttpRepo::Get(const boost::filesystem::path &path) const {
  CURLcode err = CURLE_OK;
  CurlEasyWrapper easy_handle;
  curl_easy_setopt(easy_handle.get(), CURLOPT_VERBOSE, get_curlopt_verbose());
  server_->InjectIntoCurl(path.string(), easy_handle.get());
  curl_easy_setopt(easy_handle.get(), CURLOPT_WRITEFUNCTION, &OSTreeHttpRepo::curl_handle_write);
  boost::filesystem::create_directories((root_ / path).parent_path());
  std::string filename = (root_ / path).string();
  int fp = open(filename.c_str(), O_WRONLY | O_CREAT, S_IWUSR | S_IRUSR | S_IRGRP | S_IROTH);
  if (fp == -1) {
    LOG_ERROR << "Failed to open file: " << filename;
    return false;
  }
  curl_easy_setopt(easy_handle.get(), CURLOPT_WRITEDATA, &fp);
  curl_easy_setopt(easy_handle.get(), CURLOPT_FAILONERROR, true);
  err = curl_easy_perform(easy_handle.get());
  close(fp);

  if (err == CURLE_HTTP_RETURNED_ERROR) {
    // http error (error code >= 400)
    // verbose mode will display the details
    return false;
  } else if (err != CURLE_OK) {
    // other unexpected error
    char *last_url = nullptr;
    curl_easy_getinfo(easy_handle.get(), CURLINFO_EFFECTIVE_URL, &last_url);
    LOG_ERROR << "Failed to get object:" << curl_easy_strerror(err);
    if (last_url != nullptr) {
      LOG_ERROR << "Url: " << last_url;
    }
    remove((root_ / path).c_str());
    return false;
  }

  return true;
}

OSTreeObject::ptr OSTreeHttpRepo::GetObject(const OSTreeHash hash) const {
  otable::const_iterator it;
  it = ObjectTable.find(hash);
  if (it != ObjectTable.end()) {
    return it->second;
  }

  const std::string exts[] = {".filez", ".dirtree", ".dirmeta", ".commit"};
  const std::string objpath = hash.string().insert(2, 1, '/');

  for (int i = 0; i < 3; ++i) {
    if (i > 0) {
      LOG_WARNING << "OSTree hash " << hash << " not found. Retrying (attempt " << i << " of 3)";
    }
    for (const std::string &ext : exts) {
      // NOLINTNEXTLINE(performance-inefficient-string-concatenation)
      if (Get(std::string("objects/") + objpath + ext)) {
        OSTreeObject::ptr obj(new OSTreeObject(*this, objpath + ext));
        ObjectTable[hash] = obj;
        return obj;
      }
    }
  }
  throw OSTreeObjectMissing(hash);
}

size_t OSTreeHttpRepo::curl_handle_write(void *buffer, size_t size, size_t nmemb, void *userp) {
  return static_cast<size_t>(write(*static_cast<int *>(userp), buffer, nmemb * size));
}

// vim: set tabstop=2 shiftwidth=2 expandtab:
<|endoftext|>


<commit_before>
#ifndef __ARCH_MOCK_IO_HPP__
#define __ARCH_MOCK_IO_HPP__

#include "containers/segmented_vector.hpp"
#include "utils2.hpp"
#include <stdlib.h>

struct mock_iocallback_t {
    virtual void on_io_complete(event_t *event) = 0;
    virtual ~mock_iocallback_t() {}
};

template<class inner_io_config_t>
class mock_direct_file_t
{
public:
    enum mode_t {
        mode_read = 1 << 0,
        mode_write = 1 << 1,
        mode_create = 1 << 2
    };

    mock_direct_file_t(const char *path, int mode) {
        int mode2 = 0;
        if (mode & mode_read) mode2 |= inner_io_config_t::direct_file_t::mode_read;
        // We always enable writing because the mock layer does
        // truncation on exit
        mode2 |= inner_io_config_t::direct_file_t::mode_write;
        if (mode & mode_create) mode2 |= inner_io_config_t::direct_file_t::mode_create;
        inner_file = new typename inner_io_config_t::direct_file_t(path, mode2);

        if (inner_file->is_block_device()) {
            fail_due_to_user_error(
                "Using mock_direct_file_t with a block device is a really bad idea because it "
                "reads the entire contents of the underlying file into memory, which could be "
                "a lot for a block device.");
        }

        set_size(inner_file->get_size());
        for (unsigned i = 0; i < get_size() / DEVICE_BLOCK_SIZE; i++) {
            inner_file->read_blocking(i*DEVICE_BLOCK_SIZE, DEVICE_BLOCK_SIZE, blocks[i].data);
        }
    }

    bool exists() {
        return inner_file->exists();
    }

    bool is_block_device() {
        return false;
    }

    uint64_t get_size() {
        return blocks.get_size() * DEVICE_BLOCK_SIZE;
    }

    void set_size(size_t size) {
        assert(size % DEVICE_BLOCK_SIZE == 0);
        blocks.set_size(size / DEVICE_BLOCK_SIZE);
    }

    void set_size_at_least(size_t size) {
        if (get_size() < size) {
            size_t actual_size = size + randint(10) * DEVICE_BLOCK_SIZE;
            set_size(actual_size);
        }
    }

    /* These always return 'false'; the reason they return bool instead of void
    is for consistency with other asynchronous-callback methods */
    bool read_async(size_t offset, size_t length, void *buf, mock_iocallback_t *cb) {
        read_blocking(offset, length, buf);
        random_delay(cb, &mock_iocallback_t::on_io_complete, (event_t*)NULL);
        return false;
    }

    bool write_async(size_t offset, size_t length, void *buf, mock_iocallback_t *cb) {
        write_blocking(offset, length, buf);
        random_delay(cb, &mock_iocallback_t::on_io_complete, (event_t*)NULL);
        return false;
    }

    void read_blocking(size_t offset, size_t length, void *buf) {
        verify(offset, length, buf);
        for (unsigned i = 0; i < length / DEVICE_BLOCK_SIZE; i += 1) {
            memcpy((char*)buf + i*DEVICE_BLOCK_SIZE, blocks[offset/DEVICE_BLOCK_SIZE + i].data, DEVICE_BLOCK_SIZE);
        }
    }

    void write_blocking(size_t offset, size_t length, void *buf) {
        verify(offset, length, buf);
        for (unsigned i = 0; i < length / DEVICE_BLOCK_SIZE; i += 1) {
            memcpy(blocks[offset/DEVICE_BLOCK_SIZE + i].data, (char*)buf + i*DEVICE_BLOCK_SIZE, DEVICE_BLOCK_SIZE);
        }
    }

    ~mock_direct_file_t() {
        if(exists()) {
            inner_file->set_size(get_size());
        }
        for (unsigned i = 0; i < get_size() / DEVICE_BLOCK_SIZE; i++) {
            inner_file->write_blocking(i*DEVICE_BLOCK_SIZE, DEVICE_BLOCK_SIZE, blocks[i].data);
        }

        delete inner_file;
    }

private:
    typename inner_io_config_t::direct_file_t *inner_file;

    struct block_t {
        char *data;

        block_t() {
            data = (char*)malloc_aligned(DEVICE_BLOCK_SIZE, DEVICE_BLOCK_SIZE);

            // Initialize to either random data or zeroes, choosing at random
            char d = randint(2) ? 0 : randint(0x100);
            memset(data, d, DEVICE_BLOCK_SIZE);
        }

        ~block_t() {
            free((void*)data);
        }
    };
    segmented_vector_t<block_t, 10*GIGABYTE/DEVICE_BLOCK_SIZE> blocks;

    void verify(size_t offset, size_t length, void *buf) {
        assert(buf);
        assert(offset + length <= get_size());
        assert((intptr_t)buf % DEVICE_BLOCK_SIZE == 0);
        assert(offset % DEVICE_BLOCK_SIZE == 0);
        assert(length % DEVICE_BLOCK_SIZE == 0);
    }
};

#endif /* __ARCH_MOCK_IO_HPP__ */
<commit_msg>Hopefully made mock IO a bit faster under fsck.<commit_after>
#ifndef __ARCH_MOCK_IO_HPP__
#define __ARCH_MOCK_IO_HPP__

#include "containers/segmented_vector.hpp"
#include "utils2.hpp"
#include <stdlib.h>

struct mock_iocallback_t {
    virtual void on_io_complete(event_t *event) = 0;
    virtual ~mock_iocallback_t() {}
};

template<class inner_io_config_t>
class mock_direct_file_t
{
    int mode;

public:
    enum mode_t {
        mode_read = 1 << 0,
        mode_write = 1 << 1,
        mode_create = 1 << 2
    };

    mock_direct_file_t(const char *path, int mode)
        : mode(mode)
    {
        int mode2 = 0;
        if (mode & mode_read) mode2 |= inner_io_config_t::direct_file_t::mode_read;
        // We always enable writing because the mock layer does
        // truncation on exit
        mode2 |= inner_io_config_t::direct_file_t::mode_write;
        if (mode & mode_create) mode2 |= inner_io_config_t::direct_file_t::mode_create;
        inner_file = new typename inner_io_config_t::direct_file_t(path, mode2);

        if (inner_file->is_block_device()) {
            fail_due_to_user_error(
                "Using mock_direct_file_t with a block device is a really bad idea because it "
                "reads the entire contents of the underlying file into memory, which could be "
                "a lot for a block device.");
        }

        set_size(inner_file->get_size());
        for (unsigned i = 0; i < get_size() / DEVICE_BLOCK_SIZE; i++) {
            inner_file->read_blocking(i*DEVICE_BLOCK_SIZE, DEVICE_BLOCK_SIZE, blocks[i].data);
        }
    }

    bool exists() {
        return inner_file->exists();
    }

    bool is_block_device() {
        return false;
    }

    uint64_t get_size() {
        return blocks.get_size() * DEVICE_BLOCK_SIZE;
    }

    void set_size(size_t size) {
        assert(size % DEVICE_BLOCK_SIZE == 0);
        blocks.set_size(size / DEVICE_BLOCK_SIZE);
    }

    void set_size_at_least(size_t size) {
        if (get_size() < size) {
            size_t actual_size = size + randint(10) * DEVICE_BLOCK_SIZE;
            set_size(actual_size);
        }
    }

    /* These always return 'false'; the reason they return bool instead of void
    is for consistency with other asynchronous-callback methods */
    bool read_async(size_t offset, size_t length, void *buf, mock_iocallback_t *cb) {
        assert(mode & mode_read);
        read_blocking(offset, length, buf);
        random_delay(cb, &mock_iocallback_t::on_io_complete, (event_t*)NULL);
        return false;
    }

    bool write_async(size_t offset, size_t length, void *buf, mock_iocallback_t *cb) {
        assert(mode & mode_write);
        write_blocking(offset, length, buf);
        random_delay(cb, &mock_iocallback_t::on_io_complete, (event_t*)NULL);
        return false;
    }

    void read_blocking(size_t offset, size_t length, void *buf) {
        assert(mode & mode_read);
        verify(offset, length, buf);
        for (unsigned i = 0; i < length / DEVICE_BLOCK_SIZE; i += 1) {
            memcpy((char*)buf + i*DEVICE_BLOCK_SIZE, blocks[offset/DEVICE_BLOCK_SIZE + i].data, DEVICE_BLOCK_SIZE);
        }
    }

    void write_blocking(size_t offset, size_t length, void *buf) {
        assert(mode & mode_write);
        verify(offset, length, buf);
        for (unsigned i = 0; i < length / DEVICE_BLOCK_SIZE; i += 1) {
            memcpy(blocks[offset/DEVICE_BLOCK_SIZE + i].data, (char*)buf + i*DEVICE_BLOCK_SIZE, DEVICE_BLOCK_SIZE);
        }
    }

    ~mock_direct_file_t() {
        if (mode & mode_write) {
            if(exists()) {
                inner_file->set_size(get_size());
            }
            // This is horrible. We write every single block of the file, sequentially, in a
            // blocking fashion. There must be a better way to do this...
            for (unsigned i = 0; i < get_size() / DEVICE_BLOCK_SIZE; i++) {
                inner_file->write_blocking(i*DEVICE_BLOCK_SIZE, DEVICE_BLOCK_SIZE, blocks[i].data);
            }
        }

        delete inner_file;
    }

private:
    typename inner_io_config_t::direct_file_t *inner_file;

    struct block_t {
        char *data;

        block_t() {
            data = (char*)malloc_aligned(DEVICE_BLOCK_SIZE, DEVICE_BLOCK_SIZE);

            // Initialize to either random data or zeroes, choosing at random
            char d = randint(2) ? 0 : randint(0x100);
            memset(data, d, DEVICE_BLOCK_SIZE);
        }

        ~block_t() {
            free((void*)data);
        }
    };
    segmented_vector_t<block_t, 10*GIGABYTE/DEVICE_BLOCK_SIZE> blocks;

    void verify(size_t offset, size_t length, void *buf) {
        assert(buf);
        assert(offset + length <= get_size());
        assert((intptr_t)buf % DEVICE_BLOCK_SIZE == 0);
        assert(offset % DEVICE_BLOCK_SIZE == 0);
        assert(length % DEVICE_BLOCK_SIZE == 0);
    }
};

#endif /* __ARCH_MOCK_IO_HPP__ */
<|endoftext|>


<commit_before>#include "solver.h"

#include <boost/format.hpp>
#include <cmath>
#include <fstream>
#include <functional>
#include <string>
#include <unordered_set>
#include <utility>
#include "../injector.h"
#include "../omp_hash_map/src/omp_hash_map.h"
#include "../omp_hash_map/src/reducer.h"
#include "connections.h"
#include "davidson_util.h"

#define ENERGY_FORMAT "%.12f"

class SolverImpl : public Solver {
 public:
  ~SolverImpl();

  SolverImpl(
      Session* const session,
      Connections* const connections,
      AbstractSystem* const abstrct_system);

  void setup_hf() override;

  void variation(const double eps_var) override;

  void save_variation_result(const std::string& filename) override;

  bool load_variation_result(const std::string& filename) override;

  std::vector<double> apply_hamiltonian(
      const std::vector<double>& vec) override;

  void perturbation(const double eps_pt) override;

 private:
  int n_up = 0;

  int n_dn = 0;

  double energy_hf = 0.0;

  double energy_var = 0.0;

  double energy_pt = 0.0;

  bool verbose = false;

  Session* const session;

  const std::unique_ptr<Connections> connections;

  AbstractSystem* const abstract_system;

  std::ofstream pt_result;

  void print_var_result() const;
};

SolverImpl::SolverImpl(
    Session* const session,
    Connections* const connections,
    AbstractSystem* const abstract_system)
    : session(session),
      connections(connections),
      abstract_system(abstract_system) {
  verbose = session->get_parallel()->is_master();
  Config* const config = session->get_config();
  n_up = config->get_int("n_up");
  n_dn = config->get_int("n_dn");
  if (!config->get_bool("variation_only")) {
    pt_result.open("pt_result.csv", std::ios::out | std::ios::trunc);
    pt_result << "n_orbs_pt,eps_pt,n_pt_dets,energy_corr" << std::endl;
  }
}

SolverImpl::~SolverImpl() { pt_result.close(); }

void SolverImpl::setup_hf() {
  // Create or clear wavefunction.
  if (!abstract_system->wf) {
    abstract_system->wf.reset(new data::Wavefunction());
  }
  data::Wavefunction* const wf = abstract_system->wf.get();
  wf->Clear();

  // Add a single term with coef 1.0 and no diffs.
  data::Term* term = wf->add_terms();
  term->set_coef(1.0);
  data::Determinant* det_hf = term->mutable_det();
  data::SpinDeterminant* up = det_hf->mutable_up();
  up->set_n_hf_elecs(n_up);
  data::SpinDeterminant* dn = det_hf->mutable_dn();
  dn->set_n_hf_elecs(n_dn);

  // Update HF and variational energy.
  energy_hf = energy_var = abstract_system->hamiltonian(det_hf, det_hf);

  // Clear connections.
  connections->clear();

  if (verbose) printf("HF det setup.\n");
};

void SolverImpl::variation(const double eps_var) {
  Timer* const timer = session->get_timer();

  // Contruct variational determinant hash set.
  data::Wavefunction* const wf = abstract_system->wf.get();
  const int n_start_dets = wf->terms_size();
  std::unordered_set<std::string> var_dets_set;
  std::vector<double> prev_coefs;
  var_dets_set.reserve(n_start_dets);
  prev_coefs.resize(n_start_dets, 0.0);
  for (const auto& term : wf->terms()) {
    std::string det_code;
    term.det().SerializeToString(&det_code);
    var_dets_set.insert(std::move(det_code));
  }

  // Variation iterations.
  bool converged = false;
  int iteration = 0;
  int n_new_dets = 0;
  const auto& connected_det_handler =
      [&](const data::Determinant* const connected_det) {
        std::string det_code;
        connected_det->SerializeToString(&det_code);
        if (var_dets_set.count(det_code) == 0) {
          var_dets_set.insert(std::move(det_code));
          data::Term* const term = wf->add_terms();
          term->set_coef(0.0);
          auto connected_det_copy = new data::Determinant(*connected_det);
          term->set_allocated_det(connected_det_copy);
          n_new_dets++;
        }
      };
  const auto& apply_hamiltonian_func =
      std::bind(&SolverImpl::apply_hamiltonian, this, std::placeholders::_1);
  while (!converged) {
    timer->start("Variation: " + std::to_string(iteration));

    // Find new dets.
    n_new_dets = 0;
    const int n_old_dets = wf->terms_size();
    for (int i = 0; i < n_old_dets; i++) {
      const auto& term = wf->terms(i);
      const double coef = term.coef();
      if (std::abs(coef) <= std::abs(prev_coefs[i])) continue;
      abstract_system->find_connected_dets(
          &term.det(), eps_var / std::abs(term.coef()), connected_det_handler);
    }
    const int n_total_dets = wf->terms_size();
    if (verbose) {
      printf("New / total dets: %'d / %'d\n", n_new_dets, n_total_dets);
    }
    timer->checkpoint("found new dets");

    // Update prev coefs.
    prev_coefs.resize(n_total_dets);
    for (int i = 0; i < n_total_dets; i++) prev_coefs[i] = wf->terms(i).coef();

    // Diagonalize.
    connections->update();
    std::vector<double> diagonal(n_total_dets, 0.0);
    for (int i = 0; i < n_total_dets; i++) {
      const auto& det_i = wf->terms(i).det();
      diagonal[i] = abstract_system->hamiltonian(&det_i, &det_i);
    }
    const auto& diagonalization_result = DavidsonUtil::diagonalize(
        prev_coefs, diagonal, apply_hamiltonian_func, 10, verbose);
    const double energy_var_new = diagonalization_result.first;
    const auto& new_coefs = diagonalization_result.second;
    for (int i = 0; i < n_total_dets; i++) {
      wf->mutable_terms(i)->set_coef(new_coefs[i]);
    }

    // Determine convergence.
    if (std::abs(energy_var_new - energy_var) < 1.0e-6) {
      converged = true;
    }

    energy_var = energy_var_new;
    iteration++;

    timer->end();  // iteration.
  }
};

void SolverImpl::save_variation_result(const std::string& filename) {
  std::fstream var_file(
      filename, std::ios::out | std::ios::trunc | std::ios::binary);
  data::VariationResult res;
  res.set_energy_hf(energy_hf);
  res.set_energy_var(energy_var);
  res.set_allocated_wf(abstract_system->wf.release());
  res.SerializeToOstream(&var_file);
  var_file.close();
  abstract_system->wf.reset(res.release_wf());
  if (verbose) {
    print_var_result();
    printf("Saved to: %s\n", filename.c_str());
  }
}

bool SolverImpl::load_variation_result(const std::string& filename) {
  std::fstream var_file(filename, std::ios::in | std::ios::binary);
  data::VariationResult res;
  if (!res.ParseFromIstream(&var_file)) {
    return false;
  }
  var_file.close();
  energy_hf = res.energy_hf();
  energy_var = res.energy_var();
  abstract_system->wf.reset(res.release_wf());
  if (verbose) {
    printf("Loaded from: %s\n", filename.c_str());
    print_var_result();
  }
  return true;
}

void SolverImpl::print_var_result() const {
  printf("Number of dets: %'d\n", abstract_system->wf->terms_size());
  printf("Variation energy: " ENERGY_FORMAT " Ha\n", energy_var);
  const double energy_corr = energy_var - energy_hf;
  printf("Correlation energy (var): " ENERGY_FORMAT " Ha\n", energy_corr);
}

std::vector<double> SolverImpl::apply_hamiltonian(
    const std::vector<double>& vec) {
  const int n_dets = vec.size();
  Parallel* const parallel = session->get_parallel();
  const int proc_id = parallel->get_proc_id();
  const int n_procs = parallel->get_n_procs();
  std::vector<double> res(n_dets, 0.0);

#pragma omp parallel for reduction(vec_double_plus : res) schedule(dynamic, 10)
  for (int i = proc_id; i < n_dets; i += n_procs) {
    const auto& conns = connections->get_connections(i);
    for (const auto conn : conns) {
      const int j = conn.first;
      const double H_ij = conn.second;
      res[i] += H_ij * vec[j];
      if (i != j) {
        res[j] += H_ij * vec[i];
      }
    }
  }

  parallel->reduce_to_sum(res);
  session->get_timer()->checkpoint("hamiltonian applied");

  return res;
};

void SolverImpl::perturbation(const double eps_pt) {
  // Clean variation variables.
  connections->clear();

  // Construct var dets hash set.
  std::unordered_set<std::string> var_dets_set;
  var_dets_set.reserve(abstract_system->wf->terms_size());
  for (const auto& term : abstract_system->wf->terms()) {
    var_dets_set.insert(term.det().SerializeAsString());
  }

  // Setup partial sum hash map.
  omp_hash_map<std::string, double> partial_sums;

  // Search pt dets.
  const int n_var_dets = var_dets_set.size();
  Parallel* const parallel = session->get_parallel();
  Timer* const timer = session->get_timer();
  const size_t n_procs = parallel->get_n_procs();
  const size_t proc_id = parallel->get_proc_id();
  std::hash<std::string> string_hasher;
  double target_progress = 0.25;
  timer->start("search");
#pragma omp parallel for schedule(static, 1)
  for (int i = 0; i < n_var_dets; i++) {
    const auto& term = abstract_system->wf->terms(i);
    const auto& var_det = term.det();
    const double coef = term.coef();

    const auto& pt_det_handler = [&](const auto& det_a) {
      const auto& det_a_code = det_a->SerializeAsString();
      const size_t det_a_hash = string_hasher(det_a_code);
      if (var_dets_set.count(det_a_code) == 1) return;
      if (det_a_hash % n_procs != proc_id) return;
      partial_sums.set(
          det_a_code,
          [&](double& value) {
            const double H_ai = abstract_system->hamiltonian(&var_det, det_a);
            value += H_ai * coef;
          },
          0.0);
    };

    abstract_system->find_connected_dets(
        &var_det, eps_pt / std::abs(coef), pt_det_handler);

    const double current_progress = i * 100.0 / n_var_dets;
    if (target_progress <= current_progress) {
      const int thread_id = omp_get_thread_num();
      if (thread_id == 0) {
        std::string event =
            str(boost::format("Progress: %.2f %%") % target_progress);
        timer->checkpoint(event);
        if (verbose) {
          printf("Master node PT dets: %'zu\n", partial_sums.get_n_keys());
        }
        target_progress *= 2.0;
      }
    }
  }
  timer->end();

  // Accumulate partial results.
  timer->start("accumulation");
  unsigned long long n_pt_dets = 0;
  std::vector<data::Determinant> tmp_dets(parallel->get_n_threads());
  energy_pt = partial_sums.map_reduce<double>(
      [&](const std::string& det_code, const double value) {
        const int thread_id = omp_get_thread_num();
        auto& det = tmp_dets[thread_id];
        det.ParseFromString(det_code);
        const double H_aa = abstract_system->hamiltonian(&det, &det);
#pragma omp atomic
        n_pt_dets++;
        return value * value / (energy_var - H_aa);
      },
      reducer::sum<double>,
      0.0);
  parallel->reduce_to_sum(energy_pt);

  if (verbose) {
    const int n_pt_orbs = abstract_system->get_n_orbitals();
    printf("PT orbitals: %'d\n", n_pt_orbs);
    printf("PT eps: %#.4g\n", eps_pt);
    printf("Perturbation energy: %#.15g Ha\n", energy_pt);
    const double energy_corr = energy_pt + energy_var - energy_hf;
    printf("Correlation energy (pt): %#.15g Ha\n", energy_corr);
    printf("Number of perturbation dets: %'llu\n", n_pt_dets);
    pt_result << str(boost::format("%d, %#.4g, %llu, %.12f") % n_pt_orbs %
                     eps_pt % n_pt_dets % energy_corr)
              << std::endl;
  }
  timer->end();
};

Solver* Injector::new_solver(
    Session* const session,
    Connections* const connections,
    AbstractSystem* const abstract_system) {
  return new SolverImpl(session, connections, abstract_system);
}
<commit_msg>reduce and get total number of pt dets (#20)<commit_after>#include "solver.h"

#include <boost/format.hpp>
#include <cmath>
#include <fstream>
#include <functional>
#include <string>
#include <unordered_set>
#include <utility>
#include "../injector.h"
#include "../omp_hash_map/src/omp_hash_map.h"
#include "../omp_hash_map/src/reducer.h"
#include "connections.h"
#include "davidson_util.h"

#define ENERGY_FORMAT "%.12f"

class SolverImpl : public Solver {
 public:
  ~SolverImpl();

  SolverImpl(
      Session* const session,
      Connections* const connections,
      AbstractSystem* const abstrct_system);

  void setup_hf() override;

  void variation(const double eps_var) override;

  void save_variation_result(const std::string& filename) override;

  bool load_variation_result(const std::string& filename) override;

  std::vector<double> apply_hamiltonian(
      const std::vector<double>& vec) override;

  void perturbation(const double eps_pt) override;

 private:
  int n_up = 0;

  int n_dn = 0;

  double energy_hf = 0.0;

  double energy_var = 0.0;

  double energy_pt = 0.0;

  bool verbose = false;

  Session* const session;

  const std::unique_ptr<Connections> connections;

  AbstractSystem* const abstract_system;

  std::ofstream pt_result;

  void print_var_result() const;
};

SolverImpl::SolverImpl(
    Session* const session,
    Connections* const connections,
    AbstractSystem* const abstract_system)
    : session(session),
      connections(connections),
      abstract_system(abstract_system) {
  verbose = session->get_parallel()->is_master();
  Config* const config = session->get_config();
  n_up = config->get_int("n_up");
  n_dn = config->get_int("n_dn");
  if (!config->get_bool("variation_only")) {
    pt_result.open("pt_result.csv", std::ios::out | std::ios::trunc);
    pt_result << "n_orbs_pt,eps_pt,n_pt_dets,energy_corr" << std::endl;
  }
}

SolverImpl::~SolverImpl() { pt_result.close(); }

void SolverImpl::setup_hf() {
  // Create or clear wavefunction.
  if (!abstract_system->wf) {
    abstract_system->wf.reset(new data::Wavefunction());
  }
  data::Wavefunction* const wf = abstract_system->wf.get();
  wf->Clear();

  // Add a single term with coef 1.0 and no diffs.
  data::Term* term = wf->add_terms();
  term->set_coef(1.0);
  data::Determinant* det_hf = term->mutable_det();
  data::SpinDeterminant* up = det_hf->mutable_up();
  up->set_n_hf_elecs(n_up);
  data::SpinDeterminant* dn = det_hf->mutable_dn();
  dn->set_n_hf_elecs(n_dn);

  // Update HF and variational energy.
  energy_hf = energy_var = abstract_system->hamiltonian(det_hf, det_hf);

  // Clear connections.
  connections->clear();

  if (verbose) printf("HF det setup.\n");
};

void SolverImpl::variation(const double eps_var) {
  Timer* const timer = session->get_timer();

  // Contruct variational determinant hash set.
  data::Wavefunction* const wf = abstract_system->wf.get();
  const int n_start_dets = wf->terms_size();
  std::unordered_set<std::string> var_dets_set;
  std::vector<double> prev_coefs;
  var_dets_set.reserve(n_start_dets);
  prev_coefs.resize(n_start_dets, 0.0);
  for (const auto& term : wf->terms()) {
    std::string det_code;
    term.det().SerializeToString(&det_code);
    var_dets_set.insert(std::move(det_code));
  }

  // Variation iterations.
  bool converged = false;
  int iteration = 0;
  int n_new_dets = 0;
  const auto& connected_det_handler =
      [&](const data::Determinant* const connected_det) {
        std::string det_code;
        connected_det->SerializeToString(&det_code);
        if (var_dets_set.count(det_code) == 0) {
          var_dets_set.insert(std::move(det_code));
          data::Term* const term = wf->add_terms();
          term->set_coef(0.0);
          auto connected_det_copy = new data::Determinant(*connected_det);
          term->set_allocated_det(connected_det_copy);
          n_new_dets++;
        }
      };
  const auto& apply_hamiltonian_func =
      std::bind(&SolverImpl::apply_hamiltonian, this, std::placeholders::_1);
  while (!converged) {
    timer->start("Variation: " + std::to_string(iteration));

    // Find new dets.
    n_new_dets = 0;
    const int n_old_dets = wf->terms_size();
    for (int i = 0; i < n_old_dets; i++) {
      const auto& term = wf->terms(i);
      const double coef = term.coef();
      if (std::abs(coef) <= std::abs(prev_coefs[i])) continue;
      abstract_system->find_connected_dets(
          &term.det(), eps_var / std::abs(term.coef()), connected_det_handler);
    }
    const int n_total_dets = wf->terms_size();
    if (verbose) {
      printf("New / total dets: %'d / %'d\n", n_new_dets, n_total_dets);
    }
    timer->checkpoint("found new dets");

    // Update prev coefs.
    prev_coefs.resize(n_total_dets);
    for (int i = 0; i < n_total_dets; i++) prev_coefs[i] = wf->terms(i).coef();

    // Diagonalize.
    connections->update();
    std::vector<double> diagonal(n_total_dets, 0.0);
    for (int i = 0; i < n_total_dets; i++) {
      const auto& det_i = wf->terms(i).det();
      diagonal[i] = abstract_system->hamiltonian(&det_i, &det_i);
    }
    const auto& diagonalization_result = DavidsonUtil::diagonalize(
        prev_coefs, diagonal, apply_hamiltonian_func, 10, verbose);
    const double energy_var_new = diagonalization_result.first;
    const auto& new_coefs = diagonalization_result.second;
    for (int i = 0; i < n_total_dets; i++) {
      wf->mutable_terms(i)->set_coef(new_coefs[i]);
    }

    // Determine convergence.
    if (std::abs(energy_var_new - energy_var) < 1.0e-6) {
      converged = true;
    }

    energy_var = energy_var_new;
    iteration++;

    timer->end();  // iteration.
  }
};

void SolverImpl::save_variation_result(const std::string& filename) {
  std::fstream var_file(
      filename, std::ios::out | std::ios::trunc | std::ios::binary);
  data::VariationResult res;
  res.set_energy_hf(energy_hf);
  res.set_energy_var(energy_var);
  res.set_allocated_wf(abstract_system->wf.release());
  res.SerializeToOstream(&var_file);
  var_file.close();
  abstract_system->wf.reset(res.release_wf());
  if (verbose) {
    print_var_result();
    printf("Saved to: %s\n", filename.c_str());
  }
}

bool SolverImpl::load_variation_result(const std::string& filename) {
  std::fstream var_file(filename, std::ios::in | std::ios::binary);
  data::VariationResult res;
  if (!res.ParseFromIstream(&var_file)) {
    return false;
  }
  var_file.close();
  energy_hf = res.energy_hf();
  energy_var = res.energy_var();
  abstract_system->wf.reset(res.release_wf());
  if (verbose) {
    printf("Loaded from: %s\n", filename.c_str());
    print_var_result();
  }
  return true;
}

void SolverImpl::print_var_result() const {
  printf("Number of dets: %'d\n", abstract_system->wf->terms_size());
  printf("Variation energy: " ENERGY_FORMAT " Ha\n", energy_var);
  const double energy_corr = energy_var - energy_hf;
  printf("Correlation energy (var): " ENERGY_FORMAT " Ha\n", energy_corr);
}

std::vector<double> SolverImpl::apply_hamiltonian(
    const std::vector<double>& vec) {
  const int n_dets = vec.size();
  Parallel* const parallel = session->get_parallel();
  const int proc_id = parallel->get_proc_id();
  const int n_procs = parallel->get_n_procs();
  std::vector<double> res(n_dets, 0.0);

#pragma omp parallel for reduction(vec_double_plus : res) schedule(dynamic, 10)
  for (int i = proc_id; i < n_dets; i += n_procs) {
    const auto& conns = connections->get_connections(i);
    for (const auto conn : conns) {
      const int j = conn.first;
      const double H_ij = conn.second;
      res[i] += H_ij * vec[j];
      if (i != j) {
        res[j] += H_ij * vec[i];
      }
    }
  }

  parallel->reduce_to_sum(res);
  session->get_timer()->checkpoint("hamiltonian applied");

  return res;
};

void SolverImpl::perturbation(const double eps_pt) {
  // Clean variation variables.
  connections->clear();

  // Construct var dets hash set.
  std::unordered_set<std::string> var_dets_set;
  var_dets_set.reserve(abstract_system->wf->terms_size());
  for (const auto& term : abstract_system->wf->terms()) {
    var_dets_set.insert(term.det().SerializeAsString());
  }

  // Setup partial sum hash map.
  omp_hash_map<std::string, double> partial_sums;

  // Search pt dets.
  const int n_var_dets = var_dets_set.size();
  Parallel* const parallel = session->get_parallel();
  Timer* const timer = session->get_timer();
  const size_t n_procs = parallel->get_n_procs();
  const size_t proc_id = parallel->get_proc_id();
  std::hash<std::string> string_hasher;
  double target_progress = 0.25;
  timer->start("search");
#pragma omp parallel for schedule(static, 1)
  for (int i = 0; i < n_var_dets; i++) {
    const auto& term = abstract_system->wf->terms(i);
    const auto& var_det = term.det();
    const double coef = term.coef();

    const auto& pt_det_handler = [&](const auto& det_a) {
      const auto& det_a_code = det_a->SerializeAsString();
      const size_t det_a_hash = string_hasher(det_a_code);
      if (var_dets_set.count(det_a_code) == 1) return;
      if (det_a_hash % n_procs != proc_id) return;
      partial_sums.set(
          det_a_code,
          [&](double& value) {
            const double H_ai = abstract_system->hamiltonian(&var_det, det_a);
            value += H_ai * coef;
          },
          0.0);
    };

    abstract_system->find_connected_dets(
        &var_det, eps_pt / std::abs(coef), pt_det_handler);

    const double current_progress = i * 100.0 / n_var_dets;
    if (target_progress <= current_progress) {
      const int thread_id = omp_get_thread_num();
      if (thread_id == 0) {
        std::string event =
            str(boost::format("Progress: %.2f %%") % target_progress);
        timer->checkpoint(event);
        if (verbose) {
          printf("Master node PT dets: %'zu\n", partial_sums.get_n_keys());
        }
        target_progress *= 2.0;
      }
    }
  }
  timer->end();

  // Accumulate partial results.
  timer->start("accumulation");
  unsigned long long n_pt_dets = 0;
  std::vector<data::Determinant> tmp_dets(parallel->get_n_threads());
  energy_pt = partial_sums.map_reduce<double>(
      [&](const std::string& det_code, const double value) {
        const int thread_id = omp_get_thread_num();
        auto& det = tmp_dets[thread_id];
        det.ParseFromString(det_code);
        const double H_aa = abstract_system->hamiltonian(&det, &det);
#pragma omp atomic
        n_pt_dets++;
        return value * value / (energy_var - H_aa);
      },
      reducer::sum<double>,
      0.0);
  parallel->reduce_to_sum(energy_pt);
  parallel->reduce_to_sum(n_pt_dets);

  if (verbose) {
    const int n_pt_orbs = abstract_system->get_n_orbitals();
    printf("PT orbitals: %'d\n", n_pt_orbs);
    printf("PT eps: %#.4g\n", eps_pt);
    printf("Perturbation energy: %#.15g Ha\n", energy_pt);
    const double energy_corr = energy_pt + energy_var - energy_hf;
    printf("Correlation energy (pt): %#.15g Ha\n", energy_corr);
    printf("Number of perturbation dets: %'llu\n", n_pt_dets);
    pt_result << str(boost::format("%d, %#.4g, %llu, %.12f") % n_pt_orbs %
                     eps_pt % n_pt_dets % energy_corr)
              << std::endl;
  }
  timer->end();
};

Solver* Injector::new_solver(
    Session* const session,
    Connections* const connections,
    AbstractSystem* const abstract_system) {
  return new SolverImpl(session, connections, abstract_system);
}
<|endoftext|>


<commit_before>/*

This file is part of VROOM.

Copyright (c) 2015-2019, Julien Coupey.
All rights reserved (see LICENSE).

*/

#include "structures/vroom/raw_route.h"

namespace vroom {

RawRoute::RawRoute(const Input& input, Index i)
  : fwd_peaks(2, Amount(input.amount_size())),
    bwd_peaks(2, Amount(input.amount_size())),
    vehicle_rank(i),
    has_start(input.vehicles[i].has_start()),
    has_end(input.vehicles[i].has_end()),
    capacity(input.vehicles[i].capacity) {
}

void RawRoute::set_route(const Input& input, const std::vector<Index>& r) {
  route = r;
  update_amounts(input);
}

bool RawRoute::empty() const {
  return route.empty();
}

std::size_t RawRoute::size() const {
  return route.size();
}

void RawRoute::update_amounts(const Input& input) {
  auto step_size = route.size() + 2;
  fwd_pickups.resize(route.size());
  bwd_deliveries.resize(route.size());
  current_loads.resize(step_size);
  fwd_peaks.resize(step_size);
  bwd_peaks.resize(step_size);

  if (route.empty()) {
    // So that check in is_valid_addition_for_capacity is consistent
    // with empty routes.
    std::fill(fwd_peaks.begin(), fwd_peaks.end(), Amount(input.amount_size()));
    std::fill(bwd_peaks.begin(), bwd_peaks.end(), Amount(input.amount_size()));
    return;
  }

  Amount current_pickups(input.amount_size());

  for (std::size_t i = 0; i < route.size(); ++i) {
    current_pickups += input.jobs[route[i]].pickup;
    fwd_pickups[i] = current_pickups;
  }

  Amount current_deliveries(input.amount_size());

  current_loads.back() = fwd_pickups.back();

  for (std::size_t i = 0; i < route.size(); ++i) {
    auto bwd_i = route.size() - i - 1;

    bwd_deliveries[bwd_i] = current_deliveries;
    current_loads[bwd_i + 1] = fwd_pickups[bwd_i] + current_deliveries;
    current_deliveries += input.jobs[route[bwd_i]].delivery;
  }
  current_loads[0] = current_deliveries;

  auto peak = current_loads[0];
  fwd_peaks[0] = peak;
  for (std::size_t s = 1; s < fwd_peaks.size(); ++s) {
    // Handle max component-wise.
    for (std::size_t r = 0; r < input.amount_size(); ++r) {
      peak[r] = std::max(peak[r], current_loads[s][r]);
    }
    fwd_peaks[s] = peak;
  }

  peak = current_loads.back();
  bwd_peaks.back() = peak;
  for (std::size_t s = 1; s < bwd_peaks.size(); ++s) {
    auto bwd_s = bwd_peaks.size() - s - 1;
    // Handle max component-wise.
    for (std::size_t r = 0; r < input.amount_size(); ++r) {
      peak[r] = std::max(peak[r], current_loads[bwd_s][r]);
    }
    bwd_peaks[bwd_s] = peak;
  }
}

bool RawRoute::is_valid_addition_for_capacity(const Input&,
                                              const Amount& pickup,
                                              const Amount& delivery,
                                              const Index rank) const {
  assert(rank <= route.size());

  assert(fwd_peaks.size() == route.size() + 2);

  return (fwd_peaks[rank] + delivery <= capacity) and
         (bwd_peaks[rank] + pickup <= capacity);
}

bool RawRoute::is_valid_addition_for_capacity(const Input& input,
                                              const Amount& pickup,
                                              const Amount& delivery,
                                              const Index first_rank,
                                              const Index last_rank) const {
  assert(first_rank < last_rank);
  assert(last_rank <= route.size() + 1);

  auto first_pickups = (first_rank == 0) ? Amount(input.amount_size())
                                         : fwd_pickups[first_rank - 1];
  auto first_deliveries =
    (first_rank == 0) ? current_loads[0] : bwd_deliveries[first_rank - 1];
  auto new_capacity =
    Amount(capacity + fwd_pickups[last_rank - 1] - first_pickups +
           first_deliveries - bwd_deliveries[last_rank - 1]);

  return (fwd_peaks[first_rank] + delivery <= new_capacity) and
         (bwd_peaks[last_rank] + pickup <= new_capacity);
}

Amount RawRoute::get_load(Index s) const {
  return current_loads[s];
}

void RawRoute::add(const Input& input, const Index job_rank, const Index rank) {
  route.insert(route.begin() + rank, job_rank);
  update_amounts(input);
}

void RawRoute::remove(const Input& input,
                      const Index rank,
                      const unsigned count) {
  route.erase(route.begin() + rank, route.begin() + rank + count);
  update_amounts(input);
}

} // namespace vroom
<commit_msg>Fix pickup/delivery separation in capacity check.<commit_after>/*

This file is part of VROOM.

Copyright (c) 2015-2019, Julien Coupey.
All rights reserved (see LICENSE).

*/

#include "structures/vroom/raw_route.h"

namespace vroom {

RawRoute::RawRoute(const Input& input, Index i)
  : fwd_peaks(2, Amount(input.amount_size())),
    bwd_peaks(2, Amount(input.amount_size())),
    vehicle_rank(i),
    has_start(input.vehicles[i].has_start()),
    has_end(input.vehicles[i].has_end()),
    capacity(input.vehicles[i].capacity) {
}

void RawRoute::set_route(const Input& input, const std::vector<Index>& r) {
  route = r;
  update_amounts(input);
}

bool RawRoute::empty() const {
  return route.empty();
}

std::size_t RawRoute::size() const {
  return route.size();
}

void RawRoute::update_amounts(const Input& input) {
  auto step_size = route.size() + 2;
  fwd_pickups.resize(route.size());
  bwd_deliveries.resize(route.size());
  current_loads.resize(step_size);
  fwd_peaks.resize(step_size);
  bwd_peaks.resize(step_size);

  if (route.empty()) {
    // So that check in is_valid_addition_for_capacity is consistent
    // with empty routes.
    std::fill(fwd_peaks.begin(), fwd_peaks.end(), Amount(input.amount_size()));
    std::fill(bwd_peaks.begin(), bwd_peaks.end(), Amount(input.amount_size()));
    return;
  }

  Amount current_pickups(input.amount_size());

  for (std::size_t i = 0; i < route.size(); ++i) {
    current_pickups += input.jobs[route[i]].pickup;
    fwd_pickups[i] = current_pickups;
  }

  Amount current_deliveries(input.amount_size());

  current_loads.back() = fwd_pickups.back();

  for (std::size_t i = 0; i < route.size(); ++i) {
    auto bwd_i = route.size() - i - 1;

    bwd_deliveries[bwd_i] = current_deliveries;
    current_loads[bwd_i + 1] = fwd_pickups[bwd_i] + current_deliveries;
    current_deliveries += input.jobs[route[bwd_i]].delivery;
  }
  current_loads[0] = current_deliveries;

  auto peak = current_loads[0];
  fwd_peaks[0] = peak;
  for (std::size_t s = 1; s < fwd_peaks.size(); ++s) {
    // Handle max component-wise.
    for (std::size_t r = 0; r < input.amount_size(); ++r) {
      peak[r] = std::max(peak[r], current_loads[s][r]);
    }
    fwd_peaks[s] = peak;
  }

  peak = current_loads.back();
  bwd_peaks.back() = peak;
  for (std::size_t s = 1; s < bwd_peaks.size(); ++s) {
    auto bwd_s = bwd_peaks.size() - s - 1;
    // Handle max component-wise.
    for (std::size_t r = 0; r < input.amount_size(); ++r) {
      peak[r] = std::max(peak[r], current_loads[bwd_s][r]);
    }
    bwd_peaks[bwd_s] = peak;
  }
}

bool RawRoute::is_valid_addition_for_capacity(const Input&,
                                              const Amount& pickup,
                                              const Amount& delivery,
                                              const Index rank) const {
  assert(rank <= route.size());

  assert(fwd_peaks.size() == route.size() + 2);

  return (fwd_peaks[rank] + delivery <= capacity) and
         (bwd_peaks[rank] + pickup <= capacity);
}

bool RawRoute::is_valid_addition_for_capacity(const Input& input,
                                              const Amount& pickup,
                                              const Amount& delivery,
                                              const Index first_rank,
                                              const Index last_rank) const {
  assert(first_rank < last_rank);
  assert(last_rank <= route.size() + 1);

  auto first_deliveries =
    (first_rank == 0) ? current_loads[0] : bwd_deliveries[first_rank - 1];

  auto first_pickups = (first_rank == 0) ? Amount(input.amount_size())
                                         : fwd_pickups[first_rank - 1];

  return (fwd_peaks[first_rank] + delivery <=
          capacity + first_deliveries - bwd_deliveries[last_rank - 1]) and
         (bwd_peaks[last_rank] + pickup <=
          capacity + fwd_pickups[last_rank - 1] - first_pickups);
}

Amount RawRoute::get_load(Index s) const {
  return current_loads[s];
}

void RawRoute::add(const Input& input, const Index job_rank, const Index rank) {
  route.insert(route.begin() + rank, job_rank);
  update_amounts(input);
}

void RawRoute::remove(const Input& input,
                      const Index rank,
                      const unsigned count) {
  route.erase(route.begin() + rank, route.begin() + rank + count);
  update_amounts(input);
}

} // namespace vroom
<|endoftext|>


<commit_before>#include "sotarviclient.h"

#include <boost/bind.hpp>
#include <boost/function.hpp>
#include <boost/make_shared.hpp>
#include <iostream>
#include <stdexcept>

#include "logger.h"
#include "types.h"

char service_names[6][12] = {"notify", "start", "chunk", "finish", "getpackages", "abort"};

void callbackWrapper(int fd, void *service_data, const char *parameters) {
  (void)fd;
  SotaRVIClient *client = (SotaRVIClient *)((void **)service_data)[0];
  std::string service((char *)((void **)service_data)[1]);
  Json::Reader reader;
  Json::Value json;
  reader.parse(parameters, json);

  std::cout << "got callback " << service << "\n";
  if (service == "notify") {
    client->sendEvent(boost::make_shared<event::UpdateAvailable>(
        event::UpdateAvailable(data::UpdateAvailable::fromJson(json[0]["update_available"]))));
  } else if (service == "start") {
    client->invokeAck(json[0]["update_id"].asString());
  } else if (service == "chunk") {
    client->saveChunk(json[0]);
  } else if (service == "finish") {
    client->downloadComplete(json[0]);
  } else if (service == "getpackages") {
    LOGGER_LOG(LVL_error, "service " << service << "not implemented yet");
  }
}

SotaRVIClient::SotaRVIClient(const Config &config_in, event::Channel *events_channel_in)
    : config(config_in), events_channel(events_channel_in) {
  std::string client_config(config.rvi.client_config);
  rvi = rviInitLogs(const_cast<char *>(client_config.c_str()), !loggerGetSeverity());
  if (!rvi) {
    throw std::runtime_error("cannot initialize rvi with config file " + client_config);
  }
  connection = rviConnect(rvi, config.rvi.node_host.c_str(), config.rvi.node_port.c_str());

  if (connection <= 0) {
    LOGGER_LOG(LVL_error, "cannot connect to rvi node");
  }
  unsigned int services_count = sizeof(service_names) / sizeof(service_names[0]);
  void *pointers[2];
  pointers[0] = (void *)(this);
  for (unsigned int i = 0; i < services_count; ++i) {
    pointers[1] = (void *)service_names[i];
    int stat = rviRegisterService(rvi, (std::string("sota/") + service_names[i]).c_str(), callbackWrapper, pointers,
                                  sizeof(pointers));
    if (stat) {
      LOGGER_LOG(LVL_error, "unable to register " << service_names[i]);
    }
    services[service_names[i]] = std::string("genivi.org/device/" + config.device.uuid + "/sota/") + service_names[i];
  }
  LOGGER_LOG(LVL_info, "rvi services registered!!");
  boost::thread(boost::bind(&SotaRVIClient::run, this));
}

void SotaRVIClient::sendEvent(const boost::shared_ptr<event::BaseEvent> &event) { *events_channel << event; }

void SotaRVIClient::run() {
  while (true) {
    int result = rviProcessInput(rvi, &connection, 1);
    if (result != 0) {
      LOGGER_LOG(LVL_error, "rviProcessInput error: " << result);
      sleep(5);
    }
  }
}

void SotaRVIClient::startDownload(const data::UpdateRequestId &update_id) {
  Json::Value value;
  value["device"] = config.device.uuid;
  value["update_id"] = update_id;
  value["services"] = services;
  Json::Value params(Json::arrayValue);
  params.append(value);
  int stat = rviInvokeService(rvi, "genivi.org/backend/sota/start", Json::FastWriter().write(params).c_str());
  LOGGER_LOG(LVL_info, "invoked genivi.org/backend/sota/start, return code is " << stat);
}

void SotaRVIClient::invokeAck(const std::string &update_id) {
  std::vector<int> chunks = chunks_map[update_id];

  Json::Value value;
  value["device"] = config.device.uuid;
  value["update_id"] = update_id;
  value["chunks"] = Json::arrayValue;

  for (std::vector<int>::iterator it = chunks.begin(); it != chunks.end(); ++it) {
    value["chunks"].append(*it);
  }
  Json::Value params(Json::arrayValue);
  params.append(value);
  int stat = rviInvokeService(rvi, "genivi.org/backend/sota/ack", Json::FastWriter().write(params).c_str());
  LOGGER_LOG(LVL_info, "invoked genivi.org/backend/sota/ack, return code is " << stat);
}

void SotaRVIClient::saveChunk(const Json::Value &chunk_json) {
  std::string b64_text = chunk_json["bytes"].asString();

  unsigned long long paddChars = std::count(b64_text.begin(), b64_text.end(), '=');
  std::replace(b64_text.begin(), b64_text.end(), '=', 'A');
  std::string output(base64_to_bin(b64_text.begin()), base64_to_bin(b64_text.end()));
  output.erase(output.end() - static_cast<unsigned int>(paddChars), output.end());

  std::ofstream update_file((config.device.packages_dir + chunk_json["update_id"].asString()).c_str(),
                            std::ios::out | std::ios::app | std::ios::binary);
  update_file << output;
  update_file.close();

  chunks_map[chunk_json["update_id"].asString()].push_back(chunk_json["index"].asInt());
  invokeAck(chunk_json["update_id"].asString());
}

void SotaRVIClient::downloadComplete(const Json::Value &parameters) {
  data::DownloadComplete download_complete;
  download_complete.update_id = parameters["update_id"].asString();
  download_complete.signature = parameters["signature"].asString();
  download_complete.update_image = config.device.packages_dir + download_complete.update_id;
  sendEvent(boost::make_shared<event::DownloadComplete>(download_complete));
}

void SotaRVIClient::sendUpdateReport(data::UpdateReport update_report) {
  Json::Value params(Json::arrayValue);
  params.append(update_report.toJson());

  int stat = rviInvokeService(rvi, "genivi.org/backend/sota/report", Json::FastWriter().write(params).c_str());
  LOGGER_LOG(LVL_info, "invoked genivi.org/backend/sota/report, return code is " << stat);
  chunks_map.erase(update_report.update_id);
}
<commit_msg>Fixed issue with wrong update report params, that causes abort on second update<commit_after>#include "sotarviclient.h"

#include <boost/bind.hpp>
#include <boost/function.hpp>
#include <boost/make_shared.hpp>
#include <iostream>
#include <stdexcept>

#include "logger.h"
#include "types.h"

char service_names[6][12] = {"notify", "start", "chunk", "finish", "getpackages", "abort"};

void callbackWrapper(int fd, void *service_data, const char *parameters) {
  (void)fd;
  SotaRVIClient *client = (SotaRVIClient *)((void **)service_data)[0];
  std::string service((char *)((void **)service_data)[1]);
  Json::Reader reader;
  Json::Value json;
  reader.parse(parameters, json);

  std::cout << "got callback " << service << "\n";
  if (service == "notify") {
    client->sendEvent(boost::make_shared<event::UpdateAvailable>(
        event::UpdateAvailable(data::UpdateAvailable::fromJson(json[0]["update_available"]))));
  } else if (service == "start") {
    client->invokeAck(json[0]["update_id"].asString());
  } else if (service == "chunk") {
    client->saveChunk(json[0]);
  } else if (service == "finish") {
    client->downloadComplete(json[0]);
  } else if (service == "getpackages") {
    client->sendEvent(boost::make_shared<event::InstalledSoftwareNeeded>());
  }
}

SotaRVIClient::SotaRVIClient(const Config &config_in, event::Channel *events_channel_in)
    : config(config_in), events_channel(events_channel_in) {
  std::string client_config(config.rvi.client_config);
  rvi = rviInitLogs(const_cast<char *>(client_config.c_str()), !loggerGetSeverity());
  if (!rvi) {
    throw std::runtime_error("cannot initialize rvi with config file " + client_config);
  }
  connection = rviConnect(rvi, config.rvi.node_host.c_str(), config.rvi.node_port.c_str());

  if (connection <= 0) {
    LOGGER_LOG(LVL_error, "cannot connect to rvi node");
  }
  unsigned int services_count = sizeof(service_names) / sizeof(service_names[0]);
  void *pointers[2];
  pointers[0] = (void *)(this);
  for (unsigned int i = 0; i < services_count; ++i) {
    pointers[1] = (void *)service_names[i];
    int stat = rviRegisterService(rvi, (std::string("sota/") + service_names[i]).c_str(), callbackWrapper, pointers,
                                  sizeof(pointers));
    if (stat) {
      LOGGER_LOG(LVL_error, "unable to register " << service_names[i]);
    }
    services[service_names[i]] = std::string("genivi.org/device/" + config.device.uuid + "/sota/") + service_names[i];
  }
  LOGGER_LOG(LVL_info, "rvi services registered!!");
  boost::thread(boost::bind(&SotaRVIClient::run, this));
}

void SotaRVIClient::sendEvent(const boost::shared_ptr<event::BaseEvent> &event) { *events_channel << event; }

void SotaRVIClient::run() {
  while (true) {
    int result = rviProcessInput(rvi, &connection, 1);
    if (result != 0) {
      LOGGER_LOG(LVL_error, "rviProcessInput error: " << result);
      sleep(5);
    }
  }
}

void SotaRVIClient::startDownload(const data::UpdateRequestId &update_id) {
  Json::Value value;
  value["device"] = config.device.uuid;
  value["update_id"] = update_id;
  value["services"] = services;
  Json::Value params(Json::arrayValue);
  params.append(value);
  int stat = rviInvokeService(rvi, "genivi.org/backend/sota/start", Json::FastWriter().write(params).c_str());
  LOGGER_LOG(LVL_info, "invoked genivi.org/backend/sota/start, return code is " << stat);
}

void SotaRVIClient::invokeAck(const std::string &update_id) {
  std::vector<int> chunks = chunks_map[update_id];

  Json::Value value;
  value["device"] = config.device.uuid;
  value["update_id"] = update_id;
  value["chunks"] = Json::arrayValue;

  for (std::vector<int>::iterator it = chunks.begin(); it != chunks.end(); ++it) {
    value["chunks"].append(*it);
  }
  Json::Value params(Json::arrayValue);
  params.append(value);
  int stat = rviInvokeService(rvi, "genivi.org/backend/sota/ack", Json::FastWriter().write(params).c_str());
  LOGGER_LOG(LVL_info, "invoked genivi.org/backend/sota/ack, return code is " << stat);
}

void SotaRVIClient::saveChunk(const Json::Value &chunk_json) {
  std::string b64_text = chunk_json["bytes"].asString();

  unsigned long long paddChars = std::count(b64_text.begin(), b64_text.end(), '=');
  std::replace(b64_text.begin(), b64_text.end(), '=', 'A');
  std::string output(base64_to_bin(b64_text.begin()), base64_to_bin(b64_text.end()));
  output.erase(output.end() - static_cast<unsigned int>(paddChars), output.end());

  std::ofstream update_file((config.device.packages_dir + chunk_json["update_id"].asString()).c_str(),
                            std::ios::out | std::ios::app | std::ios::binary);
  update_file << output;
  update_file.close();

  chunks_map[chunk_json["update_id"].asString()].push_back(chunk_json["index"].asInt());
  invokeAck(chunk_json["update_id"].asString());
}

void SotaRVIClient::downloadComplete(const Json::Value &parameters) {
  data::DownloadComplete download_complete;
  download_complete.update_id = parameters["update_id"].asString();
  download_complete.signature = parameters["signature"].asString();
  download_complete.update_image = config.device.packages_dir + download_complete.update_id;
  sendEvent(boost::make_shared<event::DownloadComplete>(download_complete));
}

void SotaRVIClient::sendUpdateReport(data::UpdateReport update_report) {
  Json::Value params(Json::arrayValue);
  Json::Value report;
  report["device"] = config.device.uuid;
  report["update_report"] = update_report.toJson();
  params.append(report);

  int stat = rviInvokeService(rvi, "genivi.org/backend/sota/report", Json::FastWriter().write(params).c_str());
  LOGGER_LOG(LVL_info, "invoked genivi.org/backend/sota/report, return code is " << stat);
  chunks_map.erase(update_report.update_id);
}
<|endoftext|>