--- /dev/null
+/*
+ * 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"
+ "context"
+ "encoding/binary"
+ "errors"
+ "fmt"
+ "io"
+ "io/ioutil"
+)
+
+// Size in bytes for 32-bit ints.
+const size32 = 4
+
+type headerMeta struct {
+ MagicFlags uint32
+ SequenceID int32
+ HeaderLength uint16
+}
+
+const headerMetaSize = 10
+
+type clientType int
+
+const (
+ clientUnknown clientType = iota
+ clientHeaders
+ clientFramedBinary
+ clientUnframedBinary
+ clientFramedCompact
+ clientUnframedCompact
+)
+
+// Constants defined in THeader format:
+// https://github.com/apache/thrift/blob/master/doc/specs/HeaderFormat.md
+const (
+ THeaderHeaderMagic uint32 = 0x0fff0000
+ THeaderHeaderMask uint32 = 0xffff0000
+ THeaderFlagsMask uint32 = 0x0000ffff
+ THeaderMaxFrameSize uint32 = 0x3fffffff
+)
+
+// THeaderMap is the type of the header map in THeader transport.
+type THeaderMap map[string]string
+
+// THeaderProtocolID is the wrapped protocol id used in THeader.
+type THeaderProtocolID int32
+
+// Supported THeaderProtocolID values.
+const (
+ THeaderProtocolBinary THeaderProtocolID = 0x00
+ THeaderProtocolCompact THeaderProtocolID = 0x02
+ THeaderProtocolDefault = THeaderProtocolBinary
+)
+
+// GetProtocol gets the corresponding TProtocol from the wrapped protocol id.
+func (id THeaderProtocolID) GetProtocol(trans TTransport) (TProtocol, error) {
+ switch id {
+ default:
+ return nil, NewTApplicationException(
+ INVALID_PROTOCOL,
+ fmt.Sprintf("THeader protocol id %d not supported", id),
+ )
+ case THeaderProtocolBinary:
+ return NewTBinaryProtocolFactoryDefault().GetProtocol(trans), nil
+ case THeaderProtocolCompact:
+ return NewTCompactProtocol(trans), nil
+ }
+}
+
+// THeaderTransformID defines the numeric id of the transform used.
+type THeaderTransformID int32
+
+// THeaderTransformID values
+const (
+ TransformNone THeaderTransformID = iota // 0, no special handling
+ TransformZlib // 1, zlib
+ // Rest of the values are not currently supported, namely HMAC and Snappy.
+)
+
+var supportedTransformIDs = map[THeaderTransformID]bool{
+ TransformNone: true,
+ TransformZlib: true,
+}
+
+// TransformReader is an io.ReadCloser that handles transforms reading.
+type TransformReader struct {
+ io.Reader
+
+ closers []io.Closer
+}
+
+var _ io.ReadCloser = (*TransformReader)(nil)
+
+// NewTransformReaderWithCapacity initializes a TransformReader with expected
+// closers capacity.
+//
+// If you don't know the closers capacity beforehand, just use
+//
+// &TransformReader{Reader: baseReader}
+//
+// instead would be sufficient.
+func NewTransformReaderWithCapacity(baseReader io.Reader, capacity int) *TransformReader {
+ return &TransformReader{
+ Reader: baseReader,
+ closers: make([]io.Closer, 0, capacity),
+ }
+}
+
+// Close calls the underlying closers in appropriate order,
+// stops at and returns the first error encountered.
+func (tr *TransformReader) Close() error {
+ // Call closers in reversed order
+ for i := len(tr.closers) - 1; i >= 0; i-- {
+ if err := tr.closers[i].Close(); err != nil {
+ return err
+ }
+ }
+ return nil
+}
+
+// AddTransform adds a transform.
+func (tr *TransformReader) AddTransform(id THeaderTransformID) error {
+ switch id {
+ default:
+ return NewTApplicationException(
+ INVALID_TRANSFORM,
+ fmt.Sprintf("THeaderTransformID %d not supported", id),
+ )
+ case TransformNone:
+ // no-op
+ case TransformZlib:
+ readCloser, err := zlib.NewReader(tr.Reader)
+ if err != nil {
+ return err
+ }
+ tr.Reader = readCloser
+ tr.closers = append(tr.closers, readCloser)
+ }
+ return nil
+}
+
+// TransformWriter is an io.WriteCloser that handles transforms writing.
+type TransformWriter struct {
+ io.Writer
+
+ closers []io.Closer
+}
+
+var _ io.WriteCloser = (*TransformWriter)(nil)
+
+// NewTransformWriter creates a new TransformWriter with base writer and transforms.
+func NewTransformWriter(baseWriter io.Writer, transforms []THeaderTransformID) (io.WriteCloser, error) {
+ writer := &TransformWriter{
+ Writer: baseWriter,
+ closers: make([]io.Closer, 0, len(transforms)),
+ }
+ for _, id := range transforms {
+ if err := writer.AddTransform(id); err != nil {
+ return nil, err
+ }
+ }
+ return writer, nil
+}
+
+// Close calls the underlying closers in appropriate order,
+// stops at and returns the first error encountered.
+func (tw *TransformWriter) Close() error {
+ // Call closers in reversed order
+ for i := len(tw.closers) - 1; i >= 0; i-- {
+ if err := tw.closers[i].Close(); err != nil {
+ return err
+ }
+ }
+ return nil
+}
+
+// AddTransform adds a transform.
+func (tw *TransformWriter) AddTransform(id THeaderTransformID) error {
+ switch id {
+ default:
+ return NewTApplicationException(
+ INVALID_TRANSFORM,
+ fmt.Sprintf("THeaderTransformID %d not supported", id),
+ )
+ case TransformNone:
+ // no-op
+ case TransformZlib:
+ writeCloser := zlib.NewWriter(tw.Writer)
+ tw.Writer = writeCloser
+ tw.closers = append(tw.closers, writeCloser)
+ }
+ return nil
+}
+
+// THeaderInfoType is the type id of the info headers.
+type THeaderInfoType int32
+
+// Supported THeaderInfoType values.
+const (
+ _ THeaderInfoType = iota // Skip 0
+ InfoKeyValue // 1
+ // Rest of the info types are not supported.
+)
+
+// THeaderTransport is a Transport mode that implements THeader.
+//
+// Note that THeaderTransport handles frame and zlib by itself,
+// so the underlying transport should be a raw socket transports (TSocket or TSSLSocket),
+// instead of rich transports like TZlibTransport or TFramedTransport.
+type THeaderTransport struct {
+ SequenceID int32
+ Flags uint32
+
+ transport TTransport
+
+ // THeaderMap for read and write
+ readHeaders THeaderMap
+ writeHeaders THeaderMap
+
+ // Reading related variables.
+ reader *bufio.Reader
+ // When frame is detected, we read the frame fully into frameBuffer.
+ frameBuffer bytes.Buffer
+ // When it's non-nil, Read should read from frameReader instead of
+ // reader, and EOF error indicates end of frame instead of end of all
+ // transport.
+ frameReader io.ReadCloser
+
+ // Writing related variables
+ writeBuffer bytes.Buffer
+ writeTransforms []THeaderTransformID
+
+ clientType clientType
+ protocolID THeaderProtocolID
+
+ // buffer is used in the following scenarios to avoid repetitive
+ // allocations, while 4 is big enough for all those scenarios:
+ //
+ // * header padding (max size 4)
+ // * write the frame size (size 4)
+ buffer [4]byte
+}
+
+var _ TTransport = (*THeaderTransport)(nil)
+
+// NewTHeaderTransport creates THeaderTransport from the underlying transport.
+//
+// Please note that THeaderTransport handles framing and zlib by itself,
+// so the underlying transport should be the raw socket transports (TSocket or TSSLSocket),
+// instead of rich transports like TZlibTransport or TFramedTransport.
+//
+// If trans is already a *THeaderTransport, it will be returned as is.
+func NewTHeaderTransport(trans TTransport) *THeaderTransport {
+ if ht, ok := trans.(*THeaderTransport); ok {
+ return ht
+ }
+ return &THeaderTransport{
+ transport: trans,
+ reader: bufio.NewReader(trans),
+ writeHeaders: make(THeaderMap),
+ protocolID: THeaderProtocolDefault,
+ }
+}
+
+// Open calls the underlying transport's Open function.
+func (t *THeaderTransport) Open() error {
+ return t.transport.Open()
+}
+
+// IsOpen calls the underlying transport's IsOpen function.
+func (t *THeaderTransport) IsOpen() bool {
+ return t.transport.IsOpen()
+}
+
+// ReadFrame tries to read the frame header, guess the client type, and handle
+// unframed clients.
+func (t *THeaderTransport) ReadFrame() error {
+ if !t.needReadFrame() {
+ // No need to read frame, skipping.
+ return nil
+ }
+ // Peek and handle the first 32 bits.
+ // They could either be the length field of a framed message,
+ // or the first bytes of an unframed message.
+ buf, err := t.reader.Peek(size32)
+ if err != nil {
+ return err
+ }
+ frameSize := binary.BigEndian.Uint32(buf)
+ if frameSize&VERSION_MASK == VERSION_1 {
+ t.clientType = clientUnframedBinary
+ return nil
+ }
+ if buf[0] == COMPACT_PROTOCOL_ID && buf[1]&COMPACT_VERSION_MASK == COMPACT_VERSION {
+ t.clientType = clientUnframedCompact
+ return nil
+ }
+
+ // At this point it should be a framed message,
+ // sanity check on frameSize then discard the peeked part.
+ if frameSize > THeaderMaxFrameSize {
+ return NewTProtocolExceptionWithType(
+ SIZE_LIMIT,
+ errors.New("frame too large"),
+ )
+ }
+ t.reader.Discard(size32)
+
+ // Read the frame fully into frameBuffer.
+ _, err = io.Copy(
+ &t.frameBuffer,
+ io.LimitReader(t.reader, int64(frameSize)),
+ )
+ if err != nil {
+ return err
+ }
+ t.frameReader = ioutil.NopCloser(&t.frameBuffer)
+
+ // Peek and handle the next 32 bits.
+ buf = t.frameBuffer.Bytes()[:size32]
+ version := binary.BigEndian.Uint32(buf)
+ if version&THeaderHeaderMask == THeaderHeaderMagic {
+ t.clientType = clientHeaders
+ return t.parseHeaders(frameSize)
+ }
+ if version&VERSION_MASK == VERSION_1 {
+ t.clientType = clientFramedBinary
+ return nil
+ }
+ if buf[0] == COMPACT_PROTOCOL_ID && buf[1]&COMPACT_VERSION_MASK == COMPACT_VERSION {
+ t.clientType = clientFramedCompact
+ return nil
+ }
+ if err := t.endOfFrame(); err != nil {
+ return err
+ }
+ return NewTProtocolExceptionWithType(
+ NOT_IMPLEMENTED,
+ errors.New("unsupported client transport type"),
+ )
+}
+
+// endOfFrame does end of frame handling.
+//
+// It closes frameReader, and also resets frame related states.
+func (t *THeaderTransport) endOfFrame() error {
+ defer func() {
+ t.frameBuffer.Reset()
+ t.frameReader = nil
+ }()
+ return t.frameReader.Close()
+}
+
+func (t *THeaderTransport) parseHeaders(frameSize uint32) error {
+ if t.clientType != clientHeaders {
+ return nil
+ }
+
+ var err error
+ var meta headerMeta
+ if err = binary.Read(&t.frameBuffer, binary.BigEndian, &meta); err != nil {
+ return err
+ }
+ frameSize -= headerMetaSize
+ t.Flags = meta.MagicFlags & THeaderFlagsMask
+ t.SequenceID = meta.SequenceID
+ headerLength := int64(meta.HeaderLength) * 4
+ if int64(frameSize) < headerLength {
+ return NewTProtocolExceptionWithType(
+ SIZE_LIMIT,
+ errors.New("header size is larger than the whole frame"),
+ )
+ }
+ headerBuf := NewTMemoryBuffer()
+ _, err = io.Copy(headerBuf, io.LimitReader(&t.frameBuffer, headerLength))
+ if err != nil {
+ return err
+ }
+ hp := NewTCompactProtocol(headerBuf)
+
+ // At this point the header is already read into headerBuf,
+ // and t.frameBuffer starts from the actual payload.
+ protoID, err := hp.readVarint32()
+ if err != nil {
+ return err
+ }
+ t.protocolID = THeaderProtocolID(protoID)
+ var transformCount int32
+ transformCount, err = hp.readVarint32()
+ if err != nil {
+ return err
+ }
+ if transformCount > 0 {
+ reader := NewTransformReaderWithCapacity(
+ &t.frameBuffer,
+ int(transformCount),
+ )
+ t.frameReader = reader
+ transformIDs := make([]THeaderTransformID, transformCount)
+ for i := 0; i < int(transformCount); i++ {
+ id, err := hp.readVarint32()
+ if err != nil {
+ return err
+ }
+ transformIDs[i] = THeaderTransformID(id)
+ }
+ // The transform IDs on the wire was added based on the order of
+ // writing, so on the reading side we need to reverse the order.
+ for i := transformCount - 1; i >= 0; i-- {
+ id := transformIDs[i]
+ if err := reader.AddTransform(id); err != nil {
+ return err
+ }
+ }
+ }
+
+ // The info part does not use the transforms yet, so it's
+ // important to continue using headerBuf.
+ headers := make(THeaderMap)
+ for {
+ infoType, err := hp.readVarint32()
+ if err == io.EOF {
+ break
+ }
+ if err != nil {
+ return err
+ }
+ if THeaderInfoType(infoType) == InfoKeyValue {
+ count, err := hp.readVarint32()
+ if err != nil {
+ return err
+ }
+ for i := 0; i < int(count); i++ {
+ key, err := hp.ReadString()
+ if err != nil {
+ return err
+ }
+ value, err := hp.ReadString()
+ if err != nil {
+ return err
+ }
+ headers[key] = value
+ }
+ } else {
+ // Skip reading info section on the first
+ // unsupported info type.
+ break
+ }
+ }
+ t.readHeaders = headers
+
+ return nil
+}
+
+func (t *THeaderTransport) needReadFrame() bool {
+ if t.clientType == clientUnknown {
+ // This is a new connection that's never read before.
+ return true
+ }
+ if t.isFramed() && t.frameReader == nil {
+ // We just finished the last frame.
+ return true
+ }
+ return false
+}
+
+func (t *THeaderTransport) Read(p []byte) (read int, err error) {
+ err = t.ReadFrame()
+ if err != nil {
+ return
+ }
+ if t.frameReader != nil {
+ read, err = t.frameReader.Read(p)
+ if err == io.EOF {
+ err = t.endOfFrame()
+ if err != nil {
+ return
+ }
+ if read < len(p) {
+ var nextRead int
+ nextRead, err = t.Read(p[read:])
+ read += nextRead
+ }
+ }
+ return
+ }
+ return t.reader.Read(p)
+}
+
+// Write writes data to the write buffer.
+//
+// You need to call Flush to actually write them to the transport.
+func (t *THeaderTransport) Write(p []byte) (int, error) {
+ return t.writeBuffer.Write(p)
+}
+
+// Flush writes the appropriate header and the write buffer to the underlying transport.
+func (t *THeaderTransport) Flush(ctx context.Context) error {
+ if t.writeBuffer.Len() == 0 {
+ return nil
+ }
+
+ defer t.writeBuffer.Reset()
+
+ switch t.clientType {
+ default:
+ fallthrough
+ case clientUnknown:
+ t.clientType = clientHeaders
+ fallthrough
+ case clientHeaders:
+ headers := NewTMemoryBuffer()
+ hp := NewTCompactProtocol(headers)
+ if _, err := hp.writeVarint32(int32(t.protocolID)); err != nil {
+ return NewTTransportExceptionFromError(err)
+ }
+ if _, err := hp.writeVarint32(int32(len(t.writeTransforms))); err != nil {
+ return NewTTransportExceptionFromError(err)
+ }
+ for _, transform := range t.writeTransforms {
+ if _, err := hp.writeVarint32(int32(transform)); err != nil {
+ return NewTTransportExceptionFromError(err)
+ }
+ }
+ if len(t.writeHeaders) > 0 {
+ if _, err := hp.writeVarint32(int32(InfoKeyValue)); err != nil {
+ return NewTTransportExceptionFromError(err)
+ }
+ if _, err := hp.writeVarint32(int32(len(t.writeHeaders))); err != nil {
+ return NewTTransportExceptionFromError(err)
+ }
+ for key, value := range t.writeHeaders {
+ if err := hp.WriteString(key); err != nil {
+ return NewTTransportExceptionFromError(err)
+ }
+ if err := hp.WriteString(value); err != nil {
+ return NewTTransportExceptionFromError(err)
+ }
+ }
+ }
+ padding := 4 - headers.Len()%4
+ if padding < 4 {
+ buf := t.buffer[:padding]
+ for i := range buf {
+ buf[i] = 0
+ }
+ if _, err := headers.Write(buf); err != nil {
+ return NewTTransportExceptionFromError(err)
+ }
+ }
+
+ var payload bytes.Buffer
+ meta := headerMeta{
+ MagicFlags: THeaderHeaderMagic + t.Flags&THeaderFlagsMask,
+ SequenceID: t.SequenceID,
+ HeaderLength: uint16(headers.Len() / 4),
+ }
+ if err := binary.Write(&payload, binary.BigEndian, meta); err != nil {
+ return NewTTransportExceptionFromError(err)
+ }
+ if _, err := io.Copy(&payload, headers); err != nil {
+ return NewTTransportExceptionFromError(err)
+ }
+
+ writer, err := NewTransformWriter(&payload, t.writeTransforms)
+ if err != nil {
+ return NewTTransportExceptionFromError(err)
+ }
+ if _, err := io.Copy(writer, &t.writeBuffer); err != nil {
+ return NewTTransportExceptionFromError(err)
+ }
+ if err := writer.Close(); err != nil {
+ return NewTTransportExceptionFromError(err)
+ }
+
+ // First write frame length
+ buf := t.buffer[:size32]
+ binary.BigEndian.PutUint32(buf, uint32(payload.Len()))
+ if _, err := t.transport.Write(buf); err != nil {
+ return NewTTransportExceptionFromError(err)
+ }
+ // Then write the payload
+ if _, err := io.Copy(t.transport, &payload); err != nil {
+ return NewTTransportExceptionFromError(err)
+ }
+
+ case clientFramedBinary, clientFramedCompact:
+ buf := t.buffer[:size32]
+ binary.BigEndian.PutUint32(buf, uint32(t.writeBuffer.Len()))
+ if _, err := t.transport.Write(buf); err != nil {
+ return NewTTransportExceptionFromError(err)
+ }
+ fallthrough
+ case clientUnframedBinary, clientUnframedCompact:
+ if _, err := io.Copy(t.transport, &t.writeBuffer); err != nil {
+ return NewTTransportExceptionFromError(err)
+ }
+ }
+
+ select {
+ default:
+ case <-ctx.Done():
+ return NewTTransportExceptionFromError(ctx.Err())
+ }
+
+ return t.transport.Flush(ctx)
+}
+
+// Close closes the transport, along with its underlying transport.
+func (t *THeaderTransport) Close() error {
+ if err := t.Flush(context.Background()); err != nil {
+ return err
+ }
+ return t.transport.Close()
+}
+
+// RemainingBytes calls underlying transport's RemainingBytes.
+//
+// Even in framed cases, because of all the possible compression transforms
+// involved, the remaining frame size is likely to be different from the actual
+// remaining readable bytes, so we don't bother to keep tracking the remaining
+// frame size by ourselves and just use the underlying transport's
+// RemainingBytes directly.
+func (t *THeaderTransport) RemainingBytes() uint64 {
+ return t.transport.RemainingBytes()
+}
+
+// GetReadHeaders returns the THeaderMap read from transport.
+func (t *THeaderTransport) GetReadHeaders() THeaderMap {
+ return t.readHeaders
+}
+
+// SetWriteHeader sets a header for write.
+func (t *THeaderTransport) SetWriteHeader(key, value string) {
+ t.writeHeaders[key] = value
+}
+
+// ClearWriteHeaders clears all write headers previously set.
+func (t *THeaderTransport) ClearWriteHeaders() {
+ t.writeHeaders = make(THeaderMap)
+}
+
+// AddTransform add a transform for writing.
+func (t *THeaderTransport) AddTransform(transform THeaderTransformID) error {
+ if !supportedTransformIDs[transform] {
+ return NewTProtocolExceptionWithType(
+ NOT_IMPLEMENTED,
+ fmt.Errorf("THeaderTransformID %d not supported", transform),
+ )
+ }
+ t.writeTransforms = append(t.writeTransforms, transform)
+ return nil
+}
+
+// Protocol returns the wrapped protocol id used in this THeaderTransport.
+func (t *THeaderTransport) Protocol() THeaderProtocolID {
+ switch t.clientType {
+ default:
+ return t.protocolID
+ case clientFramedBinary, clientUnframedBinary:
+ return THeaderProtocolBinary
+ case clientFramedCompact, clientUnframedCompact:
+ return THeaderProtocolCompact
+ }
+}
+
+func (t *THeaderTransport) isFramed() bool {
+ switch t.clientType {
+ default:
+ return false
+ case clientHeaders, clientFramedBinary, clientFramedCompact:
+ return true
+ }
+}
+
+// THeaderTransportFactory is a TTransportFactory implementation to create
+// THeaderTransport.
+type THeaderTransportFactory struct {
+ // The underlying factory, could be nil.
+ Factory TTransportFactory
+}
+
+// NewTHeaderTransportFactory creates a new *THeaderTransportFactory.
+func NewTHeaderTransportFactory(factory TTransportFactory) TTransportFactory {
+ return &THeaderTransportFactory{
+ Factory: factory,
+ }
+}
+
+// GetTransport implements TTransportFactory.
+func (f *THeaderTransportFactory) GetTransport(trans TTransport) (TTransport, error) {
+ if f.Factory != nil {
+ t, err := f.Factory.GetTransport(trans)
+ if err != nil {
+ return nil, err
+ }
+ return NewTHeaderTransport(t), nil
+ }
+ return NewTHeaderTransport(trans), nil
+}
projects / ceph.git / blobdiff