all: tests, tweaks for lazy extension decoding

Add a test to confirm that extensions are lazily decoded when we expect.

Drop the UnmarshalDefaultResolver flag. I added it thinking for some
reason that internal/impl couldn't depend on protoregistry; since it can
(and does), it's simpler to just test if the resolver is the expected
value.

Use a default set of options when lazily unmarshaling extensions.

Change-Id: Ied7666ffdc3bf90630260a80c9568d9a945048bc
Reviewed-on: https://go-review.googlesource.com/c/protobuf/+/218038
Reviewed-by: Joe Tsai <thebrokentoaster@gmail.com>
diff --git a/internal/impl/codec_extension.go b/internal/impl/codec_extension.go
index 1098159..efeea2f 100644
--- a/internal/impl/codec_extension.go
+++ b/internal/impl/codec_extension.go
@@ -136,7 +136,7 @@
 			wtyp := wire.Type(tag & 7)
 			var out unmarshalOutput
 			var err error
-			val, out, err = f.lazy.xi.funcs.unmarshal(b, val, num, wtyp, unmarshalOptions{})
+			val, out, err = f.lazy.xi.funcs.unmarshal(b, val, num, wtyp, lazyUnmarshalOptions)
 			if err != nil {
 				panic(errors.New("decode failure in lazy extension decoding: %v", err))
 			}
@@ -227,3 +227,34 @@
 		return f.typ.ValueOf(fn())
 	})
 }
+
+// IsLazy reports whether a field is lazily encoded.
+// It is exported for testing.
+func IsLazy(m pref.Message, fd pref.FieldDescriptor) bool {
+	var mi *MessageInfo
+	var p pointer
+	switch m := m.(type) {
+	case *messageState:
+		mi = m.messageInfo()
+		p = m.pointer()
+	case *messageReflectWrapper:
+		mi = m.messageInfo()
+		p = m.pointer()
+	default:
+		return false
+	}
+	xd, ok := fd.(pref.ExtensionTypeDescriptor)
+	if !ok {
+		return false
+	}
+	xt := xd.Type()
+	ext := mi.extensionMap(p)
+	if ext == nil {
+		return false
+	}
+	f, ok := (*ext)[int32(fd.Number())]
+	if !ok {
+		return false
+	}
+	return f.typ == xt && f.lazy != nil && atomic.LoadUint32(&f.lazy.atomicOnce) == 0
+}
diff --git a/internal/impl/decode.go b/internal/impl/decode.go
index 0fffad3..a2993d4 100644
--- a/internal/impl/decode.go
+++ b/internal/impl/decode.go
@@ -30,9 +30,11 @@
 func (o unmarshalOptions) DiscardUnknown() bool { return o.Flags&piface.UnmarshalDiscardUnknown != 0 }
 
 func (o unmarshalOptions) IsDefault() bool {
-	// The UnmarshalDefaultResolver flag indicates that we're using the default resolver.
-	// No other flag bit should be set.
-	return o.Flags == piface.UnmarshalDefaultResolver
+	return o.Flags == 0 && o.Resolver == preg.GlobalTypes
+}
+
+var lazyUnmarshalOptions = unmarshalOptions{
+	Resolver: preg.GlobalTypes,
 }
 
 type unmarshalOutput struct {
diff --git a/internal/impl/lazy_test.go b/internal/impl/lazy_test.go
new file mode 100644
index 0000000..3a52f87
--- /dev/null
+++ b/internal/impl/lazy_test.go
@@ -0,0 +1,52 @@
+// Copyright 2020 The Go Authors. All rights reserved.
+// Use of this source code is governed by a BSD-style
+// license that can be found in the LICENSE file.
+
+package impl_test
+
+import (
+	"testing"
+
+	"google.golang.org/protobuf/internal/flags"
+	"google.golang.org/protobuf/internal/impl"
+	"google.golang.org/protobuf/internal/protobuild"
+	"google.golang.org/protobuf/proto"
+
+	testpb "google.golang.org/protobuf/internal/testprotos/test"
+)
+
+func TestLazyExtensions(t *testing.T) {
+	checkLazy := func(when string, m *testpb.TestAllExtensions, want bool) {
+		xd := testpb.E_OptionalNestedMessage.TypeDescriptor()
+		if got := impl.IsLazy(m.ProtoReflect(), xd); got != want {
+			t.Errorf("%v: m.optional_nested_message lazy=%v, want %v", when, got, want)
+		}
+		e := proto.GetExtension(m, testpb.E_OptionalNestedMessage).(*testpb.TestAllExtensions_NestedMessage).Corecursive
+		if got := impl.IsLazy(e.ProtoReflect(), xd); got != want {
+			t.Errorf("%v: m.optional_nested_message.corecursive.optional_nested_message lazy=%v, want %v", when, got, want)
+		}
+	}
+
+	m1 := &testpb.TestAllExtensions{}
+	protobuild.Message{
+		"optional_nested_message": protobuild.Message{
+			"a": 1,
+			"corecursive": protobuild.Message{
+				"optional_nested_message": protobuild.Message{
+					"a": 2,
+				},
+			},
+		},
+	}.Build(m1.ProtoReflect())
+	checkLazy("before unmarshal", m1, false)
+
+	w, err := proto.Marshal(m1)
+	if err != nil {
+		t.Fatal(err)
+	}
+	m := &testpb.TestAllExtensions{}
+	if err := proto.Unmarshal(w, m); err != nil {
+		t.Fatal(err)
+	}
+	checkLazy("after unmarshal", m, flags.LazyUnmarshalExtensions)
+}