1 /*
2 * Copyright (c) 2018-present, Facebook, Inc.
3 * All rights reserved.
4 *
5 * This source code is licensed under the BSD-style license found in the
6 * LICENSE file in the root directory of this source tree.
7 */
8
9 #pragma once
10
11 #include <folly/Range.h>
12 #include <folly/io/Cursor.h>
13 #include <folly/io/IOBuf.h>
14 #include <algorithm>
15
16 namespace fizz {
17
18 /**
19 * Trims trimmed.size() bytes at the end of a chained IOBuf and fills in
20 * trimmed. The caller is responsible for making sure that buf has enough
21 * bytes to trim trimmed.size().
22 */
23 void trimBytes(folly::IOBuf& buf, folly::MutableByteRange trimmed);
24
25 /**
26 * Trim bytes from the start of an IOBuf potentially spanning multiple
27 * buffers in the chain.
28 */
29 void trimStart(folly::IOBuf& buf, size_t toTrim);
30
31 /**
32 * XOR first and second and store the result in second.
33 */
34 void XOR(folly::ByteRange first, folly::MutableByteRange second);
35
36 /**
37 * Useful when we need to run a function that performs some operation on
38 * an IOBuf and then stores the result in another IOBuf.
39 * We assume that the caller ensures that out size >= in size.
40 * Apart from that you can pass in chained buffers for output or input
41 * buffers or even the same buffer. In the case you pass the same buffer,
42 * the Func needs to make sure it can safely deal with this case.
43 */
44 template <typename Func>
transformBuffer(const folly::IOBuf & in,folly::IOBuf & out,Func func)45 void transformBuffer(const folly::IOBuf& in, folly::IOBuf& out, Func func) {
46 auto current = ∈
47 folly::IOBuf* currentOut = &out;
48 size_t offset = 0;
49
50 // Iterate using the buffers instead of iterating with
51 // the iobuf iterators.
52 do {
53 size_t currentLength = current->length();
54
55 while (currentLength != 0) {
56 size_t selected = std::min(
57 static_cast<size_t>(currentOut->length() - offset), currentLength);
58 func(
59 currentOut->writableData() + offset,
60 current->data() + (current->length() - currentLength),
61 selected);
62 currentLength -= selected;
63 offset += selected;
64 if (offset == currentOut->length()) {
65 offset = 0;
66 currentOut = currentOut->next();
67 }
68 }
69 current = current->next();
70 } while (current != &in);
71 }
72
73 /**
74 * Useful when we need to run a function that performs operations in chunks
75 * and transforms data from in -> out, regardless of whether in or out is
76 * chained. We assume out size >= in size.
77 *
78 * func is expected to take data from some input buffer, do some operation on it
79 * and then place the result of the operation into an output buffer. It only
80 * should write to output in blocks of size BlockSize. Data from the input
81 * buffer must be internally buffered by func if it can not write a full block
82 * to output.
83 */
84 template <size_t BlockSize, typename Func>
85 folly::io::RWPrivateCursor
transformBufferBlocks(const folly::IOBuf & in,folly::IOBuf & out,Func func)86 transformBufferBlocks(const folly::IOBuf& in, folly::IOBuf& out, Func func) {
87 size_t internallyBuffered = 0;
88 folly::io::RWPrivateCursor output(&out);
89 folly::io::Cursor input(&in);
90
91 // block to handle writes along output buffer boundaries
92 std::array<uint8_t, BlockSize> blockBuffer = {};
93
94 // ensure buffers are fine
95 while (!input.isAtEnd()) {
96 auto inputRange = input.peekBytes();
97 auto inputLen = inputRange.size();
98 auto outputLen = output.peekBytes().size();
99 if (inputLen + internallyBuffered < BlockSize) {
100 // input doesn't have enough - we can call func and it should
101 // internally buffer.
102 // This should be safe to just internally buffer since we took into
103 // account what was existing in the internal buffer already
104 auto numWritten = func(blockBuffer.data(), inputRange.data(), inputLen);
105 DCHECK_EQ(numWritten, 0) << "expected buffering. wrote " << numWritten;
106 // only update input offsets
107 internallyBuffered += inputLen;
108 input.skip(inputLen);
109 } else if (outputLen < BlockSize) {
110 // we have at least a block to write from input + internal buffer, so
111 // output didn't have enough room in this case
112 // copy a block from input in temp and then push onto output
113 auto numWritten = func(
114 blockBuffer.data(),
115 inputRange.data(),
116 // only provide it the amount needed for one block
117 BlockSize - internallyBuffered);
118 DCHECK_EQ(numWritten, BlockSize)
119 << "did not write full block bs=" << BlockSize
120 << " wrote=" << numWritten;
121
122 // push the block onto output
123 output.push(blockBuffer.data(), BlockSize);
124
125 // advance input
126 input.skip(BlockSize - internallyBuffered);
127 internallyBuffered = 0;
128 } else {
129 // we have at least one block that can be written from input to output
130 // calculate shared bytes while taking into account internal buffer
131 auto numSharedBytes = std::min(outputLen, inputLen + internallyBuffered);
132
133 // this is the number we can safely (and expect) to write to output
134 auto numBlockBytes = numSharedBytes - (numSharedBytes % BlockSize);
135
136 // try to grab as much from input - we can grab up to BlockSize - 1 more
137 auto maxToTake = (numBlockBytes - internallyBuffered) + (BlockSize - 1);
138 auto numToTake = std::min(inputLen, maxToTake);
139 auto numWritten =
140 func(output.writableData(), inputRange.data(), numToTake);
141
142 DCHECK_EQ(numWritten, numBlockBytes);
143
144 // advance cursors
145 input.skip(numToTake);
146 output.skip(numWritten);
147 // recalculate internal buffer state
148 internallyBuffered = (internallyBuffered + numToTake) % BlockSize;
149 }
150 }
151 return output;
152 }
153 } // namespace fizz
154