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