You can not select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.

156 lines
5.6 KiB

  1. // Copyright (C) MongoDB, Inc. 2017-present.
  2. //
  3. // Licensed under the Apache License, Version 2.0 (the "License"); you may
  4. // not use this file except in compliance with the License. You may obtain
  5. // a copy of the License at http://www.apache.org/licenses/LICENSE-2.0
  6. package bson
  7. import (
  8. "go.mongodb.org/mongo-driver/bson/bsoncodec"
  9. "go.mongodb.org/mongo-driver/bson/bsonrw"
  10. "go.mongodb.org/mongo-driver/bson/bsontype"
  11. )
  12. const defaultDstCap = 256
  13. var bvwPool = bsonrw.NewBSONValueWriterPool()
  14. var extjPool = bsonrw.NewExtJSONValueWriterPool()
  15. // Marshaler is an interface implemented by types that can marshal themselves
  16. // into a BSON document represented as bytes. The bytes returned must be a valid
  17. // BSON document if the error is nil.
  18. type Marshaler interface {
  19. MarshalBSON() ([]byte, error)
  20. }
  21. // ValueMarshaler is an interface implemented by types that can marshal
  22. // themselves into a BSON value as bytes. The type must be the valid type for
  23. // the bytes returned. The bytes and byte type together must be valid if the
  24. // error is nil.
  25. type ValueMarshaler interface {
  26. MarshalBSONValue() (bsontype.Type, []byte, error)
  27. }
  28. // Marshal returns the BSON encoding of val.
  29. //
  30. // Marshal will use the default registry created by NewRegistry to recursively
  31. // marshal val into a []byte. Marshal will inspect struct tags and alter the
  32. // marshaling process accordingly.
  33. func Marshal(val interface{}) ([]byte, error) {
  34. return MarshalWithRegistry(DefaultRegistry, val)
  35. }
  36. // MarshalAppend will append the BSON encoding of val to dst. If dst is not
  37. // large enough to hold the BSON encoding of val, dst will be grown.
  38. func MarshalAppend(dst []byte, val interface{}) ([]byte, error) {
  39. return MarshalAppendWithRegistry(DefaultRegistry, dst, val)
  40. }
  41. // MarshalWithRegistry returns the BSON encoding of val using Registry r.
  42. func MarshalWithRegistry(r *bsoncodec.Registry, val interface{}) ([]byte, error) {
  43. dst := make([]byte, 0, 256) // TODO: make the default cap a constant
  44. return MarshalAppendWithRegistry(r, dst, val)
  45. }
  46. // MarshalWithContext returns the BSON encoding of val using EncodeContext ec.
  47. func MarshalWithContext(ec bsoncodec.EncodeContext, val interface{}) ([]byte, error) {
  48. dst := make([]byte, 0, 256) // TODO: make the default cap a constant
  49. return MarshalAppendWithContext(ec, dst, val)
  50. }
  51. // MarshalAppendWithRegistry will append the BSON encoding of val to dst using
  52. // Registry r. If dst is not large enough to hold the BSON encoding of val, dst
  53. // will be grown.
  54. func MarshalAppendWithRegistry(r *bsoncodec.Registry, dst []byte, val interface{}) ([]byte, error) {
  55. return MarshalAppendWithContext(bsoncodec.EncodeContext{Registry: r}, dst, val)
  56. }
  57. // MarshalAppendWithContext will append the BSON encoding of val to dst using
  58. // EncodeContext ec. If dst is not large enough to hold the BSON encoding of val, dst
  59. // will be grown.
  60. func MarshalAppendWithContext(ec bsoncodec.EncodeContext, dst []byte, val interface{}) ([]byte, error) {
  61. sw := new(bsonrw.SliceWriter)
  62. *sw = dst
  63. vw := bvwPool.Get(sw)
  64. defer bvwPool.Put(vw)
  65. enc := encPool.Get().(*Encoder)
  66. defer encPool.Put(enc)
  67. err := enc.Reset(vw)
  68. if err != nil {
  69. return nil, err
  70. }
  71. err = enc.SetContext(ec)
  72. if err != nil {
  73. return nil, err
  74. }
  75. err = enc.Encode(val)
  76. if err != nil {
  77. return nil, err
  78. }
  79. return *sw, nil
  80. }
  81. // MarshalExtJSON returns the extended JSON encoding of val.
  82. func MarshalExtJSON(val interface{}, canonical, escapeHTML bool) ([]byte, error) {
  83. return MarshalExtJSONWithRegistry(DefaultRegistry, val, canonical, escapeHTML)
  84. }
  85. // MarshalExtJSONAppend will append the extended JSON encoding of val to dst.
  86. // If dst is not large enough to hold the extended JSON encoding of val, dst
  87. // will be grown.
  88. func MarshalExtJSONAppend(dst []byte, val interface{}, canonical, escapeHTML bool) ([]byte, error) {
  89. return MarshalExtJSONAppendWithRegistry(DefaultRegistry, dst, val, canonical, escapeHTML)
  90. }
  91. // MarshalExtJSONWithRegistry returns the extended JSON encoding of val using Registry r.
  92. func MarshalExtJSONWithRegistry(r *bsoncodec.Registry, val interface{}, canonical, escapeHTML bool) ([]byte, error) {
  93. dst := make([]byte, 0, defaultDstCap)
  94. return MarshalExtJSONAppendWithContext(bsoncodec.EncodeContext{Registry: r}, dst, val, canonical, escapeHTML)
  95. }
  96. // MarshalExtJSONWithContext returns the extended JSON encoding of val using Registry r.
  97. func MarshalExtJSONWithContext(ec bsoncodec.EncodeContext, val interface{}, canonical, escapeHTML bool) ([]byte, error) {
  98. dst := make([]byte, 0, defaultDstCap)
  99. return MarshalExtJSONAppendWithContext(ec, dst, val, canonical, escapeHTML)
  100. }
  101. // MarshalExtJSONAppendWithRegistry will append the extended JSON encoding of
  102. // val to dst using Registry r. If dst is not large enough to hold the BSON
  103. // encoding of val, dst will be grown.
  104. func MarshalExtJSONAppendWithRegistry(r *bsoncodec.Registry, dst []byte, val interface{}, canonical, escapeHTML bool) ([]byte, error) {
  105. return MarshalExtJSONAppendWithContext(bsoncodec.EncodeContext{Registry: r}, dst, val, canonical, escapeHTML)
  106. }
  107. // MarshalExtJSONAppendWithContext will append the extended JSON encoding of
  108. // val to dst using Registry r. If dst is not large enough to hold the BSON
  109. // encoding of val, dst will be grown.
  110. func MarshalExtJSONAppendWithContext(ec bsoncodec.EncodeContext, dst []byte, val interface{}, canonical, escapeHTML bool) ([]byte, error) {
  111. sw := new(bsonrw.SliceWriter)
  112. *sw = dst
  113. ejvw := extjPool.Get(sw, canonical, escapeHTML)
  114. defer extjPool.Put(ejvw)
  115. enc := encPool.Get().(*Encoder)
  116. defer encPool.Put(enc)
  117. err := enc.Reset(ejvw)
  118. if err != nil {
  119. return nil, err
  120. }
  121. err = enc.SetContext(ec)
  122. if err != nil {
  123. return nil, err
  124. }
  125. err = enc.Encode(val)
  126. if err != nil {
  127. return nil, err
  128. }
  129. return *sw, nil
  130. }