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 bsonrw |
| 8 | |
| 9 | import ( |
| 10 | "fmt" |
| 11 | "io" |
| 12 | |
| 13 | "github.com/mongodb/mongo-go-driver/bson/bsontype" |
| 14 | "github.com/mongodb/mongo-go-driver/bson/primitive" |
| 15 | "github.com/mongodb/mongo-go-driver/x/bsonx/bsoncore" |
| 16 | ) |
| 17 | |
| 18 | // Copier is a type that allows copying between ValueReaders, ValueWriters, and |
| 19 | // []byte values. |
| 20 | type Copier struct{} |
| 21 | |
| 22 | // NewCopier creates a new copier with the given registry. If a nil registry is provided |
| 23 | // a default registry is used. |
| 24 | func NewCopier() Copier { |
| 25 | return Copier{} |
| 26 | } |
| 27 | |
| 28 | // CopyDocument handles copying a document from src to dst. |
| 29 | func CopyDocument(dst ValueWriter, src ValueReader) error { |
| 30 | return Copier{}.CopyDocument(dst, src) |
| 31 | } |
| 32 | |
| 33 | // CopyDocument handles copying one document from the src to the dst. |
| 34 | func (c Copier) CopyDocument(dst ValueWriter, src ValueReader) error { |
| 35 | dr, err := src.ReadDocument() |
| 36 | if err != nil { |
| 37 | return err |
| 38 | } |
| 39 | |
| 40 | dw, err := dst.WriteDocument() |
| 41 | if err != nil { |
| 42 | return err |
| 43 | } |
| 44 | |
| 45 | return c.copyDocumentCore(dw, dr) |
| 46 | } |
| 47 | |
| 48 | // CopyDocumentFromBytes copies the values from a BSON document represented as a |
| 49 | // []byte to a ValueWriter. |
| 50 | func (c Copier) CopyDocumentFromBytes(dst ValueWriter, src []byte) error { |
| 51 | dw, err := dst.WriteDocument() |
| 52 | if err != nil { |
| 53 | return err |
| 54 | } |
| 55 | |
| 56 | err = c.CopyBytesToDocumentWriter(dw, src) |
| 57 | if err != nil { |
| 58 | return err |
| 59 | } |
| 60 | |
| 61 | return dw.WriteDocumentEnd() |
| 62 | } |
| 63 | |
| 64 | // CopyBytesToDocumentWriter copies the values from a BSON document represented as a []byte to a |
| 65 | // DocumentWriter. |
| 66 | func (c Copier) CopyBytesToDocumentWriter(dst DocumentWriter, src []byte) error { |
| 67 | // TODO(skriptble): Create errors types here. Anything thats a tag should be a property. |
| 68 | length, rem, ok := bsoncore.ReadLength(src) |
| 69 | if !ok { |
| 70 | return fmt.Errorf("couldn't read length from src, not enough bytes. length=%d", len(src)) |
| 71 | } |
| 72 | if len(src) < int(length) { |
| 73 | return fmt.Errorf("length read exceeds number of bytes available. length=%d bytes=%d", len(src), length) |
| 74 | } |
| 75 | rem = rem[:length-4] |
| 76 | |
| 77 | var t bsontype.Type |
| 78 | var key string |
| 79 | var val bsoncore.Value |
| 80 | for { |
| 81 | t, rem, ok = bsoncore.ReadType(rem) |
| 82 | if !ok { |
| 83 | return io.EOF |
| 84 | } |
| 85 | if t == bsontype.Type(0) { |
| 86 | if len(rem) != 0 { |
| 87 | return fmt.Errorf("document end byte found before end of document. remaining bytes=%v", rem) |
| 88 | } |
| 89 | break |
| 90 | } |
| 91 | |
| 92 | key, rem, ok = bsoncore.ReadKey(rem) |
| 93 | if !ok { |
| 94 | return fmt.Errorf("invalid key found. remaining bytes=%v", rem) |
| 95 | } |
| 96 | dvw, err := dst.WriteDocumentElement(key) |
| 97 | if err != nil { |
| 98 | return err |
| 99 | } |
| 100 | val, rem, ok = bsoncore.ReadValue(rem, t) |
| 101 | if !ok { |
| 102 | return fmt.Errorf("not enough bytes available to read type. bytes=%d type=%s", len(rem), t) |
| 103 | } |
| 104 | err = c.CopyValueFromBytes(dvw, t, val.Data) |
| 105 | if err != nil { |
| 106 | return err |
| 107 | } |
| 108 | } |
| 109 | return nil |
| 110 | } |
| 111 | |
| 112 | // CopyDocumentToBytes copies an entire document from the ValueReader and |
| 113 | // returns it as bytes. |
| 114 | func (c Copier) CopyDocumentToBytes(src ValueReader) ([]byte, error) { |
| 115 | return c.AppendDocumentBytes(nil, src) |
| 116 | } |
| 117 | |
| 118 | // AppendDocumentBytes functions the same as CopyDocumentToBytes, but will |
| 119 | // append the result to dst. |
| 120 | func (c Copier) AppendDocumentBytes(dst []byte, src ValueReader) ([]byte, error) { |
| 121 | if br, ok := src.(BytesReader); ok { |
| 122 | _, dst, err := br.ReadValueBytes(dst) |
| 123 | return dst, err |
| 124 | } |
| 125 | |
| 126 | vw := vwPool.Get().(*valueWriter) |
| 127 | defer vwPool.Put(vw) |
| 128 | |
| 129 | vw.reset(dst) |
| 130 | |
| 131 | err := c.CopyDocument(vw, src) |
| 132 | dst = vw.buf |
| 133 | return dst, err |
| 134 | } |
| 135 | |
| 136 | // CopyValueFromBytes will write the value represtend by t and src to dst. |
| 137 | func (c Copier) CopyValueFromBytes(dst ValueWriter, t bsontype.Type, src []byte) error { |
| 138 | if wvb, ok := dst.(BytesWriter); ok { |
| 139 | return wvb.WriteValueBytes(t, src) |
| 140 | } |
| 141 | |
| 142 | vr := vrPool.Get().(*valueReader) |
| 143 | defer vrPool.Put(vr) |
| 144 | |
| 145 | vr.reset(src) |
| 146 | vr.pushElement(t) |
| 147 | |
| 148 | return c.CopyValue(dst, vr) |
| 149 | } |
| 150 | |
| 151 | // CopyValueToBytes copies a value from src and returns it as a bsontype.Type and a |
| 152 | // []byte. |
| 153 | func (c Copier) CopyValueToBytes(src ValueReader) (bsontype.Type, []byte, error) { |
| 154 | return c.AppendValueBytes(nil, src) |
| 155 | } |
| 156 | |
| 157 | // AppendValueBytes functions the same as CopyValueToBytes, but will append the |
| 158 | // result to dst. |
| 159 | func (c Copier) AppendValueBytes(dst []byte, src ValueReader) (bsontype.Type, []byte, error) { |
| 160 | if br, ok := src.(BytesReader); ok { |
| 161 | return br.ReadValueBytes(dst) |
| 162 | } |
| 163 | |
| 164 | vw := vwPool.Get().(*valueWriter) |
| 165 | defer vwPool.Put(vw) |
| 166 | |
| 167 | start := len(dst) |
| 168 | |
| 169 | vw.reset(dst) |
| 170 | vw.push(mElement) |
| 171 | |
| 172 | err := c.CopyValue(vw, src) |
| 173 | if err != nil { |
| 174 | return 0, dst, err |
| 175 | } |
| 176 | |
| 177 | return bsontype.Type(vw.buf[start]), vw.buf[start+2:], nil |
| 178 | } |
| 179 | |
| 180 | // CopyValue will copy a single value from src to dst. |
| 181 | func (c Copier) CopyValue(dst ValueWriter, src ValueReader) error { |
| 182 | var err error |
| 183 | switch src.Type() { |
| 184 | case bsontype.Double: |
| 185 | var f64 float64 |
| 186 | f64, err = src.ReadDouble() |
| 187 | if err != nil { |
| 188 | break |
| 189 | } |
| 190 | err = dst.WriteDouble(f64) |
| 191 | case bsontype.String: |
| 192 | var str string |
| 193 | str, err = src.ReadString() |
| 194 | if err != nil { |
| 195 | return err |
| 196 | } |
| 197 | err = dst.WriteString(str) |
| 198 | case bsontype.EmbeddedDocument: |
| 199 | err = c.CopyDocument(dst, src) |
| 200 | case bsontype.Array: |
| 201 | err = c.copyArray(dst, src) |
| 202 | case bsontype.Binary: |
| 203 | var data []byte |
| 204 | var subtype byte |
| 205 | data, subtype, err = src.ReadBinary() |
| 206 | if err != nil { |
| 207 | break |
| 208 | } |
| 209 | err = dst.WriteBinaryWithSubtype(data, subtype) |
| 210 | case bsontype.Undefined: |
| 211 | err = src.ReadUndefined() |
| 212 | if err != nil { |
| 213 | break |
| 214 | } |
| 215 | err = dst.WriteUndefined() |
| 216 | case bsontype.ObjectID: |
| 217 | var oid primitive.ObjectID |
| 218 | oid, err = src.ReadObjectID() |
| 219 | if err != nil { |
| 220 | break |
| 221 | } |
| 222 | err = dst.WriteObjectID(oid) |
| 223 | case bsontype.Boolean: |
| 224 | var b bool |
| 225 | b, err = src.ReadBoolean() |
| 226 | if err != nil { |
| 227 | break |
| 228 | } |
| 229 | err = dst.WriteBoolean(b) |
| 230 | case bsontype.DateTime: |
| 231 | var dt int64 |
| 232 | dt, err = src.ReadDateTime() |
| 233 | if err != nil { |
| 234 | break |
| 235 | } |
| 236 | err = dst.WriteDateTime(dt) |
| 237 | case bsontype.Null: |
| 238 | err = src.ReadNull() |
| 239 | if err != nil { |
| 240 | break |
| 241 | } |
| 242 | err = dst.WriteNull() |
| 243 | case bsontype.Regex: |
| 244 | var pattern, options string |
| 245 | pattern, options, err = src.ReadRegex() |
| 246 | if err != nil { |
| 247 | break |
| 248 | } |
| 249 | err = dst.WriteRegex(pattern, options) |
| 250 | case bsontype.DBPointer: |
| 251 | var ns string |
| 252 | var pointer primitive.ObjectID |
| 253 | ns, pointer, err = src.ReadDBPointer() |
| 254 | if err != nil { |
| 255 | break |
| 256 | } |
| 257 | err = dst.WriteDBPointer(ns, pointer) |
| 258 | case bsontype.JavaScript: |
| 259 | var js string |
| 260 | js, err = src.ReadJavascript() |
| 261 | if err != nil { |
| 262 | break |
| 263 | } |
| 264 | err = dst.WriteJavascript(js) |
| 265 | case bsontype.Symbol: |
| 266 | var symbol string |
| 267 | symbol, err = src.ReadSymbol() |
| 268 | if err != nil { |
| 269 | break |
| 270 | } |
| 271 | err = dst.WriteSymbol(symbol) |
| 272 | case bsontype.CodeWithScope: |
| 273 | var code string |
| 274 | var srcScope DocumentReader |
| 275 | code, srcScope, err = src.ReadCodeWithScope() |
| 276 | if err != nil { |
| 277 | break |
| 278 | } |
| 279 | |
| 280 | var dstScope DocumentWriter |
| 281 | dstScope, err = dst.WriteCodeWithScope(code) |
| 282 | if err != nil { |
| 283 | break |
| 284 | } |
| 285 | err = c.copyDocumentCore(dstScope, srcScope) |
| 286 | case bsontype.Int32: |
| 287 | var i32 int32 |
| 288 | i32, err = src.ReadInt32() |
| 289 | if err != nil { |
| 290 | break |
| 291 | } |
| 292 | err = dst.WriteInt32(i32) |
| 293 | case bsontype.Timestamp: |
| 294 | var t, i uint32 |
| 295 | t, i, err = src.ReadTimestamp() |
| 296 | if err != nil { |
| 297 | break |
| 298 | } |
| 299 | err = dst.WriteTimestamp(t, i) |
| 300 | case bsontype.Int64: |
| 301 | var i64 int64 |
| 302 | i64, err = src.ReadInt64() |
| 303 | if err != nil { |
| 304 | break |
| 305 | } |
| 306 | err = dst.WriteInt64(i64) |
| 307 | case bsontype.Decimal128: |
| 308 | var d128 primitive.Decimal128 |
| 309 | d128, err = src.ReadDecimal128() |
| 310 | if err != nil { |
| 311 | break |
| 312 | } |
| 313 | err = dst.WriteDecimal128(d128) |
| 314 | case bsontype.MinKey: |
| 315 | err = src.ReadMinKey() |
| 316 | if err != nil { |
| 317 | break |
| 318 | } |
| 319 | err = dst.WriteMinKey() |
| 320 | case bsontype.MaxKey: |
| 321 | err = src.ReadMaxKey() |
| 322 | if err != nil { |
| 323 | break |
| 324 | } |
| 325 | err = dst.WriteMaxKey() |
| 326 | default: |
| 327 | err = fmt.Errorf("Cannot copy unknown BSON type %s", src.Type()) |
| 328 | } |
| 329 | |
| 330 | return err |
| 331 | } |
| 332 | |
| 333 | func (c Copier) copyArray(dst ValueWriter, src ValueReader) error { |
| 334 | ar, err := src.ReadArray() |
| 335 | if err != nil { |
| 336 | return err |
| 337 | } |
| 338 | |
| 339 | aw, err := dst.WriteArray() |
| 340 | if err != nil { |
| 341 | return err |
| 342 | } |
| 343 | |
| 344 | for { |
| 345 | vr, err := ar.ReadValue() |
| 346 | if err == ErrEOA { |
| 347 | break |
| 348 | } |
| 349 | if err != nil { |
| 350 | return err |
| 351 | } |
| 352 | |
| 353 | vw, err := aw.WriteArrayElement() |
| 354 | if err != nil { |
| 355 | return err |
| 356 | } |
| 357 | |
| 358 | err = c.CopyValue(vw, vr) |
| 359 | if err != nil { |
| 360 | return err |
| 361 | } |
| 362 | } |
| 363 | |
| 364 | return aw.WriteArrayEnd() |
| 365 | } |
| 366 | |
| 367 | func (c Copier) copyDocumentCore(dw DocumentWriter, dr DocumentReader) error { |
| 368 | for { |
| 369 | key, vr, err := dr.ReadElement() |
| 370 | if err == ErrEOD { |
| 371 | break |
| 372 | } |
| 373 | if err != nil { |
| 374 | return err |
| 375 | } |
| 376 | |
| 377 | vw, err := dw.WriteDocumentElement(key) |
| 378 | if err != nil { |
| 379 | return err |
| 380 | } |
| 381 | |
| 382 | err = c.CopyValue(vw, vr) |
| 383 | if err != nil { |
| 384 | return err |
| 385 | } |
| 386 | } |
| 387 | |
| 388 | return dw.WriteDocumentEnd() |
| 389 | } |