From 46724e29848fc0799dfb9c0ccc6c68d3da419707 Mon Sep 17 00:00:00 2001
From: Gus Class <gguuss@gmail.com>
Date: Fri, 17 Nov 2017 12:53:44 -0800
Subject: [PATCH] Adds credential refresh to MQTT example (#917)

---
 .../cloud/iot/examples/MqttExample.java       | 20 ++++++++++++++++++-
 .../iot/examples/MqttExampleOptions.java      | 12 +++++++++++
 2 files changed, 31 insertions(+), 1 deletion(-)

diff --git a/iot/api-client/mqtt_example/src/main/java/com/google/cloud/iot/examples/MqttExample.java b/iot/api-client/mqtt_example/src/main/java/com/google/cloud/iot/examples/MqttExample.java
index 5c21f7aceda..524a519abf1 100644
--- a/iot/api-client/mqtt_example/src/main/java/com/google/cloud/iot/examples/MqttExample.java
+++ b/iot/api-client/mqtt_example/src/main/java/com/google/cloud/iot/examples/MqttExample.java
@@ -127,7 +127,7 @@ public static void main(String[] args) throws Exception {
     // to authorize the device.
     connectOptions.setUserName("unused");
 
-    System.out.println(options.algorithm);
+    DateTime iat = new DateTime();
     if (options.algorithm.equals("RS256")) {
       connectOptions.setPassword(
           createJwtRsa(options.projectId, options.privateKeyFile).toCharArray());
@@ -161,6 +161,24 @@ public static void main(String[] args) throws Exception {
             "Publishing %s message %d/%d: '%s'\n",
             options.messageType, i, options.numMessages, payload);
 
+        // Refresh the connection credentials before the JWT expires.
+        long secsSinceRefresh = ((new DateTime()).getMillis() - iat.getMillis()) / 1000;
+        if (secsSinceRefresh > (options.tokenExpMins * 60)) {
+          System.out.format("\tRefreshing token after: %d seconds\n", secsSinceRefresh);
+          iat = new DateTime();
+          if (options.algorithm.equals("RS256")) {
+            connectOptions.setPassword(
+                createJwtRsa(options.projectId, options.privateKeyFile).toCharArray());
+          } else if (options.algorithm.equals("ES256")) {
+            connectOptions.setPassword(
+                createJwtEs(options.projectId, options.privateKeyFile).toCharArray());
+          } else {
+            throw new IllegalArgumentException(
+                "Invalid algorithm " + options.algorithm
+                    + ". Should be one of 'RS256' or 'ES256'.");
+          }
+        }
+
         // Publish "payload" to the MQTT topic. qos=1 means at least once delivery. Cloud IoT Core
         // also supports qos=0 for at most once delivery.
         MqttMessage message = new MqttMessage(payload.getBytes());
diff --git a/iot/api-client/mqtt_example/src/main/java/com/google/cloud/iot/examples/MqttExampleOptions.java b/iot/api-client/mqtt_example/src/main/java/com/google/cloud/iot/examples/MqttExampleOptions.java
index 102368acb7a..2dd3c0b547a 100644
--- a/iot/api-client/mqtt_example/src/main/java/com/google/cloud/iot/examples/MqttExampleOptions.java
+++ b/iot/api-client/mqtt_example/src/main/java/com/google/cloud/iot/examples/MqttExampleOptions.java
@@ -30,6 +30,7 @@ public class MqttExampleOptions {
   String algorithm;
   String cloudRegion = "us-central1";
   int numMessages = 100;
+  int tokenExpMins = 20;
   String mqttBridgeHostname = "mqtt.googleapis.com";
   short mqttBridgePort = 8883;
   String messageType = "event";
@@ -101,6 +102,13 @@ public static MqttExampleOptions fromFlags(String[] args) {
             .hasArg()
             .desc("MQTT bridge hostname.")
             .build());
+    options.addOption(
+        Option.builder()
+            .type(Number.class)
+            .longOpt("token_exp_minutes")
+            .hasArg()
+            .desc("Minutes to JWT token refresh (token expiration time).")
+            .build());
     options.addOption(
         Option.builder()
             .type(Number.class)
@@ -133,6 +141,10 @@ public static MqttExampleOptions fromFlags(String[] args) {
       if (commandLine.hasOption("num_messages")) {
         res.numMessages = ((Number) commandLine.getParsedOptionValue("num_messages")).intValue();
       }
+      if (commandLine.hasOption("token_exp_minutes")) {
+        res.tokenExpMins =
+            ((Number) commandLine.getParsedOptionValue("token_exp_minutes")).intValue();
+      }
       if (commandLine.hasOption("mqtt_bridge_hostname")) {
         res.mqttBridgeHostname = commandLine.getOptionValue("mqtt_bridge_hostname");
       }