What: add spec for array ops with []uint8 or []byte type (#128)

Why: * mongodb aquires array for $in/$nin/$all ops How: * encode with array for spec
globalsign · Apr 3, 2018 · 69bef6a · 69bef6a
1 parent 76ea203
commit 69bef6a
Show file tree

Hide file tree

Showing 2 changed files with 56 additions and 12 deletions.
diff --git a/bson/bson_test.go b/bson/bson_test.go
@@ -34,9 +34,9 @@ import (
 	"errors"
 	"net/url"
 	"reflect"
+	"strings"
 	"testing"
 	"time"
-	"strings"
 
 	"github.com/globalsign/mgo/bson"
 	. "gopkg.in/check.v1"
@@ -111,6 +111,10 @@ var sampleItems = []testItemType{
 	{bson.M{"BSON": []interface{}{"awesome", float64(5.05), 1986}},
 		"1\x00\x00\x00\x04BSON\x00&\x00\x00\x00\x020\x00\x08\x00\x00\x00" +
 			"awesome\x00\x011\x00333333\x14@\x102\x00\xc2\x07\x00\x00\x00\x00"},
+	{bson.M{"slice": []uint8{1, 2}},
+		"\x13\x00\x00\x00\x05slice\x00\x02\x00\x00\x00\x00\x01\x02\x00"},
+	{bson.M{"slice": []byte{1, 2}},
+		"\x13\x00\x00\x00\x05slice\x00\x02\x00\x00\x00\x00\x01\x02\x00"},
 }
 
 func (s *S) TestMarshalSampleItems(c *C) {
@@ -343,6 +347,27 @@ func (s *S) TestOneWayMarshalItems(c *C) {
 	}
 }
 
+// --------------------------------------------------------------------------
+// Some ops marshaling operations which would encode []uint8 or []byte in array.
+
+var arrayOpsMarshalItems = []testItemType{
+	{bson.M{"_": bson.M{"$in": []uint8{1, 2}}},
+		"\x03_\x00\x1d\x00\x00\x00\x04$in\x00\x13\x00\x00\x00\x100\x00\x01\x00\x00\x00\x101\x00\x02\x00\x00\x00\x00\x00"},
+	{bson.M{"_": bson.M{"$nin": []uint8{1, 2}}},
+		"\x03_\x00\x1e\x00\x00\x00\x04$nin\x00\x13\x00\x00\x00\x100\x00\x01\x00\x00\x00\x101\x00\x02\x00\x00\x00\x00\x00"},
+	{bson.M{"_": bson.M{"$all": []uint8{1, 2}}},
+		"\x03_\x00\x1e\x00\x00\x00\x04$all\x00\x13\x00\x00\x00\x100\x00\x01\x00\x00\x00\x101\x00\x02\x00\x00\x00\x00\x00"},
+}
+
+func (s *S) TestArrayOpsMarshalItems(c *C) {
+	for i, item := range arrayOpsMarshalItems {
+		data, err := bson.Marshal(item.obj)
+		c.Assert(err, IsNil)
+		c.Assert(string(data), Equals, wrapInDoc(item.data),
+			Commentf("Failed on item %d", i))
+	}
+}
+
 // --------------------------------------------------------------------------
 // Two-way tests for user-defined structures using the samples
 // from bsonspec.org.

diff --git a/bson/encode.go b/bson/encode.go
@@ -60,6 +60,15 @@ var (
 	typeTimeDuration   = reflect.TypeOf(time.Duration(0))
 )
 
+var (
+	// spec for []uint8 or []byte encoding
+	arrayOps = map[string]bool{
+		"$in":  true,
+		"$nin": true,
+		"$all": true,
+	}
+)
+
 const itoaCacheSize = 32
 
 const (
@@ -423,8 +432,13 @@ func (e *encoder) addElem(name string, v reflect.Value, minSize bool) {
 		vt := v.Type()
 		et := vt.Elem()
 		if et.Kind() == reflect.Uint8 {
-			e.addElemName(0x05, name)
-			e.addBinary(0x00, v.Bytes())
+			if arrayOps[name] {
+				e.addElemName(0x04, name)
+				e.addDoc(v)
+			} else {
+				e.addElemName(0x05, name)
+				e.addBinary(0x00, v.Bytes())
+			}
 		} else if et == typeDocElem || et == typeRawDocElem {
 			e.addElemName(0x03, name)
 			e.addDoc(v)
@@ -436,16 +450,21 @@ func (e *encoder) addElem(name string, v reflect.Value, minSize bool) {
 	case reflect.Array:
 		et := v.Type().Elem()
 		if et.Kind() == reflect.Uint8 {
-			e.addElemName(0x05, name)
-			if v.CanAddr() {
-				e.addBinary(0x00, v.Slice(0, v.Len()).Interface().([]byte))
+			if arrayOps[name] {
+				e.addElemName(0x04, name)
+				e.addDoc(v)
 			} else {
-				n := v.Len()
-				e.addInt32(int32(n))
-				e.addBytes(0x00)
-				for i := 0; i < n; i++ {
-					el := v.Index(i)
-					e.addBytes(byte(el.Uint()))
+				e.addElemName(0x05, name)
+				if v.CanAddr() {
+					e.addBinary(0x00, v.Slice(0, v.Len()).Interface().([]byte))
+				} else {
+					n := v.Len()
+					e.addInt32(int32(n))
+					e.addBytes(0x00)
+					for i := 0; i < n; i++ {
+						el := v.Index(i)
+						e.addBytes(byte(el.Uint()))
+					}
 				}
 			}
 		} else {