Xiaomi: Refactor XiaomiCharacteristic to improve logging and ordering

Fixes a potential race condition on initialization, since the chunked
commands were being scheduled on a separate transaction builder, which
would be scheduled to be written before the initialization.

app/src/main/java/nodomain/freeyourgadget/gadgetbridge/service/devices/xiaomi/XiaomiAuthService.java

@@ -66,6 +66,8 @@ public class XiaomiAuthService extends AbstractXiaomiService {
     public static final int CMD_NONCE = 26;
     public static final int CMD_AUTH = 27;
 
+    private boolean encryptionInitialized = false;
+
     private final byte[] secretKey = new byte[16];
     private final byte[] nonce = new byte[16];
     private final byte[] encryptionKey = new byte[16];
@@ -77,17 +79,19 @@ public class XiaomiAuthService extends AbstractXiaomiService {
         super(support);
     }
 
+    public boolean isEncryptionInitialized() {
+        return encryptionInitialized;
+    }
+
     protected void startEncryptedHandshake(final TransactionBuilder builder) {
+        encryptionInitialized = false;
+
         builder.add(new SetDeviceStateAction(getSupport().getDevice(), GBDevice.State.AUTHENTICATING, getSupport().getContext()));
 
         System.arraycopy(getSecretKey(getSupport().getDevice()), 0, secretKey, 0, 16);
         new SecureRandom().nextBytes(nonce);
 
-        // TODO use sendCommand
+        getSupport().sendCommand(builder, buildNonceCommand(nonce));
-        builder.write(
-                getSupport().getCharacteristic(XiaomiEncryptedSupport.UUID_CHARACTERISTIC_XIAOMI_COMMAND_WRITE),
-                ArrayUtils.addAll(PAYLOAD_HEADER_AUTH, buildNonceCommand(nonce))
-        );
     }
 
     protected void startClearTextHandshake(final TransactionBuilder builder, String userId) {
@@ -103,7 +107,7 @@ public class XiaomiAuthService extends AbstractXiaomiService {
                 .setAuth(auth)
                 .build();
 
-        getSupport().sendCommand("start clear text handshake", command);
+        getSupport().sendCommand(builder, command);
     }
 
     @Override
@@ -138,6 +142,8 @@ public class XiaomiAuthService extends AbstractXiaomiService {
                 if (cmd.getSubtype() == CMD_AUTH || cmd.getAuth().getStatus() == 1) {
                     LOG.info("Authenticated!");
 
+                    encryptionInitialized = cmd.getSubtype() == CMD_AUTH;
+
                     final TransactionBuilder builder = getSupport().createTransactionBuilder("phase 2 initialize");
                     builder.add(new SetDeviceStateAction(getSupport().getDevice(), GBDevice.State.INITIALIZED, getSupport().getContext()));
                     getSupport().phase2Initialize();
@@ -227,7 +233,7 @@ public class XiaomiAuthService extends AbstractXiaomiService {
         return cmd.setAuth(auth.build()).build();
     }
 
-    public static byte[] buildNonceCommand(final byte[] nonce) {
+    public static XiaomiProto.Command buildNonceCommand(final byte[] nonce) {
         final XiaomiProto.PhoneNonce.Builder phoneNonce = XiaomiProto.PhoneNonce.newBuilder();
         phoneNonce.setNonce(ByteString.copyFrom(nonce));
 
@@ -238,7 +244,7 @@ public class XiaomiAuthService extends AbstractXiaomiService {
         command.setType(COMMAND_TYPE);
         command.setSubtype(CMD_NONCE);
         command.setAuth(auth.build());
-        return command.build().toByteArray();
+        return command.build();
     }
 
     public static byte[] computeAuthStep3Hmac(final byte[] secretKey,

app/src/main/java/nodomain/freeyourgadget/gadgetbridge/service/devices/xiaomi/XiaomiCharacteristic.java

@@ -50,8 +50,8 @@ public class XiaomiCharacteristic {
     private final UUID characteristicUUID;
 
     // Encryption
-    private XiaomiAuthService authService = null;
+    private final XiaomiAuthService authService;
-    private boolean isEncrypted = false;
+    private boolean isEncrypted;
     private short encryptedIndex = 0;
 
     // Chunking
@@ -61,10 +61,10 @@ public class XiaomiCharacteristic {
 
     // Scheduling
     // TODO timeouts
-    private final Queue<byte[]> payloadQueue = new LinkedList<>();
+    private final Queue<Payload> payloadQueue = new LinkedList<>();
     private boolean waitingAck = false;
     private boolean sendingChunked = false;
-    private byte[] currentSending = null;
+    private Payload currentPayload = null;
 
     private Handler handler = null;
 
@@ -99,30 +99,35 @@ public class XiaomiCharacteristic {
         this.payloadQueue.clear();
         this.waitingAck = false;
         this.sendingChunked = false;
-        this.currentSending = null;
+        this.currentPayload = null;
     }
 
     /**
      * Write bytes to this characteristic, encrypting and splitting it into chunks if necessary.
      */
-    public void write(final byte[] value) {
+    public void write(final String taskName, final byte[] value) {
-        payloadQueue.add(value);
+        write(null, new Payload(taskName, value));
-        sendNext();
     }
 
     /**
-     * Write bytes to this characteristic directly.
+     * Write bytes to this characteristic, encrypting and splitting it into chunks if necessary. Uses
+     * the provided if we need to schedule something, otherwise it will be queued as other commands.
      */
-    public void writeDirect(final TransactionBuilder builder, final byte[] value) {
+    public void write(final TransactionBuilder builder, final byte[] value) {
-        builder.write(bluetoothGattCharacteristic, value);
+        write(builder, new Payload(builder.getTaskName(), value));
+    }
+
+    private void write(final TransactionBuilder builder, final Payload payload) {
+        payloadQueue.add(payload);
+        sendNext(builder);
     }
 
     public void onCharacteristicChanged(final byte[] value) {
         if (Arrays.equals(value, PAYLOAD_ACK)) {
             LOG.debug("Got ack");
-            currentSending = null;
+            currentPayload = null;
             waitingAck = false;
-            sendNext();
+            sendNext(null);
             return;
         }
 
@@ -178,34 +183,35 @@ public class XiaomiCharacteristic {
                     switch (subtype) {
                         case 0:
                             LOG.debug("Got chunked ack end");
-                            currentSending = null;
+                            currentPayload = null;
                             sendingChunked = false;
-                            sendNext();
+                            sendNext(null);
                             return;
                         case 1:
                             LOG.debug("Got chunked ack start");
-                            final TransactionBuilder builder = mSupport.createTransactionBuilder("send chunks");
+                            final TransactionBuilder builder = mSupport.createTransactionBuilder("send chunks for " + currentPayload.getTaskName());
-                            for (int i = 0; i * MAX_WRITE_SIZE < currentSending.length; i++) {
+                            final byte[] payload = currentPayload.getBytesToSend();
+                            for (int i = 0; i * MAX_WRITE_SIZE < payload.length; i++) {
                                 final int startIndex = i * MAX_WRITE_SIZE;
-                                final int endIndex = Math.min((i + 1) * MAX_WRITE_SIZE, currentSending.length);
+                                final int endIndex = Math.min((i + 1) * MAX_WRITE_SIZE, payload.length);
-                                LOG.debug("Sending chunk {} from {} to {}", i, startIndex, endIndex);
+                                LOG.debug("Sending chunk {} from {} to {} for {}", i, startIndex, endIndex, currentPayload.getTaskName());
                                 final byte[] chunkToSend = new byte[2 + endIndex - startIndex];
                                 BLETypeConversions.writeUint16(chunkToSend, 0, i + 1);
-                                System.arraycopy(currentSending, startIndex, chunkToSend, 2, endIndex - startIndex);
+                                System.arraycopy(payload, startIndex, chunkToSend, 2, endIndex - startIndex);
                                 builder.write(bluetoothGattCharacteristic, chunkToSend);
                             }
 
                             builder.queue(mSupport.getQueue());
                             return;
                         case 2:
-                            LOG.warn("Got chunked nack");
+                            LOG.warn("Got chunked nack for {}", currentPayload.getTaskName());
-                            currentSending = null;
+                            currentPayload = null;
                             sendingChunked = false;
-                            sendNext();
+                            sendNext(null);
                             return;
                     }
 
-                    LOG.warn("Unknown chunked ack subtype {}", subtype);
+                    LOG.warn("Unknown chunked ack subtype {} for {}", subtype, currentPayload.getTaskName());
 
                     return;
                 case 2:
@@ -233,66 +239,72 @@ public class XiaomiCharacteristic {
         }
     }
 
-    private void sendNext() {
+    private void sendNext(@Nullable final TransactionBuilder b) {
         if (waitingAck || sendingChunked) {
             LOG.debug("Already sending something");
             return;
         }
 
-        final byte[] payload = payloadQueue.poll();
+        currentPayload = payloadQueue.poll();
-        if (payload == null) {
+        if (currentPayload == null) {
             LOG.debug("Nothing to send");
             return;
         }
 
-        if (isEncrypted) {
+        final boolean encrypt = isEncrypted && authService.isEncryptionInitialized();
-            currentSending = authService.encrypt(payload, encryptedIndex);
+
-        } else {
+        if (encrypt) {
-            currentSending = payload;
+            currentPayload.setBytesToSend(authService.encrypt(currentPayload.getBytesToSend(), encryptedIndex));
         }
 
-        if (shouldWriteChunked(currentSending)) {
+        if (shouldWriteChunked(currentPayload.getBytesToSend())) {
-            if (isEncrypted) {
+            if (encrypt) {
                 // Prepend encrypted index for the nonce
-                currentSending = ByteBuffer.allocate(2 + currentSending.length).order(ByteOrder.LITTLE_ENDIAN)
+                currentPayload.setBytesToSend(
-                        .putShort(encryptedIndex++)
+                        ByteBuffer.allocate(2 + currentPayload.getBytesToSend().length).order(ByteOrder.LITTLE_ENDIAN)
-                        .put(currentSending)
+                                .putShort(encryptedIndex++)
-                        .array();
+                                .put(currentPayload.getBytesToSend())
+                                .array()
+                );
             }
 
-            LOG.debug("Sending next - chunked");
+            LOG.debug("Sending {} - chunked", currentPayload.getTaskName());
 
             sendingChunked = true;
 
             final ByteBuffer buf = ByteBuffer.allocate(6).order(ByteOrder.LITTLE_ENDIAN);
             buf.putShort((short) 0);
             buf.put((byte) 0);
-            buf.put((byte) (isEncrypted ? 1 : 0));
+            buf.put((byte) (encrypt ? 1 : 0));
-            buf.putShort((short) Math.ceil(currentSending.length / (float) MAX_WRITE_SIZE));
+            buf.putShort((short) Math.ceil(currentPayload.getBytesToSend().length / (float) MAX_WRITE_SIZE));
 
-            final TransactionBuilder builder = mSupport.createTransactionBuilder("send chunked start");
+            final TransactionBuilder builder = b == null ? mSupport.createTransactionBuilder("send chunked start for " + currentPayload.getTaskName()) : b;
             builder.write(bluetoothGattCharacteristic, buf.array());
-            builder.queue(mSupport.getQueue());
+            if (b == null) {
+                builder.queue(mSupport.getQueue());
+            }
         } else {
-            LOG.debug("Sending next - single");
+            LOG.debug("Sending {} - single", currentPayload.getTaskName());
 
             // Encrypt single command
-            final int commandLength = (isEncrypted ? 6 : 4) + currentSending.length;
+            final int commandLength = (encrypt ? 6 : 4) + currentPayload.getBytesToSend().length;
 
             final ByteBuffer buf = ByteBuffer.allocate(commandLength).order(ByteOrder.LITTLE_ENDIAN);
             buf.putShort((short) 0);
             buf.put((byte) 2); // 2 for command
-            buf.put((byte) (isEncrypted ? 1 : 2));
+            buf.put((byte) (encrypt ? 1 : 2));
-            if (isEncrypted) {
+            if (encrypt) {
                 buf.putShort(encryptedIndex++);
             }
-            buf.put(currentSending); // it's already encrypted
+            buf.put(currentPayload.getBytesToSend()); // it's already encrypted
 
             waitingAck = true;
 
-            final TransactionBuilder builder = mSupport.createTransactionBuilder("send single command");
+            final TransactionBuilder builder = b == null ? mSupport.createTransactionBuilder("send single command for " + currentPayload.getTaskName()) : b;
             builder.write(bluetoothGattCharacteristic, buf.array());
-            builder.queue(mSupport.getQueue());
+            if (b == null) {
+                builder.queue(mSupport.getQueue());
+            }
         }
     }
 
@@ -327,4 +339,29 @@ public class XiaomiCharacteristic {
     public interface Handler {
         void handle(final byte[] payload);
     }
+
+    private static class Payload {
+        private final String taskName;
+        private final byte[] bytes;
+
+        // Bytes that will actually be sent (might be encrypted)
+        private byte[] bytesToSend;
+
+        public Payload(final String taskName, final byte[] bytes) {
+            this.taskName = taskName;
+            this.bytes = bytes;
+        }
+
+        public String getTaskName() {
+            return taskName;
+        }
+
+        public void setBytesToSend(final byte[] bytesToSend) {
+            this.bytesToSend = bytesToSend;
+        }
+
+        public byte[] getBytesToSend() {
+            return bytesToSend != null ? bytesToSend : bytes;
+        }
+    }
 }

app/src/main/java/nodomain/freeyourgadget/gadgetbridge/service/devices/xiaomi/XiaomiSupport.java

@@ -415,10 +415,22 @@ public abstract class XiaomiSupport extends AbstractBTLEDeviceSupport {
             return;
         }
 
-        // FIXME builder is ignored
+        this.characteristicCommandWrite.write(taskName, command.toByteArray());
-        final byte[] commandBytes = command.toByteArray();
+    }
-        LOG.debug("Sending command {}", GB.hexdump(commandBytes));
+
-        this.characteristicCommandWrite.write(commandBytes);
+    /**
+     * Realistically, this function should only be used during auth, as we must schedule the command after
+     * notifications were enabled on the characteristics, and for that we need the builder to guarantee the
+     * order.
+     */
+    public void sendCommand(final TransactionBuilder builder, final XiaomiProto.Command command) {
+        if (this.characteristicCommandWrite == null) {
+            // Can sometimes happen in race conditions when connecting + receiving calendar event or weather updates
+            LOG.warn("characteristicCommandWrite is null!");
+            return;
+        }
+
+        this.characteristicCommandWrite.write(builder, command.toByteArray());
     }
 
     public void sendCommand(final String taskName, final int type, final int subtype) {