]>
Commit | Line | Data |
---|---|---|
f67539c2 TL |
1 | /* |
2 | * Licensed to the Apache Software Foundation (ASF) under one | |
3 | * or more contributor license agreements. See the NOTICE file | |
4 | * distributed with this work for additional information | |
5 | * regarding copyright ownership. The ASF licenses this file | |
6 | * to you under the Apache License, Version 2.0 (the | |
7 | * "License"); you may not use this file except in compliance | |
8 | * with the License. You may obtain a copy of the License at | |
9 | * | |
10 | * http://www.apache.org/licenses/LICENSE-2.0 | |
11 | * | |
12 | * Unless required by applicable law or agreed to in writing, | |
13 | * software distributed under the License is distributed on an | |
14 | * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY | |
15 | * KIND, either express or implied. See the License for the | |
16 | * specific language governing permissions and limitations | |
17 | * under the License. | |
18 | */ | |
19 | ||
20 | package thrift | |
21 | ||
22 | import ( | |
23 | "bufio" | |
24 | "bytes" | |
25 | "context" | |
26 | "encoding/binary" | |
27 | "fmt" | |
28 | "io" | |
29 | ) | |
30 | ||
31 | const DEFAULT_MAX_LENGTH = 16384000 | |
32 | ||
33 | type TFramedTransport struct { | |
34 | transport TTransport | |
35 | buf bytes.Buffer | |
36 | reader *bufio.Reader | |
37 | frameSize uint32 //Current remaining size of the frame. if ==0 read next frame header | |
38 | buffer [4]byte | |
39 | maxLength uint32 | |
40 | } | |
41 | ||
42 | type tFramedTransportFactory struct { | |
43 | factory TTransportFactory | |
44 | maxLength uint32 | |
45 | } | |
46 | ||
47 | func NewTFramedTransportFactory(factory TTransportFactory) TTransportFactory { | |
48 | return &tFramedTransportFactory{factory: factory, maxLength: DEFAULT_MAX_LENGTH} | |
49 | } | |
50 | ||
51 | func NewTFramedTransportFactoryMaxLength(factory TTransportFactory, maxLength uint32) TTransportFactory { | |
52 | return &tFramedTransportFactory{factory: factory, maxLength: maxLength} | |
53 | } | |
54 | ||
55 | func (p *tFramedTransportFactory) GetTransport(base TTransport) (TTransport, error) { | |
56 | tt, err := p.factory.GetTransport(base) | |
57 | if err != nil { | |
58 | return nil, err | |
59 | } | |
60 | return NewTFramedTransportMaxLength(tt, p.maxLength), nil | |
61 | } | |
62 | ||
63 | func NewTFramedTransport(transport TTransport) *TFramedTransport { | |
64 | return &TFramedTransport{transport: transport, reader: bufio.NewReader(transport), maxLength: DEFAULT_MAX_LENGTH} | |
65 | } | |
66 | ||
67 | func NewTFramedTransportMaxLength(transport TTransport, maxLength uint32) *TFramedTransport { | |
68 | return &TFramedTransport{transport: transport, reader: bufio.NewReader(transport), maxLength: maxLength} | |
69 | } | |
70 | ||
71 | func (p *TFramedTransport) Open() error { | |
72 | return p.transport.Open() | |
73 | } | |
74 | ||
75 | func (p *TFramedTransport) IsOpen() bool { | |
76 | return p.transport.IsOpen() | |
77 | } | |
78 | ||
79 | func (p *TFramedTransport) Close() error { | |
80 | return p.transport.Close() | |
81 | } | |
82 | ||
83 | func (p *TFramedTransport) Read(buf []byte) (l int, err error) { | |
84 | if p.frameSize == 0 { | |
85 | p.frameSize, err = p.readFrameHeader() | |
86 | if err != nil { | |
87 | return | |
88 | } | |
89 | } | |
90 | if p.frameSize < uint32(len(buf)) { | |
91 | frameSize := p.frameSize | |
92 | tmp := make([]byte, p.frameSize) | |
93 | l, err = p.Read(tmp) | |
94 | copy(buf, tmp) | |
95 | if err == nil { | |
96 | // Note: It's important to only return an error when l | |
97 | // is zero. | |
98 | // In io.Reader.Read interface, it's perfectly fine to | |
99 | // return partial data and nil error, which means | |
100 | // "This is all the data we have right now without | |
101 | // blocking. If you need the full data, call Read again | |
102 | // or use io.ReadFull instead". | |
103 | // Returning partial data with an error actually means | |
104 | // there's no more data after the partial data just | |
105 | // returned, which is not true in this case | |
106 | // (it might be that the other end just haven't written | |
107 | // them yet). | |
108 | if l == 0 { | |
109 | err = NewTTransportExceptionFromError(fmt.Errorf("Not enough frame size %d to read %d bytes", frameSize, len(buf))) | |
110 | } | |
111 | return | |
112 | } | |
113 | } | |
114 | got, err := p.reader.Read(buf) | |
115 | p.frameSize = p.frameSize - uint32(got) | |
116 | //sanity check | |
117 | if p.frameSize < 0 { | |
118 | return 0, NewTTransportException(UNKNOWN_TRANSPORT_EXCEPTION, "Negative frame size") | |
119 | } | |
120 | return got, NewTTransportExceptionFromError(err) | |
121 | } | |
122 | ||
123 | func (p *TFramedTransport) ReadByte() (c byte, err error) { | |
124 | if p.frameSize == 0 { | |
125 | p.frameSize, err = p.readFrameHeader() | |
126 | if err != nil { | |
127 | return | |
128 | } | |
129 | } | |
130 | if p.frameSize < 1 { | |
131 | return 0, NewTTransportExceptionFromError(fmt.Errorf("Not enough frame size %d to read %d bytes", p.frameSize, 1)) | |
132 | } | |
133 | c, err = p.reader.ReadByte() | |
134 | if err == nil { | |
135 | p.frameSize-- | |
136 | } | |
137 | return | |
138 | } | |
139 | ||
140 | func (p *TFramedTransport) Write(buf []byte) (int, error) { | |
141 | n, err := p.buf.Write(buf) | |
142 | return n, NewTTransportExceptionFromError(err) | |
143 | } | |
144 | ||
145 | func (p *TFramedTransport) WriteByte(c byte) error { | |
146 | return p.buf.WriteByte(c) | |
147 | } | |
148 | ||
149 | func (p *TFramedTransport) WriteString(s string) (n int, err error) { | |
150 | return p.buf.WriteString(s) | |
151 | } | |
152 | ||
153 | func (p *TFramedTransport) Flush(ctx context.Context) error { | |
154 | size := p.buf.Len() | |
155 | buf := p.buffer[:4] | |
156 | binary.BigEndian.PutUint32(buf, uint32(size)) | |
157 | _, err := p.transport.Write(buf) | |
158 | if err != nil { | |
159 | p.buf.Truncate(0) | |
160 | return NewTTransportExceptionFromError(err) | |
161 | } | |
162 | if size > 0 { | |
163 | if n, err := p.buf.WriteTo(p.transport); err != nil { | |
164 | print("Error while flushing write buffer of size ", size, " to transport, only wrote ", n, " bytes: ", err.Error(), "\n") | |
165 | p.buf.Truncate(0) | |
166 | return NewTTransportExceptionFromError(err) | |
167 | } | |
168 | } | |
169 | err = p.transport.Flush(ctx) | |
170 | return NewTTransportExceptionFromError(err) | |
171 | } | |
172 | ||
173 | func (p *TFramedTransport) readFrameHeader() (uint32, error) { | |
174 | buf := p.buffer[:4] | |
175 | if _, err := io.ReadFull(p.reader, buf); err != nil { | |
176 | return 0, err | |
177 | } | |
178 | size := binary.BigEndian.Uint32(buf) | |
179 | if size < 0 || size > p.maxLength { | |
180 | return 0, NewTTransportException(UNKNOWN_TRANSPORT_EXCEPTION, fmt.Sprintf("Incorrect frame size (%d)", size)) | |
181 | } | |
182 | return size, nil | |
183 | } | |
184 | ||
185 | func (p *TFramedTransport) RemainingBytes() (num_bytes uint64) { | |
186 | return uint64(p.frameSize) | |
187 | } |