encode.go 15 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646
  1. // BSON library for Go
  2. //
  3. // Copyright (c) 2010-2012 - Gustavo Niemeyer <gustavo@niemeyer.net>
  4. //
  5. // All rights reserved.
  6. //
  7. // Redistribution and use in source and binary forms, with or without
  8. // modification, are permitted provided that the following conditions are met:
  9. //
  10. // 1. Redistributions of source code must retain the above copyright notice, this
  11. // list of conditions and the following disclaimer.
  12. // 2. Redistributions in binary form must reproduce the above copyright notice,
  13. // this list of conditions and the following disclaimer in the documentation
  14. // and/or other materials provided with the distribution.
  15. //
  16. // THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND
  17. // ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED
  18. // WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
  19. // DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR CONTRIBUTORS BE LIABLE FOR
  20. // ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES
  21. // (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;
  22. // LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND
  23. // ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
  24. // (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS
  25. // SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
  26. // gobson - BSON library for Go.
  27. package bson
  28. import (
  29. "encoding/json"
  30. "fmt"
  31. "math"
  32. "net/url"
  33. "reflect"
  34. "sort"
  35. "strconv"
  36. "sync"
  37. "time"
  38. )
  39. // --------------------------------------------------------------------------
  40. // Some internal infrastructure.
  41. var (
  42. typeBinary = reflect.TypeOf(Binary{})
  43. typeObjectId = reflect.TypeOf(ObjectId(""))
  44. typeDBPointer = reflect.TypeOf(DBPointer{"", ObjectId("")})
  45. typeSymbol = reflect.TypeOf(Symbol(""))
  46. typeMongoTimestamp = reflect.TypeOf(MongoTimestamp(0))
  47. typeOrderKey = reflect.TypeOf(MinKey)
  48. typeDocElem = reflect.TypeOf(DocElem{})
  49. typeRawDocElem = reflect.TypeOf(RawDocElem{})
  50. typeRaw = reflect.TypeOf(Raw{})
  51. typeRawPtr = reflect.PtrTo(reflect.TypeOf(Raw{}))
  52. typeURL = reflect.TypeOf(url.URL{})
  53. typeTime = reflect.TypeOf(time.Time{})
  54. typeString = reflect.TypeOf("")
  55. typeJSONNumber = reflect.TypeOf(json.Number(""))
  56. typeTimeDuration = reflect.TypeOf(time.Duration(0))
  57. )
  58. var (
  59. // spec for []uint8 or []byte encoding
  60. arrayOps = map[string]bool{
  61. "$in": true,
  62. "$nin": true,
  63. "$all": true,
  64. }
  65. )
  66. const itoaCacheSize = 32
  67. const (
  68. getterUnknown = iota
  69. getterNone
  70. getterTypeVal
  71. getterTypePtr
  72. getterAddr
  73. )
  74. var itoaCache []string
  75. var getterStyles map[reflect.Type]int
  76. var getterIface reflect.Type
  77. var getterMutex sync.RWMutex
  78. func init() {
  79. itoaCache = make([]string, itoaCacheSize)
  80. for i := 0; i != itoaCacheSize; i++ {
  81. itoaCache[i] = strconv.Itoa(i)
  82. }
  83. var iface Getter
  84. getterIface = reflect.TypeOf(&iface).Elem()
  85. getterStyles = make(map[reflect.Type]int)
  86. }
  87. func itoa(i int) string {
  88. if i < itoaCacheSize {
  89. return itoaCache[i]
  90. }
  91. return strconv.Itoa(i)
  92. }
  93. func getterStyle(outt reflect.Type) int {
  94. getterMutex.RLock()
  95. style := getterStyles[outt]
  96. getterMutex.RUnlock()
  97. if style != getterUnknown {
  98. return style
  99. }
  100. getterMutex.Lock()
  101. defer getterMutex.Unlock()
  102. if outt.Implements(getterIface) {
  103. vt := outt
  104. for vt.Kind() == reflect.Ptr {
  105. vt = vt.Elem()
  106. }
  107. if vt.Implements(getterIface) {
  108. style = getterTypeVal
  109. } else {
  110. style = getterTypePtr
  111. }
  112. } else if reflect.PtrTo(outt).Implements(getterIface) {
  113. style = getterAddr
  114. } else {
  115. style = getterNone
  116. }
  117. getterStyles[outt] = style
  118. return style
  119. }
  120. func getGetter(outt reflect.Type, out reflect.Value) Getter {
  121. style := getterStyle(outt)
  122. if style == getterNone {
  123. return nil
  124. }
  125. if style == getterAddr {
  126. if !out.CanAddr() {
  127. return nil
  128. }
  129. return out.Addr().Interface().(Getter)
  130. }
  131. if style == getterTypeVal && out.Kind() == reflect.Ptr && out.IsNil() {
  132. return nil
  133. }
  134. return out.Interface().(Getter)
  135. }
  136. // --------------------------------------------------------------------------
  137. // Marshaling of the document value itself.
  138. type encoder struct {
  139. out []byte
  140. }
  141. func (e *encoder) addDoc(v reflect.Value) {
  142. for {
  143. if vi, ok := v.Interface().(Getter); ok {
  144. getv, err := vi.GetBSON()
  145. if err != nil {
  146. panic(err)
  147. }
  148. v = reflect.ValueOf(getv)
  149. continue
  150. }
  151. if v.Kind() == reflect.Ptr {
  152. v = v.Elem()
  153. continue
  154. }
  155. break
  156. }
  157. if v.Type() == typeRaw {
  158. raw := v.Interface().(Raw)
  159. if raw.Kind != 0x03 && raw.Kind != 0x00 {
  160. panic("Attempted to marshal Raw kind " + strconv.Itoa(int(raw.Kind)) + " as a document")
  161. }
  162. if len(raw.Data) == 0 {
  163. panic("Attempted to marshal empty Raw document")
  164. }
  165. e.addBytes(raw.Data...)
  166. return
  167. }
  168. start := e.reserveInt32()
  169. switch v.Kind() {
  170. case reflect.Map:
  171. e.addMap(v)
  172. case reflect.Struct:
  173. e.addStruct(v)
  174. case reflect.Array, reflect.Slice:
  175. e.addSlice(v)
  176. default:
  177. panic("Can't marshal " + v.Type().String() + " as a BSON document")
  178. }
  179. e.addBytes(0)
  180. e.setInt32(start, int32(len(e.out)-start))
  181. }
  182. func (e *encoder) addMap(v reflect.Value) {
  183. for _, k := range v.MapKeys() {
  184. e.addElem(fmt.Sprint(k), v.MapIndex(k), false)
  185. }
  186. }
  187. func (e *encoder) addStruct(v reflect.Value) {
  188. sinfo, err := getStructInfo(v.Type())
  189. if err != nil {
  190. panic(err)
  191. }
  192. var value reflect.Value
  193. if sinfo.InlineMap >= 0 {
  194. m := v.Field(sinfo.InlineMap)
  195. if m.Len() > 0 {
  196. for _, k := range m.MapKeys() {
  197. ks := k.String()
  198. if _, found := sinfo.FieldsMap[ks]; found {
  199. panic(fmt.Sprintf("Can't have key %q in inlined map; conflicts with struct field", ks))
  200. }
  201. e.addElem(ks, m.MapIndex(k), false)
  202. }
  203. }
  204. }
  205. for _, info := range sinfo.FieldsList {
  206. if info.Inline == nil {
  207. value = v.Field(info.Num)
  208. } else {
  209. // as pointers to struct are allowed here,
  210. // there is no guarantee that pointer won't be nil.
  211. //
  212. // It is expected allowed behaviour
  213. // so info.Inline MAY consist index to a nil pointer
  214. // and that is why we safely call v.FieldByIndex and just continue on panic
  215. field, errField := safeFieldByIndex(v, info.Inline)
  216. if errField != nil {
  217. continue
  218. }
  219. value = field
  220. }
  221. if info.OmitEmpty && isZero(value) {
  222. continue
  223. }
  224. if useRespectNilValues &&
  225. (value.Kind() == reflect.Slice || value.Kind() == reflect.Map) &&
  226. value.IsNil() {
  227. e.addElem(info.Key, reflect.ValueOf(nil), info.MinSize)
  228. continue
  229. }
  230. e.addElem(info.Key, value, info.MinSize)
  231. }
  232. }
  233. func safeFieldByIndex(v reflect.Value, index []int) (result reflect.Value, err error) {
  234. defer func() {
  235. if recovered := recover(); recovered != nil {
  236. switch r := recovered.(type) {
  237. case string:
  238. err = fmt.Errorf("%s", r)
  239. case error:
  240. err = r
  241. }
  242. }
  243. }()
  244. result = v.FieldByIndex(index)
  245. return
  246. }
  247. func isZero(v reflect.Value) bool {
  248. switch v.Kind() {
  249. case reflect.String:
  250. return len(v.String()) == 0
  251. case reflect.Ptr, reflect.Interface:
  252. return v.IsNil()
  253. case reflect.Slice:
  254. return v.Len() == 0
  255. case reflect.Map:
  256. return v.Len() == 0
  257. case reflect.Int, reflect.Int8, reflect.Int16, reflect.Int32, reflect.Int64:
  258. return v.Int() == 0
  259. case reflect.Uint, reflect.Uint8, reflect.Uint16, reflect.Uint32, reflect.Uint64, reflect.Uintptr:
  260. return v.Uint() == 0
  261. case reflect.Float32, reflect.Float64:
  262. return v.Float() == 0
  263. case reflect.Bool:
  264. return !v.Bool()
  265. case reflect.Struct:
  266. vt := v.Type()
  267. if vt == typeTime {
  268. return v.Interface().(time.Time).IsZero()
  269. }
  270. for i := 0; i < v.NumField(); i++ {
  271. if vt.Field(i).PkgPath != "" && !vt.Field(i).Anonymous {
  272. continue // Private field
  273. }
  274. if !isZero(v.Field(i)) {
  275. return false
  276. }
  277. }
  278. return true
  279. }
  280. return false
  281. }
  282. func (e *encoder) addSlice(v reflect.Value) {
  283. vi := v.Interface()
  284. if d, ok := vi.(D); ok {
  285. for _, elem := range d {
  286. e.addElem(elem.Name, reflect.ValueOf(elem.Value), false)
  287. }
  288. return
  289. }
  290. if d, ok := vi.(RawD); ok {
  291. for _, elem := range d {
  292. e.addElem(elem.Name, reflect.ValueOf(elem.Value), false)
  293. }
  294. return
  295. }
  296. l := v.Len()
  297. et := v.Type().Elem()
  298. if et == typeDocElem {
  299. for i := 0; i < l; i++ {
  300. elem := v.Index(i).Interface().(DocElem)
  301. e.addElem(elem.Name, reflect.ValueOf(elem.Value), false)
  302. }
  303. return
  304. }
  305. if et == typeRawDocElem {
  306. for i := 0; i < l; i++ {
  307. elem := v.Index(i).Interface().(RawDocElem)
  308. e.addElem(elem.Name, reflect.ValueOf(elem.Value), false)
  309. }
  310. return
  311. }
  312. for i := 0; i < l; i++ {
  313. e.addElem(itoa(i), v.Index(i), false)
  314. }
  315. }
  316. // --------------------------------------------------------------------------
  317. // Marshaling of elements in a document.
  318. func (e *encoder) addElemName(kind byte, name string) {
  319. e.addBytes(kind)
  320. e.addBytes([]byte(name)...)
  321. e.addBytes(0)
  322. }
  323. func (e *encoder) addElem(name string, v reflect.Value, minSize bool) {
  324. if !v.IsValid() {
  325. e.addElemName(0x0A, name)
  326. return
  327. }
  328. if getter := getGetter(v.Type(), v); getter != nil {
  329. getv, err := getter.GetBSON()
  330. if err != nil {
  331. panic(err)
  332. }
  333. e.addElem(name, reflect.ValueOf(getv), minSize)
  334. return
  335. }
  336. switch v.Kind() {
  337. case reflect.Interface:
  338. e.addElem(name, v.Elem(), minSize)
  339. case reflect.Ptr:
  340. e.addElem(name, v.Elem(), minSize)
  341. case reflect.String:
  342. s := v.String()
  343. switch v.Type() {
  344. case typeObjectId:
  345. if len(s) != 12 {
  346. panic("ObjectIDs must be exactly 12 bytes long (got " +
  347. strconv.Itoa(len(s)) + ")")
  348. }
  349. e.addElemName(0x07, name)
  350. e.addBytes([]byte(s)...)
  351. case typeSymbol:
  352. e.addElemName(0x0E, name)
  353. e.addStr(s)
  354. case typeJSONNumber:
  355. n := v.Interface().(json.Number)
  356. if i, err := n.Int64(); err == nil {
  357. e.addElemName(0x12, name)
  358. e.addInt64(i)
  359. } else if f, err := n.Float64(); err == nil {
  360. e.addElemName(0x01, name)
  361. e.addFloat64(f)
  362. } else {
  363. panic("failed to convert json.Number to a number: " + s)
  364. }
  365. default:
  366. e.addElemName(0x02, name)
  367. e.addStr(s)
  368. }
  369. case reflect.Float32, reflect.Float64:
  370. e.addElemName(0x01, name)
  371. e.addFloat64(v.Float())
  372. case reflect.Uint, reflect.Uint8, reflect.Uint16, reflect.Uint32, reflect.Uint64, reflect.Uintptr:
  373. u := v.Uint()
  374. if int64(u) < 0 {
  375. panic("BSON has no uint64 type, and value is too large to fit correctly in an int64")
  376. } else if u <= math.MaxInt32 && (minSize || v.Kind() <= reflect.Uint32) {
  377. e.addElemName(0x10, name)
  378. e.addInt32(int32(u))
  379. } else {
  380. e.addElemName(0x12, name)
  381. e.addInt64(int64(u))
  382. }
  383. case reflect.Int, reflect.Int8, reflect.Int16, reflect.Int32, reflect.Int64:
  384. switch v.Type() {
  385. case typeMongoTimestamp:
  386. e.addElemName(0x11, name)
  387. e.addInt64(v.Int())
  388. case typeOrderKey:
  389. if v.Int() == int64(MaxKey) {
  390. e.addElemName(0x7F, name)
  391. } else {
  392. e.addElemName(0xFF, name)
  393. }
  394. case typeTimeDuration:
  395. // Stored as int64
  396. e.addElemName(0x12, name)
  397. e.addInt64(int64(v.Int() / 1e6))
  398. default:
  399. i := v.Int()
  400. if (minSize || v.Type().Kind() != reflect.Int64) && i >= math.MinInt32 && i <= math.MaxInt32 {
  401. // It fits into an int32, encode as such.
  402. e.addElemName(0x10, name)
  403. e.addInt32(int32(i))
  404. } else {
  405. e.addElemName(0x12, name)
  406. e.addInt64(i)
  407. }
  408. }
  409. case reflect.Bool:
  410. e.addElemName(0x08, name)
  411. if v.Bool() {
  412. e.addBytes(1)
  413. } else {
  414. e.addBytes(0)
  415. }
  416. case reflect.Map:
  417. e.addElemName(0x03, name)
  418. e.addDoc(v)
  419. case reflect.Slice:
  420. vt := v.Type()
  421. et := vt.Elem()
  422. if et.Kind() == reflect.Uint8 {
  423. if arrayOps[name] {
  424. e.addElemName(0x04, name)
  425. e.addDoc(v)
  426. } else {
  427. e.addElemName(0x05, name)
  428. e.addBinary(0x00, v.Bytes())
  429. }
  430. } else if et == typeDocElem || et == typeRawDocElem {
  431. e.addElemName(0x03, name)
  432. e.addDoc(v)
  433. } else {
  434. e.addElemName(0x04, name)
  435. e.addDoc(v)
  436. }
  437. case reflect.Array:
  438. et := v.Type().Elem()
  439. if et.Kind() == reflect.Uint8 {
  440. if arrayOps[name] {
  441. e.addElemName(0x04, name)
  442. e.addDoc(v)
  443. } else {
  444. e.addElemName(0x05, name)
  445. if v.CanAddr() {
  446. e.addBinary(0x00, v.Slice(0, v.Len()).Interface().([]byte))
  447. } else {
  448. n := v.Len()
  449. e.addInt32(int32(n))
  450. e.addBytes(0x00)
  451. for i := 0; i < n; i++ {
  452. el := v.Index(i)
  453. e.addBytes(byte(el.Uint()))
  454. }
  455. }
  456. }
  457. } else {
  458. e.addElemName(0x04, name)
  459. e.addDoc(v)
  460. }
  461. case reflect.Struct:
  462. switch s := v.Interface().(type) {
  463. case Raw:
  464. kind := s.Kind
  465. if kind == 0x00 {
  466. kind = 0x03
  467. }
  468. if len(s.Data) == 0 && kind != 0x06 && kind != 0x0A && kind != 0xFF && kind != 0x7F {
  469. panic("Attempted to marshal empty Raw document")
  470. }
  471. e.addElemName(kind, name)
  472. e.addBytes(s.Data...)
  473. case Binary:
  474. e.addElemName(0x05, name)
  475. e.addBinary(s.Kind, s.Data)
  476. case Decimal128:
  477. e.addElemName(0x13, name)
  478. e.addInt64(int64(s.l))
  479. e.addInt64(int64(s.h))
  480. case DBPointer:
  481. e.addElemName(0x0C, name)
  482. e.addStr(s.Namespace)
  483. if len(s.Id) != 12 {
  484. panic("ObjectIDs must be exactly 12 bytes long (got " +
  485. strconv.Itoa(len(s.Id)) + ")")
  486. }
  487. e.addBytes([]byte(s.Id)...)
  488. case RegEx:
  489. e.addElemName(0x0B, name)
  490. e.addCStr(s.Pattern)
  491. options := runes(s.Options)
  492. sort.Sort(options)
  493. e.addCStr(string(options))
  494. case JavaScript:
  495. if s.Scope == nil {
  496. e.addElemName(0x0D, name)
  497. e.addStr(s.Code)
  498. } else {
  499. e.addElemName(0x0F, name)
  500. start := e.reserveInt32()
  501. e.addStr(s.Code)
  502. e.addDoc(reflect.ValueOf(s.Scope))
  503. e.setInt32(start, int32(len(e.out)-start))
  504. }
  505. case time.Time:
  506. // MongoDB handles timestamps as milliseconds.
  507. e.addElemName(0x09, name)
  508. e.addInt64(s.Unix()*1000 + int64(s.Nanosecond()/1e6))
  509. case url.URL:
  510. e.addElemName(0x02, name)
  511. e.addStr(s.String())
  512. case undefined:
  513. e.addElemName(0x06, name)
  514. default:
  515. e.addElemName(0x03, name)
  516. e.addDoc(v)
  517. }
  518. default:
  519. panic("Can't marshal " + v.Type().String() + " in a BSON document")
  520. }
  521. }
  522. // -------------
  523. // Helper method for sorting regex options
  524. type runes []rune
  525. func (a runes) Len() int { return len(a) }
  526. func (a runes) Swap(i, j int) { a[i], a[j] = a[j], a[i] }
  527. func (a runes) Less(i, j int) bool { return a[i] < a[j] }
  528. // --------------------------------------------------------------------------
  529. // Marshaling of base types.
  530. func (e *encoder) addBinary(subtype byte, v []byte) {
  531. if subtype == 0x02 {
  532. // Wonder how that brilliant idea came to life. Obsolete, luckily.
  533. e.addInt32(int32(len(v) + 4))
  534. e.addBytes(subtype)
  535. e.addInt32(int32(len(v)))
  536. } else {
  537. e.addInt32(int32(len(v)))
  538. e.addBytes(subtype)
  539. }
  540. e.addBytes(v...)
  541. }
  542. func (e *encoder) addStr(v string) {
  543. e.addInt32(int32(len(v) + 1))
  544. e.addCStr(v)
  545. }
  546. func (e *encoder) addCStr(v string) {
  547. e.addBytes([]byte(v)...)
  548. e.addBytes(0)
  549. }
  550. func (e *encoder) reserveInt32() (pos int) {
  551. pos = len(e.out)
  552. e.addBytes(0, 0, 0, 0)
  553. return pos
  554. }
  555. func (e *encoder) setInt32(pos int, v int32) {
  556. e.out[pos+0] = byte(v)
  557. e.out[pos+1] = byte(v >> 8)
  558. e.out[pos+2] = byte(v >> 16)
  559. e.out[pos+3] = byte(v >> 24)
  560. }
  561. func (e *encoder) addInt32(v int32) {
  562. u := uint32(v)
  563. e.addBytes(byte(u), byte(u>>8), byte(u>>16), byte(u>>24))
  564. }
  565. func (e *encoder) addInt64(v int64) {
  566. u := uint64(v)
  567. e.addBytes(byte(u), byte(u>>8), byte(u>>16), byte(u>>24),
  568. byte(u>>32), byte(u>>40), byte(u>>48), byte(u>>56))
  569. }
  570. func (e *encoder) addFloat64(v float64) {
  571. e.addInt64(int64(math.Float64bits(v)))
  572. }
  573. func (e *encoder) addBytes(v ...byte) {
  574. e.out = append(e.out, v...)
  575. }