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 #ifndef _THRIFT_TRANSPORT_TSHORTREADTRANSPORT_H_
21 #define _THRIFT_TRANSPORT_TSHORTREADTRANSPORT_H_ 1
22 
23 #include <cstdlib>
24 
25 #include <thrift/transport/TTransport.h>
26 #include <thrift/transport/TVirtualTransport.h>
27 
28 namespace apache {
29 namespace thrift {
30 namespace transport {
31 namespace test {
32 
33 /**
34  * This class is only meant for testing.  It wraps another transport.
35  * Calls to read are passed through with some probability.  Otherwise,
36  * the read amount is randomly reduced before being passed through.
37  *
38  */
39 class TShortReadTransport : public TVirtualTransport<TShortReadTransport> {
40 public:
41   TShortReadTransport(std::shared_ptr<TTransport> transport, double full_prob,
42                      std::shared_ptr<TConfiguration> config = nullptr)
TVirtualTransport(config)43     : TVirtualTransport(config), transport_(transport), fullProb_(full_prob) {
44     }
45 
isOpen()46   bool isOpen() const override { return transport_->isOpen(); }
47 
peek()48   bool peek() override { return transport_->peek(); }
49 
open()50   void open() override { transport_->open(); }
51 
close()52   void close() override { transport_->close(); }
53 
read(uint8_t * buf,uint32_t len)54   uint32_t read(uint8_t* buf, uint32_t len) {
55     checkReadBytesAvailable(len);
56     if (len == 0) {
57       return 0;
58     }
59 
60     if (rand() / (double)RAND_MAX >= fullProb_) {
61       len = 1 + rand() % len;
62     }
63     return transport_->read(buf, len);
64   }
65 
write(const uint8_t * buf,uint32_t len)66   void write(const uint8_t* buf, uint32_t len) { transport_->write(buf, len); }
67 
flush()68   void flush() override {
69     resetConsumedMessageSize();
70     transport_->flush();
71   }
72 
borrow(uint8_t * buf,uint32_t * len)73   const uint8_t* borrow(uint8_t* buf, uint32_t* len) { return transport_->borrow(buf, len); }
74 
consume(uint32_t len)75   void consume(uint32_t len) {
76     countConsumedMessageBytes(len);
77     return transport_->consume(len);
78   }
79 
getUnderlyingTransport()80   std::shared_ptr<TTransport> getUnderlyingTransport() { return transport_; }
81 
82 protected:
83   std::shared_ptr<TTransport> transport_;
84   double fullProb_;
85 };
86 }
87 }
88 }
89 } // apache::thrift::transport::test
90 
91 #endif // #ifndef _THRIFT_TRANSPORT_TSHORTREADTRANSPORT_H_
92