Adds JWT refresh to http example. (GoogleCloudPlatform#906)

gguuss · web-flow · commit 476565a9e770 · 2017-11-02T15:58:20.000-07:00
diff --git a/iot/api-client/http_example/src/main/java/com/google/cloud/iot/examples/HttpExample.java b/iot/api-client/http_example/src/main/java/com/google/cloud/iot/examples/HttpExample.java
@@ -145,6 +145,7 @@ public static void main(String[] args) throws Exception {
 
     // Create the corresponding JWT depending on the selected algorithm.
     String token;
+    DateTime iat = new DateTime();
     if (options.algorithm.equals("RS256")) {
       token = createJwtRsa(options.projectId, options.privateKeyFile);
     } else if (options.algorithm.equals("ES256")) {
@@ -154,15 +155,28 @@ public static void main(String[] args) throws Exception {
           "Invalid algorithm " + options.algorithm + ". Should be one of 'RS256' or 'ES256'.");
     }
 
+    String urlPath = String.format("%s/%s/", options.httpBridgeAddress, options.apiVersion);
+    System.out.format("Using URL: '%s'\n", urlPath);
+
     // Publish numMessages messages to the HTTP bridge.
     for (int i = 1; i <= options.numMessages; ++i) {
       String payload = String.format("%s/%s-payload-%d", options.registryId, options.deviceId, i);
       System.out.format(
           "Publishing %s message %d/%d: '%s'\n",
           options.messageType, i, options.numMessages, payload);
 
-      String urlPath = String.format("%s/%s/", options.httpBridgeAddress, options.apiVersion);
-      System.out.format("Using URL: '%s'\n", urlPath);
+      // Refresh the authentication token if the token has expired.
+      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")) {
+          token = createJwtRsa(options.projectId, options.privateKeyFile);
+        } else if (options.algorithm.equals("ES256")) {
+          token = createJwtEs(options.projectId, options.privateKeyFile);
+        }
+      }
 
       publishMessage(payload, urlPath, options.messageType, token, options.projectId,
               options.cloudRegion, options.registryId, options.deviceId);
diff --git a/iot/api-client/http_example/src/main/java/com/google/cloud/iot/examples/HttpExampleOptions.java b/iot/api-client/http_example/src/main/java/com/google/cloud/iot/examples/HttpExampleOptions.java
@@ -30,6 +30,7 @@ public class HttpExampleOptions {
   String algorithm;
   String cloudRegion = "us-central1";
   int numMessages = 100;
+  int tokenExpMins = 20;
   String httpBridgeAddress = "https://cloudiot-device.googleapis.com";
   String apiVersion = "v1beta1";
   String messageType = "event";
@@ -94,6 +95,13 @@ public static HttpExampleOptions fromFlags(String[] args) {
             .hasArg()
             .desc("Number of messages to publish.")
             .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(String.class)
@@ -133,6 +141,10 @@ public static HttpExampleOptions 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("http_bridge_address")) {
         res.httpBridgeAddress = commandLine.getOptionValue("http_bridge_address");
       }
