/* Copyright 2016 The Roughtime Authors.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *   http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License. */

#ifndef SECURITY_ROUGHTIME_UDP_PROCESSOR_H_
#define SECURITY_ROUGHTIME_UDP_PROCESSOR_H_

#include <arpa/inet.h>
#include <netinet/in.h>

#include <memory>
#include <vector>

#include "server.h"

namespace roughtime {

// UdpProcessor manages a set of receive buffers for processing UDP requests.
class UdpProcessor {
 public:
  // Stats contains various counters for a single batch.
  struct Stats {
    size_t bytes_in;
    size_t bytes_out;
    unsigned packets_in;
    unsigned packets_out;
    unsigned requests_invalid;
    unsigned requests_truncated;
  };

  UdpProcessor();
  virtual ~UdpProcessor();

  // AddBrokenReplyGenerator adds a BrokenReplyGenerator that will be used for
  // a fraction of single-request batches. It can answer the request with
  // replies that are designed to test edge cases in clients. The sum of
  // probabilities of all generators passed to |AddBrokenReplyGenerator| must
  // be ≤ 1024.
  bool AddBrokenReplyGenerator(std::unique_ptr<BrokenReplyGenerator> generator);

  // ProcessBatch reads zero or more requests from |fd|, has |server| process
  // them and sends out the replies. It returns false if there was an
  // unexpected error during processing and true otherwise. (Reading zero
  // packets, finding invalid requests etc are not counted as unexpected
  // errors.)
  virtual bool ProcessBatch(int fd, Server* server, Stats* out_stats);

  // MakeSocket sets |*out_sock| to a UDP socket bound to the given port and
  // sets |*out_port| to the bound port number. If |port| is zero then a free
  // port number is used. It returns true on success or false on error.
  static bool MakeSocket(int port, int* out_sock, uint16_t* out_port);

 protected:
  // See recvmmsg(2) for help understanding these.  Note that indices in the
  // sending arrays don't correspond 1:1 with indices in the receiving arrays,
  // due to the possibility of invalid requests.
  mmsghdr recv_mmsghdrs_[kBatchSize];
  iovec recv_iov_[kBatchSize];
  uint8_t recv_buf_[kBatchSize][kMaxRecvPacketSize];

  mmsghdr send_mmsghdrs_[kBatchSize];
  iovec send_iov_[kBatchSize];
  uint8_t send_buf_[kBatchSize][kMaxResponseSize];

  // PrepareResponse sets up |send_mmsghdrs[index]|.
  virtual void PrepareResponse(struct msghdr* out_send_header,
                               const struct msghdr& recv_header);

  // Reset clears state in order to prepare for recvmmsg(2).
  virtual void Reset();

 private:
  // MaybeBreakResponse takes a valid request in |request| and the standard
  // response in |normal_response|. If any element of
  // |broken_reply_generators_| should be applied then it does so. It returns
  // true if the response in |out| should be sent and false if not.
  bool MaybeBreakResponse(uint8_t* out, size_t* out_len, size_t max_out_len,
                          const uint8_t* normal_response,
                          size_t normal_response_len, const uint8_t* request,
                          size_t request_len);

  sockaddr_storage sockaddrs_[kBatchSize];

  // requests_processed_ is the total number of valid requests that have been
  // processed.
  uint64_t requests_processed_ = 0;

  std::vector<std::unique_ptr<BrokenReplyGenerator>> broken_reply_generators_;

  // broken_reply_generator_sum_ is the sum of the probabilities of
  // |broken_reply_generators_|.
  uint16_t broken_reply_generator_sum_ = 0;
};

}  // namespace roughtime

#endif  // SECURITY_ROUGHTIME_UDP_PROCESSOR_H_