| OLD | NEW |
|---|
| (Empty) |
| 1 // Copyright 2015 The Chromium Authors. All rights reserved. | |
| 2 // Use of this source code is governed by a BSD-style license that can be | |
| 3 // found in the LICENSE file. | |
| 4 | |
| 5 package helper | |
| 6 | |
| 7 import ( | |
| 8 "bytes" | |
| 9 "errors" | |
| 10 "fmt" | |
| 11 "sort" | |
| 12 "time" | |
| 13 | |
| 14 "github.com/luci/gae" | |
| 15 "github.com/luci/luci-go/common/cmpbin" | |
| 16 ) | |
| 17 | |
| 18 // WriteDSPropertyMapDeterministic allows tests to make WriteDSPropertyMap | |
| 19 // deterministic. | |
| 20 var WriteDSPropertyMapDeterministic = false | |
| 21 | |
| 22 // ReadDSPropertyMapReasonableLimit sets a limit on the number of rows and | |
| 23 // number of properties per row which can be read by ReadDSPropertyMap. The | |
| 24 // total number of Property objects readable by this method is this number | |
| 25 // squared (e.g. Limit rows * Limit properties) | |
| 26 const ReadDSPropertyMapReasonableLimit uint64 = 30000 | |
| 27 | |
| 28 // ReadKeyNumToksReasonableLimit is the maximum number of DSKey tokens that | |
| 29 // ReadDSKey is willing to read for a single key. | |
| 30 const ReadKeyNumToksReasonableLimit uint64 = 50 | |
| 31 | |
| 32 // DSKeyContext controls whether the various Write and Read serializtion | |
| 33 // routines should encode the context of DSKeys (read: the appid and namespace). | |
| 34 // Frequently the appid and namespace of keys are known in advance and so there'
s | |
| 35 // no reason to redundantly encode them. | |
| 36 type DSKeyContext bool | |
| 37 | |
| 38 // With- and WithoutContext indicate if the serialization method should include | |
| 39 // context for DSKeys. See DSKeyContext for more information. | |
| 40 const ( | |
| 41 WithContext DSKeyContext = true | |
| 42 WithoutContext = false | |
| 43 ) | |
| 44 | |
| 45 // WriteDSKey encodes a key to the buffer. If context is WithContext, then this | |
| 46 // encoded value will include the appid and namespace of the key. | |
| 47 func WriteDSKey(buf *bytes.Buffer, context DSKeyContext, k gae.DSKey) { | |
| 48 // [appid ++ namespace]? ++ #tokens ++ tokens* | |
| 49 appid, namespace, toks := DSKeySplit(k) | |
| 50 if context == WithContext { | |
| 51 buf.WriteByte(1) | |
| 52 cmpbin.WriteString(buf, appid) | |
| 53 cmpbin.WriteString(buf, namespace) | |
| 54 } else { | |
| 55 buf.WriteByte(0) | |
| 56 } | |
| 57 cmpbin.WriteUint(buf, uint64(len(toks))) | |
| 58 for _, tok := range toks { | |
| 59 WriteDSKeyTok(buf, tok) | |
| 60 } | |
| 61 } | |
| 62 | |
| 63 // ReadDSKey deserializes a key from the buffer. The value of context must match | |
| 64 // the value of context that was passed to WriteDSKey when the key was encoded. | |
| 65 // If context == WithoutContext, then the appid and namespace parameters are | |
| 66 // used in the decoded DSKey. Otherwise they're ignored. | |
| 67 func ReadDSKey(buf *bytes.Buffer, context DSKeyContext, appid, namespace string)
(ret gae.DSKey, err error) { | |
| 68 actualCtx, err := buf.ReadByte() | |
| 69 if err != nil { | |
| 70 return | |
| 71 } | |
| 72 | |
| 73 actualAid, actualNS := "", "" | |
| 74 if actualCtx == 1 { | |
| 75 if actualAid, _, err = cmpbin.ReadString(buf); err != nil { | |
| 76 return | |
| 77 } | |
| 78 if actualNS, _, err = cmpbin.ReadString(buf); err != nil { | |
| 79 return | |
| 80 } | |
| 81 } else if actualCtx != 0 { | |
| 82 err = fmt.Errorf("helper: expected actualCtx to be 0 or 1, got %
d", actualCtx) | |
| 83 return | |
| 84 } | |
| 85 | |
| 86 if context == WithoutContext { | |
| 87 // overrwrite with the supplied ones | |
| 88 actualAid = appid | |
| 89 actualNS = namespace | |
| 90 } | |
| 91 | |
| 92 numToks := uint64(0) | |
| 93 if numToks, _, err = cmpbin.ReadUint(buf); err != nil { | |
| 94 return | |
| 95 } | |
| 96 if numToks > ReadKeyNumToksReasonableLimit { | |
| 97 err = fmt.Errorf("helper: tried to decode huge key of length %d"
, numToks) | |
| 98 return | |
| 99 } | |
| 100 | |
| 101 toks := make([]gae.DSKeyTok, numToks) | |
| 102 for i := uint64(0); i < numToks; i++ { | |
| 103 if toks[i], err = ReadDSKeyTok(buf); err != nil { | |
| 104 return | |
| 105 } | |
| 106 } | |
| 107 | |
| 108 return NewDSKeyToks(actualAid, actualNS, toks), nil | |
| 109 } | |
| 110 | |
| 111 // WriteDSKeyTok writes a DSKeyTok to the buffer. You usually want WriteDSKey | |
| 112 // instead of this. | |
| 113 func WriteDSKeyTok(buf *bytes.Buffer, tok gae.DSKeyTok) { | |
| 114 // tok.kind ++ typ ++ [tok.stringID || tok.intID] | |
| 115 cmpbin.WriteString(buf, tok.Kind) | |
| 116 if tok.StringID != "" { | |
| 117 buf.WriteByte(byte(gae.DSPTString)) | |
| 118 cmpbin.WriteString(buf, tok.StringID) | |
| 119 } else { | |
| 120 buf.WriteByte(byte(gae.DSPTInt)) | |
| 121 cmpbin.WriteInt(buf, tok.IntID) | |
| 122 } | |
| 123 } | |
| 124 | |
| 125 // ReadDSKeyTok reads a DSKeyTok from the buffer. You usually want ReadDSKey | |
| 126 // instead of this. | |
| 127 func ReadDSKeyTok(buf *bytes.Buffer) (ret gae.DSKeyTok, err error) { | |
| 128 if ret.Kind, _, err = cmpbin.ReadString(buf); err != nil { | |
| 129 return | |
| 130 } | |
| 131 typ, err := buf.ReadByte() | |
| 132 if err != nil { | |
| 133 return | |
| 134 } | |
| 135 switch gae.DSPropertyType(typ) { | |
| 136 case gae.DSPTString: | |
| 137 ret.StringID, _, err = cmpbin.ReadString(buf) | |
| 138 case gae.DSPTInt: | |
| 139 ret.IntID, _, err = cmpbin.ReadInt(buf) | |
| 140 if err == nil && ret.IntID <= 0 { | |
| 141 err = errors.New("helper: decoded key with empty stringI
D and zero/negative intID") | |
| 142 } | |
| 143 default: | |
| 144 err = fmt.Errorf("helper: invalid type %s", gae.DSPropertyType(t
yp)) | |
| 145 } | |
| 146 return | |
| 147 } | |
| 148 | |
| 149 // WriteDSGeoPoint writes a DSGeoPoint to the buffer. | |
| 150 func WriteDSGeoPoint(buf *bytes.Buffer, gp gae.DSGeoPoint) { | |
| 151 cmpbin.WriteFloat64(buf, gp.Lat) | |
| 152 cmpbin.WriteFloat64(buf, gp.Lng) | |
| 153 } | |
| 154 | |
| 155 // ReadDSGeoPoint reads a DSGeoPoint from the buffer. | |
| 156 func ReadDSGeoPoint(buf *bytes.Buffer) (gp gae.DSGeoPoint, err error) { | |
| 157 if gp.Lat, _, err = cmpbin.ReadFloat64(buf); err != nil { | |
| 158 return | |
| 159 } | |
| 160 gp.Lng, _, err = cmpbin.ReadFloat64(buf) | |
| 161 if err == nil && !gp.Valid() { | |
| 162 err = fmt.Errorf("helper: decoded invalid DSGeoPoint: %v", gp) | |
| 163 } | |
| 164 return | |
| 165 } | |
| 166 | |
| 167 // WriteTime writes a time.Time in a byte-sortable way. | |
| 168 // | |
| 169 // This method truncates the time to microseconds and drops the timezone, | |
| 170 // because that's the (undocumented) way that the appengine SDK does it. | |
| 171 func WriteTime(buf *bytes.Buffer, t time.Time) { | |
| 172 name, off := t.Zone() | |
| 173 if name != "UTC" || off != 0 { | |
| 174 panic(fmt.Errorf("helper: UTC OR DEATH: %s", t)) | |
| 175 } | |
| 176 cmpbin.WriteUint(buf, uint64(t.Unix())*1e6+uint64(t.Nanosecond()/1e3)) | |
| 177 } | |
| 178 | |
| 179 // ReadTime reads a time.Time from the buffer. | |
| 180 func ReadTime(buf *bytes.Buffer) (time.Time, error) { | |
| 181 v, _, err := cmpbin.ReadUint(buf) | |
| 182 if err != nil { | |
| 183 return time.Time{}, err | |
| 184 } | |
| 185 return time.Unix(int64(v/1e6), int64((v%1e6)*1e3)).UTC(), nil | |
| 186 } | |
| 187 | |
| 188 // WriteDSProperty writes a DSProperty to the buffer. `context` behaves the same | |
| 189 // way that it does for WriteDSKey, but only has an effect if `p` contains a | |
| 190 // DSKey as its Value. | |
| 191 func WriteDSProperty(buf *bytes.Buffer, p gae.DSProperty, context DSKeyContext)
{ | |
| 192 typb := byte(p.Type()) | |
| 193 if p.IndexSetting() == gae.NoIndex { | |
| 194 typb |= 0x80 | |
| 195 } | |
| 196 buf.WriteByte(typb) | |
| 197 switch p.Type() { | |
| 198 case gae.DSPTNull, gae.DSPTBoolTrue, gae.DSPTBoolFalse: | |
| 199 case gae.DSPTInt: | |
| 200 cmpbin.WriteInt(buf, p.Value().(int64)) | |
| 201 case gae.DSPTFloat: | |
| 202 cmpbin.WriteFloat64(buf, p.Value().(float64)) | |
| 203 case gae.DSPTString: | |
| 204 cmpbin.WriteString(buf, p.Value().(string)) | |
| 205 case gae.DSPTBytes: | |
| 206 if p.IndexSetting() == gae.NoIndex { | |
| 207 cmpbin.WriteBytes(buf, p.Value().([]byte)) | |
| 208 } else { | |
| 209 cmpbin.WriteBytes(buf, p.Value().(gae.DSByteString)) | |
| 210 } | |
| 211 case gae.DSPTTime: | |
| 212 WriteTime(buf, p.Value().(time.Time)) | |
| 213 case gae.DSPTGeoPoint: | |
| 214 WriteDSGeoPoint(buf, p.Value().(gae.DSGeoPoint)) | |
| 215 case gae.DSPTKey: | |
| 216 WriteDSKey(buf, context, p.Value().(gae.DSKey)) | |
| 217 case gae.DSPTBlobKey: | |
| 218 cmpbin.WriteString(buf, string(p.Value().(gae.BSKey))) | |
| 219 } | |
| 220 } | |
| 221 | |
| 222 // ReadDSProperty reads a DSProperty from the buffer. `context`, `appid`, and | |
| 223 // `namespace` behave the same way they do for ReadDSKey, but only have an | |
| 224 // effect if the decoded property has a DSKey value. | |
| 225 func ReadDSProperty(buf *bytes.Buffer, context DSKeyContext, appid, namespace st
ring) (p gae.DSProperty, err error) { | |
| 226 val := interface{}(nil) | |
| 227 typb, err := buf.ReadByte() | |
| 228 if err != nil { | |
| 229 return | |
| 230 } | |
| 231 is := gae.ShouldIndex | |
| 232 if (typb & 0x80) != 0 { | |
| 233 is = gae.NoIndex | |
| 234 } | |
| 235 switch gae.DSPropertyType(typb & 0x7f) { | |
| 236 case gae.DSPTNull: | |
| 237 case gae.DSPTBoolTrue: | |
| 238 val = true | |
| 239 case gae.DSPTBoolFalse: | |
| 240 val = false | |
| 241 case gae.DSPTInt: | |
| 242 val, _, err = cmpbin.ReadInt(buf) | |
| 243 case gae.DSPTFloat: | |
| 244 val, _, err = cmpbin.ReadFloat64(buf) | |
| 245 case gae.DSPTString: | |
| 246 val, _, err = cmpbin.ReadString(buf) | |
| 247 case gae.DSPTBytes: | |
| 248 b := []byte(nil) | |
| 249 if b, _, err = cmpbin.ReadBytes(buf); err != nil { | |
| 250 break | |
| 251 } | |
| 252 if is == gae.NoIndex { | |
| 253 val = b | |
| 254 } else { | |
| 255 val = gae.DSByteString(b) | |
| 256 } | |
| 257 case gae.DSPTTime: | |
| 258 val, err = ReadTime(buf) | |
| 259 case gae.DSPTGeoPoint: | |
| 260 val, err = ReadDSGeoPoint(buf) | |
| 261 case gae.DSPTKey: | |
| 262 val, err = ReadDSKey(buf, context, appid, namespace) | |
| 263 case gae.DSPTBlobKey: | |
| 264 s := "" | |
| 265 if s, _, err = cmpbin.ReadString(buf); err != nil { | |
| 266 break | |
| 267 } | |
| 268 val = gae.BSKey(s) | |
| 269 default: | |
| 270 err = fmt.Errorf("read: unknown type! %v", typb) | |
| 271 } | |
| 272 if err == nil { | |
| 273 err = p.SetValue(val, is) | |
| 274 } | |
| 275 return | |
| 276 } | |
| 277 | |
| 278 // WriteDSPropertyMap writes an entire DSPropertyMap to the buffer. `context` | |
| 279 // behaves the same way that it does for WriteDSKey. If | |
| 280 // WriteDSPropertyMapDeterministic is true, then the rows will be sorted by | |
| 281 // property name before they're serialized to buf (mostly useful for testing, | |
| 282 // but also potentially useful if you need to make a hash of the property data). | |
| 283 func WriteDSPropertyMap(buf *bytes.Buffer, propMap gae.DSPropertyMap, context DS
KeyContext) { | |
| 284 rows := make(sort.StringSlice, 0, len(propMap)) | |
| 285 tmpBuf := &bytes.Buffer{} | |
| 286 for name, vals := range propMap { | |
| 287 tmpBuf.Reset() | |
| 288 cmpbin.WriteString(tmpBuf, name) | |
| 289 cmpbin.WriteUint(tmpBuf, uint64(len(vals))) | |
| 290 for _, p := range vals { | |
| 291 WriteDSProperty(tmpBuf, p, context) | |
| 292 } | |
| 293 rows = append(rows, tmpBuf.String()) | |
| 294 } | |
| 295 | |
| 296 if WriteDSPropertyMapDeterministic { | |
| 297 rows.Sort() | |
| 298 } | |
| 299 | |
| 300 cmpbin.WriteUint(buf, uint64(len(propMap))) | |
| 301 for _, r := range rows { | |
| 302 buf.WriteString(r) | |
| 303 } | |
| 304 } | |
| 305 | |
| 306 // ReadDSPropertyMap reads a DSPropertyMap from the buffer. `context` and | |
| 307 // friends behave the same way that they do for ReadDSKey. | |
| 308 func ReadDSPropertyMap(buf *bytes.Buffer, context DSKeyContext, appid, namespace
string) (propMap gae.DSPropertyMap, err error) { | |
| 309 numRows := uint64(0) | |
| 310 if numRows, _, err = cmpbin.ReadUint(buf); err != nil { | |
| 311 return | |
| 312 } | |
| 313 if numRows > ReadDSPropertyMapReasonableLimit { | |
| 314 err = fmt.Errorf("helper: tried to decode map with huge number o
f rows %d", numRows) | |
| 315 return | |
| 316 } | |
| 317 | |
| 318 name, prop := "", gae.DSProperty{} | |
| 319 propMap = make(gae.DSPropertyMap, numRows) | |
| 320 for i := uint64(0); i < numRows; i++ { | |
| 321 if name, _, err = cmpbin.ReadString(buf); err != nil { | |
| 322 return | |
| 323 } | |
| 324 | |
| 325 numProps := uint64(0) | |
| 326 if numProps, _, err = cmpbin.ReadUint(buf); err != nil { | |
| 327 return | |
| 328 } | |
| 329 if numProps > ReadDSPropertyMapReasonableLimit { | |
| 330 err = fmt.Errorf("helper: tried to decode map with huge
number of properties %d", numProps) | |
| 331 return | |
| 332 } | |
| 333 props := make([]gae.DSProperty, 0, numProps) | |
| 334 for j := uint64(0); j < numProps; j++ { | |
| 335 if prop, err = ReadDSProperty(buf, context, appid, names
pace); err != nil { | |
| 336 return | |
| 337 } | |
| 338 props = append(props, prop) | |
| 339 } | |
| 340 propMap[name] = props | |
| 341 } | |
| 342 return | |
| 343 } | |
| OLD | NEW |
|---|
Chromium Code Reviews
