extensions_test.go 20 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650
  1. // Copyright 2014 The Go Authors. All rights reserved.
  2. // Use of this source code is governed by a BSD-style
  3. // license that can be found in the LICENSE file.
  4. package proto_test
  5. import (
  6. "bytes"
  7. "fmt"
  8. "reflect"
  9. "sort"
  10. "strings"
  11. "sync"
  12. "testing"
  13. "github.com/golang/protobuf/proto"
  14. pb "github.com/golang/protobuf/proto/test_proto"
  15. )
  16. func TestGetExtensionsWithMissingExtensions(t *testing.T) {
  17. msg := &pb.MyMessage{}
  18. ext1 := &pb.Ext{}
  19. if err := proto.SetExtension(msg, pb.E_Ext_More, ext1); err != nil {
  20. t.Fatalf("Could not set ext1: %s", err)
  21. }
  22. exts, err := proto.GetExtensions(msg, []*proto.ExtensionDesc{
  23. pb.E_Ext_More,
  24. pb.E_Ext_Text,
  25. })
  26. if err != nil {
  27. t.Fatalf("GetExtensions() failed: %s", err)
  28. }
  29. if exts[0] != ext1 {
  30. t.Errorf("ext1 not in returned extensions: %T %v", exts[0], exts[0])
  31. }
  32. if exts[1] != nil {
  33. t.Errorf("ext2 in returned extensions: %T %v", exts[1], exts[1])
  34. }
  35. }
  36. func TestGetExtensionForIncompleteDesc(t *testing.T) {
  37. msg := &pb.MyMessage{Count: proto.Int32(0)}
  38. extdesc1 := &proto.ExtensionDesc{
  39. ExtendedType: (*pb.MyMessage)(nil),
  40. ExtensionType: (*bool)(nil),
  41. Field: 123456789,
  42. Name: "a.b",
  43. Tag: "varint,123456789,opt",
  44. }
  45. ext1 := proto.Bool(true)
  46. if err := proto.SetExtension(msg, extdesc1, ext1); err != nil {
  47. t.Fatalf("Could not set ext1: %s", err)
  48. }
  49. extdesc2 := &proto.ExtensionDesc{
  50. ExtendedType: (*pb.MyMessage)(nil),
  51. ExtensionType: ([]byte)(nil),
  52. Field: 123456790,
  53. Name: "a.c",
  54. Tag: "bytes,123456790,opt",
  55. }
  56. ext2 := []byte{0, 1, 2, 3, 4, 5, 6, 7}
  57. if err := proto.SetExtension(msg, extdesc2, ext2); err != nil {
  58. t.Fatalf("Could not set ext2: %s", err)
  59. }
  60. extdesc3 := &proto.ExtensionDesc{
  61. ExtendedType: (*pb.MyMessage)(nil),
  62. ExtensionType: (*pb.Ext)(nil),
  63. Field: 123456791,
  64. Name: "a.d",
  65. Tag: "bytes,123456791,opt",
  66. }
  67. ext3 := &pb.Ext{Data: proto.String("foo")}
  68. if err := proto.SetExtension(msg, extdesc3, ext3); err != nil {
  69. t.Fatalf("Could not set ext3: %s", err)
  70. }
  71. b, err := proto.Marshal(msg)
  72. if err != nil {
  73. t.Fatalf("Could not marshal msg: %v", err)
  74. }
  75. if err := proto.Unmarshal(b, msg); err != nil {
  76. t.Fatalf("Could not unmarshal into msg: %v", err)
  77. }
  78. var expected proto.Buffer
  79. if err := expected.EncodeVarint(uint64((extdesc1.Field << 3) | proto.WireVarint)); err != nil {
  80. t.Fatalf("failed to compute expected prefix for ext1: %s", err)
  81. }
  82. if err := expected.EncodeVarint(1 /* bool true */); err != nil {
  83. t.Fatalf("failed to compute expected value for ext1: %s", err)
  84. }
  85. if b, err := proto.GetExtension(msg, &proto.ExtensionDesc{Field: extdesc1.Field}); err != nil {
  86. t.Fatalf("Failed to get raw value for ext1: %s", err)
  87. } else if !reflect.DeepEqual(b, expected.Bytes()) {
  88. t.Fatalf("Raw value for ext1: got %v, want %v", b, expected.Bytes())
  89. }
  90. expected = proto.Buffer{} // reset
  91. if err := expected.EncodeVarint(uint64((extdesc2.Field << 3) | proto.WireBytes)); err != nil {
  92. t.Fatalf("failed to compute expected prefix for ext2: %s", err)
  93. }
  94. if err := expected.EncodeRawBytes(ext2); err != nil {
  95. t.Fatalf("failed to compute expected value for ext2: %s", err)
  96. }
  97. if b, err := proto.GetExtension(msg, &proto.ExtensionDesc{Field: extdesc2.Field}); err != nil {
  98. t.Fatalf("Failed to get raw value for ext2: %s", err)
  99. } else if !reflect.DeepEqual(b, expected.Bytes()) {
  100. t.Fatalf("Raw value for ext2: got %v, want %v", b, expected.Bytes())
  101. }
  102. expected = proto.Buffer{} // reset
  103. if err := expected.EncodeVarint(uint64((extdesc3.Field << 3) | proto.WireBytes)); err != nil {
  104. t.Fatalf("failed to compute expected prefix for ext3: %s", err)
  105. }
  106. if b, err := proto.Marshal(ext3); err != nil {
  107. t.Fatalf("failed to compute expected value for ext3: %s", err)
  108. } else if err := expected.EncodeRawBytes(b); err != nil {
  109. t.Fatalf("failed to compute expected value for ext3: %s", err)
  110. }
  111. if b, err := proto.GetExtension(msg, &proto.ExtensionDesc{Field: extdesc3.Field}); err != nil {
  112. t.Fatalf("Failed to get raw value for ext3: %s", err)
  113. } else if !reflect.DeepEqual(b, expected.Bytes()) {
  114. t.Fatalf("Raw value for ext3: got %v, want %v", b, expected.Bytes())
  115. }
  116. }
  117. func TestExtensionDescsWithUnregisteredExtensions(t *testing.T) {
  118. msg := &pb.MyMessage{Count: proto.Int32(0)}
  119. extdesc1 := pb.E_Ext_More
  120. if descs, err := proto.ExtensionDescs(msg); len(descs) != 0 || err != nil {
  121. t.Errorf("proto.ExtensionDescs: got %d descs, error %v; want 0, nil", len(descs), err)
  122. }
  123. ext1 := &pb.Ext{}
  124. if err := proto.SetExtension(msg, extdesc1, ext1); err != nil {
  125. t.Fatalf("Could not set ext1: %s", err)
  126. }
  127. extdesc2 := &proto.ExtensionDesc{
  128. ExtendedType: (*pb.MyMessage)(nil),
  129. ExtensionType: (*bool)(nil),
  130. Field: 123456789,
  131. Name: "a.b",
  132. Tag: "varint,123456789,opt",
  133. }
  134. ext2 := proto.Bool(false)
  135. if err := proto.SetExtension(msg, extdesc2, ext2); err != nil {
  136. t.Fatalf("Could not set ext2: %s", err)
  137. }
  138. b, err := proto.Marshal(msg)
  139. if err != nil {
  140. t.Fatalf("Could not marshal msg: %v", err)
  141. }
  142. if err := proto.Unmarshal(b, msg); err != nil {
  143. t.Fatalf("Could not unmarshal into msg: %v", err)
  144. }
  145. descs, err := proto.ExtensionDescs(msg)
  146. if err != nil {
  147. t.Fatalf("proto.ExtensionDescs: got error %v", err)
  148. }
  149. sortExtDescs(descs)
  150. wantDescs := []*proto.ExtensionDesc{extdesc1, {Field: extdesc2.Field}}
  151. if !reflect.DeepEqual(descs, wantDescs) {
  152. t.Errorf("proto.ExtensionDescs(msg) sorted extension ids: got %+v, want %+v", descs, wantDescs)
  153. }
  154. }
  155. type ExtensionDescSlice []*proto.ExtensionDesc
  156. func (s ExtensionDescSlice) Len() int { return len(s) }
  157. func (s ExtensionDescSlice) Less(i, j int) bool { return s[i].Field < s[j].Field }
  158. func (s ExtensionDescSlice) Swap(i, j int) { s[i], s[j] = s[j], s[i] }
  159. func sortExtDescs(s []*proto.ExtensionDesc) {
  160. sort.Sort(ExtensionDescSlice(s))
  161. }
  162. func TestGetExtensionStability(t *testing.T) {
  163. check := func(m *pb.MyMessage) bool {
  164. ext1, err := proto.GetExtension(m, pb.E_Ext_More)
  165. if err != nil {
  166. t.Fatalf("GetExtension() failed: %s", err)
  167. }
  168. ext2, err := proto.GetExtension(m, pb.E_Ext_More)
  169. if err != nil {
  170. t.Fatalf("GetExtension() failed: %s", err)
  171. }
  172. return ext1 == ext2
  173. }
  174. msg := &pb.MyMessage{Count: proto.Int32(4)}
  175. ext0 := &pb.Ext{}
  176. if err := proto.SetExtension(msg, pb.E_Ext_More, ext0); err != nil {
  177. t.Fatalf("Could not set ext1: %s", ext0)
  178. }
  179. if !check(msg) {
  180. t.Errorf("GetExtension() not stable before marshaling")
  181. }
  182. bb, err := proto.Marshal(msg)
  183. if err != nil {
  184. t.Fatalf("Marshal() failed: %s", err)
  185. }
  186. msg1 := &pb.MyMessage{}
  187. err = proto.Unmarshal(bb, msg1)
  188. if err != nil {
  189. t.Fatalf("Unmarshal() failed: %s", err)
  190. }
  191. if !check(msg1) {
  192. t.Errorf("GetExtension() not stable after unmarshaling")
  193. }
  194. }
  195. func TestGetExtensionDefaults(t *testing.T) {
  196. var setFloat64 float64 = 1
  197. var setFloat32 float32 = 2
  198. var setInt32 int32 = 3
  199. var setInt64 int64 = 4
  200. var setUint32 uint32 = 5
  201. var setUint64 uint64 = 6
  202. var setBool = true
  203. var setBool2 = false
  204. var setString = "Goodnight string"
  205. var setBytes = []byte("Goodnight bytes")
  206. var setEnum = pb.DefaultsMessage_TWO
  207. type testcase struct {
  208. ext *proto.ExtensionDesc // Extension we are testing.
  209. want interface{} // Expected value of extension, or nil (meaning that GetExtension will fail).
  210. def interface{} // Expected value of extension after ClearExtension().
  211. }
  212. tests := []testcase{
  213. {pb.E_NoDefaultDouble, setFloat64, nil},
  214. {pb.E_NoDefaultFloat, setFloat32, nil},
  215. {pb.E_NoDefaultInt32, setInt32, nil},
  216. {pb.E_NoDefaultInt64, setInt64, nil},
  217. {pb.E_NoDefaultUint32, setUint32, nil},
  218. {pb.E_NoDefaultUint64, setUint64, nil},
  219. {pb.E_NoDefaultSint32, setInt32, nil},
  220. {pb.E_NoDefaultSint64, setInt64, nil},
  221. {pb.E_NoDefaultFixed32, setUint32, nil},
  222. {pb.E_NoDefaultFixed64, setUint64, nil},
  223. {pb.E_NoDefaultSfixed32, setInt32, nil},
  224. {pb.E_NoDefaultSfixed64, setInt64, nil},
  225. {pb.E_NoDefaultBool, setBool, nil},
  226. {pb.E_NoDefaultBool, setBool2, nil},
  227. {pb.E_NoDefaultString, setString, nil},
  228. {pb.E_NoDefaultBytes, setBytes, nil},
  229. {pb.E_NoDefaultEnum, setEnum, nil},
  230. {pb.E_DefaultDouble, setFloat64, float64(3.1415)},
  231. {pb.E_DefaultFloat, setFloat32, float32(3.14)},
  232. {pb.E_DefaultInt32, setInt32, int32(42)},
  233. {pb.E_DefaultInt64, setInt64, int64(43)},
  234. {pb.E_DefaultUint32, setUint32, uint32(44)},
  235. {pb.E_DefaultUint64, setUint64, uint64(45)},
  236. {pb.E_DefaultSint32, setInt32, int32(46)},
  237. {pb.E_DefaultSint64, setInt64, int64(47)},
  238. {pb.E_DefaultFixed32, setUint32, uint32(48)},
  239. {pb.E_DefaultFixed64, setUint64, uint64(49)},
  240. {pb.E_DefaultSfixed32, setInt32, int32(50)},
  241. {pb.E_DefaultSfixed64, setInt64, int64(51)},
  242. {pb.E_DefaultBool, setBool, true},
  243. {pb.E_DefaultBool, setBool2, true},
  244. {pb.E_DefaultString, setString, "Hello, string,def=foo"},
  245. {pb.E_DefaultBytes, setBytes, []byte("Hello, bytes")},
  246. {pb.E_DefaultEnum, setEnum, pb.DefaultsMessage_ONE},
  247. }
  248. checkVal := func(t *testing.T, name string, test testcase, msg *pb.DefaultsMessage, valWant interface{}) {
  249. t.Run(name, func(t *testing.T) {
  250. val, err := proto.GetExtension(msg, test.ext)
  251. if err != nil {
  252. if valWant != nil {
  253. t.Errorf("GetExtension(): %s", err)
  254. return
  255. }
  256. if want := proto.ErrMissingExtension; err != want {
  257. t.Errorf("Unexpected error: got %v, want %v", err, want)
  258. return
  259. }
  260. return
  261. }
  262. // All proto2 extension values are either a pointer to a value or a slice of values.
  263. ty := reflect.TypeOf(val)
  264. tyWant := reflect.TypeOf(test.ext.ExtensionType)
  265. if got, want := ty, tyWant; got != want {
  266. t.Errorf("unexpected reflect.TypeOf(): got %v want %v", got, want)
  267. return
  268. }
  269. tye := ty.Elem()
  270. tyeWant := tyWant.Elem()
  271. if got, want := tye, tyeWant; got != want {
  272. t.Errorf("unexpected reflect.TypeOf().Elem(): got %v want %v", got, want)
  273. return
  274. }
  275. // Check the name of the type of the value.
  276. // If it is an enum it will be type int32 with the name of the enum.
  277. if got, want := tye.Name(), tye.Name(); got != want {
  278. t.Errorf("unexpected reflect.TypeOf().Elem().Name(): got %v want %v", got, want)
  279. return
  280. }
  281. // Check that value is what we expect.
  282. // If we have a pointer in val, get the value it points to.
  283. valExp := val
  284. if ty.Kind() == reflect.Ptr {
  285. valExp = reflect.ValueOf(val).Elem().Interface()
  286. }
  287. if got, want := valExp, valWant; !reflect.DeepEqual(got, want) {
  288. t.Errorf("unexpected reflect.DeepEqual(): got %v want %v", got, want)
  289. return
  290. }
  291. })
  292. }
  293. setTo := func(test testcase) interface{} {
  294. setTo := reflect.ValueOf(test.want)
  295. if typ := reflect.TypeOf(test.ext.ExtensionType); typ.Kind() == reflect.Ptr {
  296. setTo = reflect.New(typ).Elem()
  297. setTo.Set(reflect.New(setTo.Type().Elem()))
  298. setTo.Elem().Set(reflect.ValueOf(test.want))
  299. }
  300. return setTo.Interface()
  301. }
  302. for _, test := range tests {
  303. msg := &pb.DefaultsMessage{}
  304. name := test.ext.Name
  305. // Check the initial value.
  306. checkVal(t, name+"/initial", test, msg, test.def)
  307. // Set the per-type value and check value.
  308. if err := proto.SetExtension(msg, test.ext, setTo(test)); err != nil {
  309. t.Errorf("%s: SetExtension(): %v", name, err)
  310. continue
  311. }
  312. checkVal(t, name+"/set", test, msg, test.want)
  313. // Set and check the value.
  314. proto.ClearExtension(msg, test.ext)
  315. checkVal(t, name+"/cleared", test, msg, test.def)
  316. }
  317. }
  318. func TestNilMessage(t *testing.T) {
  319. name := "nil interface"
  320. if got, err := proto.GetExtension(nil, pb.E_Ext_More); err == nil {
  321. t.Errorf("%s: got %T %v, expected to fail", name, got, got)
  322. } else if !strings.Contains(err.Error(), "extendable") {
  323. t.Errorf("%s: got error %v, expected not-extendable error", name, err)
  324. }
  325. // Regression tests: all functions of the Extension API
  326. // used to panic when passed (*M)(nil), where M is a concrete message
  327. // type. Now they handle this gracefully as a no-op or reported error.
  328. var nilMsg *pb.MyMessage
  329. desc := pb.E_Ext_More
  330. isNotExtendable := func(err error) bool {
  331. return strings.Contains(fmt.Sprint(err), "not an extendable")
  332. }
  333. if proto.HasExtension(nilMsg, desc) {
  334. t.Error("HasExtension(nil) = true")
  335. }
  336. if _, err := proto.GetExtensions(nilMsg, []*proto.ExtensionDesc{desc}); !isNotExtendable(err) {
  337. t.Errorf("GetExtensions(nil) = %q (wrong error)", err)
  338. }
  339. if _, err := proto.ExtensionDescs(nilMsg); !isNotExtendable(err) {
  340. t.Errorf("ExtensionDescs(nil) = %q (wrong error)", err)
  341. }
  342. if err := proto.SetExtension(nilMsg, desc, nil); !isNotExtendable(err) {
  343. t.Errorf("SetExtension(nil) = %q (wrong error)", err)
  344. }
  345. proto.ClearExtension(nilMsg, desc) // no-op
  346. proto.ClearAllExtensions(nilMsg) // no-op
  347. }
  348. func TestExtensionsRoundTrip(t *testing.T) {
  349. msg := &pb.MyMessage{}
  350. ext1 := &pb.Ext{
  351. Data: proto.String("hi"),
  352. }
  353. ext2 := &pb.Ext{
  354. Data: proto.String("there"),
  355. }
  356. exists := proto.HasExtension(msg, pb.E_Ext_More)
  357. if exists {
  358. t.Error("Extension More present unexpectedly")
  359. }
  360. if err := proto.SetExtension(msg, pb.E_Ext_More, ext1); err != nil {
  361. t.Error(err)
  362. }
  363. if err := proto.SetExtension(msg, pb.E_Ext_More, ext2); err != nil {
  364. t.Error(err)
  365. }
  366. e, err := proto.GetExtension(msg, pb.E_Ext_More)
  367. if err != nil {
  368. t.Error(err)
  369. }
  370. x, ok := e.(*pb.Ext)
  371. if !ok {
  372. t.Errorf("e has type %T, expected test_proto.Ext", e)
  373. } else if *x.Data != "there" {
  374. t.Errorf("SetExtension failed to overwrite, got %+v, not 'there'", x)
  375. }
  376. proto.ClearExtension(msg, pb.E_Ext_More)
  377. if _, err = proto.GetExtension(msg, pb.E_Ext_More); err != proto.ErrMissingExtension {
  378. t.Errorf("got %v, expected ErrMissingExtension", e)
  379. }
  380. if _, err := proto.GetExtension(msg, pb.E_X215); err == nil {
  381. t.Error("expected bad extension error, got nil")
  382. }
  383. if err := proto.SetExtension(msg, pb.E_X215, 12); err == nil {
  384. t.Error("expected extension err")
  385. }
  386. if err := proto.SetExtension(msg, pb.E_Ext_More, 12); err == nil {
  387. t.Error("expected some sort of type mismatch error, got nil")
  388. }
  389. }
  390. func TestNilExtension(t *testing.T) {
  391. msg := &pb.MyMessage{
  392. Count: proto.Int32(1),
  393. }
  394. if err := proto.SetExtension(msg, pb.E_Ext_Text, proto.String("hello")); err != nil {
  395. t.Fatal(err)
  396. }
  397. if err := proto.SetExtension(msg, pb.E_Ext_More, (*pb.Ext)(nil)); err == nil {
  398. t.Error("expected SetExtension to fail due to a nil extension")
  399. } else if want := fmt.Sprintf("proto: SetExtension called with nil value of type %T", new(pb.Ext)); err.Error() != want {
  400. t.Errorf("expected error %v, got %v", want, err)
  401. }
  402. // Note: if the behavior of Marshal is ever changed to ignore nil extensions, update
  403. // this test to verify that E_Ext_Text is properly propagated through marshal->unmarshal.
  404. }
  405. func TestMarshalUnmarshalRepeatedExtension(t *testing.T) {
  406. // Add a repeated extension to the result.
  407. tests := []struct {
  408. name string
  409. ext []*pb.ComplexExtension
  410. }{
  411. {
  412. "two fields",
  413. []*pb.ComplexExtension{
  414. {First: proto.Int32(7)},
  415. {Second: proto.Int32(11)},
  416. },
  417. },
  418. {
  419. "repeated field",
  420. []*pb.ComplexExtension{
  421. {Third: []int32{1000}},
  422. {Third: []int32{2000}},
  423. },
  424. },
  425. {
  426. "two fields and repeated field",
  427. []*pb.ComplexExtension{
  428. {Third: []int32{1000}},
  429. {First: proto.Int32(9)},
  430. {Second: proto.Int32(21)},
  431. {Third: []int32{2000}},
  432. },
  433. },
  434. }
  435. for _, test := range tests {
  436. // Marshal message with a repeated extension.
  437. msg1 := new(pb.OtherMessage)
  438. err := proto.SetExtension(msg1, pb.E_RComplex, test.ext)
  439. if err != nil {
  440. t.Fatalf("[%s] Error setting extension: %v", test.name, err)
  441. }
  442. b, err := proto.Marshal(msg1)
  443. if err != nil {
  444. t.Fatalf("[%s] Error marshaling message: %v", test.name, err)
  445. }
  446. // Unmarshal and read the merged proto.
  447. msg2 := new(pb.OtherMessage)
  448. err = proto.Unmarshal(b, msg2)
  449. if err != nil {
  450. t.Fatalf("[%s] Error unmarshaling message: %v", test.name, err)
  451. }
  452. e, err := proto.GetExtension(msg2, pb.E_RComplex)
  453. if err != nil {
  454. t.Fatalf("[%s] Error getting extension: %v", test.name, err)
  455. }
  456. ext := e.([]*pb.ComplexExtension)
  457. if ext == nil {
  458. t.Fatalf("[%s] Invalid extension", test.name)
  459. }
  460. if len(ext) != len(test.ext) {
  461. t.Errorf("[%s] Wrong length of ComplexExtension: got: %v want: %v\n", test.name, len(ext), len(test.ext))
  462. }
  463. for i := range test.ext {
  464. if !proto.Equal(ext[i], test.ext[i]) {
  465. t.Errorf("[%s] Wrong value for ComplexExtension[%d]: got: %v want: %v\n", test.name, i, ext[i], test.ext[i])
  466. }
  467. }
  468. }
  469. }
  470. func TestUnmarshalRepeatingNonRepeatedExtension(t *testing.T) {
  471. // We may see multiple instances of the same extension in the wire
  472. // format. For example, the proto compiler may encode custom options in
  473. // this way. Here, we verify that we merge the extensions together.
  474. tests := []struct {
  475. name string
  476. ext []*pb.ComplexExtension
  477. }{
  478. {
  479. "two fields",
  480. []*pb.ComplexExtension{
  481. {First: proto.Int32(7)},
  482. {Second: proto.Int32(11)},
  483. },
  484. },
  485. {
  486. "repeated field",
  487. []*pb.ComplexExtension{
  488. {Third: []int32{1000}},
  489. {Third: []int32{2000}},
  490. },
  491. },
  492. {
  493. "two fields and repeated field",
  494. []*pb.ComplexExtension{
  495. {Third: []int32{1000}},
  496. {First: proto.Int32(9)},
  497. {Second: proto.Int32(21)},
  498. {Third: []int32{2000}},
  499. },
  500. },
  501. }
  502. for _, test := range tests {
  503. var buf bytes.Buffer
  504. var want pb.ComplexExtension
  505. // Generate a serialized representation of a repeated extension
  506. // by catenating bytes together.
  507. for i, e := range test.ext {
  508. // Merge to create the wanted proto.
  509. proto.Merge(&want, e)
  510. // serialize the message
  511. msg := new(pb.OtherMessage)
  512. err := proto.SetExtension(msg, pb.E_Complex, e)
  513. if err != nil {
  514. t.Fatalf("[%s] Error setting extension %d: %v", test.name, i, err)
  515. }
  516. b, err := proto.Marshal(msg)
  517. if err != nil {
  518. t.Fatalf("[%s] Error marshaling message %d: %v", test.name, i, err)
  519. }
  520. buf.Write(b)
  521. }
  522. // Unmarshal and read the merged proto.
  523. msg2 := new(pb.OtherMessage)
  524. err := proto.Unmarshal(buf.Bytes(), msg2)
  525. if err != nil {
  526. t.Fatalf("[%s] Error unmarshaling message: %v", test.name, err)
  527. }
  528. e, err := proto.GetExtension(msg2, pb.E_Complex)
  529. if err != nil {
  530. t.Fatalf("[%s] Error getting extension: %v", test.name, err)
  531. }
  532. ext := e.(*pb.ComplexExtension)
  533. if ext == nil {
  534. t.Fatalf("[%s] Invalid extension", test.name)
  535. }
  536. if !proto.Equal(ext, &want) {
  537. t.Errorf("[%s] Wrong value for ComplexExtension: got: %s want: %s\n", test.name, ext, &want)
  538. }
  539. }
  540. }
  541. func TestClearAllExtensions(t *testing.T) {
  542. // unregistered extension
  543. desc := &proto.ExtensionDesc{
  544. ExtendedType: (*pb.MyMessage)(nil),
  545. ExtensionType: (*bool)(nil),
  546. Field: 101010100,
  547. Name: "emptyextension",
  548. Tag: "varint,0,opt",
  549. }
  550. m := &pb.MyMessage{}
  551. if proto.HasExtension(m, desc) {
  552. t.Errorf("proto.HasExtension(%s): got true, want false", proto.MarshalTextString(m))
  553. }
  554. if err := proto.SetExtension(m, desc, proto.Bool(true)); err != nil {
  555. t.Errorf("proto.SetExtension(m, desc, true): got error %q, want nil", err)
  556. }
  557. if !proto.HasExtension(m, desc) {
  558. t.Errorf("proto.HasExtension(%s): got false, want true", proto.MarshalTextString(m))
  559. }
  560. proto.ClearAllExtensions(m)
  561. if proto.HasExtension(m, desc) {
  562. t.Errorf("proto.HasExtension(%s): got true, want false", proto.MarshalTextString(m))
  563. }
  564. }
  565. func TestMarshalRace(t *testing.T) {
  566. ext := &pb.Ext{}
  567. m := &pb.MyMessage{Count: proto.Int32(4)}
  568. if err := proto.SetExtension(m, pb.E_Ext_More, ext); err != nil {
  569. t.Fatalf("proto.SetExtension(m, desc, true): got error %q, want nil", err)
  570. }
  571. b, err := proto.Marshal(m)
  572. if err != nil {
  573. t.Fatalf("Could not marshal message: %v", err)
  574. }
  575. if err := proto.Unmarshal(b, m); err != nil {
  576. t.Fatalf("Could not unmarshal message: %v", err)
  577. }
  578. // after Unmarshal, the extension is in undecoded form.
  579. // GetExtension will decode it lazily. Make sure this does
  580. // not race against Marshal.
  581. wg := sync.WaitGroup{}
  582. errs := make(chan error, 3)
  583. for n := 3; n > 0; n-- {
  584. wg.Add(1)
  585. go func() {
  586. defer wg.Done()
  587. _, err := proto.Marshal(m)
  588. errs <- err
  589. }()
  590. }
  591. wg.Wait()
  592. close(errs)
  593. for err = range errs {
  594. if err != nil {
  595. t.Fatal(err)
  596. }
  597. }
  598. }