+++ /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