text: // Copyright (C) 2017-2019 Baidu, Inc. All Rights Reserved.
  //
  // Redistribution and use in source and binary forms, with or without
  // modification, are permitted provided that the following conditions
  // are met:
  //
  //  * Redistributions of source code must retain the above copyright
  //    notice, this list of conditions and the following disclaimer.
  //  * Redistributions in binary form must reproduce the above copyright
  //    notice, this list of conditions and the following disclaimer in
  //    the documentation and/or other materials provided with the
  //    distribution.
  //  * Neither the name of Baidu, Inc., nor the names of its
  //    contributors may be used to endorse or promote products derived
  //    from this software without specific prior written permission.
  //
  // THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
  // "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
  // LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR
  // A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT
  // OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,
  // SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT
  // LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
  // DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY
  // THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
  // (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
  // OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
  
  #include "MessageHandler.h"
  
  using namespace util;
  
  MessageHandler::MessageHandler(int port) {
      this->nm = NetworkManagerServer::getInstance(port);
  }
  
  MessageHandler::~MessageHandler() {
      delete this->enclave;
  }
  
  
  int MessageHandler::init() {
      this->nm->Init();
      this->nm->connectCallbackHandler([this](string v, int type) {
          return this->incomingHandler(v, type);
      });
  }
  
  void MessageHandler::start() {
  
      sgx_status_t ret = this->initEnclave();
      if (SGX_SUCCESS != ret) {
          Log("Error, call initEnclave fail", log::error);
          return;
      }
  
      sgx_status_t status;
      ret = initialize(this->enclave->getID(), &status);
      if ((SGX_SUCCESS != ret) || (SGX_SUCCESS != status)) {
          Log("Error, call generate_salt fail", log::error);
          return;
      }
  
      Log("Call initEnclave success");
      this->nm->startService();
  }
  
  
  sgx_status_t MessageHandler::initEnclave() {
      this->enclave = Enclave::getInstance();
      return this->enclave->createEnclave();
  }
  
  uint32_t MessageHandler::getExtendedEPID_GID() {
      uint32_t extended_epid_group_id = 0;
      int ret = sgx_get_extended_epid_group_id(&extended_epid_group_id);
  
      if (SGX_SUCCESS != ret) {
          ret = -1;
          Log("Error, call sgx_get_extended_epid_group_id fail");
          return ret;
      }
  
      Log("Call sgx_get_extended_epid_group_id success");
  
      return extended_epid_group_id;
  }
  
  
  string MessageHandler::generateMSG0() {
      Log("Call MSG0 generate");
  
      uint32_t extended_epid_group_id = this->getExtendedEPID_GID();
  
      Messages::MessageMsg0 msg;
      msg.set_type(RA_MSG0);
      msg.set_epid(extended_epid_group_id);
  
      return nm->serialize(msg);
  }
  
  
  string MessageHandler::generateMSG1() {
      int retGIDStatus = 0;
      int count = 0;
      sgx_status_t ret;
      sgx_ra_context_t context = INT_MAX;
      sgx_ra_msg1_t sgxMsg1Obj;
  
      ret = this->enclave->raInit(&context);
      if (SGX_SUCCESS != ret) {
          Log("Error, call enclave_init_ra fail", log::error);
          return "";
      }
  
      while (1) {
          retGIDStatus = sgx_ra_get_msg1(context,
                                         this->enclave->getID(),
                                         sgx_ra_get_ga,
                                         &sgxMsg1Obj);
  
          if (retGIDStatus == SGX_SUCCESS) {
              break;
          } else if (retGIDStatus == SGX_ERROR_BUSY) {
              if (count == 5) { //retried 5 times, so fail out
                  Log("Error, sgx_ra_get_msg1 is busy - 5 retries failed", log::error);
                  break;;
              } else {
                  sleep(3);
                  count++;
              }
          } else {    //error other than busy
              Log("Error, failed to generate MSG1", log::error);
              break;
          }
      }
  
  
      if (SGX_SUCCESS == retGIDStatus) {
          Log("MSG1 generated Successfully");
  
          Messages::MessageMSG1 msg;
          msg.set_type(RA_MSG1);
          msg.set_context(context);
  
          for (auto x : sgxMsg1Obj.g_a.gx)
              msg.add_gax(x);
  
          for (auto x : sgxMsg1Obj.g_a.gy)
              msg.add_gay(x);
  
          for (auto x : sgxMsg1Obj.gid) {
              msg.add_gid(x);
          }
  
          return nm->serialize(msg);
      }
  
      return "";
  }
  
  
  void MessageHandler::assembleMSG2(Messages::MessageMSG2 msg, sgx_ra_msg2_t **pp_msg2) {
      uint32_t size = msg.size();
  
      sgx_ra_msg2_t *p_msg2 = NULL;
      p_msg2 = (sgx_ra_msg2_t*) malloc(size + sizeof(sgx_ra_msg2_t));
  
      uint8_t pub_key_gx[32];
      uint8_t pub_key_gy[32];
  
      sgx_ec256_signature_t sign_gb_ga;
      sgx_spid_t spid;
  
      for (int i; i<32; i++) {
          pub_key_gx[i] = msg.public_key_gx(i);
          pub_key_gy[i] = msg.public_key_gy(i);
      }
  
      for (int i=0; i<16; i++) {
          spid.id[i] = msg.spid(i);
      }
  
      for (int i=0; i<8; i++) {
          sign_gb_ga.x[i] = msg.signature_x(i);
          sign_gb_ga.y[i] = msg.signature_y(i);
      }
  
      memcpy(&p_msg2->g_b.gx, &pub_key_gx, sizeof(pub_key_gx));
      memcpy(&p_msg2->g_b.gy, &pub_key_gy, sizeof(pub_key_gy));
      memcpy(&p_msg2->sign_gb_ga, &sign_gb_ga, sizeof(sign_gb_ga));
      memcpy(&p_msg2->spid, &spid, sizeof(spid));
  
      p_msg2->quote_type = (uint16_t)msg.quote_type();
      p_msg2->kdf_id = msg.cmac_kdf_id();
  
      uint8_t smac[16];
      for (int i=0; i<16; i++)
          smac[i] = msg.smac(i);
  
      memcpy(&p_msg2->mac, &smac, sizeof(smac));
  
      p_msg2->sig_rl_size = msg.size_sigrl();
  
      for (int i=0; i<msg.size_sigrl(); i++)
          p_msg2->sig_rl[i] = msg.sigrl(i);
  
      *pp_msg2 = p_msg2;
  }
  
  
  string MessageHandler::handleMSG2(Messages::MessageMSG2 msg) {
      Log("Received MSG2");
  
      uint32_t size = msg.size();
      sgx_ra_context_t context = msg.context();
  
      sgx_ra_msg2_t *p_msg2;
      this->assembleMSG2(msg, &p_msg2);
  
      sgx_ra_msg3_t *p_msg3 = NULL;
      uint32_t msg3_size;
      int ret = 0;
  
      do {
          ret = sgx_ra_proc_msg2(context,
                                 this->enclave->getID(),
                                 sgx_ra_proc_msg2_trusted,
                                 sgx_ra_get_msg3_trusted,
                                 p_msg2,
                                 size,
                                 &p_msg3,
                                 &msg3_size);
      } while (SGX_ERROR_BUSY == ret && busy_retry_time--);
  
      SafeFree(p_msg2);
  
      if (SGX_SUCCESS != (sgx_status_t)ret) {
          Log("Error, call sgx_ra_proc_msg2 fail, error code: 0x%x", ret);
      } else {
          Log("Call sgx_ra_proc_msg2 success");
  
          Messages::MessageMSG3 msg3;
  
          msg3.set_type(RA_MSG3);
          msg3.set_size(msg3_size);
          msg3.set_context(context);
  
          for (int i=0; i<SGX_MAC_SIZE; i++)
              msg3.add_sgx_mac(p_msg3->mac[i]);
  
          for (int i=0; i<SGX_ECP256_KEY_SIZE; i++) {
              msg3.add_gax_msg3(p_msg3->g_a.gx[i]);
              msg3.add_gay_msg3(p_msg3->g_a.gy[i]);
          }
  
          for (int i=0; i<256; i++) {
              msg3.add_sec_property(p_msg3->ps_sec_prop.sgx_ps_sec_prop_desc[i]);
          }
  
  
          for (int i=0; i<1116; i++) {
              msg3.add_quote(p_msg3->quote[i]);
          }
  
          SafeFree(p_msg3);
  
          return nm->serialize(msg3);
      }
  
      SafeFree(p_msg3);
  
      return "";
  }
  
  
  void MessageHandler::assembleAttestationMSG(Messages::AttestationMessage msg, ra_samp_response_header_t **pp_att_msg) {
      sample_ra_att_result_msg_t *p_att_result_msg = NULL;
      ra_samp_response_header_t* p_att_result_msg_full = NULL;
  
      int total_size = msg.size() + sizeof(ra_samp_response_header_t) + msg.result_size();
      p_att_result_msg_full = (ra_samp_response_header_t*) malloc(total_size);
  
      memset(p_att_result_msg_full, 0, total_size);
      p_att_result_msg_full->type = RA_ATT_RESULT;
      p_att_result_msg_full->size = msg.size();
  
      p_att_result_msg = (sample_ra_att_result_msg_t *) p_att_result_msg_full->body;
  
      p_att_result_msg->platform_info_blob.sample_epid_group_status = msg.epid_group_status();
      p_att_result_msg->platform_info_blob.sample_tcb_evaluation_status = msg.tcb_evaluation_status();
      p_att_result_msg->platform_info_blob.pse_evaluation_status = msg.pse_evaluation_status();
  
      for (int i=0; i<PSVN_SIZE; i++)
          p_att_result_msg->platform_info_blob.latest_equivalent_tcb_psvn[i] = msg.latest_equivalent_tcb_psvn(i);
  
      for (int i=0; i<ISVSVN_SIZE; i++)
          p_att_result_msg->platform_info_blob.latest_pse_isvsvn[i] = msg.latest_pse_isvsvn(i);
  
      for (int i=0; i<PSDA_SVN_SIZE; i++)
          p_att_result_msg->platform_info_blob.latest_psda_svn[i] = msg.latest_psda_svn(i);
  
      for (int i=0; i<GID_SIZE; i++)
          p_att_result_msg->platform_info_blob.performance_rekey_gid[i] = msg.performance_rekey_gid(i);
  
      for (int i=0; i<SAMPLE_NISTP256_KEY_SIZE; i++) {
          p_att_result_msg->platform_info_blob.signature.x[i] = msg.ec_sign256_x(i);
          p_att_result_msg->platform_info_blob.signature.y[i] = msg.ec_sign256_y(i);
      }
  
      for (int i=0; i<SAMPLE_MAC_SIZE; i++)
          p_att_result_msg->mac[i] = msg.mac_smk(i);
  
  
      p_att_result_msg->secret.payload_size = msg.result_size();
  
      for (int i=0; i<12; i++)
          p_att_result_msg->secret.reserved[i] = msg.reserved(i);
  
      for (int i=0; i<SAMPLE_SP_TAG_SIZE; i++)
          p_att_result_msg->secret.payload_tag[i] = msg.payload_tag(i);
  
      for (int i=0; i<SAMPLE_SP_TAG_SIZE; i++)
          p_att_result_msg->secret.payload_tag[i] = msg.payload_tag(i);
  
      for (int i=0; i<msg.result_size(); i++) {
          p_att_result_msg->secret.payload[i] = (uint8_t)msg.payload(i);
      }
  
      *pp_att_msg = p_att_result_msg_full;
  }
  
  string MessageHandler::generateAttestationFailed(uint32_t id, sgx_ra_context_t context) {
  
      Messages::MessagePsiSalt msg;
  
      msg.set_type(RA_PSI_SLAT);
      msg.set_size(0);
      msg.set_state(0);
      msg.set_context(context);
      msg.add_salt(0);
      msg.add_mac(0);
      msg.set_id(id);
  
      return nm->serialize(msg);
  }
  
  string MessageHandler::handleAttestationResult(Messages::AttestationMessage msg) {
      Log("Received Attestation result");
  
      ra_samp_response_header_t *p_att_result_msg_full = NULL;
      this->assembleAttestationMSG(msg, &p_att_result_msg_full);
      sample_ra_att_result_msg_t *p_att_result_msg_body = (sample_ra_att_result_msg_t *) ((uint8_t*) p_att_result_msg_full + sizeof(ra_samp_response_header_t));
  
      sgx_status_t status;
      sgx_status_t ret;
      sgx_ra_context_t context = msg.context();
      uint32_t id = 0;
      uint8_t salt[SALT_SIZE];
      uint8_t mac[SGX_MAC_SIZE];
  
      ret = verify_att_result_mac(this->enclave->getID(),
                                  &status,
                                  context,
                                  (uint8_t*)&p_att_result_msg_body->platform_info_blob,
                                  sizeof(ias_platform_info_blob_t),
                                  (uint8_t*)&p_att_result_msg_body->mac);
  
  
      if ((SGX_SUCCESS != ret) || (SGX_SUCCESS != status)) {
          Log("Error: INTEGRITY FAILED - attestation result message MK based cmac failed", log::error);
          SafeFree(p_att_result_msg_full);
          return generateAttestationFailed(context, id);
      }
  
      if (0 != p_att_result_msg_full->status[0] || 0 != p_att_result_msg_full->status[1]) {
          Log("Error, attestation mac result message MK based cmac failed", log::error);
          SafeFree(p_att_result_msg_full);
          return generateAttestationFailed(context, id);
      } else {
          ret = verify_secret_data(this->enclave->getID(),
                                   &status,
                                   context,
                                   p_att_result_msg_body->secret.payload,
                                   p_att_result_msg_body->secret.payload_size,
                                   p_att_result_msg_body->secret.payload_tag,
                                   MAX_VERIFICATION_RESULT,
                                   salt,
                                   mac,
                                   &id);
  
          SafeFree(p_att_result_msg_full);
  
          if (SGX_SUCCESS != ret) {
              Log("Error, attestation result message secret using SK based AESGCM failed", log::error);
              Log("Error  on ret , code : %08X\n",ret);
              print_error_message(ret);
  
              return generateAttestationFailed(context, id);
  
          } else if (SGX_SUCCESS != status) {
              Log("Error, attestation result message secret using SK based AESGCM failed", log::error);
              Log("Error  on status, code : %08X\n",status);
              print_error_message(status);
  
              return generateAttestationFailed(context, id);
  
          } else {
              Log("Send attestation okay");
  
              Messages::MessagePsiSalt msg;
              msg.set_type(RA_PSI_SLAT);
              msg.set_size(0);
              msg.set_state(1);
              msg.set_context(context);
              msg.set_id(id);
  
              for (int i = 0; i < SALT_SIZE; i++) {
                  msg.add_salt(salt[i]);
              }
  
              for (int i = 0; i < SGX_MAC_SIZE; i++) {
                  msg.add_mac(mac[i]);
              }
  
              return nm->serialize(msg);
          }
      }
  
      return "";
  }
  
  
  string MessageHandler::handleMSG0(Messages::MessageMsg0 msg) {
      Log("MSG0 response received");
  
      if (msg.status() == TYPE_OK) {
          Log("Sending msg1 to remote attestation service provider. Expecting msg2 back");
          auto ret = this->generateMSG1();
          return ret;
      } else {
          Log("MSG0 response status was not OK", log::error);
      }
  
      return "";
  }
  
  
  string MessageHandler::handleVerification() {
      Log("Verification request received");
      return this->generateMSG0();
  }
  
  string MessageHandler::createInitMsg(int type, string msg) {
      Messages::InitialMessage init_msg;
      init_msg.set_type(type);
      init_msg.set_size(0);
  
      return nm->serialize(init_msg);
  }
  
  string MessageHandler::handlePsiHashData(Messages::MessagePsiHashData msg) {
      Log("[PSI] Received hash data");
  
      uint8_t mac[SGX_MAC_SIZE] = {0};
      sgx_ra_context_t context = msg.context();
      uint32_t id = msg.id();
      int data_size = msg.data_size();
      uint8_t* data = (uint8_t*)malloc(data_size);
  
      for (int i = 0; i < SGX_MAC_SIZE; i++) {
          mac[i] = (uint8_t)msg.mac(i);
      }
  
      for (int i = 0; i < data_size; i++) {
          data[i] = (uint8_t)msg.data(i);
      }
  
      sgx_status_t status;
      sgx_status_t ret = add_hash_data(this->enclave->getID(),
                                      &status,
                                      id,
                                      context,
                                      data,
                                      data_size,
                                      mac);
      if (SGX_SUCCESS != ret || SGX_SUCCESS != status) {
          Log("[PSI] add_hash_data failed, %d, %d!", ret, status);
          SafeFree(data);
          return "";
      }
  
      SafeFree(data);
  
      Messages::MessagePsiResult result;
      result.set_type(RA_PSI_RESULT);
      result.set_size(0);
      result.set_state(0);
      result.set_context(msg.context());
      result.set_id(msg.id());
  
      return nm->serialize(result);
  }
  
  string MessageHandler::handlePsiHashDataFinished(Messages::MessagePsiHashDataFinished msg, bool* again) {
  
      sgx_ra_context_t context = msg.context();
      uint32_t id = msg.id();
      uint8_t mac[SGX_MAC_SIZE] = {0};
      uint8_t * data = NULL;
      size_t data_size = 0;
      sgx_status_t status;
      sgx_status_t ret;
  
      Log("[PSI] Received hash data finished, %d", id);
  
      ret = get_result_size(this->enclave->getID(), &status, id, &data_size);
      if (SGX_SUCCESS != ret) {
          Log("[PSI] get_result_size failed, %d", ret);
          return "";
      }
  
      if (SGX_SUCCESS != status) {
          if (status == SGX_ERROR_INVALID_STATE) {
              *again = true;
  
              Messages::MessagePsiResult result;
              result.set_type(RA_PSI_RESULT);
              result.set_size(0);
              result.set_state(1); //tell sp req result again
              result.set_context(msg.context());
              result.set_id(msg.id());
  
              Log("[PSI] has not calc result success");
  
              return nm->serialize(result);
          } else {
              Log("[PSI] get_result_size failed, %d, %d", ret, status);
              return "";
          }
      }
  
      if (data_size > 0) {
          data = (uint8_t*)malloc(data_size);
          if (data == NULL) {
              Log("[PSI] alloc buffer for intersect data failed!");
              return "";
          }
  
          ret = get_result(this->enclave->getID(),
                          &status,
                          id,
                          context,
                          data,
                          data_size,
                          mac);
          if (SGX_SUCCESS != ret || SGX_SUCCESS != status) {
              Log("[PSI] get_result failed, %d, %d", ret, status);
              return "";
          }
      }
  
      Messages::MessagePsiIntersect intersect;
      intersect.set_type(RA_PSI_INTERSECT);
      intersect.set_size(0);
      intersect.set_id(msg.id());
      intersect.set_context(context);
  
      for (int i = 0; i < SGX_MAC_SIZE; i++) {
          intersect.add_mac(mac[i]);
      }
  
      if (data && data_size) {
          for (int i = 0; i < data_size; i++) {
              intersect.add_data(data[i]);
          }
      }
      if (data) {
          SafeFree(data);
      }
  
      enclave_ra_close(this->enclave->getID(), &status, context);
  
      Log("[PSI] get result success, %d", data_size/SGX_HASH_SIZE);
  
      return nm->serialize(intersect);
  }
  
  vector<string> MessageHandler::incomingHandler(string v, int type) {
      vector<string> res;
      string s;
      bool ret;
  
      switch (type) {
      case RA_VERIFICATION: {	//Verification request
          Messages::InitialMessage init_msg;
          ret = init_msg.ParseFromString(v);
          if (ret && init_msg.type() == RA_VERIFICATION) {
              s = this->handleVerification();
              res.push_back(to_string(RA_MSG0));
          }
      }
      break;
      case RA_MSG0: {		//Reply to MSG0
          Messages::MessageMsg0 msg0;
          ret = msg0.ParseFromString(v);
          if (ret && (msg0.type() == RA_MSG0)) {
              s = this->handleMSG0(msg0);
              res.push_back(to_string(RA_MSG1));
          }
      }
      break;
      case RA_MSG2: {		//MSG2
          Messages::MessageMSG2 msg2;
          ret = msg2.ParseFromString(v);
          if (ret && (msg2.type() == RA_MSG2)) {
              s = this->handleMSG2(msg2);
              res.push_back(to_string(RA_MSG3));
          }
      }
      break;
      case RA_ATT_RESULT: {	//Reply to MSG3
          Messages::AttestationMessage att_msg;
          ret = att_msg.ParseFromString(v);
          if (ret && att_msg.type() == RA_ATT_RESULT) {
              s = this->handleAttestationResult(att_msg);
              res.push_back(to_string(RA_PSI_SLAT));
          }
      }
      break;
      case RA_PSI_HASHDATA: {
          Messages::MessagePsiHashData data_msg;
          ret = data_msg.ParseFromString(v);
          if (ret && data_msg.type() == RA_PSI_HASHDATA) {
              s = this->handlePsiHashData(data_msg);
              res.push_back(to_string(RA_PSI_RESULT));
          }
      }
      break;
      case RA_PSI_HASHDATA_FINISHED: {
          Messages::MessagePsiHashDataFinished finish_msg;
          ret = finish_msg.ParseFromString(v);
          if (ret && finish_msg.type() == RA_PSI_HASHDATA_FINISHED) {
              bool again = false;
              s = this->handlePsiHashDataFinished(finish_msg, &again);
              if (again) {
                  res.push_back(to_string(RA_PSI_RESULT));
              } else {
                  res.push_back(to_string(RA_PSI_INTERSECT));
              }
          }
      }
      break;
      default:
          Log("Unknown type: %d", type, log::error);
          break;
      }
  
      res.push_back(s);
  
      return res;
  }
  