folly/wangle/codec/LengthFieldPrepender.cpp

   1 /*
   2  * Copyright 2015 Facebook, Inc.
   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 #include <folly/wangle/codec/LengthFieldPrepender.h>
  17
  18 namespace folly { namespace wangle {
  19
  20 LengthFieldPrepender::LengthFieldPrepender(
  21     int lengthFieldLength,
  22     int lengthAdjustment,
  23     bool lengthIncludesLengthField,
  24     bool networkByteOrder)
  25     : lengthFieldLength_(lengthFieldLength)
  26     , lengthAdjustment_(lengthAdjustment)
  27     , lengthIncludesLengthField_(lengthIncludesLengthField)
  28     , networkByteOrder_(networkByteOrder) {
  29     CHECK(lengthFieldLength == 1 ||
  30           lengthFieldLength == 2 ||
  31           lengthFieldLength == 4 ||
  32           lengthFieldLength == 8 );
  33   }
  34
  35 Future<Unit> LengthFieldPrepender::write(
  36     Context* ctx, std::unique_ptr<IOBuf> buf) {
  37   int length = lengthAdjustment_ + buf->computeChainDataLength();
  38   if (lengthIncludesLengthField_) {
  39     length += lengthFieldLength_;
  40   }
  41
  42   if (length < 0) {
  43     throw std::runtime_error("Length field < 0");
  44   }
  45
  46   auto len = IOBuf::create(lengthFieldLength_);
  47   len->append(lengthFieldLength_);
  48   folly::io::RWPrivateCursor c(len.get());
  49
  50   switch (lengthFieldLength_) {
  51     case 1: {
  52       if (length >= 256) {
  53         throw std::runtime_error("length does not fit byte");
  54       }
  55       if (networkByteOrder_) {
  56         c.writeBE((uint8_t)length);
  57       } else {
  58         c.writeLE((uint8_t)length);
  59       }
  60       break;
  61     }
  62     case 2: {
  63       if (length >= 65536) {
  64         throw std::runtime_error("length does not fit byte");
  65       }
  66       if (networkByteOrder_) {
  67         c.writeBE((uint16_t)length);
  68       } else {
  69         c.writeLE((uint16_t)length);
  70       }
  71       break;
  72     }
  73     case 4: {
  74       if (networkByteOrder_) {
  75         c.writeBE((uint32_t)length);
  76       } else {
  77         c.writeLE((uint32_t)length);
  78       }
  79       break;
  80     }
  81     case 8: {
  82       if (networkByteOrder_) {
  83         c.writeBE((uint64_t)length);
  84       } else {
  85         c.writeLE((uint64_t)length);
  86       }
  87       break;
  88     }
  89     default: {
  90       throw std::runtime_error("Invalid lengthFieldLength");
  91     }
  92   }
  93
  94   len->prependChain(std::move(buf));
  95   return ctx->fireWrite(std::move(len));
  96 }
  97
  98
  99 }} // Namespace