roughtime.googlesource.com/roughtime.git@v0.0.0-20201210012726-dd529367052d/udp_processor.h (about)

     1  /* Copyright 2016 The Roughtime Authors.
     2   *
     3   * Licensed under the Apache License, Version 2.0 (the "License");
     4   * you may not use this file except in compliance with the License.
     5   * You may obtain a copy of the License at
     6   *
     7   *   http://www.apache.org/licenses/LICENSE-2.0
     8   *
     9   * Unless required by applicable law or agreed to in writing, software
    10   * distributed under the License is distributed on an "AS IS" BASIS,
    11   * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    12   * See the License for the specific language governing permissions and
    13   * limitations under the License. */
    14  
    15  #ifndef SECURITY_ROUGHTIME_UDP_PROCESSOR_H_
    16  #define SECURITY_ROUGHTIME_UDP_PROCESSOR_H_
    17  
    18  #include <arpa/inet.h>
    19  #include <netinet/in.h>
    20  
    21  #include <memory>
    22  #include <vector>
    23  
    24  #include "server.h"
    25  
    26  namespace roughtime {
    27  
    28  #if defined(__MACH__) || defined(__Fuchsia__)
    29  struct mmsghdr {
    30    uint8_t *iov_base;
    31    size_t msg_len;
    32    msghdr msg_hdr;
    33  };
    34  #endif
    35  
    36  // UdpProcessor manages a set of receive buffers for processing UDP requests.
    37  class UdpProcessor {
    38   public:
    39    // Stats contains various counters for a single batch.
    40    struct Stats {
    41      size_t bytes_in;
    42      size_t bytes_out;
    43      unsigned packets_in;
    44      unsigned packets_out;
    45      unsigned requests_invalid;
    46      unsigned requests_truncated;
    47    };
    48  
    49    UdpProcessor();
    50    virtual ~UdpProcessor();
    51  
    52    // AddBrokenReplyGenerator adds a BrokenReplyGenerator that will be used for
    53    // a fraction of single-request batches. It can answer the request with
    54    // replies that are designed to test edge cases in clients. The sum of
    55    // probabilities of all generators passed to |AddBrokenReplyGenerator| must
    56    // be ≤ 1024.
    57    bool AddBrokenReplyGenerator(std::unique_ptr<BrokenReplyGenerator> generator);
    58  
    59    // ProcessBatch reads zero or more requests from |fd|, has |server| process
    60    // them and sends out the replies. It returns false if there was an
    61    // unexpected error during processing and true otherwise. (Reading zero
    62    // packets, finding invalid requests etc are not counted as unexpected
    63    // errors.)
    64    virtual bool ProcessBatch(int fd, Server* server, Stats* out_stats);
    65  
    66    // HandleOne processes a single |packet| of length |len| received from |from|.
    67    // It should return true if the message was accepted by |server|.
    68    virtual bool HandleOne(const struct msghdr* from, const uint8_t* packet,
    69                           size_t len, Server* server);
    70  
    71    // MakeSocket sets |*out_sock| to a UDP socket bound to the given port and
    72    // sets |*out_port| to the bound port number. If |port| is zero then a free
    73    // port number is used. It returns true on success or false on error.
    74    static bool MakeSocket(int port, int* out_sock, uint16_t* out_port);
    75  
    76   protected:
    77    // See recvmmsg(2) for help understanding these.  Note that indices in the
    78    // sending arrays don't correspond 1:1 with indices in the receiving arrays,
    79    // due to the possibility of invalid requests.
    80    mmsghdr recv_mmsghdrs_[kBatchSize];
    81    iovec recv_iov_[kBatchSize];
    82    uint8_t recv_buf_[kBatchSize][kMaxRecvPacketSize];
    83  
    84    mmsghdr send_mmsghdrs_[kBatchSize];
    85    iovec send_iov_[kBatchSize];
    86    uint8_t send_buf_[kBatchSize][kMaxResponseSize];
    87  
    88    // PrepareResponse sets up |send_mmsghdrs[index]|.
    89    virtual void PrepareResponse(struct msghdr* out_send_header,
    90                                 const struct msghdr& recv_header);
    91  
    92    // Reset clears state in order to prepare for recvmmsg(2).
    93    virtual void Reset();
    94  
    95   private:
    96    // MaybeBreakResponse takes a valid request in |request| and the standard
    97    // response in |normal_response|. If any element of
    98    // |broken_reply_generators_| should be applied then it does so. It returns
    99    // true if the response in |out| should be sent and false if not.
   100    bool MaybeBreakResponse(uint8_t* out, size_t* out_len, size_t max_out_len,
   101                            const uint8_t* normal_response,
   102                            size_t normal_response_len, const uint8_t* request,
   103                            size_t request_len);
   104  
   105    sockaddr_storage sockaddrs_[kBatchSize];
   106  
   107    // requests_processed_ is the total number of valid requests that have been
   108    // processed.
   109    uint64_t requests_processed_ = 0;
   110  
   111    std::vector<std::unique_ptr<BrokenReplyGenerator>> broken_reply_generators_;
   112  
   113    // broken_reply_generator_sum_ is the sum of the probabilities of
   114    // |broken_reply_generators_|.
   115    uint16_t broken_reply_generator_sum_ = 0;
   116  };
   117  
   118  }  // namespace roughtime
   119  
   120  #endif  // SECURITY_ROUGHTIME_UDP_PROCESSOR_H_