1 files changed, 85 insertions, 42 deletions

diff --git a/src/device_trezor/device_trezor_base.cpp b/src/device_trezor/device_trezor_base.cpp
index 3a98bba5a..38c20c30b 100644
--- a/src/device_trezor/device_trezor_base.cpp
+++ b/src/device_trezor/device_trezor_base.cpp
@@ -32,43 +32,11 @@
 namespace hw {
 namespace trezor {
 
-#if WITH_DEVICE_TREZOR
+#ifdef WITH_DEVICE_TREZOR
 
 #undef MONERO_DEFAULT_LOG_CATEGORY
 #define MONERO_DEFAULT_LOG_CATEGORY "device.trezor"
 
-    std::shared_ptr<google::protobuf::Message> trezor_protocol_callback::on_button_request(const messages::common::ButtonRequest * msg){
-      MDEBUG("on_button_request");
-      device.on_button_request();
-      return std::make_shared<messages::common::ButtonAck>();
-    }
-
-    std::shared_ptr<google::protobuf::Message> trezor_protocol_callback::on_pin_matrix_request(const messages::common::PinMatrixRequest * msg){
-      MDEBUG("on_pin_request");
-      epee::wipeable_string pin;
-      device.on_pin_request(pin);
-      auto resp = std::make_shared<messages::common::PinMatrixAck>();
-      resp->set_pin(pin.data(), pin.size());
-      return resp;
-    }
-
-    std::shared_ptr<google::protobuf::Message> trezor_protocol_callback::on_passphrase_request(const messages::common::PassphraseRequest * msg){
-      MDEBUG("on_passhprase_request");
-      epee::wipeable_string passphrase;
-      device.on_passphrase_request(msg->on_device(), passphrase);
-      auto resp = std::make_shared<messages::common::PassphraseAck>();
-      if (!msg->on_device()){
-        resp->set_passphrase(passphrase.data(), passphrase.size());
-      }
-      return resp;
-    }
-
-    std::shared_ptr<google::protobuf::Message> trezor_protocol_callback::on_passphrase_state_request(const messages::common::PassphraseStateRequest * msg){
-      MDEBUG("on_passhprase_state_request");
-      device.on_passphrase_state_request(msg->state());
-      return std::make_shared<messages::common::PassphraseStateAck>();
-    }
-
     const uint32_t device_trezor_base::DEFAULT_BIP44_PATH[] = {0x8000002c, 0x80000080, 0x80000000};
 
     device_trezor_base::device_trezor_base() {
@@ -117,9 +85,6 @@ namespace trezor {
         return false;
       }
 
-      if (!m_protocol_callback){
-        m_protocol_callback = std::make_shared<trezor_protocol_callback>(*this);
-      }
       return true;
     }
 
@@ -245,6 +210,49 @@ namespace trezor {
       }
     }
 
+    void device_trezor_base::write_raw(const google::protobuf::Message * msg){
+      require_connected();
+      CHECK_AND_ASSERT_THROW_MES(msg, "Empty message");
+      this->getTransport()->write(*msg);
+    }
+
+    GenericMessage device_trezor_base::read_raw(){
+      require_connected();
+      std::shared_ptr<google::protobuf::Message> msg_resp;
+      hw::trezor::messages::MessageType msg_resp_type;
+
+      this->getTransport()->read(msg_resp, &msg_resp_type);
+      return GenericMessage(msg_resp_type, msg_resp);
+    }
+
+    GenericMessage device_trezor_base::call_raw(const google::protobuf::Message * msg) {
+      write_raw(msg);
+      return read_raw();
+    }
+
+    bool device_trezor_base::message_handler(GenericMessage & input){
+      // Later if needed this generic message handler can be replaced by a pointer to
+      // a protocol message handler which by default points to the device class which implements
+      // the default handler.
+      switch(input.m_type){
+        case messages::MessageType_ButtonRequest:
+          on_button_request(input, dynamic_cast<const messages::common::ButtonRequest*>(input.m_msg.get()));
+          return true;
+        case messages::MessageType_PassphraseRequest:
+          on_passphrase_request(input, dynamic_cast<const messages::common::PassphraseRequest*>(input.m_msg.get()));
+          return true;
+        case messages::MessageType_PassphraseStateRequest:
+          on_passphrase_state_request(input, dynamic_cast<const messages::common::PassphraseStateRequest*>(input.m_msg.get()));
+          return true;
+        case messages::MessageType_PinMatrixRequest:
+          on_pin_request(input, dynamic_cast<const messages::common::PinMatrixRequest*>(input.m_msg.get()));
+          return true;
+        default:
+          return false;
+      }
+    }
+
+
     /* ======================================================================= */
     /*                              TREZOR PROTOCOL                            */
     /* ======================================================================= */
@@ -269,32 +277,67 @@ namespace trezor {
       return false;
     }
 
-    void device_trezor_base::on_button_request()
+    void device_trezor_base::on_button_request(GenericMessage & resp, const messages::common::ButtonRequest * msg)
     {
+      CHECK_AND_ASSERT_THROW_MES(msg, "Empty message");
+      MDEBUG("on_button_request, code: " << msg->code());
+
+      messages::common::ButtonAck ack;
+      write_raw(&ack);
+
       if (m_callback){
         m_callback->on_button_request();
       }
+
+      resp = read_raw();
     }
 
-    void device_trezor_base::on_pin_request(epee::wipeable_string & pin)
+    void device_trezor_base::on_pin_request(GenericMessage & resp, const messages::common::PinMatrixRequest * msg)
     {
+      MDEBUG("on_pin_request");
+      CHECK_AND_ASSERT_THROW_MES(msg, "Empty message");
+
+      epee::wipeable_string pin;
+
       if (m_callback){
         m_callback->on_pin_request(pin);
       }
+
+      // TODO: remove PIN from memory
+      messages::common::PinMatrixAck m;
+      m.set_pin(pin.data(), pin.size());
+      resp = call_raw(&m);
     }
 
-    void device_trezor_base::on_passphrase_request(bool on_device, epee::wipeable_string & passphrase)
+    void device_trezor_base::on_passphrase_request(GenericMessage & resp, const messages::common::PassphraseRequest * msg)
     {
+      CHECK_AND_ASSERT_THROW_MES(msg, "Empty message");
+      MDEBUG("on_passhprase_request, on device: " << msg->on_device());
+      epee::wipeable_string passphrase;
+
       if (m_callback){
-        m_callback->on_passphrase_request(on_device, passphrase);
+        m_callback->on_passphrase_request(msg->on_device(), passphrase);
       }
+
+      messages::common::PassphraseAck m;
+      if (!msg->on_device()){
+        // TODO: remove passphrase from memory
+        m.set_passphrase(passphrase.data(), passphrase.size());
+      }
+      resp = call_raw(&m);
     }
 
-    void device_trezor_base::on_passphrase_state_request(const std::string & state)
+    void device_trezor_base::on_passphrase_state_request(GenericMessage & resp, const messages::common::PassphraseStateRequest * msg)
     {
+      MDEBUG("on_passhprase_state_request");
+      CHECK_AND_ASSERT_THROW_MES(msg, "Empty message");
+
       if (m_callback){
-        m_callback->on_passphrase_state_request(state);
+        m_callback->on_passphrase_state_request(msg->state());
       }
+
+      messages::common::PassphraseStateAck m;
+      resp = call_raw(&m);
     }
 
 #endif //WITH_DEVICE_TREZOR