1 /*
2 * Copyright (c) Facebook, Inc. and its affiliates.
3 *
4 * Licensed under the Apache License, Version 2.0 (the "License");
5 * you may not use this file except in compliance with the License.
6 * You may obtain a copy of the License at
7 *
8 * http://www.apache.org/licenses/LICENSE-2.0
9 *
10 * Unless required by applicable law or agreed to in writing, software
11 * distributed under the License is distributed on an "AS IS" BASIS,
12 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13 * See the License for the specific language governing permissions and
14 * limitations under the License.
15 */
16
17 #include <thrift/lib/cpp2/transport/core/testutil/TAsyncSocketIntercepted.h>
18
19 namespace apache {
20 namespace thrift {
21 namespace async {
22
performWrite(const iovec * vec,uint32_t count,folly::WriteFlags flags,uint32_t * countWritten,uint32_t * partialWritten)23 folly::AsyncSocket::WriteResult TAsyncSocketIntercepted::performWrite(
24 const iovec* vec,
25 uint32_t count,
26 folly::WriteFlags flags,
27 uint32_t* countWritten,
28 uint32_t* partialWritten) {
29 std::vector<iovec> newiov;
30 std::string corruptedDataHolder;
31 if (params_.get() && params_->corruptLastWriteByte_) {
32 newiov.resize(count);
33 memcpy(newiov.data(), vec, sizeof(iovec) * count);
34 int toModify = static_cast<int>(count) - 1;
35 while (toModify >= 0 && vec[toModify].iov_len == 0) {
36 --toModify;
37 }
38 if (toModify >= 0) {
39 corruptedDataHolder.assign(
40 static_cast<const char*>(newiov[toModify].iov_base),
41 newiov[toModify].iov_len);
42 ++corruptedDataHolder[corruptedDataHolder.size() - 1];
43 newiov[toModify].iov_base = const_cast<void*>(
44 static_cast<const void*>(corruptedDataHolder.data()));
45 }
46 vec = newiov.data();
47 }
48 WriteResult writeRes = folly::AsyncSocket::performWrite(
49 vec, count, flags, countWritten, partialWritten);
50 totalBytesWritten_ += writeRes.writeReturn;
51 return writeRes;
52 }
53
performRead(void ** buf,size_t * buflen,size_t * offset)54 folly::AsyncSocket::ReadResult TAsyncSocketIntercepted::performRead(
55 void** buf, size_t* buflen, size_t* offset) {
56 ReadResult res = folly::AsyncSocket::performRead(buf, buflen, offset);
57
58 if (params_.get() && params_->corruptLastReadByte_ && !res.exception &&
59 res.readReturn > 0 &&
60 res.readReturn >= params_->corruptLastReadByteMinSize_) {
61 static_cast<char*>(*buf)[res.readReturn - 1]++;
62 }
63 totalBytesRead_ += res.readReturn;
64 return res;
65 }
66
67 } // namespace async
68 } // namespace thrift
69 } // namespace apache
70