1 files changed, 117 insertions, 33 deletions

diff --git a/src/device_trezor/trezor/transport.cpp b/src/device_trezor/trezor/transport.cpp
index cd66e59e8..9149e0168 100644
--- a/src/device_trezor/trezor/transport.cpp
+++ b/src/device_trezor/trezor/transport.cpp
@@ -84,6 +84,17 @@ namespace trezor{
     return std::string(in.GetString());
   }
 
+  uint64_t pack_version(uint32_t major, uint32_t minor, uint32_t patch)
+  {
+    // packing (major, minor, patch) to 64 B: 16 B | 24 B | 24 B
+    const unsigned bits_1 = 16;
+    const unsigned bits_2 = 24;
+    const uint32_t mask_1 = (1 << bits_1) - 1;
+    const uint32_t mask_2 = (1 << bits_2) - 1;
+    CHECK_AND_ASSERT_THROW_MES(major <= mask_1 && minor <= mask_2 && patch <= mask_2, "Version numbers overflow packing scheme");
+    return patch | (((uint64_t)minor) << bits_2) | (((uint64_t)major) << (bits_1 + bits_2));
+  }
+
   //
   // Helpers
   //
@@ -212,6 +223,40 @@ namespace trezor{
     msg = msg_wrap;
   }
 
+  Transport::Transport(): m_open_counter(0) {
+
+  }
+
+  bool Transport::pre_open(){
+    if (m_open_counter > 0){
+      MTRACE("Already opened, count: " << m_open_counter);
+      m_open_counter += 1;
+      return false;
+
+    } else if (m_open_counter < 0){
+      MTRACE("Negative open value: " << m_open_counter);
+
+    }
+
+    // Caller should set m_open_counter to 1 after open
+    m_open_counter = 0;
+    return true;
+  }
+
+  bool Transport::pre_close(){
+    m_open_counter -= 1;
+
+    if (m_open_counter < 0){
+      MDEBUG("Already closed. Counter " << m_open_counter);
+
+    } else if (m_open_counter == 0) {
+      return true;
+
+    }
+
+    return false;
+  }
+
   //
   // Bridge transport
   //
@@ -246,6 +291,10 @@ namespace trezor{
   }
 
   void BridgeTransport::open() {
+    if (!pre_open()){
+      return;
+    }
+
     if (!m_device_path){
       throw exc::CommunicationException("Coud not open, empty device path");
     }
@@ -259,9 +308,15 @@ namespace trezor{
     }
 
     m_session = boost::make_optional(json_get_string(bridge_res["session"]));
+    m_open_counter = 1;
   }
 
   void BridgeTransport::close() {
+    if (!pre_close()){
+      return;
+    }
+
+    MTRACE("Closing Trezor:BridgeTransport");
     if (!m_device_path || !m_session){
       throw exc::CommunicationException("Device not open");
     }
@@ -423,6 +478,10 @@ namespace trezor{
   }
 
   void UdpTransport::open() {
+    if (!pre_open()){
+      return;
+    }
+
     udp::resolver resolver(m_io_service);
     udp::resolver::query query(udp::v4(), m_device_host, std::to_string(m_device_port));
     m_endpoint = *resolver.resolve(query);
@@ -434,10 +493,16 @@ namespace trezor{
     check_deadline();
 
     m_proto->session_begin(*this);
+    m_open_counter = 1;
   }
 
   void UdpTransport::close() {
-    if (!m_socket){
+    if (!pre_close()){
+      return;
+    }
+
+    MTRACE("Closing Trezor:UdpTransport");
+    if (!m_socket) {
       throw exc::CommunicationException("Socket is already closed");
     }
 
@@ -446,6 +511,19 @@ namespace trezor{
     m_socket = nullptr;
   }
 
+  std::shared_ptr<Transport> UdpTransport::find_debug() {
+#ifdef WITH_TREZOR_DEBUGGING
+    std::shared_ptr<UdpTransport> t = std::make_shared<UdpTransport>();
+    t->m_proto = std::make_shared<ProtocolV1>();
+    t->m_device_host = m_device_host;
+    t->m_device_port = m_device_port + 1;
+    return t;
+#else
+    MINFO("Debug link is disabled in production");
+    return nullptr;
+#endif
+  }
+
   void UdpTransport::write_chunk(const void * buff, size_t size){
     require_socket();
 
@@ -660,8 +738,7 @@ namespace trezor{
   WebUsbTransport::WebUsbTransport(
       boost::optional<libusb_device_descriptor*> descriptor,
       boost::optional<std::shared_ptr<Protocol>> proto
-  ): m_conn_count(0),
-     m_usb_session(nullptr), m_usb_device(nullptr), m_usb_device_handle(nullptr),
+  ): m_usb_session(nullptr), m_usb_device(nullptr), m_usb_device_handle(nullptr),
      m_bus_id(-1), m_device_addr(-1)
   {
     if (descriptor){
@@ -672,7 +749,7 @@ namespace trezor{
 
     m_proto = proto ? proto.get() : std::make_shared<ProtocolV1>();
 
-#ifdef WITH_TREZOR_DEBUG
+#ifdef WITH_TREZOR_DEBUGGING
     m_debug_mode = false;
 #endif
   }
@@ -757,12 +834,10 @@ namespace trezor{
   };
 
   void WebUsbTransport::open() {
-    const int interface = get_interface();
-    if (m_conn_count > 0){
-      MTRACE("Already opened, count: " << m_conn_count);
-      m_conn_count += 1;
+    if (!pre_open()){
       return;
     }
+    const int interface = get_interface();
 
 #define TREZOR_DESTROY_SESSION() do { libusb_exit(m_usb_session); m_usb_session = nullptr; } while(0)
 
@@ -840,45 +915,55 @@ namespace trezor{
       throw exc::DeviceAcquireException("Unable to claim libusb device");
     }
 
-    m_conn_count = 1;
+    m_open_counter = 1;
     m_proto->session_begin(*this);
     
 #undef TREZOR_DESTROY_SESSION
   };
 
   void WebUsbTransport::close() {
-    m_conn_count -= 1;
-
-    if (m_conn_count < 0){
-      MERROR("Close counter is negative: " << m_conn_count);
-
-    } else if (m_conn_count == 0){
-      MTRACE("Closing webusb device");
+    if (!pre_close()){
+      return;
+    }
 
-      m_proto->session_end(*this);
+    MTRACE("Closing Trezor:WebUsbTransport");
+    m_proto->session_end(*this);
 
-      int r = libusb_release_interface(m_usb_device_handle, get_interface());
-      if (r != 0){
-        MERROR("Could not release libusb interface: " << r);
-      }
+    int r = libusb_release_interface(m_usb_device_handle, get_interface());
+    if (r != 0){
+      MERROR("Could not release libusb interface: " << r);
+    }
 
-      m_usb_device = nullptr;
-      if (m_usb_device_handle) {
-        libusb_close(m_usb_device_handle);
-        m_usb_device_handle = nullptr;
-      }
+    m_usb_device = nullptr;
+    if (m_usb_device_handle) {
+      libusb_close(m_usb_device_handle);
+      m_usb_device_handle = nullptr;
+    }
 
-      if (m_usb_session) {
-        libusb_exit(m_usb_session);
-        m_usb_session = nullptr;
-      }
+    if (m_usb_session) {
+      libusb_exit(m_usb_session);
+      m_usb_session = nullptr;
     }
   };
 
+  std::shared_ptr<Transport> WebUsbTransport::find_debug() {
+#ifdef WITH_TREZOR_DEBUGGING
+    require_device();
+    auto t = std::make_shared<WebUsbTransport>(boost::make_optional(m_usb_device_desc.get()));
+    t->m_bus_id = m_bus_id;
+    t->m_device_addr = m_device_addr;
+    t->m_port_numbers = m_port_numbers;
+    t->m_debug_mode = true;
+    return t;
+#else
+      MINFO("Debug link is disabled in production");
+      return nullptr;
+#endif
+    }
 
   int WebUsbTransport::get_interface() const{
     const int INTERFACE_NORMAL = 0;
-#ifdef WITH_TREZOR_DEBUG
+#ifdef WITH_TREZOR_DEBUGGING
     const int INTERFACE_DEBUG = 1;
     return m_debug_mode ? INTERFACE_DEBUG : INTERFACE_NORMAL;
 #else
@@ -888,7 +973,7 @@ namespace trezor{
 
   unsigned char WebUsbTransport::get_endpoint() const{
     const unsigned char ENDPOINT_NORMAL = 1;
-#ifdef WITH_TREZOR_DEBUG
+#ifdef WITH_TREZOR_DEBUGGING
     const unsigned char ENDPOINT_DEBUG = 2;
     return m_debug_mode ? ENDPOINT_DEBUG : ENDPOINT_NORMAL;
 #else
@@ -1047,4 +1132,3 @@ namespace trezor{
 }
 }
 
-