forked from apache/thrift
/
header.go
709 lines (630 loc) · 17.9 KB
/
header.go
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528
529
530
531
532
533
534
535
536
537
538
539
540
541
542
543
544
545
546
547
548
549
550
551
552
553
554
555
556
557
558
559
560
561
562
563
564
565
566
567
568
569
570
571
572
573
574
575
576
577
578
579
580
581
582
583
584
585
586
587
588
589
590
591
592
593
594
595
596
597
598
599
600
601
602
603
604
605
606
607
608
609
610
611
612
613
614
615
616
617
618
619
620
621
622
623
624
625
626
627
628
629
630
631
632
633
634
635
636
637
638
639
640
641
642
643
644
645
646
647
648
649
650
651
652
653
654
655
656
657
658
659
660
661
662
663
664
665
666
667
668
669
670
671
672
673
674
675
676
677
678
679
680
681
682
683
684
685
686
687
688
689
690
691
692
693
694
695
696
697
698
699
700
701
702
703
704
705
706
707
708
709
/*
* Licensed to the Apache Software Foundation (ASF) under one
* or more contributor license agreements. See the NOTICE file
* distributed with this work for additional information
* regarding copyright ownership. The ASF licenses this file
* to you under the Apache License, Version 2.0 (the
* "License"); you may not use this file except in compliance
* with the License. You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing,
* software distributed under the License is distributed on an
* "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
* KIND, either express or implied. See the License for the
* specific language governing permissions and limitations
* under the License.
*/
package thrift
import (
"bufio"
"bytes"
"compress/zlib"
"encoding/binary"
"fmt"
"io"
)
// Header keys
const (
IdentityHeader string = "identity"
IDVersionHeader string = "id_version"
IDVersion string = "1"
PriorityHeader string = "thrift_priority"
ClientTimeoutHeader string = "client_timeout"
QueueTimeoutHeader string = "queue_timeout"
// Header Magicks
// 0 and 16th bits must be 0 to differentiate from framed & unframed
HeaderMagic uint32 = 0x0FFF0000
HeaderMask uint32 = 0xFFFF0000
FlagsMask uint32 = 0x0000FFFF
HTTPServerMagic uint32 = 0x504F5354 // POST
HTTPClientMagic uint32 = 0x48545450 // HTTP
HTTPGetClientMagic uint32 = 0x47455420 // GET
HTTPHeadClientMagic uint32 = 0x48454144 // HEAD
BigFrameMagic uint32 = 0x42494746 // BIGF
MaxFrameSize uint32 = 0x3FFFFFFF
CommonHeaderSize uint64 = 10
MaxHeaderSize uint32 = 131071
)
type ClientType int64
const (
HeaderClientType ClientType = iota
FramedDeprecated
UnframedDeprecated
HTTPServerType
HTTPClientType
FramedCompact
HeaderSASLClientType
HTTPGetClientType
UnknownClientType
UnframedCompactDeprecated
)
func (c ClientType) String() string {
switch c {
case HeaderClientType:
return "Header"
case FramedDeprecated:
return "FramedDeprecated"
case UnframedDeprecated:
return "UnframedDeprecated"
case HTTPServerType:
return "HTTPServer"
case HTTPClientType:
return "HTTPClient"
case FramedCompact:
return "FramedCompact"
case HeaderSASLClientType:
return "HeaderSASL"
case HTTPGetClientType:
return "HTTPGet"
case UnframedCompactDeprecated:
return "UnframedCompactDeprecated"
case UnknownClientType:
fallthrough
default:
return "Unknown"
}
}
type HeaderFlags uint32
const (
HeaderFlagSupportOutOfOrder HeaderFlags = 0x01
HeaderFlagDuplexReverse HeaderFlags = 0x08
HeaderFlagSASL HeaderFlags = 0x10
)
type InfoIDType uint32
const (
InfoIDPadding InfoIDType = 0
InfoIDKeyValue InfoIDType = 1
InfoIDPKeyValue InfoIDType = 2
)
// TransformID Numerical ID of transform function
type TransformID uint32
const (
// TransformNone Default null transform
TransformNone TransformID = 0
// TransformZlib Apply zlib compression
TransformZlib TransformID = 1
// TransformHMAC Deprecated and no longer supported
TransformHMAC TransformID = 2
// TransformSnappy Apply snappy compression
TransformSnappy TransformID = 3
// TransformQLZ Deprecated and no longer supported
TransformQLZ TransformID = 4
// TransformZstd Apply zstd compression
TransformZstd TransformID = 5
)
func (c TransformID) String() string {
switch c {
case TransformNone:
return "none"
case TransformZlib:
return "zlib"
case TransformHMAC:
return "hmac"
case TransformSnappy:
return "snappy"
case TransformQLZ:
return "qlz"
case TransformZstd:
return "zstd"
default:
return "unknown"
}
}
var supportedTransforms = map[TransformID]bool{
TransformNone: true,
TransformZlib: true,
TransformHMAC: false,
TransformSnappy: false,
TransformQLZ: false,
TransformZstd: false,
}
// Untransformer will find a transform function to wrap a reader with to transformed the data.
func (c TransformID) Untransformer() (func(byteReader) (byteReader, error), error) {
switch c {
case TransformNone:
return func(rd byteReader) (byteReader, error) {
return rd, nil
}, nil
case TransformZlib:
return func(rd byteReader) (byteReader, error) {
zlrd, err := zlib.NewReader(rd)
if err != nil {
return nil, err
}
return ensureByteReader(zlrd), nil
}, nil
default:
return nil, NewTProtocolExceptionWithType(
NOT_IMPLEMENTED, fmt.Errorf("Header transform %s not supported", c.String()),
)
}
}
type tHeader struct {
length uint64
flags uint16
seq uint32
headerLen uint16
payloadLen uint64
protoID ProtocolID
transforms []TransformID
// Map to use for headers
headers map[string]string
pHeaders map[string]string
// clientType Negotiated client type
clientType ClientType
}
// byteReader Combined interface to expose original ReadByte calls
type byteReader interface {
io.Reader
io.ByteReader
}
// ensureByteReader If a reader does not implement ReadByte, wrap it with a
// buffer that can. Needed for most thrift interfaces.
func ensureByteReader(rd io.Reader) byteReader {
if brr, ok := rd.(byteReader); ok {
return brr
}
return bufio.NewReader(rd)
}
// limitedByteReader Keep the ByteReader interface when wrapping with a limit
type limitedByteReader struct {
io.LimitedReader
// Copy of the original interface given to us that implemented ByteReader
orig byteReader
}
func newLimitedByteReader(rd byteReader, n int64) *limitedByteReader {
return &limitedByteReader{
LimitedReader: io.LimitedReader{R: rd, N: n}, orig: rd,
}
}
func (r *limitedByteReader) ReadByte() (byte, error) {
if r.N <= 0 {
return '0', io.EOF
}
b, err := r.orig.ReadByte()
r.N--
return b, err
}
func readVarString(buf byteReader) (string, error) {
strlen, err := binary.ReadUvarint(buf)
if err != nil {
return "", fmt.Errorf("tHeader: error reading len of kv string: %s", err.Error())
}
strbuf := make([]byte, strlen)
_, err = io.ReadFull(buf, strbuf)
if err != nil {
return "", fmt.Errorf("tHeader: error reading kv string: %s", err.Error())
}
return string(strbuf), nil
}
// readHeaderMaps Consume a set of key/value pairs from the buffer
func readInfoHeaderSet(buf byteReader) (map[string]string, error) {
headers := map[string]string{}
numkvs, err := binary.ReadUvarint(buf)
if err != nil {
return nil, fmt.Errorf("tHeader: error reading number of keyvalues: %s", err.Error())
}
for i := uint64(0); i < numkvs; i++ {
key, err := readVarString(buf)
if err != nil {
return nil, fmt.Errorf("tHeader: error reading keyvalue key: %s", err.Error())
}
val, err := readVarString(buf)
if err != nil {
return nil, fmt.Errorf("tHeader: error reading keyvalue val: %s", err.Error())
}
headers[key] = val
}
return headers, nil
}
// readTransforms Consume a size delimited transform set from the buffer
// If the there is an unknown or unsupported transform we will bail out.
func readTransforms(buf byteReader) ([]TransformID, error) {
transforms := []TransformID{}
numtransforms, err := binary.ReadUvarint(buf)
if err != nil {
return nil, NewTTransportExceptionFromError(
fmt.Errorf("tHeader: error reading number of transforms: %s", err.Error()),
)
}
// Read transforms
for i := uint64(0); i < numtransforms; i++ {
transformID, err := binary.ReadUvarint(buf)
if err != nil {
return nil, NewTTransportExceptionFromError(
fmt.Errorf("tHeader: error reading transforms: %s", err.Error()),
)
}
tid := TransformID(transformID)
if supported, ok := supportedTransforms[tid]; ok {
if supported {
transforms = append(transforms, tid)
} else {
return nil, NewTTransportExceptionFromError(
fmt.Errorf("tHeader: unsupported transform: %s", tid.String()),
)
}
} else {
return nil, NewTTransportExceptionFromError(
fmt.Errorf("tHeader: unknown transform ID: %#x", tid),
)
}
}
return transforms, nil
}
// readInfoHeaders Read the K/V headers at the end of the header
// This will keep consuming bytes until the buffer returns EOF
func readInfoHeaders(buf byteReader) (map[string]string, map[string]string, error) {
// var err error
infoheaders := map[string]string{}
infopHeaders := map[string]string{}
for {
infoID, err := binary.ReadUvarint(buf)
// this is the last field, read until there is no more padding
if err == io.EOF {
break
}
if err != nil {
return nil, nil, NewTTransportExceptionFromError(
fmt.Errorf("tHeader: error reading infoID: %s", err.Error()),
)
}
switch InfoIDType(infoID) {
case InfoIDPadding:
continue
case InfoIDKeyValue:
hdrs, err := readInfoHeaderSet(buf)
if err != nil {
return nil, nil, err
}
for k, v := range hdrs {
infoheaders[k] = v
}
case InfoIDPKeyValue:
hdrs, err := readInfoHeaderSet(buf)
if err != nil {
return nil, nil, err
}
for k, v := range hdrs {
infopHeaders[k] = v
}
default:
return nil, nil, NewTTransportExceptionFromError(
fmt.Errorf("tHeader: error reading infoIDType: %#x", infoID),
)
}
}
return infoheaders, infopHeaders, nil
}
// readVarHeader Read the variable-length trailing header
func (hdr *tHeader) readVarHeader(buf byteReader) error {
// Read protocol ID
protoID, err := binary.ReadUvarint(buf)
if err != nil {
return NewTTransportExceptionFromError(
fmt.Errorf("tHeader: error reading protocol ID: %s", err.Error()),
)
}
hdr.protoID = ProtocolID(protoID)
hdr.transforms, err = readTransforms(buf)
if err != nil {
return err
}
hdr.headers, hdr.pHeaders, err = readInfoHeaders(buf)
if err != nil {
return err
}
return nil
}
// isCompactFramed Check if the magic value corresponds to compact proto
func isCompactFramed(magic uint32) bool {
protocolID := int8(magic >> 24)
protocolVersion := int8((magic >> 16) & uint32(COMPACT_VERSION_MASK))
return uint8(protocolID) == uint8(COMPACT_PROTOCOL_ID) && (protocolVersion == int8(COMPACT_VERSION) ||
protocolVersion == int8(COMPACT_VERSION_BE))
}
// analyzeFirst32Bit Guess client type from the first 4 bytes
func analyzeFirst32Bit(word uint32) ClientType {
if (word & BinaryVersionMask) == BinaryVersion1 {
return UnframedDeprecated
} else if isCompactFramed(word) {
return UnframedCompactDeprecated
} else if word == HTTPServerMagic ||
word == HTTPGetClientMagic ||
word == HTTPHeadClientMagic {
return HTTPServerType
} else if word == HTTPClientMagic {
return HTTPClientType
}
return UnknownClientType
}
// analyzeSecond32Bit Find the header client type from the 4-8th bytes of header
func analyzeSecond32Bit(word uint32) ClientType {
if (word & BinaryVersionMask) == BinaryVersion1 {
return FramedDeprecated
}
if isCompactFramed(word) {
return FramedCompact
}
if (word & HeaderMask) == HeaderMagic {
if (word & uint32(HeaderFlagSASL)) != 0 {
return HeaderSASLClientType
}
return HeaderClientType
}
return UnknownClientType
}
// checkFramed If the client type is framed, set appropriate protocolID in
// the header. Otherwise, return an unknown transport error.
func checkFramed(hdr *tHeader, clientType ClientType) error {
switch clientType {
case FramedDeprecated:
hdr.protoID = ProtocolIDBinary
hdr.clientType = clientType
hdr.payloadLen = hdr.length
return nil
case FramedCompact:
hdr.protoID = ProtocolIDCompact
hdr.clientType = clientType
hdr.payloadLen = hdr.length
return nil
default:
return NewTProtocolExceptionWithType(
NOT_IMPLEMENTED, fmt.Errorf("Transport %s not supported on tHeader", clientType),
)
}
}
// readHeaderInfo Consume header information from the buffer
func (hdr *tHeader) Read(buf *bufio.Reader) error {
var (
err error
firstword uint32
secondword uint32
wordbuf []byte
)
if wordbuf, err = buf.Peek(4); err != nil {
return NewTTransportExceptionFromError(err)
}
firstword = binary.BigEndian.Uint32(wordbuf)
// Check the first word if it matches http/unframed signatures
// We don't support non-framed protocols, so bail out
switch clientType := analyzeFirst32Bit(firstword); clientType {
case UnframedDeprecated:
hdr.clientType = clientType
hdr.protoID = ProtocolIDBinary
return nil
case UnframedCompactDeprecated:
hdr.clientType = clientType
hdr.protoID = ProtocolIDCompact
return nil
case UnknownClientType:
break
default:
return NewTTransportExceptionFromError(
fmt.Errorf("Transport %s not supported on tHeader (word=%#x)", clientType, firstword),
)
}
// From here on out, all protocols supported are frame-based. First word is length.
hdr.length = uint64(firstword)
if firstword > MaxFrameSize {
return NewTTransportExceptionFromError(
fmt.Errorf("BigFrames not supported: got size %d", firstword),
)
}
// First word is always length, discard.
_, err = buf.Discard(4)
if err != nil {
// Shouldn't be possible to fail here, but check anyways
return NewTTransportExceptionFromError(err)
}
// Only peek here. If it was framed transport, we are now reading payload.
if wordbuf, err = buf.Peek(4); err != nil {
return NewTTransportExceptionFromError(err)
}
secondword = binary.BigEndian.Uint32(wordbuf)
// Check if we can detect a framed proto, and bail out if we do.
if clientType := analyzeSecond32Bit(secondword); clientType != HeaderClientType {
return checkFramed(hdr, clientType)
}
// It was not framed proto, assume header and discard that word.
_, err = buf.Discard(4)
if err != nil {
// Shouldn't be possible to fail here, but check anyways
return NewTTransportExceptionFromError(err)
}
// Assume header protocol from here on in, parse rest of header
hdr.flags = uint16(secondword & FlagsMask)
err = binary.Read(buf, binary.BigEndian, &hdr.seq)
if err != nil {
return NewTTransportExceptionFromError(err)
}
err = binary.Read(buf, binary.BigEndian, &hdr.headerLen)
if err != nil {
return NewTTransportExceptionFromError(err)
}
if uint32(hdr.headerLen*4) > MaxHeaderSize {
return NewTTransportExceptionFromError(
fmt.Errorf("invalid header length: %d", int64(hdr.headerLen*4)),
)
}
// The length of the payload without the header (fixed is 10)
hdr.payloadLen = hdr.length - 10 - uint64(hdr.headerLen*4)
// Limit the reader for the header so we can't overrun
limbuf := newLimitedByteReader(buf, int64(hdr.headerLen*4))
hdr.clientType = HeaderClientType
return hdr.readVarHeader(limbuf)
}
func writeTransforms(transforms []TransformID, buf io.Writer) (int, error) {
size := 0
n, err := writeUvarint(uint64(len(transforms)), buf)
size += n
if err != nil {
return size, err
}
if transforms == nil {
return size, nil
}
for _, trans := range transforms {
// FIXME: We should only write supported xforms
n, err = writeUvarint(uint64(trans), buf)
size += n
if err != nil {
return size, err
}
}
return size, nil
}
func writeUvarint(v uint64, buf io.Writer) (int, error) {
var b [10]byte
n := binary.PutUvarint(b[:], v)
return buf.Write(b[:n])
}
func writeVarString(s string, buf io.Writer) (int, error) {
n, err := writeUvarint(uint64(len(s)), buf)
if err != nil {
return n, err
}
n2, err := buf.Write([]byte(s))
return n + n2, err
}
func writeInfoHeaders(headers map[string]string, infoidtype InfoIDType, buf io.Writer) (int, error) {
cnt := len(headers)
size := 0
if cnt < 1 {
return 0, nil
}
n, err := writeUvarint(uint64(infoidtype), buf)
size += n
if err != nil {
return 0, err
}
n, err = writeUvarint(uint64(cnt), buf)
size += n
if err != nil {
return 0, err
}
for k, v := range headers {
n, err = writeVarString(k, buf)
size += n
if err != nil {
return 0, err
}
n, err = writeVarString(v, buf)
size += n
if err != nil {
return 0, err
}
}
return size, nil
}
func (hdr *tHeader) writeVarHeader(buf io.Writer) (int, error) {
size := 0
n, err := writeUvarint(uint64(hdr.protoID), buf)
size += n
if err != nil {
return size, err
}
n, err = writeTransforms(hdr.transforms, buf)
size += n
if err != nil {
return size, err
}
n, err = writeInfoHeaders(hdr.pHeaders, InfoIDPKeyValue, buf)
size += n
if err != nil {
return size, err
}
n, err = writeInfoHeaders(hdr.headers, InfoIDKeyValue, buf)
size += n
if err != nil {
return size, err
}
padding := (4 - size%4) % 4
for i := 0; i < padding; i++ {
buf.Write([]byte{byte(0)})
size++
}
return size, err
}
func (hdr *tHeader) calcLenFromPayload() error {
fixedlen := uint64(0)
switch hdr.clientType {
case FramedCompact:
hdr.length = hdr.payloadLen
return nil
case FramedDeprecated:
hdr.length = hdr.payloadLen
return nil
case HeaderClientType:
// TODO: Changes with bigframes
fixedlen = 10
default:
return NewTApplicationException(
UNKNOWN_TRANSPORT_EXCEPTION,
fmt.Sprintf("cannot get length of non-framed transport %s", hdr.clientType.String()),
)
}
framesize := uint64(hdr.payloadLen + fixedlen + uint64(hdr.headerLen)*4)
// FIXME: support bigframes
if framesize > uint64(MaxFrameSize) {
return NewTTransportException(
INVALID_FRAME_SIZE,
fmt.Sprintf("cannot send bigframe of size %d", framesize),
)
}
hdr.length = framesize
return nil
}
// Write Write out the header, requires payloadLen be set.
func (hdr *tHeader) Write(buf io.Writer) error {
// Make a reasonably sized temp buffer for the variable header
hdrbuf := bytes.NewBuffer(nil)
_, err := hdr.writeVarHeader(hdrbuf)
if err != nil {
return err
}
if (hdrbuf.Len() % 4) > 0 {
return NewTTransportException(
INVALID_FRAME_SIZE, fmt.Sprintf("unable to write header of size %d (must be multiple of 4)", hdr.headerLen),
)
}
if hdrbuf.Len() > int(MaxHeaderSize) {
return NewTApplicationException(
INVALID_FRAME_SIZE, fmt.Sprintf("unable to write header of size %d (max is %d)", hdrbuf.Len(), MaxHeaderSize),
)
}
hdr.headerLen = uint16(hdrbuf.Len() / 4)
err = hdr.calcLenFromPayload()
if err != nil {
return err
}
binary.Write(buf, binary.BigEndian, uint32(hdr.length))
binary.Write(buf, binary.BigEndian, uint16(HeaderMagic>>16))
binary.Write(buf, binary.BigEndian, hdr.flags)
binary.Write(buf, binary.BigEndian, hdr.seq)
binary.Write(buf, binary.BigEndian, hdr.headerLen)
hdrbuf.WriteTo(buf)
return nil
}