serialize.c++ 9.19 KB
Newer Older
Kenton Varda's avatar
Kenton Varda committed
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24
// Copyright (c) 2013, Kenton Varda <temporal@gmail.com>
// All rights reserved.
//
// Redistribution and use in source and binary forms, with or without
// modification, are permitted provided that the following conditions are met:
//
// 1. Redistributions of source code must retain the above copyright notice, this
//    list of conditions and the following disclaimer.
// 2. Redistributions in binary form must reproduce the above copyright notice,
//    this list of conditions and the following disclaimer in the documentation
//    and/or other materials provided with the distribution.
//
// THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND
// ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED
// WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
// DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR CONTRIBUTORS BE LIABLE FOR
// ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES
// (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;
// LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND
// ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
// (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS
// SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

#include "serialize.h"
25
#include "layout.h"
Kenton Varda's avatar
Kenton Varda committed
26
#include <kj/debug.h>
27
#include <exception>
Kenton Varda's avatar
Kenton Varda committed
28

29
namespace capnp {
Kenton Varda's avatar
Kenton Varda committed
30

31 32
FlatArrayMessageReader::FlatArrayMessageReader(
    kj::ArrayPtr<const word> array, ReaderOptions options)
Kenton Varda's avatar
Kenton Varda committed
33 34 35 36 37 38
    : MessageReader(options) {
  if (array.size() < 1) {
    // Assume empty message.
    return;
  }

39 40
  const _::WireValue<uint32_t>* table =
      reinterpret_cast<const _::WireValue<uint32_t>*>(array.begin());
Kenton Varda's avatar
Kenton Varda committed
41

42
  uint segmentCount = table[0].get() + 1;
Kenton Varda's avatar
Kenton Varda committed
43 44
  size_t offset = segmentCount / 2u + 1u;

45
  KJ_REQUIRE(array.size() >= offset, "Message ends prematurely in segment table.") {
Kenton Varda's avatar
Kenton Varda committed
46 47 48 49 50 51 52 53 54
    return;
  }

  if (segmentCount == 0) {
    return;
  }

  uint segmentSize = table[1].get();

55 56
  KJ_REQUIRE(array.size() >= offset + segmentSize,
             "Message ends prematurely in first segment.") {
Kenton Varda's avatar
Kenton Varda committed
57 58 59 60 61 62 63
    return;
  }

  segment0 = array.slice(offset, offset + segmentSize);
  offset += segmentSize;

  if (segmentCount > 1) {
Kenton Varda's avatar
Kenton Varda committed
64
    moreSegments = kj::heapArray<kj::ArrayPtr<const word>>(segmentCount - 1);
Kenton Varda's avatar
Kenton Varda committed
65 66 67 68

    for (uint i = 1; i < segmentCount; i++) {
      uint segmentSize = table[i + 1].get();

69
      KJ_REQUIRE(array.size() >= offset + segmentSize, "Message ends prematurely.") {
Kenton Varda's avatar
Kenton Varda committed
70 71 72 73 74 75 76 77 78 79
        moreSegments = nullptr;
        return;
      }

      moreSegments[i - 1] = array.slice(offset, offset + segmentSize);
      offset += segmentSize;
    }
  }
}

80
kj::ArrayPtr<const word> FlatArrayMessageReader::getSegment(uint id) {
Kenton Varda's avatar
Kenton Varda committed
81 82 83 84 85 86 87 88 89
  if (id == 0) {
    return segment0;
  } else if (id <= moreSegments.size()) {
    return moreSegments[id - 1];
  } else {
    return nullptr;
  }
}

90
kj::Array<word> messageToFlatArray(kj::ArrayPtr<const kj::ArrayPtr<const word>> segments) {
91
  KJ_REQUIRE(segments.size() > 0, "Tried to serialize uninitialized message.");
92

Kenton Varda's avatar
Kenton Varda committed
93 94 95 96 97 98
  size_t totalSize = segments.size() / 2 + 1;

  for (auto& segment: segments) {
    totalSize += segment.size();
  }

Kenton Varda's avatar
Kenton Varda committed
99
  kj::Array<word> result = kj::heapArray<word>(totalSize);
Kenton Varda's avatar
Kenton Varda committed
100

101 102
  _::WireValue<uint32_t>* table =
      reinterpret_cast<_::WireValue<uint32_t>*>(result.begin());
Kenton Varda's avatar
Kenton Varda committed
103

104 105 106
  // We write the segment count - 1 because this makes the first word zero for single-segment
  // messages, improving compression.  We don't bother doing this with segment sizes because
  // one-word segments are rare anyway.
107
  table[0].set(segments.size() - 1);
Kenton Varda's avatar
Kenton Varda committed
108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124

  for (uint i = 0; i < segments.size(); i++) {
    table[i + 1].set(segments[i].size());
  }

  if (segments.size() % 2 == 0) {
    // Set padding byte.
    table[segments.size() + 1].set(0);
  }

  word* dst = result.begin() + segments.size() / 2 + 1;

  for (auto& segment: segments) {
    memcpy(dst, segment.begin(), segment.size() * sizeof(word));
    dst += segment.size();
  }

125
  KJ_DASSERT(dst == result.end(), "Buffer overrun/underrun bug in code above.");
Kenton Varda's avatar
Kenton Varda committed
126

Kenton Varda's avatar
Kenton Varda committed
127
  return kj::mv(result);
Kenton Varda's avatar
Kenton Varda committed
128 129 130 131
}

// =======================================================================================

132
InputStreamMessageReader::InputStreamMessageReader(
133
    kj::InputStream& inputStream, ReaderOptions options, kj::ArrayPtr<word> scratchSpace)
134
    : MessageReader(options), inputStream(inputStream), readPos(nullptr) {
135
  _::WireValue<uint32_t> firstWord[2];
Kenton Varda's avatar
Kenton Varda committed
136

Kenton Varda's avatar
Kenton Varda committed
137
  inputStream.read(firstWord, sizeof(firstWord));
Kenton Varda's avatar
Kenton Varda committed
138

139
  uint segmentCount = firstWord[0].get() + 1;
140
  uint segment0Size = segmentCount == 0 ? 0 : firstWord[1].get();
Kenton Varda's avatar
Kenton Varda committed
141

142
  size_t totalWords = segment0Size;
Kenton Varda's avatar
Kenton Varda committed
143

144
  // Reject messages with too many segments for security reasons.
145
  KJ_REQUIRE(segmentCount < 512, "Message has too many segments.") {
146 147
    segmentCount = 1;
    segment0Size = 1;
148
    break;
149
  }
150

151
  // Read sizes for all segments except the first.  Include padding if necessary.
152
  _::WireValue<uint32_t> moreSizes[segmentCount & ~1];
153
  if (segmentCount > 1) {
Kenton Varda's avatar
Kenton Varda committed
154
    inputStream.read(moreSizes, sizeof(moreSizes));
155 156
    for (uint i = 0; i < segmentCount - 1; i++) {
      totalWords += moreSizes[i].get();
Kenton Varda's avatar
Kenton Varda committed
157 158 159
    }
  }

160 161 162
  // Don't accept a message which the receiver couldn't possibly traverse without hitting the
  // traversal limit.  Without this check, a malicious client could transmit a very large segment
  // size to make the receiver allocate excessive space and possibly crash.
163 164
  KJ_REQUIRE(totalWords <= options.traversalLimitInWords,
             "Message is too large.  To increase the limit on the receiving end, see "
165
             "capnp::ReaderOptions.") {
166
    segmentCount = 1;
167
    segment0Size = kj::min(segment0Size, options.traversalLimitInWords);
168
    totalWords = segment0Size;
169
    break;
170
  }
171

172
  if (scratchSpace.size() < totalWords) {
173 174
    // TODO(perf):  Consider allocating each segment as a separate chunk to reduce memory
    //   fragmentation.
Kenton Varda's avatar
Kenton Varda committed
175
    ownedSpace = kj::heapArray<word>(totalWords);
176
    scratchSpace = ownedSpace;
Kenton Varda's avatar
Kenton Varda committed
177 178
  }

179
  segment0 = scratchSpace.slice(0, segment0Size);
Kenton Varda's avatar
Kenton Varda committed
180

181
  if (segmentCount > 1) {
Kenton Varda's avatar
Kenton Varda committed
182
    moreSegments = kj::heapArray<kj::ArrayPtr<const word>>(segmentCount - 1);
183
    size_t offset = segment0Size;
Kenton Varda's avatar
Kenton Varda committed
184

185 186 187 188
    for (uint i = 0; i < segmentCount - 1; i++) {
      uint segmentSize = moreSizes[i].get();
      moreSegments[i] = scratchSpace.slice(offset, offset + segmentSize);
      offset += segmentSize;
Kenton Varda's avatar
Kenton Varda committed
189 190 191
    }
  }

192
  if (segmentCount == 1) {
Kenton Varda's avatar
Kenton Varda committed
193
    inputStream.read(scratchSpace.begin(), totalWords * sizeof(word));
194 195 196
  } else if (segmentCount > 1) {
    readPos = reinterpret_cast<byte*>(scratchSpace.begin());
    readPos += inputStream.read(readPos, segment0Size * sizeof(word), totalWords * sizeof(word));
Kenton Varda's avatar
Kenton Varda committed
197
  }
198
}
Kenton Varda's avatar
Kenton Varda committed
199

200
InputStreamMessageReader::~InputStreamMessageReader() noexcept(false) {
201
  if (readPos != nullptr) {
202 203 204 205
    unwindDetector.catchExceptionsIfUnwinding([&]() {
      // Note that lazy reads only happen when we have multiple segments, so moreSegments.back() is
      // valid.
      const byte* allEnd = reinterpret_cast<const byte*>(moreSegments.back().end());
Kenton Varda's avatar
Kenton Varda committed
206
      inputStream.skip(allEnd - readPos);
207
    });
Kenton Varda's avatar
Kenton Varda committed
208 209 210
  }
}

211
kj::ArrayPtr<const word> InputStreamMessageReader::getSegment(uint id) {
Kenton Varda's avatar
Kenton Varda committed
212 213 214 215
  if (id > moreSegments.size()) {
    return nullptr;
  }

216
  kj::ArrayPtr<const word> segment = id == 0 ? segment0 : moreSegments[id - 1];
Kenton Varda's avatar
Kenton Varda committed
217

218 219 220 221 222 223 224 225
  if (readPos != nullptr) {
    // May need to lazily read more data.
    const byte* segmentEnd = reinterpret_cast<const byte*>(segment.end());
    if (readPos < segmentEnd) {
      // Note that lazy reads only happen when we have multiple segments, so moreSegments.back() is
      // valid.
      const byte* allEnd = reinterpret_cast<const byte*>(moreSegments.back().end());
      readPos += inputStream.read(readPos, segmentEnd - readPos, allEnd - readPos);
Kenton Varda's avatar
Kenton Varda committed
226 227 228
    }
  }

229
  return segment;
Kenton Varda's avatar
Kenton Varda committed
230 231 232 233
}

// -------------------------------------------------------------------

234
void writeMessage(kj::OutputStream& output, kj::ArrayPtr<const kj::ArrayPtr<const word>> segments) {
235
  KJ_REQUIRE(segments.size() > 0, "Tried to serialize uninitialized message.");
236

237
  _::WireValue<uint32_t> table[(segments.size() + 2) & ~size_t(1)];
Kenton Varda's avatar
Kenton Varda committed
238

239 240 241
  // We write the segment count - 1 because this makes the first word zero for single-segment
  // messages, improving compression.  We don't bother doing this with segment sizes because
  // one-word segments are rare anyway.
242
  table[0].set(segments.size() - 1);
Kenton Varda's avatar
Kenton Varda committed
243 244 245 246 247 248 249 250
  for (uint i = 0; i < segments.size(); i++) {
    table[i + 1].set(segments[i].size());
  }
  if (segments.size() % 2 == 0) {
    // Set padding byte.
    table[segments.size() + 1].set(0);
  }

251 252
  KJ_STACK_ARRAY(kj::ArrayPtr<const byte>, pieces, segments.size() + 1, 4, 32);
  pieces[0] = kj::arrayPtr(reinterpret_cast<byte*>(table), sizeof(table));
Kenton Varda's avatar
Kenton Varda committed
253

254
  for (uint i = 0; i < segments.size(); i++) {
255 256
    pieces[i + 1] = kj::arrayPtr(reinterpret_cast<const byte*>(segments[i].begin()),
                                 reinterpret_cast<const byte*>(segments[i].end()));
Kenton Varda's avatar
Kenton Varda committed
257
  }
258

Kenton Varda's avatar
Kenton Varda committed
259
  output.write(pieces);
Kenton Varda's avatar
Kenton Varda committed
260 261 262
}

// =======================================================================================
263
StreamFdMessageReader::~StreamFdMessageReader() noexcept(false) {}
Kenton Varda's avatar
Kenton Varda committed
264

265
void writeMessageToFd(int fd, kj::ArrayPtr<const kj::ArrayPtr<const word>> segments) {
266
  kj::FdOutputStream stream(fd);
Kenton Varda's avatar
Kenton Varda committed
267 268 269
  writeMessage(stream, segments);
}

270
}  // namespace capnp