1# encoding: ascii-8bit
2#
3# Licensed to the Apache Software Foundation (ASF) under one
4# or more contributor license agreements. See the NOTICE file
5# distributed with this work for additional information
6# regarding copyright ownership. The ASF licenses this file
7# to you under the Apache License, Version 2.0 (the
8# "License"); you may not use this file except in compliance
9# with the License. You may obtain a copy of the License at
10#
11#   http://www.apache.org/licenses/LICENSE-2.0
12#
13# Unless required by applicable law or agreed to in writing,
14# software distributed under the License is distributed on an
15# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
16# KIND, either express or implied. See the License for the
17# specific language governing permissions and limitations
18# under the License.
19#
20
21module Thrift
22  class FramedTransport < BaseTransport
23    def initialize(transport, read=true, write=true)
24      @transport = transport
25      @rbuf      = Bytes.empty_byte_buffer
26      @wbuf      = Bytes.empty_byte_buffer
27      @read      = read
28      @write     = write
29      @index      = 0
30    end
31
32    def open?
33      @transport.open?
34    end
35
36    def open
37      @transport.open
38    end
39
40    def close
41      @transport.close
42    end
43
44    def read(sz)
45      return @transport.read(sz) unless @read
46
47      return Bytes.empty_byte_buffer if sz <= 0
48
49      read_frame if @index >= @rbuf.length
50
51      @index += sz
52      @rbuf.slice(@index - sz, sz) || Bytes.empty_byte_buffer
53    end
54
55    def read_byte
56      return @transport.read_byte() unless @read
57
58      read_frame if @index >= @rbuf.length
59
60      # The read buffer has some data now, read a single byte. Using get_string_byte() avoids
61      # allocating a temp string of size 1 unnecessarily.
62      @index += 1
63      return Bytes.get_string_byte(@rbuf, @index - 1)
64    end
65
66    def read_into_buffer(buffer, size)
67      i = 0
68      while i < size
69        read_frame if @index >= @rbuf.length
70
71        # The read buffer has some data now, so copy bytes over to the output buffer.
72        byte = Bytes.get_string_byte(@rbuf, @index)
73        Bytes.set_string_byte(buffer, i, byte)
74        @index += 1
75        i += 1
76      end
77      i
78    end
79
80    def write(buf, sz=nil)
81      return @transport.write(buf) unless @write
82
83      buf = Bytes.force_binary_encoding(buf)
84      @wbuf << (sz ? buf[0...sz] : buf)
85    end
86
87    #
88    # Writes the output buffer to the stream in the format of a 4-byte length
89    # followed by the actual data.
90    #
91    def flush
92      return @transport.flush unless @write
93
94      out = [@wbuf.length].pack('N')
95      # Array#pack should return a BINARY encoded String, so it shouldn't be necessary to force encoding
96      out << @wbuf
97      @transport.write(out)
98      @transport.flush
99      @wbuf = Bytes.empty_byte_buffer
100    end
101
102    def to_s
103      "framed(#{@transport.to_s})"
104    end
105
106    private
107
108    def read_frame
109      sz = @transport.read_all(4).unpack('N').first
110
111      @index = 0
112      @rbuf = @transport.read_all(sz)
113    end
114  end
115
116  class FramedTransportFactory < BaseTransportFactory
117    def get_transport(transport)
118      return FramedTransport.new(transport)
119    end
120
121    def to_s
122      "framed"
123    end
124  end
125end
126