value_writer.go raw
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 "errors"
11 "fmt"
12 "io"
13 "math"
14 "strconv"
15 "strings"
16 "sync"
17
18 "go.mongodb.org/mongo-driver/bson/bsontype"
19 "go.mongodb.org/mongo-driver/bson/primitive"
20 "go.mongodb.org/mongo-driver/x/bsonx/bsoncore"
21 )
22
23 var _ ValueWriter = (*valueWriter)(nil)
24
25 var vwPool = sync.Pool{
26 New: func() interface{} {
27 return new(valueWriter)
28 },
29 }
30
31 func putValueWriter(vw *valueWriter) {
32 if vw != nil {
33 vw.w = nil // don't leak the writer
34 vwPool.Put(vw)
35 }
36 }
37
38 // BSONValueWriterPool is a pool for BSON ValueWriters.
39 //
40 // Deprecated: BSONValueWriterPool will not be supported in Go Driver 2.0.
41 type BSONValueWriterPool struct {
42 pool sync.Pool
43 }
44
45 // NewBSONValueWriterPool creates a new pool for ValueWriter instances that write to BSON.
46 //
47 // Deprecated: BSONValueWriterPool will not be supported in Go Driver 2.0.
48 func NewBSONValueWriterPool() *BSONValueWriterPool {
49 return &BSONValueWriterPool{
50 pool: sync.Pool{
51 New: func() interface{} {
52 return new(valueWriter)
53 },
54 },
55 }
56 }
57
58 // Get retrieves a BSON ValueWriter from the pool and resets it to use w as the destination.
59 //
60 // Deprecated: BSONValueWriterPool will not be supported in Go Driver 2.0.
61 func (bvwp *BSONValueWriterPool) Get(w io.Writer) ValueWriter {
62 vw := bvwp.pool.Get().(*valueWriter)
63
64 // TODO: Having to call reset here with the same buffer doesn't really make sense.
65 vw.reset(vw.buf)
66 vw.buf = vw.buf[:0]
67 vw.w = w
68 return vw
69 }
70
71 // GetAtModeElement retrieves a ValueWriterFlusher from the pool and resets it to use w as the destination.
72 //
73 // Deprecated: BSONValueWriterPool will not be supported in Go Driver 2.0.
74 func (bvwp *BSONValueWriterPool) GetAtModeElement(w io.Writer) ValueWriterFlusher {
75 vw := bvwp.Get(w).(*valueWriter)
76 vw.push(mElement)
77 return vw
78 }
79
80 // Put inserts a ValueWriter into the pool. If the ValueWriter is not a BSON ValueWriter, nothing
81 // happens and ok will be false.
82 //
83 // Deprecated: BSONValueWriterPool will not be supported in Go Driver 2.0.
84 func (bvwp *BSONValueWriterPool) Put(vw ValueWriter) (ok bool) {
85 bvw, ok := vw.(*valueWriter)
86 if !ok {
87 return false
88 }
89
90 bvwp.pool.Put(bvw)
91 return true
92 }
93
94 // This is here so that during testing we can change it and not require
95 // allocating a 4GB slice.
96 var maxSize = math.MaxInt32
97
98 var errNilWriter = errors.New("cannot create a ValueWriter from a nil io.Writer")
99
100 type errMaxDocumentSizeExceeded struct {
101 size int64
102 }
103
104 func (mdse errMaxDocumentSizeExceeded) Error() string {
105 return fmt.Sprintf("document size (%d) is larger than the max int32", mdse.size)
106 }
107
108 type vwMode int
109
110 const (
111 _ vwMode = iota
112 vwTopLevel
113 vwDocument
114 vwArray
115 vwValue
116 vwElement
117 vwCodeWithScope
118 )
119
120 func (vm vwMode) String() string {
121 var str string
122
123 switch vm {
124 case vwTopLevel:
125 str = "TopLevel"
126 case vwDocument:
127 str = "DocumentMode"
128 case vwArray:
129 str = "ArrayMode"
130 case vwValue:
131 str = "ValueMode"
132 case vwElement:
133 str = "ElementMode"
134 case vwCodeWithScope:
135 str = "CodeWithScopeMode"
136 default:
137 str = "UnknownMode"
138 }
139
140 return str
141 }
142
143 type vwState struct {
144 mode mode
145 key string
146 arrkey int
147 start int32
148 }
149
150 type valueWriter struct {
151 w io.Writer
152 buf []byte
153
154 stack []vwState
155 frame int64
156 }
157
158 func (vw *valueWriter) advanceFrame() {
159 vw.frame++
160 if vw.frame >= int64(len(vw.stack)) {
161 vw.stack = append(vw.stack, vwState{})
162 }
163 }
164
165 func (vw *valueWriter) push(m mode) {
166 vw.advanceFrame()
167
168 // Clean the stack
169 vw.stack[vw.frame] = vwState{mode: m}
170
171 switch m {
172 case mDocument, mArray, mCodeWithScope:
173 vw.reserveLength() // WARN: this is not needed
174 }
175 }
176
177 func (vw *valueWriter) reserveLength() {
178 vw.stack[vw.frame].start = int32(len(vw.buf))
179 vw.buf = append(vw.buf, 0x00, 0x00, 0x00, 0x00)
180 }
181
182 func (vw *valueWriter) pop() {
183 switch vw.stack[vw.frame].mode {
184 case mElement, mValue:
185 vw.frame--
186 case mDocument, mArray, mCodeWithScope:
187 vw.frame -= 2 // we pop twice to jump over the mElement: mDocument -> mElement -> mDocument/mTopLevel/etc...
188 }
189 }
190
191 // NewBSONValueWriter creates a ValueWriter that writes BSON to w.
192 //
193 // This ValueWriter will only write entire documents to the io.Writer and it
194 // will buffer the document as it is built.
195 func NewBSONValueWriter(w io.Writer) (ValueWriter, error) {
196 if w == nil {
197 return nil, errNilWriter
198 }
199 return newValueWriter(w), nil
200 }
201
202 func newValueWriter(w io.Writer) *valueWriter {
203 vw := new(valueWriter)
204 stack := make([]vwState, 1, 5)
205 stack[0] = vwState{mode: mTopLevel}
206 vw.w = w
207 vw.stack = stack
208
209 return vw
210 }
211
212 // TODO: only used in tests
213 func newValueWriterFromSlice(buf []byte) *valueWriter {
214 vw := new(valueWriter)
215 stack := make([]vwState, 1, 5)
216 stack[0] = vwState{mode: mTopLevel}
217 vw.stack = stack
218 vw.buf = buf
219
220 return vw
221 }
222
223 func (vw *valueWriter) reset(buf []byte) {
224 if vw.stack == nil {
225 vw.stack = make([]vwState, 1, 5)
226 }
227 vw.stack = vw.stack[:1]
228 vw.stack[0] = vwState{mode: mTopLevel}
229 vw.buf = buf
230 vw.frame = 0
231 vw.w = nil
232 }
233
234 func (vw *valueWriter) invalidTransitionError(destination mode, name string, modes []mode) error {
235 te := TransitionError{
236 name: name,
237 current: vw.stack[vw.frame].mode,
238 destination: destination,
239 modes: modes,
240 action: "write",
241 }
242 if vw.frame != 0 {
243 te.parent = vw.stack[vw.frame-1].mode
244 }
245 return te
246 }
247
248 func (vw *valueWriter) writeElementHeader(t bsontype.Type, destination mode, callerName string, addmodes ...mode) error {
249 frame := &vw.stack[vw.frame]
250 switch frame.mode {
251 case mElement:
252 key := frame.key
253 if !isValidCString(key) {
254 return errors.New("BSON element key cannot contain null bytes")
255 }
256 vw.appendHeader(t, key)
257 case mValue:
258 vw.appendIntHeader(t, frame.arrkey)
259 default:
260 modes := []mode{mElement, mValue}
261 if addmodes != nil {
262 modes = append(modes, addmodes...)
263 }
264 return vw.invalidTransitionError(destination, callerName, modes)
265 }
266
267 return nil
268 }
269
270 func (vw *valueWriter) WriteValueBytes(t bsontype.Type, b []byte) error {
271 if err := vw.writeElementHeader(t, mode(0), "WriteValueBytes"); err != nil {
272 return err
273 }
274 vw.buf = append(vw.buf, b...)
275 vw.pop()
276 return nil
277 }
278
279 func (vw *valueWriter) WriteArray() (ArrayWriter, error) {
280 if err := vw.writeElementHeader(bsontype.Array, mArray, "WriteArray"); err != nil {
281 return nil, err
282 }
283
284 vw.push(mArray)
285
286 return vw, nil
287 }
288
289 func (vw *valueWriter) WriteBinary(b []byte) error {
290 return vw.WriteBinaryWithSubtype(b, 0x00)
291 }
292
293 func (vw *valueWriter) WriteBinaryWithSubtype(b []byte, btype byte) error {
294 if err := vw.writeElementHeader(bsontype.Binary, mode(0), "WriteBinaryWithSubtype"); err != nil {
295 return err
296 }
297
298 vw.buf = bsoncore.AppendBinary(vw.buf, btype, b)
299 vw.pop()
300 return nil
301 }
302
303 func (vw *valueWriter) WriteBoolean(b bool) error {
304 if err := vw.writeElementHeader(bsontype.Boolean, mode(0), "WriteBoolean"); err != nil {
305 return err
306 }
307
308 vw.buf = bsoncore.AppendBoolean(vw.buf, b)
309 vw.pop()
310 return nil
311 }
312
313 func (vw *valueWriter) WriteCodeWithScope(code string) (DocumentWriter, error) {
314 if err := vw.writeElementHeader(bsontype.CodeWithScope, mCodeWithScope, "WriteCodeWithScope"); err != nil {
315 return nil, err
316 }
317
318 // CodeWithScope is a different than other types because we need an extra
319 // frame on the stack. In the EndDocument code, we write the document
320 // length, pop, write the code with scope length, and pop. To simplify the
321 // pop code, we push a spacer frame that we'll always jump over.
322 vw.push(mCodeWithScope)
323 vw.buf = bsoncore.AppendString(vw.buf, code)
324 vw.push(mSpacer)
325 vw.push(mDocument)
326
327 return vw, nil
328 }
329
330 func (vw *valueWriter) WriteDBPointer(ns string, oid primitive.ObjectID) error {
331 if err := vw.writeElementHeader(bsontype.DBPointer, mode(0), "WriteDBPointer"); err != nil {
332 return err
333 }
334
335 vw.buf = bsoncore.AppendDBPointer(vw.buf, ns, oid)
336 vw.pop()
337 return nil
338 }
339
340 func (vw *valueWriter) WriteDateTime(dt int64) error {
341 if err := vw.writeElementHeader(bsontype.DateTime, mode(0), "WriteDateTime"); err != nil {
342 return err
343 }
344
345 vw.buf = bsoncore.AppendDateTime(vw.buf, dt)
346 vw.pop()
347 return nil
348 }
349
350 func (vw *valueWriter) WriteDecimal128(d128 primitive.Decimal128) error {
351 if err := vw.writeElementHeader(bsontype.Decimal128, mode(0), "WriteDecimal128"); err != nil {
352 return err
353 }
354
355 vw.buf = bsoncore.AppendDecimal128(vw.buf, d128)
356 vw.pop()
357 return nil
358 }
359
360 func (vw *valueWriter) WriteDouble(f float64) error {
361 if err := vw.writeElementHeader(bsontype.Double, mode(0), "WriteDouble"); err != nil {
362 return err
363 }
364
365 vw.buf = bsoncore.AppendDouble(vw.buf, f)
366 vw.pop()
367 return nil
368 }
369
370 func (vw *valueWriter) WriteInt32(i32 int32) error {
371 if err := vw.writeElementHeader(bsontype.Int32, mode(0), "WriteInt32"); err != nil {
372 return err
373 }
374
375 vw.buf = bsoncore.AppendInt32(vw.buf, i32)
376 vw.pop()
377 return nil
378 }
379
380 func (vw *valueWriter) WriteInt64(i64 int64) error {
381 if err := vw.writeElementHeader(bsontype.Int64, mode(0), "WriteInt64"); err != nil {
382 return err
383 }
384
385 vw.buf = bsoncore.AppendInt64(vw.buf, i64)
386 vw.pop()
387 return nil
388 }
389
390 func (vw *valueWriter) WriteJavascript(code string) error {
391 if err := vw.writeElementHeader(bsontype.JavaScript, mode(0), "WriteJavascript"); err != nil {
392 return err
393 }
394
395 vw.buf = bsoncore.AppendJavaScript(vw.buf, code)
396 vw.pop()
397 return nil
398 }
399
400 func (vw *valueWriter) WriteMaxKey() error {
401 if err := vw.writeElementHeader(bsontype.MaxKey, mode(0), "WriteMaxKey"); err != nil {
402 return err
403 }
404
405 vw.pop()
406 return nil
407 }
408
409 func (vw *valueWriter) WriteMinKey() error {
410 if err := vw.writeElementHeader(bsontype.MinKey, mode(0), "WriteMinKey"); err != nil {
411 return err
412 }
413
414 vw.pop()
415 return nil
416 }
417
418 func (vw *valueWriter) WriteNull() error {
419 if err := vw.writeElementHeader(bsontype.Null, mode(0), "WriteNull"); err != nil {
420 return err
421 }
422
423 vw.pop()
424 return nil
425 }
426
427 func (vw *valueWriter) WriteObjectID(oid primitive.ObjectID) error {
428 if err := vw.writeElementHeader(bsontype.ObjectID, mode(0), "WriteObjectID"); err != nil {
429 return err
430 }
431
432 vw.buf = bsoncore.AppendObjectID(vw.buf, oid)
433 vw.pop()
434 return nil
435 }
436
437 func (vw *valueWriter) WriteRegex(pattern string, options string) error {
438 if !isValidCString(pattern) || !isValidCString(options) {
439 return errors.New("BSON regex values cannot contain null bytes")
440 }
441 if err := vw.writeElementHeader(bsontype.Regex, mode(0), "WriteRegex"); err != nil {
442 return err
443 }
444
445 vw.buf = bsoncore.AppendRegex(vw.buf, pattern, sortStringAlphebeticAscending(options))
446 vw.pop()
447 return nil
448 }
449
450 func (vw *valueWriter) WriteString(s string) error {
451 if err := vw.writeElementHeader(bsontype.String, mode(0), "WriteString"); err != nil {
452 return err
453 }
454
455 vw.buf = bsoncore.AppendString(vw.buf, s)
456 vw.pop()
457 return nil
458 }
459
460 func (vw *valueWriter) WriteDocument() (DocumentWriter, error) {
461 if vw.stack[vw.frame].mode == mTopLevel {
462 vw.reserveLength()
463 return vw, nil
464 }
465 if err := vw.writeElementHeader(bsontype.EmbeddedDocument, mDocument, "WriteDocument", mTopLevel); err != nil {
466 return nil, err
467 }
468
469 vw.push(mDocument)
470 return vw, nil
471 }
472
473 func (vw *valueWriter) WriteSymbol(symbol string) error {
474 if err := vw.writeElementHeader(bsontype.Symbol, mode(0), "WriteSymbol"); err != nil {
475 return err
476 }
477
478 vw.buf = bsoncore.AppendSymbol(vw.buf, symbol)
479 vw.pop()
480 return nil
481 }
482
483 func (vw *valueWriter) WriteTimestamp(t uint32, i uint32) error {
484 if err := vw.writeElementHeader(bsontype.Timestamp, mode(0), "WriteTimestamp"); err != nil {
485 return err
486 }
487
488 vw.buf = bsoncore.AppendTimestamp(vw.buf, t, i)
489 vw.pop()
490 return nil
491 }
492
493 func (vw *valueWriter) WriteUndefined() error {
494 if err := vw.writeElementHeader(bsontype.Undefined, mode(0), "WriteUndefined"); err != nil {
495 return err
496 }
497
498 vw.pop()
499 return nil
500 }
501
502 func (vw *valueWriter) WriteDocumentElement(key string) (ValueWriter, error) {
503 switch vw.stack[vw.frame].mode {
504 case mTopLevel, mDocument:
505 default:
506 return nil, vw.invalidTransitionError(mElement, "WriteDocumentElement", []mode{mTopLevel, mDocument})
507 }
508
509 vw.push(mElement)
510 vw.stack[vw.frame].key = key
511
512 return vw, nil
513 }
514
515 func (vw *valueWriter) WriteDocumentEnd() error {
516 switch vw.stack[vw.frame].mode {
517 case mTopLevel, mDocument:
518 default:
519 return fmt.Errorf("incorrect mode to end document: %s", vw.stack[vw.frame].mode)
520 }
521
522 vw.buf = append(vw.buf, 0x00)
523
524 err := vw.writeLength()
525 if err != nil {
526 return err
527 }
528
529 if vw.stack[vw.frame].mode == mTopLevel {
530 if err = vw.Flush(); err != nil {
531 return err
532 }
533 }
534
535 vw.pop()
536
537 if vw.stack[vw.frame].mode == mCodeWithScope {
538 // We ignore the error here because of the guarantee of writeLength.
539 // See the docs for writeLength for more info.
540 _ = vw.writeLength()
541 vw.pop()
542 }
543 return nil
544 }
545
546 func (vw *valueWriter) Flush() error {
547 if vw.w == nil {
548 return nil
549 }
550
551 if _, err := vw.w.Write(vw.buf); err != nil {
552 return err
553 }
554 // reset buffer
555 vw.buf = vw.buf[:0]
556 return nil
557 }
558
559 func (vw *valueWriter) WriteArrayElement() (ValueWriter, error) {
560 if vw.stack[vw.frame].mode != mArray {
561 return nil, vw.invalidTransitionError(mValue, "WriteArrayElement", []mode{mArray})
562 }
563
564 arrkey := vw.stack[vw.frame].arrkey
565 vw.stack[vw.frame].arrkey++
566
567 vw.push(mValue)
568 vw.stack[vw.frame].arrkey = arrkey
569
570 return vw, nil
571 }
572
573 func (vw *valueWriter) WriteArrayEnd() error {
574 if vw.stack[vw.frame].mode != mArray {
575 return fmt.Errorf("incorrect mode to end array: %s", vw.stack[vw.frame].mode)
576 }
577
578 vw.buf = append(vw.buf, 0x00)
579
580 err := vw.writeLength()
581 if err != nil {
582 return err
583 }
584
585 vw.pop()
586 return nil
587 }
588
589 // NOTE: We assume that if we call writeLength more than once the same function
590 // within the same function without altering the vw.buf that this method will
591 // not return an error. If this changes ensure that the following methods are
592 // updated:
593 //
594 // - WriteDocumentEnd
595 func (vw *valueWriter) writeLength() error {
596 length := len(vw.buf)
597 if length > maxSize {
598 return errMaxDocumentSizeExceeded{size: int64(len(vw.buf))}
599 }
600 frame := &vw.stack[vw.frame]
601 length = length - int(frame.start)
602 start := frame.start
603
604 _ = vw.buf[start+3] // BCE
605 vw.buf[start+0] = byte(length)
606 vw.buf[start+1] = byte(length >> 8)
607 vw.buf[start+2] = byte(length >> 16)
608 vw.buf[start+3] = byte(length >> 24)
609 return nil
610 }
611
612 func isValidCString(cs string) bool {
613 // Disallow the zero byte in a cstring because the zero byte is used as the
614 // terminating character.
615 //
616 // It's safe to check bytes instead of runes because all multibyte UTF-8
617 // code points start with (binary) 11xxxxxx or 10xxxxxx, so 00000000 (i.e.
618 // 0) will never be part of a multibyte UTF-8 code point. This logic is the
619 // same as the "r < utf8.RuneSelf" case in strings.IndexRune but can be
620 // inlined.
621 //
622 // https://cs.opensource.google/go/go/+/refs/tags/go1.21.1:src/strings/strings.go;l=127
623 return strings.IndexByte(cs, 0) == -1
624 }
625
626 // appendHeader is the same as bsoncore.AppendHeader but does not check if the
627 // key is a valid C string since the caller has already checked for that.
628 //
629 // The caller of this function must check if key is a valid C string.
630 func (vw *valueWriter) appendHeader(t bsontype.Type, key string) {
631 vw.buf = bsoncore.AppendType(vw.buf, t)
632 vw.buf = append(vw.buf, key...)
633 vw.buf = append(vw.buf, 0x00)
634 }
635
636 func (vw *valueWriter) appendIntHeader(t bsontype.Type, key int) {
637 vw.buf = bsoncore.AppendType(vw.buf, t)
638 vw.buf = strconv.AppendInt(vw.buf, int64(key), 10)
639 vw.buf = append(vw.buf, 0x00)
640 }
641