// Copyright 2018 The Chromium Authors. All rights reserved.
// Use of this source code is governed by a BSD-style license that can be
// found in the LICENSE file.

#ifndef SERVICES_NETWORK_MDNS_RESPONDER_H_
#define SERVICES_NETWORK_MDNS_RESPONDER_H_

#include <map>
#include <memory>
#include <set>
#include <string>
#include <vector>

#include "base/containers/flat_set.h"
#include "base/containers/unique_ptr_adapters.h"
#include "mojo/public/cpp/bindings/binding_set.h"
#include "net/dns/dns_query.h"
#include "net/dns/dns_response.h"
#include "services/network/public/mojom/mdns_responder.mojom.h"

namespace base {
class TickClock;
}  // namespace base

namespace net {
class IOBufferWithSize;
class IPAddress;
class MDnsSocketFactory;
}  // namespace net

namespace network {

class MdnsResponder;

namespace mdns_helper {

// Creates an mDNS response, of which the Answer section contains the address
// records for |name_addr_map|, and the Additional section contains the
// corresponding NSEC records that assert the existence of only address
// records in the Answer section.
COMPONENT_EXPORT(NETWORK_SERVICE)
scoped_refptr<net::IOBufferWithSize> CreateResolutionResponse(
    const base::TimeDelta& ttl,
    const std::map<std::string, net::IPAddress>& name_addr_map);
// Creates an mDNS response, of which the Answer section contains NSEC records
// that assert the existence of only address records for |name_addr_map|, and
// the Additional section contains the corresponding address records.
COMPONENT_EXPORT(NETWORK_SERVICE)
scoped_refptr<net::IOBufferWithSize> CreateNegativeResponse(
    const std::map<std::string, net::IPAddress>& name_addr_map);

}  // namespace mdns_helper

// Options to configure the transmission of mDNS responses.
struct COMPONENT_EXPORT(NETWORK_SERVICE) MdnsResponseSendOption
    : public base::RefCounted<MdnsResponseSendOption> {
 public:
  enum class ResponseClass {
    UNSPECIFIED,
    ANNOUNCEMENT,
    PROBE_RESOLUTION,
    REGULAR_RESOLUTION,
    NEGATIVE,
    GOODBYE,
  };

  MdnsResponseSendOption();
  // As a shorthand, an empty set denotes all interfaces.
  std::set<uint16_t> send_socket_handler_ids;
  // Used for rate limiting.
  base::flat_set<std::string> names_for_rate_limit;
  // Used for retry after send failure.
  ResponseClass klass = ResponseClass::UNSPECIFIED;
  // The number of retries done for the same response due to send failure.
  uint8_t num_send_retries_done = 0;

 private:
  friend class RefCounted<MdnsResponseSendOption>;

  ~MdnsResponseSendOption();
};

// The responder manager creates and manages responder instances spawned for
// each Mojo binding. It also manages the underlying network IO by delegating
// the IO task to socket handlers of each interface. When there is a network
// stack error or a Mojo binding error, the manager also offers the
// corresponding error handling.
class COMPONENT_EXPORT(NETWORK_SERVICE) MdnsResponderManager {
 public:
  // Wraps a name generation method that can be configured in tests via
  // SetNameGeneratorForTesting below.
  class NameGenerator {
   public:
    virtual ~NameGenerator() = default;
    virtual std::string CreateName() = 0;
  };

  MdnsResponderManager();
  explicit MdnsResponderManager(net::MDnsSocketFactory* socket_factory);
  ~MdnsResponderManager();

  // Creates an instance of MdnsResponder for the binding request from an
  // InterfacePtr.
  void CreateMdnsResponder(mojom::MdnsResponderRequest request);
#ifdef DEBUG
  // The methods below are only used for extra uniqueness validation of names
  // owned by responders. By default, we use the RandomUuidNameGenerator (see
  // mdns_responder.cc), which probabilistically guarantees the uniqueness of
  // generated names.
  //
  // Adds a name to the set of all existing names generated by all responders
  // (i.e., names owned by an instance of responder). Return true if the name is
  // not in the set before the addition; false otherwise.
  bool AddName(const std::string& name) {
    auto result = names_.insert(name);
    return result.second;
  }
  // Removes a name from the set of all existing names generated by all
  // responders. Return true if the name exists in the set before the removal;
  // false otherwise.
  bool RemoveName(const std::string& name) { return names_.erase(name) == 1; }
#endif
  // Sends an mDNS response in the wire format given by |buf|. See
  // MdnsResponseSendOption for configurable options in |option|.
  //
  // Sending responses is rate-limited, and this method returns true if the
  // response is successfully scheduled to send on all successfully bound
  // interfaces specified in |option.send_socket_handler_ids|, and false
  // otherwise.
  bool Send(scoped_refptr<net::IOBufferWithSize> buf,
            scoped_refptr<MdnsResponseSendOption> option);
  // The error handler that is invoked when the Mojo binding of |responder| is
  // closed (e.g. the InterfacePtr on the client side is destroyed) or
  // encounters an error. It removes this responder instance, which further
  // clears the existing name-address associations owned by this responder in
  // the local network.
  void OnMojoConnectionError(MdnsResponder* responder);

  // Called when an external mDNS response is received and it contains
  // records of names generated by an owned MdnsResponder instance.
  void HandleNameConflictIfAny(
      const std::map<std::string, std::set<net::IPAddress>>& external_maps);

  NameGenerator* name_generator() const { return name_generator_.get(); }
  // Sets the name generator that is shared by all MdnsResponder instances.
  // Changing the name generator affects all existing responder instances and
  // also the ones spawned in the future.
  //
  // Used for tests only.
  void SetNameGeneratorForTesting(
      std::unique_ptr<NameGenerator> name_generator);

  // Sets the tick clock that is used for rate limiting of mDNS responses, and
  // also resets the internal schedule for rate limiting.
  //
  // Used for tests only.
  void SetTickClockForTesting(const base::TickClock* tick_clock);

 private:
  enum class SocketHandlerStartResult {
    UNSPECIFIED,
    // Handlers started for all interfaces.
    ALL_SUCCESS,
    // Handlers started for a subset of interfaces.
    PARTIAL_SUCCESS,
    // No handler started.
    ALL_FAILURE,
  };
  // Handles the underlying sending and receiving of mDNS messages on each
  // interface available. The implementation mostly resembles
  // net::MdnsConnection::SocketHandler.
  class SocketHandler;

  // Initializes socket handlers and sets |start_result_|;
  void Start();
  // Dispatches a parsed query from a socket handler to each responder instance.
  void OnMdnsQueryReceived(const net::DnsQuery& query,
                           uint16_t recv_socket_handler_id);
  void OnSocketHandlerReadError(uint16_t socket_handler_id, int result);

  std::unique_ptr<net::MDnsSocketFactory> owned_socket_factory_;
  net::MDnsSocketFactory* socket_factory_;
  // Only the socket handlers that have successfully bound and started are kept.
  std::map<uint16_t, std::unique_ptr<SocketHandler>> socket_handler_by_id_;
  SocketHandlerStartResult start_result_ =
      SocketHandlerStartResult::UNSPECIFIED;
#ifdef DEBUG
  // Used in debug only for the extra uniqueness validation of names generated
  // by responders.
  std::set<std::string> names_;
#endif
  std::unique_ptr<NameGenerator> name_generator_;
  std::set<std::unique_ptr<MdnsResponder>, base::UniquePtrComparator>
      responders_;

  DISALLOW_COPY_AND_ASSIGN(MdnsResponderManager);
};

// Implementation of the mDNS service that can provide utilities of an mDNS
// responder.
class COMPONENT_EXPORT(NETWORK_SERVICE) MdnsResponder
    : public mojom::MdnsResponder {
 public:
  MdnsResponder(mojom::MdnsResponderRequest request,
                MdnsResponderManager* manager);
  // When destroyed, clears all existing name-address associations owned by this
  // responder in the local network by sending out goodbye packets. See
  // SendGoodbyePacketForNameAddressMap below.
  ~MdnsResponder() override;

  // mojom::MdnsResponder overrides.
  void CreateNameForAddress(
      const net::IPAddress& address,
      mojom::MdnsResponder::CreateNameForAddressCallback callback) override;
  void RemoveNameForAddress(
      const net::IPAddress& address,
      mojom::MdnsResponder::RemoveNameForAddressCallback callback) override;

  // Processes the given query and generates a response if the query contains an
  // mDNS name that this responder has a mapped IP address.
  void OnMdnsQueryReceived(const net::DnsQuery& query,
                           uint16_t recv_socket_handler_id);

  bool HasConflictWithExternalResolution(
      const std::string& name,
      const std::set<net::IPAddress>& external_mapped_addreses);

  void SetNameGeneratorForTesting(
      MdnsResponderManager::NameGenerator* name_generator) {
    name_generator_ = name_generator;
  }

 private:
  // Returns true if the response is successfully scheduled to send on all
  // successfully bound interfaces after rate limiting, and false otherwise. See
  // also MdnsResponderManager::Send.
  bool SendMdnsResponse(scoped_refptr<net::IOBufferWithSize> response,
                        scoped_refptr<MdnsResponseSendOption> option);
  // RFC 6761, Section 10.1 "Goodbye Packets".
  //
  // The responder should send out an unsolicited mDNS response with an resource
  // record of zero TTL to clear the name-to-address mapping in neighboring
  // hosts, when the mapping is no longer valid.
  //
  // Returns true if the goodbye message is successfully scheduled to send on
  // all interfaces after rate limiting, and false otherwise. See also
  // MdnsResponderManager::Send.
  bool SendGoodbyePacketForNameAddressMap(
      const std::map<std::string, net::IPAddress>& name_addr_map);

  std::map<std::string, net::IPAddress>::iterator FindNameCreatedForAddress(
      const net::IPAddress& address);

  mojo::Binding<network::mojom::MdnsResponder> binding_;
  // A back pointer to the responder manager that owns this responder. The
  // responder should be destroyed before |manager_| becomes invalid or a weak
  // reference should be used to access the manager when there is no such
  // guarantee in an operation.
  MdnsResponderManager* const manager_;
  std::map<std::string, net::IPAddress> name_addr_map_;
  std::map<std::string, uint16_t> name_refcount_map_;
  MdnsResponderManager::NameGenerator* name_generator_;

  DISALLOW_COPY_AND_ASSIGN(MdnsResponder);
};

}  // namespace network

#endif  // SERVICES_NETWORK_MDNS_RESPONDER_H_