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 7package bson 8 9import ( 10 "go.mongodb.org/mongo-driver/bson/bsoncodec" 11 "go.mongodb.org/mongo-driver/bson/bsonrw" 12 "go.mongodb.org/mongo-driver/bson/bsontype" 13) 14 15const defaultDstCap = 256 16 17var bvwPool = bsonrw.NewBSONValueWriterPool() 18var extjPool = bsonrw.NewExtJSONValueWriterPool() 19 20// Marshaler is an interface implemented by types that can marshal themselves 21// into a BSON document represented as bytes. The bytes returned must be a valid 22// BSON document if the error is nil. 23type Marshaler interface { 24 MarshalBSON() ([]byte, error) 25} 26 27// ValueMarshaler is an interface implemented by types that can marshal 28// themselves into a BSON value as bytes. The type must be the valid type for 29// the bytes returned. The bytes and byte type together must be valid if the 30// error is nil. 31type ValueMarshaler interface { 32 MarshalBSONValue() (bsontype.Type, []byte, error) 33} 34 35// Marshal returns the BSON encoding of val as a BSON document. If val is not a type that can be transformed into a 36// document, MarshalValue should be used instead. 37// 38// Marshal will use the default registry created by NewRegistry to recursively 39// marshal val into a []byte. Marshal will inspect struct tags and alter the 40// marshaling process accordingly. 41func Marshal(val interface{}) ([]byte, error) { 42 return MarshalWithRegistry(DefaultRegistry, val) 43} 44 45// MarshalAppend will encode val as a BSON document and append the bytes to dst. If dst is not large enough to hold the 46// bytes, it will be grown. If val is not a type that can be transformed into a document, MarshalValueAppend should be 47// used instead. 48func MarshalAppend(dst []byte, val interface{}) ([]byte, error) { 49 return MarshalAppendWithRegistry(DefaultRegistry, dst, val) 50} 51 52// MarshalWithRegistry returns the BSON encoding of val as a BSON document. If val is not a type that can be transformed 53// into a document, MarshalValueWithRegistry should be used instead. 54func MarshalWithRegistry(r *bsoncodec.Registry, val interface{}) ([]byte, error) { 55 dst := make([]byte, 0, 256) // TODO: make the default cap a constant 56 return MarshalAppendWithRegistry(r, dst, val) 57} 58 59// MarshalWithContext returns the BSON encoding of val as a BSON document using EncodeContext ec. If val is not a type 60// that can be transformed into a document, MarshalValueWithContext should be used instead. 61func MarshalWithContext(ec bsoncodec.EncodeContext, val interface{}) ([]byte, error) { 62 dst := make([]byte, 0, 256) // TODO: make the default cap a constant 63 return MarshalAppendWithContext(ec, dst, val) 64} 65 66// MarshalAppendWithRegistry will encode val as a BSON document using Registry r and append the bytes to dst. If dst is 67// not large enough to hold the bytes, it will be grown. If val is not a type that can be transformed into a document, 68// MarshalValueAppendWithRegistry should be used instead. 69func MarshalAppendWithRegistry(r *bsoncodec.Registry, dst []byte, val interface{}) ([]byte, error) { 70 return MarshalAppendWithContext(bsoncodec.EncodeContext{Registry: r}, dst, val) 71} 72 73// MarshalAppendWithContext will encode val as a BSON document using Registry r and EncodeContext ec and append the 74// bytes to dst. If dst is not large enough to hold the bytes, it will be grown. If val is not a type that can be 75// transformed into a document, MarshalValueAppendWithContext should be used instead. 76func MarshalAppendWithContext(ec bsoncodec.EncodeContext, dst []byte, val interface{}) ([]byte, error) { 77 sw := new(bsonrw.SliceWriter) 78 *sw = dst 79 vw := bvwPool.Get(sw) 80 defer bvwPool.Put(vw) 81 82 enc := encPool.Get().(*Encoder) 83 defer encPool.Put(enc) 84 85 err := enc.Reset(vw) 86 if err != nil { 87 return nil, err 88 } 89 err = enc.SetContext(ec) 90 if err != nil { 91 return nil, err 92 } 93 94 err = enc.Encode(val) 95 if err != nil { 96 return nil, err 97 } 98 99 return *sw, nil 100} 101 102// MarshalValue returns the BSON encoding of val. 103// 104// MarshalValue will use bson.DefaultRegistry to transform val into a BSON value. If val is a struct, this function will 105// inspect struct tags and alter the marshalling process accordingly. 106func MarshalValue(val interface{}) (bsontype.Type, []byte, error) { 107 return MarshalValueWithRegistry(DefaultRegistry, val) 108} 109 110// MarshalValueAppend will append the BSON encoding of val to dst. If dst is not large enough to hold the BSON encoding 111// of val, dst will be grown. 112func MarshalValueAppend(dst []byte, val interface{}) (bsontype.Type, []byte, error) { 113 return MarshalValueAppendWithRegistry(DefaultRegistry, dst, val) 114} 115 116// MarshalValueWithRegistry returns the BSON encoding of val using Registry r. 117func MarshalValueWithRegistry(r *bsoncodec.Registry, val interface{}) (bsontype.Type, []byte, error) { 118 dst := make([]byte, 0, defaultDstCap) 119 return MarshalValueAppendWithRegistry(r, dst, val) 120} 121 122// MarshalValueWithContext returns the BSON encoding of val using EncodeContext ec. 123func MarshalValueWithContext(ec bsoncodec.EncodeContext, val interface{}) (bsontype.Type, []byte, error) { 124 dst := make([]byte, 0, defaultDstCap) 125 return MarshalValueAppendWithContext(ec, dst, val) 126} 127 128// MarshalValueAppendWithRegistry will append the BSON encoding of val to dst using Registry r. If dst is not large 129// enough to hold the BSON encoding of val, dst will be grown. 130func MarshalValueAppendWithRegistry(r *bsoncodec.Registry, dst []byte, val interface{}) (bsontype.Type, []byte, error) { 131 return MarshalValueAppendWithContext(bsoncodec.EncodeContext{Registry: r}, dst, val) 132} 133 134// MarshalValueAppendWithContext will append the BSON encoding of val to dst using EncodeContext ec. If dst is not large 135// enough to hold the BSON encoding of val, dst will be grown. 136func MarshalValueAppendWithContext(ec bsoncodec.EncodeContext, dst []byte, val interface{}) (bsontype.Type, []byte, error) { 137 // get a ValueWriter configured to write to dst 138 sw := new(bsonrw.SliceWriter) 139 *sw = dst 140 vwFlusher := bvwPool.GetAtModeElement(sw) 141 142 // get an Encoder and encode the value 143 enc := encPool.Get().(*Encoder) 144 defer encPool.Put(enc) 145 if err := enc.Reset(vwFlusher); err != nil { 146 return 0, nil, err 147 } 148 if err := enc.SetContext(ec); err != nil { 149 return 0, nil, err 150 } 151 if err := enc.Encode(val); err != nil { 152 return 0, nil, err 153 } 154 155 // flush the bytes written because we cannot guarantee that a full document has been written 156 // after the flush, *sw will be in the format 157 // [value type, 0 (null byte to indicate end of empty element name), value bytes..] 158 if err := vwFlusher.Flush(); err != nil { 159 return 0, nil, err 160 } 161 buffer := *sw 162 return bsontype.Type(buffer[0]), buffer[2:], nil 163} 164 165// MarshalExtJSON returns the extended JSON encoding of val. 166func MarshalExtJSON(val interface{}, canonical, escapeHTML bool) ([]byte, error) { 167 return MarshalExtJSONWithRegistry(DefaultRegistry, val, canonical, escapeHTML) 168} 169 170// MarshalExtJSONAppend will append the extended JSON encoding of val to dst. 171// If dst is not large enough to hold the extended JSON encoding of val, dst 172// will be grown. 173func MarshalExtJSONAppend(dst []byte, val interface{}, canonical, escapeHTML bool) ([]byte, error) { 174 return MarshalExtJSONAppendWithRegistry(DefaultRegistry, dst, val, canonical, escapeHTML) 175} 176 177// MarshalExtJSONWithRegistry returns the extended JSON encoding of val using Registry r. 178func MarshalExtJSONWithRegistry(r *bsoncodec.Registry, val interface{}, canonical, escapeHTML bool) ([]byte, error) { 179 dst := make([]byte, 0, defaultDstCap) 180 return MarshalExtJSONAppendWithContext(bsoncodec.EncodeContext{Registry: r}, dst, val, canonical, escapeHTML) 181} 182 183// MarshalExtJSONWithContext returns the extended JSON encoding of val using Registry r. 184func MarshalExtJSONWithContext(ec bsoncodec.EncodeContext, val interface{}, canonical, escapeHTML bool) ([]byte, error) { 185 dst := make([]byte, 0, defaultDstCap) 186 return MarshalExtJSONAppendWithContext(ec, dst, val, canonical, escapeHTML) 187} 188 189// MarshalExtJSONAppendWithRegistry will append the extended JSON encoding of 190// val to dst using Registry r. If dst is not large enough to hold the BSON 191// encoding of val, dst will be grown. 192func MarshalExtJSONAppendWithRegistry(r *bsoncodec.Registry, dst []byte, val interface{}, canonical, escapeHTML bool) ([]byte, error) { 193 return MarshalExtJSONAppendWithContext(bsoncodec.EncodeContext{Registry: r}, dst, val, canonical, escapeHTML) 194} 195 196// MarshalExtJSONAppendWithContext will append the extended JSON encoding of 197// val to dst using Registry r. If dst is not large enough to hold the BSON 198// encoding of val, dst will be grown. 199func MarshalExtJSONAppendWithContext(ec bsoncodec.EncodeContext, dst []byte, val interface{}, canonical, escapeHTML bool) ([]byte, error) { 200 sw := new(bsonrw.SliceWriter) 201 *sw = dst 202 ejvw := extjPool.Get(sw, canonical, escapeHTML) 203 defer extjPool.Put(ejvw) 204 205 enc := encPool.Get().(*Encoder) 206 defer encPool.Put(enc) 207 208 err := enc.Reset(ejvw) 209 if err != nil { 210 return nil, err 211 } 212 err = enc.SetContext(ec) 213 if err != nil { 214 return nil, err 215 } 216 217 err = enc.Encode(val) 218 if err != nil { 219 return nil, err 220 } 221 222 return *sw, nil 223} 224