From 9220b25043bf0a60c091dc37dafa4a95ebc6cf13 Mon Sep 17 00:00:00 2001
From: Greg V <greg@unrelenting.technology>
Date: Thu, 21 Jan 2021 20:19:20 +0300
Subject: [PATCH] Add API for (un)registering connections

At least the J-Link clone I have (serial 123456, HW 7.00, FW
"J-Link ARM-OB STM32 compiled Aug 22 2012 19:52:04") absolutely requires
a registered connection, at least for SWD usage. Nothing ever gets ACKed otherwise.

OpenOCD always calls jaylink_register (with all zeros for info)
when the REGISTER capability is present, and jaylink_unregister on exit.
---
 src/lib.rs | 71 +++++++++++++++++++++++++++++++++++++++++++++++++++++-
 1 file changed, 70 insertions(+), 1 deletion(-)

diff --git a/src/lib.rs b/src/lib.rs
index d9e8fee..557c2f4 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -93,7 +93,7 @@ pub use self::interface::{Interface, InterfaceIter, Interfaces};
 use self::bits::IteratorExt as _;
 use self::error::ResultExt as _;
 use bitflags::bitflags;
-use byteorder::{LittleEndian, ReadBytesExt};
+use byteorder::{BigEndian, LittleEndian, ReadBytesExt, WriteBytesExt};
 use io::Cursor;
 use log::{debug, trace, warn};
 use std::cell::{Cell, RefCell, RefMut};
@@ -162,6 +162,15 @@ enum Command {
 
     ReadConfig = 0xF2,
     WriteConfig = 0xF3,
+
+    Register = 0x09,
+}
+
+#[repr(u8)]
+#[allow(dead_code)]
+enum RegisterCommand {
+    Register = 0x64,
+    Unregister = 0x65,
 }
 
 #[repr(u8)]
@@ -208,6 +217,15 @@ impl SwoStatus {
     }
 }
 
+#[derive(Default, Debug)]
+pub struct Connection {
+    pub handle: u16,
+    pub pid: u32,
+    pub hid: Option<std::net::Ipv4Addr>,
+    pub iid: u8,
+    pub cid: u8,
+}
+
 /// A handle to a J-Link USB device.
 ///
 /// This is the main interface type of this library. There are multiple ways of obtaining an
@@ -882,6 +900,57 @@ impl JayLink {
         Ok(())
     }
 
+    /// Registers a connection on the device.
+    /// Updates the Connection with the returned handle number.
+    ///
+    /// This requires the [`Register`] capability.
+    ///
+    /// **Note**: This may be **required** on some devices for SWD to work at all.
+    ///
+    /// [`Register`]: Capability::Register
+    pub fn register(&self, conn: &mut Connection) -> Result<()> {
+        let handle = self.raw_register(conn, RegisterCommand::Register)?;
+        conn.handle = handle;
+        Ok(())
+    }
+
+    /// Unregisters a connection on the device.
+    ///
+    /// This requires the [`Register`] capability.
+    ///
+    /// [`Register`]: Capability::Register
+    pub fn unregister(&self, conn: &Connection) -> Result<()> {
+        self.raw_register(conn, RegisterCommand::Unregister).map(|_| ())
+    }
+
+    fn raw_register(&self, conn: &Connection, cmd: RegisterCommand) -> Result<u16> {
+        self.require_capability(Capability::Register)?;
+        let mut buf = Vec::with_capacity(14);
+        buf.push(Command::Register as u8);
+        buf.push(cmd as u8);
+        buf.write_u32::<LittleEndian>(conn.pid).unwrap();
+        buf.write_u32::<BigEndian>(conn.hid.unwrap_or(std::net::Ipv4Addr::UNSPECIFIED).into())
+            .unwrap();
+        buf.push(conn.iid);
+        buf.push(conn.cid);
+        buf.write_u16::<LittleEndian>(conn.handle).unwrap();
+        self.write_cmd(&buf)?;
+
+        // Receive connection table
+        let mut buf = [0; 76];
+        self.read(&mut buf)?;
+        let new_handle = u16::from_le_bytes([buf[0], buf[1]]);
+        let num_entries = u16::from_le_bytes([buf[2], buf[3]]);
+        let entry_size = u16::from_le_bytes([buf[4], buf[5]]);
+        let info_size = u16::from_le_bytes([buf[6], buf[7]]);
+        let all_size = 8 + (num_entries * entry_size) + info_size;
+        if all_size > 76 {
+            let mut extra_buf = vec![0; (all_size - 76) as usize];
+            self.read(&mut extra_buf)?;
+        }
+        Ok(new_handle)
+    }
+
     /// Performs a JTAG I/O operation.
     ///
     /// This will shift out data on `TMS` (pin 7) and `TDI` (pin 5), while reading data shifted