Fix memory corruption bug in impl/memory

impl/memory/datastore_data.go

 // Copyright 2015 The Chromium Authors. All rights reserved.
 // Use of this source code is governed by a BSD-style license that can be
 // found in the LICENSE file.
 
 package memory
 
 import (
         "bytes"
         "fmt"
         "strings"
         "sync"
         "sync/atomic"
 
         ds "github.com/luci/gae/service/datastore"
         "github.com/luci/gae/service/datastore/serialize"
         "github.com/luci/luci-go/common/errors"
         "golang.org/x/net/context"
 )
 
 //////////////////////////////// dataStoreData /////////////////////////////////
 
 type dataStoreData struct {
         rwlock sync.RWMutex
 
         // the 'appid' of this datastore
         aid string
 
         // See README.md for head schema.
-        head *memStore
+        head memStore
         // if snap is nil, that means that this is always-consistent, and
         // getQuerySnaps will return (head, head)
-        snap *memStore
+        snap memStore
         // For testing, see SetTransactionRetryCount.
         txnFakeRetry int
         // true means that queries with insufficent indexes will pause to add them
         // and then continue instead of failing.
         autoIndex bool
         // true means that all of the __...__ keys which are normally automatically
         // maintained will be omitted. This also means that Put with an incomplete
         // key will become an error.
         disableSpecialEntities bool
 }
@@ skipping 72 matching lines @@
         defer d.Unlock()
         d.disableSpecialEntities = true
 }
 
 func (d *dataStoreData) getDisableSpecialEntities() bool {
         d.rwlock.RLock()
         defer d.rwlock.RUnlock()
         return d.disableSpecialEntities
 }
 
-func (d *dataStoreData) getQuerySnaps(consistent bool) (idx, head *memStore) {
+func (d *dataStoreData) getQuerySnaps(consistent bool) (idx, head memStore) {
         d.rwlock.RLock()
         defer d.rwlock.RUnlock()
         if d.snap == nil {
                 // we're 'always consistent'
                 snap := d.head.Snapshot()
                 return snap, snap
         }
 
         head = d.head.Snapshot()
         if consistent {
                 idx = head
         } else {
                 idx = d.snap
         }
         return
 }
 
-func (d *dataStoreData) takeSnapshot() *memStore {
+func (d *dataStoreData) takeSnapshot() memStore {
         d.rwlock.RLock()
         defer d.rwlock.RUnlock()
         return d.head.Snapshot()
 }
 
-func (d *dataStoreData) setSnapshot(snap *memStore) {
+func (d *dataStoreData) setSnapshot(snap memStore) {
         d.rwlock.Lock()
         defer d.rwlock.Unlock()
         if d.snap == nil {
                 // we're 'always consistent'
                 return
         }
         d.snap = snap
 }
 
 func (d *dataStoreData) catchupIndexes() {
@@ skipping 20 matching lines @@
 }
 
 func groupIDsKey(key *ds.Key) []byte {
         return keyBytes(ds.NewKey("", "", "__entity_group_ids__", "", 1, key.Root()))
 }
 
 func rootIDsKey(kind string) []byte {
         return keyBytes(ds.NewKey("", "", "__entity_root_ids__", kind, 0, nil))
 }
 
-func curVersion(ents *memCollection, key []byte) int64 {
+func curVersion(ents memCollection, key []byte) int64 {
         if ents != nil {
                 if v := ents.Get(key); v != nil {
                         pm, err := rpm(v)
                         memoryCorruption(err)
 
                         pl, ok := pm["__version__"]
                         if ok && len(pl) > 0 && pl[0].Type() == ds.PTInt {
                                 return pl[0].Value().(int64)
                         }
 
                         memoryCorruption(fmt.Errorf("__version__ property missing or wrong: %v", pm))
                 }
         }
         return 0
 }
 
-func incrementLocked(ents *memCollection, key []byte, amt int) int64 {
+func incrementLocked(ents memCollection, key []byte, amt int) int64 {
         if amt <= 0 {
                 panic(fmt.Errorf("incrementLocked called with bad `amt`: %d", amt))
         }
         ret := curVersion(ents, key) + 1
         ents.Set(key, serialize.ToBytes(ds.PropertyMap{
                 "__version__": {ds.MkPropertyNI(ret + int64(amt-1))},
         }))
         return ret
 }
 
-func (d *dataStoreData) mutableEntsLocked(ns string) *memCollection {
-        coll := "ents:" + ns
-        ents := d.head.GetCollection(coll)
-        if ents == nil {
-                ents = d.head.SetCollection(coll, nil)
-        }
-        return ents
-}
-
 func (d *dataStoreData) allocateIDs(incomplete *ds.Key, n int) (int64, error) {
         d.Lock()
         defer d.Unlock()
 
-        ents := d.mutableEntsLocked(incomplete.Namespace())
+        ents := d.head.GetOrCreateCollection("ents:" + incomplete.Namespace())
         return d.allocateIDsLocked(ents, incomplete, n)
 }
 
-func (d *dataStoreData) allocateIDsLocked(ents *memCollection, incomplete *ds.Key, n int) (int64, error) {
+func (d *dataStoreData) allocateIDsLocked(ents memCollection, incomplete *ds.Key, n int) (int64, error) {
         if d.disableSpecialEntities {
                 return 0, errors.New("disableSpecialEntities is true so allocateIDs is disabled")
         }
 
         idKey := []byte(nil)
         if incomplete.Parent() == nil {
                 idKey = rootIDsKey(incomplete.Kind())
         } else {
                 idKey = groupIDsKey(incomplete)
         }
         return incrementLocked(ents, idKey, n), nil
 }
 
-func (d *dataStoreData) fixKeyLocked(ents *memCollection, key *ds.Key) (*ds.Key, error) {
+func (d *dataStoreData) fixKeyLocked(ents memCollection, key *ds.Key) (*ds.Key, error) {
         if key.Incomplete() {
                 id, err := d.allocateIDsLocked(ents, key, 1)
                 if err != nil {
                         return key, err
                 }
                 key = ds.NewKey(key.AppID(), key.Namespace(), key.Kind(), "", id, key.Parent())
         }
         return key, nil
 }
 
 func (d *dataStoreData) putMulti(keys []*ds.Key, vals []ds.PropertyMap, cb ds.PutMultiCB) error {
         ns := keys[0].Namespace()
 
         for i, k := range keys {
                 pmap, _ := vals[i].Save(false)
                 dataBytes := serialize.ToBytes(pmap)
 
                 k, err := func() (ret *ds.Key, err error) {
                         d.Lock()
                         defer d.Unlock()
 
-                        ents := d.mutableEntsLocked(ns)
+                        ents := d.head.GetOrCreateCollection("ents:" + ns)
 
                         ret, err = d.fixKeyLocked(ents, k)
                         if err != nil {
                                 return
                         }
                         if !d.disableSpecialEntities {
                                 incrementLocked(ents, groupMetaKey(ret), 1)
                         }
 
                         old := ents.Get(keyBytes(ret))
                         oldPM := ds.PropertyMap(nil)
                         if old != nil {
                                 if oldPM, err = rpm(old); err != nil {
                                         return
                                 }
                         }
                         ents.Set(keyBytes(ret), dataBytes)
                         updateIndexes(d.head, ret, oldPM, pmap)
                         return
                 }()
                 if cb != nil {
                         if err := cb(k, err); err != nil {
                                 return err
                         }
                 }
         }
         return nil
 }
 
-func getMultiInner(keys []*ds.Key, cb ds.GetMultiCB, getColl func() (*memCollection, error)) error {
+func getMultiInner(keys []*ds.Key, cb ds.GetMultiCB, getColl func() (memCollection, error)) error {
         ents, err := getColl()
         if err != nil {
                 return err
         }
         if ents == nil {
                 for range keys {
                         cb(nil, ds.ErrNoSuchEntity)
                 }
                 return nil
         }
 
         for _, k := range keys {
                 pdata := ents.Get(keyBytes(k))
                 if pdata == nil {
                         cb(nil, ds.ErrNoSuchEntity)
                         continue
                 }
                 cb(rpm(pdata))
         }
         return nil
 }
 
 func (d *dataStoreData) getMulti(keys []*ds.Key, cb ds.GetMultiCB) error {
-        return getMultiInner(keys, cb, func() (*memCollection, error) {
+        return getMultiInner(keys, cb, func() (memCollection, error) {
                 s := d.takeSnapshot()
 
                 return s.GetCollection("ents:" + keys[0].Namespace()), nil
         })
 }
 
 func (d *dataStoreData) delMulti(keys []*ds.Key, cb ds.DeleteMultiCB) error {
         ns := keys[0].Namespace()
 
         hasEntsInNS := func() bool {
                 d.Lock()
                 defer d.Unlock()
-                return d.mutableEntsLocked(ns) != nil
+                return d.head.GetOrCreateCollection("ents:"+ns) != nil
         }()
 
         if hasEntsInNS {
                 for _, k := range keys {
                         err := func() error {
                                 kb := keyBytes(k)
 
                                 d.Lock()
                                 defer d.Unlock()
 
-                                ents := d.mutableEntsLocked(ns)
+                                ents := d.head.GetOrCreateCollection("ents:" + ns)
 
                                 if !d.disableSpecialEntities {
                                         incrementLocked(ents, groupMetaKey(k), 1)
                                 }
                                 if old := ents.Get(kb); old != nil {
                                         oldPM, err := rpm(old)
                                         if err != nil {
                                                 return err
                                         }
                                         ents.Delete(kb)
@@ skipping 85 matching lines @@
 
 type txnDataStoreData struct {
         sync.Mutex
 
         parent *dataStoreData
 
         // boolean 0 or 1, use atomic.*Int32 to access.
         closed int32
         isXG   bool
 
-        snap *memStore
+        snap memStore
 
         // string is the raw-bytes encoding of the entity root incl. namespace
         muts map[string][]txnMutation
         // TODO(riannucci): account for 'transaction size' limit of 10MB by summing
         // length of encoded keys + values.
 }
 
 var _ memContextObj = (*txnDataStoreData)(nil)
 
 const xgEGLimit = 25
@@ skipping 60 matching lines @@
         return nil
 }
 
 func (td *txnDataStoreData) putMulti(keys []*ds.Key, vals []ds.PropertyMap, cb ds.PutMultiCB) {
         ns := keys[0].Namespace()
 
         for i, k := range keys {
                 err := func() (err error) {
                         td.parent.Lock()
                         defer td.parent.Unlock()
-                        ents := td.parent.mutableEntsLocked(ns)
+                        ents := td.parent.head.GetOrCreateCollection("ents:" + ns)
 
                         k, err = td.parent.fixKeyLocked(ents, k)
                         return
                 }()
                 if err == nil {
                         err = td.writeMutation(false, k, vals[i])
                 }
                 if cb != nil {
                         cb(k, err)
                 }
         }
 }
 
 func (td *txnDataStoreData) getMulti(keys []*ds.Key, cb ds.GetMultiCB) error {
-        return getMultiInner(keys, cb, func() (*memCollection, error) {
+        return getMultiInner(keys, cb, func() (memCollection, error) {
                 err := error(nil)
                 for _, key := range keys {
                         err = td.writeMutation(true, key, nil)
                         if err != nil {
                                 return nil, err
                         }
                 }
                 return td.snap.GetCollection("ents:" + keys[0].Namespace()), nil
         })
 }
@@ skipping 10 matching lines @@
 
 func keyBytes(key *ds.Key) []byte {
         return serialize.ToBytes(ds.MkProperty(key))
 }
 
 func rpm(data []byte) (ds.PropertyMap, error) {
         return serialize.ReadPropertyMap(bytes.NewBuffer(data),
                 serialize.WithContext, "", "")
 }
 
-func namespaces(store *memStore) []string {
+func namespaces(store memStore) []string {
         var namespaces []string
         for _, c := range store.GetCollectionNames() {
                 ns, has := trimPrefix(c, "ents:")
                 if !has {
                         if len(namespaces) > 0 {
                                 break
                         }
                         continue
                 }
                 namespaces = append(namespaces, ns)
         }
         return namespaces
 }
 
 func trimPrefix(v, p string) (string, bool) {
         if strings.HasPrefix(v, p) {
                 return v[len(p):], true
         }
         return v, false
 }