Don Newton | 379ae25 | 2019-04-01 12:17:06 -0400 | [diff] [blame^] | 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 | |
| 7 | package bson |
| 8 | |
| 9 | import ( |
| 10 | "github.com/mongodb/mongo-go-driver/bson/bsoncodec" |
| 11 | "github.com/mongodb/mongo-go-driver/bson/bsonrw" |
| 12 | "github.com/mongodb/mongo-go-driver/bson/bsontype" |
| 13 | ) |
| 14 | |
| 15 | const defaultDstCap = 256 |
| 16 | |
| 17 | var bvwPool = bsonrw.NewBSONValueWriterPool() |
| 18 | var 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. |
| 23 | type 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. |
| 31 | type ValueMarshaler interface { |
| 32 | MarshalBSONValue() (bsontype.Type, []byte, error) |
| 33 | } |
| 34 | |
| 35 | // Marshal returns the BSON encoding of val. |
| 36 | // |
| 37 | // Marshal will use the default registry created by NewRegistry to recursively |
| 38 | // marshal val into a []byte. Marshal will inspect struct tags and alter the |
| 39 | // marshaling process accordingly. |
| 40 | func Marshal(val interface{}) ([]byte, error) { |
| 41 | return MarshalWithRegistry(DefaultRegistry, val) |
| 42 | } |
| 43 | |
| 44 | // MarshalAppend will append the BSON encoding of val to dst. If dst is not |
| 45 | // large enough to hold the BSON encoding of val, dst will be grown. |
| 46 | func MarshalAppend(dst []byte, val interface{}) ([]byte, error) { |
| 47 | return MarshalAppendWithRegistry(DefaultRegistry, dst, val) |
| 48 | } |
| 49 | |
| 50 | // MarshalWithRegistry returns the BSON encoding of val using Registry r. |
| 51 | func MarshalWithRegistry(r *bsoncodec.Registry, val interface{}) ([]byte, error) { |
| 52 | dst := make([]byte, 0, 256) // TODO: make the default cap a constant |
| 53 | return MarshalAppendWithRegistry(r, dst, val) |
| 54 | } |
| 55 | |
| 56 | // MarshalWithContext returns the BSON encoding of val using EncodeContext ec. |
| 57 | func MarshalWithContext(ec bsoncodec.EncodeContext, val interface{}) ([]byte, error) { |
| 58 | dst := make([]byte, 0, 256) // TODO: make the default cap a constant |
| 59 | return MarshalAppendWithContext(ec, dst, val) |
| 60 | } |
| 61 | |
| 62 | // MarshalAppendWithRegistry will append the BSON encoding of val to dst using |
| 63 | // Registry r. If dst is not large enough to hold the BSON encoding of val, dst |
| 64 | // will be grown. |
| 65 | func MarshalAppendWithRegistry(r *bsoncodec.Registry, dst []byte, val interface{}) ([]byte, error) { |
| 66 | return MarshalAppendWithContext(bsoncodec.EncodeContext{Registry: r}, dst, val) |
| 67 | } |
| 68 | |
| 69 | // MarshalAppendWithContext will append the BSON encoding of val to dst using |
| 70 | // EncodeContext ec. If dst is not large enough to hold the BSON encoding of val, dst |
| 71 | // will be grown. |
| 72 | func MarshalAppendWithContext(ec bsoncodec.EncodeContext, dst []byte, val interface{}) ([]byte, error) { |
| 73 | sw := new(bsonrw.SliceWriter) |
| 74 | *sw = dst |
| 75 | vw := bvwPool.Get(sw) |
| 76 | defer bvwPool.Put(vw) |
| 77 | |
| 78 | enc := encPool.Get().(*Encoder) |
| 79 | defer encPool.Put(enc) |
| 80 | |
| 81 | err := enc.Reset(vw) |
| 82 | if err != nil { |
| 83 | return nil, err |
| 84 | } |
| 85 | err = enc.SetContext(ec) |
| 86 | if err != nil { |
| 87 | return nil, err |
| 88 | } |
| 89 | |
| 90 | err = enc.Encode(val) |
| 91 | if err != nil { |
| 92 | return nil, err |
| 93 | } |
| 94 | |
| 95 | return *sw, nil |
| 96 | } |
| 97 | |
| 98 | // MarshalExtJSON returns the extended JSON encoding of val. |
| 99 | func MarshalExtJSON(val interface{}, canonical, escapeHTML bool) ([]byte, error) { |
| 100 | return MarshalExtJSONWithRegistry(DefaultRegistry, val, canonical, escapeHTML) |
| 101 | } |
| 102 | |
| 103 | // MarshalExtJSONAppend will append the extended JSON encoding of val to dst. |
| 104 | // If dst is not large enough to hold the extended JSON encoding of val, dst |
| 105 | // will be grown. |
| 106 | func MarshalExtJSONAppend(dst []byte, val interface{}, canonical, escapeHTML bool) ([]byte, error) { |
| 107 | return MarshalExtJSONAppendWithRegistry(DefaultRegistry, dst, val, canonical, escapeHTML) |
| 108 | } |
| 109 | |
| 110 | // MarshalExtJSONWithRegistry returns the extended JSON encoding of val using Registry r. |
| 111 | func MarshalExtJSONWithRegistry(r *bsoncodec.Registry, val interface{}, canonical, escapeHTML bool) ([]byte, error) { |
| 112 | dst := make([]byte, 0, defaultDstCap) |
| 113 | return MarshalExtJSONAppendWithContext(bsoncodec.EncodeContext{Registry: r}, dst, val, canonical, escapeHTML) |
| 114 | } |
| 115 | |
| 116 | // MarshalExtJSONWithContext returns the extended JSON encoding of val using Registry r. |
| 117 | func MarshalExtJSONWithContext(ec bsoncodec.EncodeContext, val interface{}, canonical, escapeHTML bool) ([]byte, error) { |
| 118 | dst := make([]byte, 0, defaultDstCap) |
| 119 | return MarshalExtJSONAppendWithContext(ec, dst, val, canonical, escapeHTML) |
| 120 | } |
| 121 | |
| 122 | // MarshalExtJSONAppendWithRegistry will append the extended JSON encoding of |
| 123 | // val to dst using Registry r. If dst is not large enough to hold the BSON |
| 124 | // encoding of val, dst will be grown. |
| 125 | func MarshalExtJSONAppendWithRegistry(r *bsoncodec.Registry, dst []byte, val interface{}, canonical, escapeHTML bool) ([]byte, error) { |
| 126 | return MarshalExtJSONAppendWithContext(bsoncodec.EncodeContext{Registry: r}, dst, val, canonical, escapeHTML) |
| 127 | } |
| 128 | |
| 129 | // MarshalExtJSONAppendWithContext will append the extended JSON encoding of |
| 130 | // val to dst using Registry r. If dst is not large enough to hold the BSON |
| 131 | // encoding of val, dst will be grown. |
| 132 | func MarshalExtJSONAppendWithContext(ec bsoncodec.EncodeContext, dst []byte, val interface{}, canonical, escapeHTML bool) ([]byte, error) { |
| 133 | sw := new(bsonrw.SliceWriter) |
| 134 | *sw = dst |
| 135 | ejvw := extjPool.Get(sw, canonical, escapeHTML) |
| 136 | defer extjPool.Put(ejvw) |
| 137 | |
| 138 | enc := encPool.Get().(*Encoder) |
| 139 | defer encPool.Put(enc) |
| 140 | |
| 141 | err := enc.Reset(ejvw) |
| 142 | if err != nil { |
| 143 | return nil, err |
| 144 | } |
| 145 | err = enc.SetContext(ec) |
| 146 | if err != nil { |
| 147 | return nil, err |
| 148 | } |
| 149 | |
| 150 | err = enc.Encode(val) |
| 151 | if err != nil { |
| 152 | return nil, err |
| 153 | } |
| 154 | |
| 155 | return *sw, nil |
| 156 | } |