From 51730290c6cf97c4c23938186d2e1dff096e9f26 Mon Sep 17 00:00:00 2001
From: ajaaym <34161822+ajaaym@users.noreply.github.com>
Date: Tue, 7 May 2019 17:47:14 -0400
Subject: [PATCH] Example for pubsub authenticated push (#1407)

---
 appengine-java8/pubsub/README.md              |  24 ++++
 appengine-java8/pubsub/pom.xml                |   5 +
 .../appengine/pubsub/MessageRepository.java   |  24 +++-
 .../pubsub/MessageRepositoryImpl.java         |  74 +++++++++++-
 .../pubsub/PubSubAuthenticatedPush.java       | 112 ++++++++++++++++++
 .../example/appengine/pubsub/PubSubHome.java  |  99 ++++++++++++++--
 .../appengine/pubsub/PubSubPublish.java       |  11 +-
 .../example/appengine/pubsub/PubSubPush.java  |   2 +
 .../pubsub/src/main/webapp/index.jsp          |  24 +---
 9 files changed, 327 insertions(+), 48 deletions(-)
 create mode 100644 appengine-java8/pubsub/src/main/java/com/example/appengine/pubsub/PubSubAuthenticatedPush.java

diff --git a/appengine-java8/pubsub/README.md b/appengine-java8/pubsub/README.md
index c3294dcfb5e..2cb9314f3f3 100644
--- a/appengine-java8/pubsub/README.md
+++ b/appengine-java8/pubsub/README.md
@@ -53,6 +53,21 @@ gcloud beta pubsub subscriptions create <your-subscription-name> \
   --ack-deadline 30
 ```
 
+- Create a subscription for authenticated pushes to send messages to a Google Cloud Project URL such as https://<your-project-id>.appspot.com/authenticated-push.
+
+The push auth service account must have Service Account Token Creator Role assigned, which can be done in the Cloud Console [IAM & admin](https://console.cloud.google.com/iam-admin/iam) UI.
+`--push-auth-token-audience` is optional. If set, remember to modify the audience field check in [PubSubAuthenticatedPush.java](src/main/java/com/example/appengine/pubsub/PubSubAuthenticatedPush.java#L36).
+
+```
+gcloud beta pubsub subscriptions create <your-subscription-name> \
+  --topic <your-topic-name> \
+  --push-endpoint \
+    https://<your-project-id>.appspot.com/pubsub/authenticated-push?token=<your-verification-token> \
+  --ack-deadline 30 \
+  --push-auth-service-account=[your-service-account-email] \
+  --push-auth-token-audience=example.com
+```
+
 ## Run locally
 Set the following environment variables and run using shown Maven command. You can then
 direct your browser to `http://localhost:8080/`
@@ -70,6 +85,15 @@ mvn appengine:run
    "localhost:8080/pubsub/push?token=<your-token>"
 ```
 
+### Authenticated push notifications
+
+Simulating authenticated push requests will fail because requests need to contain a Cloud Pub/Sub-generated JWT in the "Authorization" header.
+
+```
+   curl -H "Content-Type: application/json" -i --data @sample_message.json
+   "localhost:8080/pubsub/authenticated-push?token=<your-token>"
+```
+
 ## Deploy
 
 Update the environment variables `PUBSUB_TOPIC` and `PUBSUB_VERIFICATION_TOKEN` in
diff --git a/appengine-java8/pubsub/pom.xml b/appengine-java8/pubsub/pom.xml
index fc40186461e..d3e9750c666 100644
--- a/appengine-java8/pubsub/pom.xml
+++ b/appengine-java8/pubsub/pom.xml
@@ -47,6 +47,11 @@
       <type>jar</type>
       <scope>provided</scope>
     </dependency>
+    <dependency>
+      <groupId>com.googlecode.jatl</groupId>
+      <artifactId>jatl</artifactId>
+      <version>0.2.2</version>
+    </dependency>
 
     <!-- [START dependencies] -->
     <dependency>
diff --git a/appengine-java8/pubsub/src/main/java/com/example/appengine/pubsub/MessageRepository.java b/appengine-java8/pubsub/src/main/java/com/example/appengine/pubsub/MessageRepository.java
index c2fc2f78761..c4defa3b2da 100644
--- a/appengine-java8/pubsub/src/main/java/com/example/appengine/pubsub/MessageRepository.java
+++ b/appengine-java8/pubsub/src/main/java/com/example/appengine/pubsub/MessageRepository.java
@@ -14,7 +14,6 @@
  * limitations under the License.
  */
 
-
 package com.example.appengine.pubsub;
 
 import java.util.List;
@@ -26,8 +25,31 @@ public interface MessageRepository {
 
   /**
    * Retrieve most recent stored messages.
+   *
    * @param limit number of messages
    * @return list of messages
    */
   List<Message> retrieve(int limit);
+
+  /** Save claim to persistent storage. */
+  void saveClaim(String claim);
+
+  /**
+   * Retrieve most recent stored claims.
+   *
+   * @param limit number of messages
+   * @return list of claims
+   */
+  List<String> retrieveClaims(int limit);
+
+  /** Save token to persistent storage. */
+  void saveToken(String token);
+
+  /**
+   * Retrieve most recent stored tokens.
+   *
+   * @param limit number of messages
+   * @return list of tokens
+   */
+  List<String> retrieveTokens(int limit);
 }
diff --git a/appengine-java8/pubsub/src/main/java/com/example/appengine/pubsub/MessageRepositoryImpl.java b/appengine-java8/pubsub/src/main/java/com/example/appengine/pubsub/MessageRepositoryImpl.java
index 521c82f13dd..85a66bdfc74 100644
--- a/appengine-java8/pubsub/src/main/java/com/example/appengine/pubsub/MessageRepositoryImpl.java
+++ b/appengine-java8/pubsub/src/main/java/com/example/appengine/pubsub/MessageRepositoryImpl.java
@@ -14,7 +14,6 @@
  * limitations under the License.
  */
 
-
 package com.example.appengine.pubsub;
 
 import com.google.cloud.datastore.Datastore;
@@ -35,6 +34,12 @@ public class MessageRepositoryImpl implements MessageRepository {
 
   private String messagesKind = "messages";
   private KeyFactory keyFactory = getDatastoreInstance().newKeyFactory().setKind(messagesKind);
+  private String claimsKind = "claims";
+  private KeyFactory claimsKindKeyFactory =
+      getDatastoreInstance().newKeyFactory().setKind(claimsKind);
+  private String tokensKind = "tokens";
+  private KeyFactory tokensKindKeyFactory =
+      getDatastoreInstance().newKeyFactory().setKind(tokensKind);
 
   @Override
   public void save(Message message) {
@@ -42,8 +47,8 @@ public void save(Message message) {
     Datastore datastore = getDatastoreInstance();
     Key key = datastore.allocateId(keyFactory.newKey());
 
-    Entity.Builder messageEntityBuilder = Entity.newBuilder(key)
-        .set("messageId", message.getMessageId());
+    Entity.Builder messageEntityBuilder =
+        Entity.newBuilder(key).set("messageId", message.getMessageId());
 
     if (message.getData() != null) {
       messageEntityBuilder = messageEntityBuilder.set("data", message.getData());
@@ -84,15 +89,72 @@ public List<Message> retrieve(int limit) {
     return messages;
   }
 
+  @Override
+  public void saveClaim(String claim) {
+    // Save message to "messages"
+    Datastore datastore = getDatastoreInstance();
+    Key key = datastore.allocateId(claimsKindKeyFactory.newKey());
+
+    Entity.Builder claimEntityBuilder = Entity.newBuilder(key).set("claim", claim);
+
+    datastore.put(claimEntityBuilder.build());
+  }
+
+  @Override
+  public List<String> retrieveClaims(int limit) {
+    // Get claim saved in Datastore
+    Datastore datastore = getDatastoreInstance();
+    Query<Entity> query = Query.newEntityQueryBuilder().setKind(claimsKind).setLimit(limit).build();
+    QueryResults<Entity> results = datastore.run(query);
+
+    List<String> claims = new ArrayList<>();
+    while (results.hasNext()) {
+      Entity entity = results.next();
+      String claim = entity.getString("claim");
+      if (claim != null) {
+        claims.add(claim);
+      }
+    }
+    return claims;
+  }
+
+  @Override
+  public void saveToken(String token) {
+    // Save message to "messages"
+    Datastore datastore = getDatastoreInstance();
+    Key key = datastore.allocateId(tokensKindKeyFactory.newKey());
+
+    Entity.Builder tokenEntityBuilder = Entity.newBuilder(key).set("token", token);
+
+    datastore.put(tokenEntityBuilder.build());
+  }
+
+  @Override
+  public List<String> retrieveTokens(int limit) {
+    // Get token saved in Datastore
+    Datastore datastore = getDatastoreInstance();
+    Query<Entity> query = Query.newEntityQueryBuilder().setKind(tokensKind).setLimit(limit).build();
+    QueryResults<Entity> results = datastore.run(query);
+
+    List<String> tokens = new ArrayList<>();
+    while (results.hasNext()) {
+      Entity entity = results.next();
+      String token = entity.getString("token");
+      if (token != null) {
+        tokens.add(token);
+      }
+    }
+    return tokens;
+  }
+
   private Datastore getDatastoreInstance() {
     return DatastoreOptions.getDefaultInstance().getService();
   }
 
-  private MessageRepositoryImpl() {
-  }
+  private MessageRepositoryImpl() {}
 
   // retrieve a singleton instance
-  public static synchronized  MessageRepositoryImpl getInstance() {
+  public static synchronized MessageRepositoryImpl getInstance() {
     if (instance == null) {
       instance = new MessageRepositoryImpl();
     }
diff --git a/appengine-java8/pubsub/src/main/java/com/example/appengine/pubsub/PubSubAuthenticatedPush.java b/appengine-java8/pubsub/src/main/java/com/example/appengine/pubsub/PubSubAuthenticatedPush.java
new file mode 100644
index 00000000000..532bae32c95
--- /dev/null
+++ b/appengine-java8/pubsub/src/main/java/com/example/appengine/pubsub/PubSubAuthenticatedPush.java
@@ -0,0 +1,112 @@
+/*
+ * Copyright 2019 Google LLC
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.example.appengine.pubsub;
+
+import com.google.api.client.googleapis.auth.oauth2.GoogleIdToken;
+import com.google.api.client.googleapis.auth.oauth2.GoogleIdTokenVerifier;
+import com.google.api.client.http.javanet.NetHttpTransport;
+import com.google.api.client.json.jackson2.JacksonFactory;
+import com.google.gson.Gson;
+import com.google.gson.JsonElement;
+import com.google.gson.JsonParser;
+import java.io.IOException;
+import java.util.Base64;
+import java.util.Collections;
+import java.util.stream.Collectors;
+import javax.servlet.ServletException;
+import javax.servlet.annotation.WebServlet;
+import javax.servlet.http.HttpServlet;
+import javax.servlet.http.HttpServletRequest;
+import javax.servlet.http.HttpServletResponse;
+
+// [START gae_standard_pubsub_auth_push]
+@WebServlet(value = "/pubsub/authenticated-push")
+public class PubSubAuthenticatedPush extends HttpServlet {
+  private final String pubsubVerificationToken = System.getenv("PUBSUB_VERIFICATION_TOKEN");
+  private final MessageRepository messageRepository;
+  private final GoogleIdTokenVerifier verifier =
+      new GoogleIdTokenVerifier.Builder(new NetHttpTransport(), new JacksonFactory())
+          /**
+           * Please change example.com to match with value you are providing while creating
+           * subscription as provided in @see <a
+           * href="https://github.com/GoogleCloudPlatform/java-docs-samples/tree/master/appengine-java8/pubsub">README</a>.
+           */
+          .setAudience(Collections.singletonList("example.com"))
+          .build();
+  private final Gson gson = new Gson();
+  private final JsonParser jsonParser = new JsonParser();
+
+  @Override
+  public void doPost(HttpServletRequest req, HttpServletResponse resp)
+      throws IOException, ServletException {
+
+    // Verify that the request originates from the application.
+    if (req.getParameter("token").compareTo(pubsubVerificationToken) != 0) {
+      resp.setStatus(HttpServletResponse.SC_BAD_REQUEST);
+      return;
+    }
+    // Get the Cloud Pub/Sub-generated JWT in the "Authorization" header.
+    String authorizationHeader = req.getHeader("Authorization");
+    if (authorizationHeader == null
+        || authorizationHeader.isEmpty()
+        || authorizationHeader.split(" ").length != 2) {
+      resp.setStatus(HttpServletResponse.SC_BAD_REQUEST);
+      return;
+    }
+    String authorization = authorizationHeader.split(" ")[1];
+
+    try {
+      // Verify and decode the JWT.
+      GoogleIdToken idToken = verifier.verify(authorization);
+      messageRepository.saveToken(authorization);
+      messageRepository.saveClaim(idToken.getPayload().toPrettyString());
+      // parse message object from "message" field in the request body json
+      // decode message data from base64
+      Message message = getMessage(req);
+      messageRepository.save(message);
+      // 200, 201, 204, 102 status codes are interpreted as success by the Pub/Sub system
+      resp.setStatus(102);
+      super.doPost(req, resp);
+    } catch (Exception e) {
+      resp.setStatus(HttpServletResponse.SC_BAD_REQUEST);
+    }
+  }
+
+  private Message getMessage(HttpServletRequest request) throws IOException {
+    String requestBody = request.getReader().lines().collect(Collectors.joining("\n"));
+    JsonElement jsonRoot = jsonParser.parse(requestBody);
+    String messageStr = jsonRoot.getAsJsonObject().get("message").toString();
+    Message message = gson.fromJson(messageStr, Message.class);
+    // decode from base64
+    String decoded = decode(message.getData());
+    message.setData(decoded);
+    return message;
+  }
+
+  private String decode(String data) {
+    return new String(Base64.getDecoder().decode(data));
+  }
+
+  PubSubAuthenticatedPush(MessageRepository messageRepository) {
+    this.messageRepository = messageRepository;
+  }
+
+  public PubSubAuthenticatedPush() {
+    this(MessageRepositoryImpl.getInstance());
+  }
+}
+// [END gae_standard_pubsub_auth_push]
diff --git a/appengine-java8/pubsub/src/main/java/com/example/appengine/pubsub/PubSubHome.java b/appengine-java8/pubsub/src/main/java/com/example/appengine/pubsub/PubSubHome.java
index be9ab2546e6..105add23235 100644
--- a/appengine-java8/pubsub/src/main/java/com/example/appengine/pubsub/PubSubHome.java
+++ b/appengine-java8/pubsub/src/main/java/com/example/appengine/pubsub/PubSubHome.java
@@ -16,6 +16,9 @@
 
 package com.example.appengine.pubsub;
 
+import com.googlecode.jatl.Html;
+import java.io.StringWriter;
+import java.io.Writer;
 import java.util.List;
 
 public class PubSubHome {
@@ -28,22 +31,92 @@ public class PubSubHome {
    *
    * @return html representation of messages (one per row)
    */
-  public static String getReceivedMessages() {
+  public static List<Message> getReceivedMessages() {
     List<Message> messageList = messageRepository.retrieve(MAX_MESSAGES);
-    return convertToHtmlTable(messageList);
+    return messageList;
   }
 
-  private static String convertToHtmlTable(List<Message> messages) {
-    StringBuilder sb = new StringBuilder();
-    for (Message message : messages) {
-      sb.append("<tr>");
-      sb.append("<td>" + message.getMessageId() + "</td>");
-      sb.append("<td>" + message.getData() + "</td>");
-      sb.append("<td>" + message.getPublishTime() + "</td>");
-      sb.append("</tr>");
-    }
-    return sb.toString();
+  /**
+   * Retrieve received claims in html.
+   *
+   * @return html representation of claims (one per row)
+   */
+  public static List<String> getReceivedClaims() {
+    List<String> claimList = messageRepository.retrieveClaims(MAX_MESSAGES);
+    return claimList;
+  }
+
+  /**
+   * Retrieve received tokens in html.
+   *
+   * @return html representation of tokens (one per row)
+   */
+  public static List<String> getReceivedTokens() {
+    List<String> tokenList = messageRepository.retrieveTokens(MAX_MESSAGES);
+    return tokenList;
+  }
+
+  public static String convertToHtml() {
+    Writer writer = new StringWriter(1024);
+    new Html(writer) {
+      {
+        html();
+        head();
+        meta().httpEquiv("refresh").content("10").end();
+        end();
+        body();
+        h3().text("Publish a message").end();
+        form().action("pubsub/publish").method("POST");
+        label().text("Message:").end();
+        input().id("payload").type("input").name("payload").end();
+        input().id("submit").type("submit").value("Send").end();
+        end();
+        h3().text("Last received tokens").end();
+        table().border("1").cellpadding("10");
+        tr();
+        th().text("Tokens").end();
+        end();
+        markupString(getReceivedTokens());
+        h3().text("Last received claims").end();
+        table().border("1").cellpadding("10");
+        tr();
+        th().text("Claims").end();
+        end();
+        markupString(getReceivedClaims());
+        h3().text("Last received messages").end();
+        table().border("1").cellpadding("10");
+        tr();
+        th().text("Id").end();
+        th().text("Data").end();
+        th().text("PublishTime").end();
+        end();
+        markupMessage(getReceivedMessages());
+        endAll();
+        done();
+      }
+
+      Html markupString(List<String> strings) {
+        for (String string : strings) {
+          tr();
+          th().text(string).end();
+          end();
+        }
+        return end();
+      }
+
+      Html markupMessage(List<Message> messages) {
+        for (Message message : messages) {
+          tr();
+          th().text(message.getMessageId()).end();
+          th().text(message.getData()).end();
+          th().text(message.getPublishTime()).end();
+          end();
+        }
+        return end();
+      }
+    };
+    return ((StringWriter) writer).getBuffer().toString();
   }
 
-  private PubSubHome() { }
+  private PubSubHome() {}
 }
diff --git a/appengine-java8/pubsub/src/main/java/com/example/appengine/pubsub/PubSubPublish.java b/appengine-java8/pubsub/src/main/java/com/example/appengine/pubsub/PubSubPublish.java
index 73d850284e2..e4915ea131d 100644
--- a/appengine-java8/pubsub/src/main/java/com/example/appengine/pubsub/PubSubPublish.java
+++ b/appengine-java8/pubsub/src/main/java/com/example/appengine/pubsub/PubSubPublish.java
@@ -40,10 +40,11 @@ public void doPost(HttpServletRequest req, HttpServletResponse resp)
       String topicId = System.getenv("PUBSUB_TOPIC");
       // create a publisher on the topic
       if (publisher == null) {
-        ProjectTopicName topicName = ProjectTopicName.newBuilder()
-            .setProject(ServiceOptions.getDefaultProjectId())
-            .setTopic(topicId)
-            .build();
+        ProjectTopicName topicName =
+            ProjectTopicName.newBuilder()
+                .setProject(ServiceOptions.getDefaultProjectId())
+                .setTopic(topicId)
+                .build();
         publisher = Publisher.newBuilder(topicName).build();
       }
       // construct a pubsub message from the payload
@@ -61,7 +62,7 @@ public void doPost(HttpServletRequest req, HttpServletResponse resp)
 
   private Publisher publisher;
 
-  public PubSubPublish() { }
+  public PubSubPublish() {}
 
   PubSubPublish(Publisher publisher) {
     this.publisher = publisher;
diff --git a/appengine-java8/pubsub/src/main/java/com/example/appengine/pubsub/PubSubPush.java b/appengine-java8/pubsub/src/main/java/com/example/appengine/pubsub/PubSubPush.java
index 0d3f5a2def6..14f6f9a58bc 100644
--- a/appengine-java8/pubsub/src/main/java/com/example/appengine/pubsub/PubSubPush.java
+++ b/appengine-java8/pubsub/src/main/java/com/example/appengine/pubsub/PubSubPush.java
@@ -28,6 +28,7 @@
 import javax.servlet.http.HttpServletRequest;
 import javax.servlet.http.HttpServletResponse;
 
+// [START gae_standard_pubsub_push]
 @WebServlet(value = "/pubsub/push")
 public class PubSubPush extends HttpServlet {
 
@@ -79,3 +80,4 @@ public PubSubPush() {
     this.messageRepository = MessageRepositoryImpl.getInstance();
   }
 }
+// [END gae_standard_pubsub_push]
diff --git a/appengine-java8/pubsub/src/main/webapp/index.jsp b/appengine-java8/pubsub/src/main/webapp/index.jsp
index 9b5bc1cfd85..7d358582809 100644
--- a/appengine-java8/pubsub/src/main/webapp/index.jsp
+++ b/appengine-java8/pubsub/src/main/webapp/index.jsp
@@ -1,24 +1,2 @@
 <%@ page import="com.example.appengine.pubsub.PubSubHome" %>
-<html>
-  <head>
-    <meta http-equiv="refresh" content="10">
-  </head>
-  <title>An example of using PubSub on App Engine Flex</title>
-  <body>
-    <h3> Publish a message </h3>
-    <form action="pubsub/publish" method="POST">
-      <label for="payload">Message:</label>
-      <input id="payload" type="input" name="payload" />
-      <input id="submit"  type="submit" value="Send" />
-    </form>
-    <h3> Last received messages </h3>
-    <table border="1" cellpadding="10">
-      <tr>
-      <th>Id</th>
-      <th>Data</th>
-      <th>PublishTime</th>
-      </tr>
-      <%= PubSubHome.getReceivedMessages() %>
-    </table>
-  </body>
-</html>
+<%= PubSubHome.convertToHtml() %>