Allow use of a global ExtensionRegistry

2 files changed, 40 insertions, 2 deletions

diff --git a/protobuf-lite/src/main/java/io/grpc/protobuf/lite/ProtoLiteUtils.java b/protobuf-lite/src/main/java/io/grpc/protobuf/lite/ProtoLiteUtils.java
index c5677ff5e..1961afe29 100644
--- a/protobuf-lite/src/main/java/io/grpc/protobuf/lite/ProtoLiteUtils.java
+++ b/protobuf-lite/src/main/java/io/grpc/protobuf/lite/ProtoLiteUtils.java
@@ -31,7 +31,10 @@
 
 package io.grpc.protobuf.lite;
 
+import static com.google.common.base.Preconditions.checkNotNull;
+
 import com.google.protobuf.CodedInputStream;
+import com.google.protobuf.ExtensionRegistryLite;
 import com.google.protobuf.InvalidProtocolBufferException;
 import com.google.protobuf.MessageLite;
 import com.google.protobuf.Parser;
@@ -49,6 +52,28 @@ import java.io.InputStream;
 @ExperimentalApi("Experimental until Lite is stable in protobuf")
 public class ProtoLiteUtils {
 
+  private static volatile ExtensionRegistryLite globalRegistry =
+      ExtensionRegistryLite.getEmptyRegistry();
+
+  /**
+   * Sets the global registry for proto marshalling shared across all servers and clients.
+   *
+   * <p>Warning:  This API will likely change over time.  It is not possible to have separate
+   * registries per Process, Server, Channel, Service, or Method.  This is intentional until there
+   * is a more appropriate API to set them.
+   *
+   * <p>Warning:  Do NOT modify the extension registry after setting it.  It is thread safe to call
+   * {@link #setExtensionRegistry}, but not to modify the underlying object.
+   *
+   * <p>If you need custom parsing behavior for protos, you will need to make your own
+   * {@code MethodDescriptor.Marhsaller} for the time being.
+   *
+   */
+  @ExperimentalApi
+  public static void setExtensionRegistry(ExtensionRegistryLite newRegistry) {
+    globalRegistry = checkNotNull(newRegistry, "newRegistry");
+  }
+
   /** Create a {@code Marshaller} for protos of the same type as {@code defaultInstance}. */
   public static <T extends MessageLite> Marshaller<T> marshaller(final T defaultInstance) {
     @SuppressWarnings("unchecked")
@@ -96,7 +121,7 @@ public class ProtoLiteUtils {
         CodedInputStream codedInput = CodedInputStream.newInstance(stream);
         codedInput.setSizeLimit(Integer.MAX_VALUE);
 
-        T message = parser.parseFrom(codedInput);
+        T message = parser.parseFrom(codedInput, globalRegistry);
         try {
           codedInput.checkLastTagWas(0);
           return message;
@@ -123,7 +148,7 @@ public class ProtoLiteUtils {
       @SuppressWarnings("unchecked")
       public T parseBytes(byte[] serialized) {
         try {
-          return (T) instance.getParserForType().parseFrom(serialized);
+          return (T) instance.getParserForType().parseFrom(serialized, globalRegistry);
         } catch (InvalidProtocolBufferException ipbe) {
           throw new IllegalArgumentException(ipbe);
         }
diff --git a/protobuf-lite/src/test/java/io/grpc/protobuf/lite/ProtoLiteUtilsTest.java b/protobuf-lite/src/test/java/io/grpc/protobuf/lite/ProtoLiteUtilsTest.java
index a5042a435..918d5657d 100644
--- a/protobuf-lite/src/test/java/io/grpc/protobuf/lite/ProtoLiteUtilsTest.java
+++ b/protobuf-lite/src/test/java/io/grpc/protobuf/lite/ProtoLiteUtilsTest.java
@@ -50,7 +50,9 @@ import io.grpc.MethodDescriptor.Marshaller;
 import io.grpc.Status;
 import io.grpc.StatusRuntimeException;
 
+import org.junit.Rule;
 import org.junit.Test;
+import org.junit.rules.ExpectedException;
 import org.junit.runner.RunWith;
 import org.junit.runners.JUnit4;
 
@@ -63,6 +65,9 @@ import java.util.Arrays;
 /** Unit tests for {@link ProtoLiteUtils}. */
 @RunWith(JUnit4.class)
 public class ProtoLiteUtilsTest {
+
+  @Rule public final ExpectedException thrown = ExpectedException.none();
+
   private Marshaller<Type> marshaller = ProtoLiteUtils.marshaller(Type.getDefaultInstance());
   private Type proto = Type.newBuilder().setName("name").build();
 
@@ -211,4 +216,12 @@ public class ProtoLiteUtilsTest {
       assertNotNull(((InvalidProtocolBufferException) ex.getCause()).getUnfinishedMessage());
     }
   }
+
+  @Test
+  public void extensionRegistry_notNull() {
+    thrown.expect(NullPointerException.class);
+    thrown.expectMessage("newRegistry");
+
+    ProtoLiteUtils.setExtensionRegistry(null);
+  }
 }