Fix MDnsClient's cache entry cleanup logic.

net/dns/mdns_client_impl.h

 // Copyright 2013 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 NET_DNS_MDNS_CLIENT_IMPL_H_
 #define NET_DNS_MDNS_CLIENT_IMPL_H_
 
 #include <map>
+#include <queue>
 #include <string>
 #include <utility>
 #include <vector>
 
 #include "base/cancelable_callback.h"
 #include "base/memory/scoped_vector.h"
 #include "base/observer_list.h"
 #include "net/base/io_buffer.h"
 #include "net/base/ip_endpoint.h"
 #include "net/dns/mdns_cache.h"
 #include "net/dns/mdns_client.h"
 #include "net/udp/datagram_server_socket.h"
 #include "net/udp/udp_server_socket.h"
 #include "net/udp/udp_socket.h"
 
+namespace base {
+class Clock;
+class Timer;
+}  // namespace base
+
 namespace net {
 
 class MDnsSocketFactoryImpl : public MDnsSocketFactory {
  public:
-  MDnsSocketFactoryImpl() {};
+  MDnsSocketFactoryImpl() {}
   ~MDnsSocketFactoryImpl() override{};
 
   void CreateSockets(ScopedVector<DatagramServerSocket>* sockets) override;
 
  private:
   DISALLOW_COPY_AND_ASSIGN(MDnsSocketFactoryImpl);
 };
 
 // A connection to the network for multicast DNS clients. It reads data into
 // DnsResponse objects and alerts the delegate that a packet has been received.
@@ skipping 63 matching lines @@
 class MDnsListenerImpl;
 
 class NET_EXPORT_PRIVATE MDnsClientImpl : public MDnsClient {
  public:
   // The core object exists while the MDnsClient is listening, and is deleted
   // whenever the number of listeners reaches zero. The deletion happens
   // asychronously, so destroying the last listener does not immediately
   // invalidate the core.
   class Core : public base::SupportsWeakPtr<Core>, MDnsConnection::Delegate {
    public:
-    Core();
+    Core(base::Clock* clock, base::Timer* timer);
     ~Core() override;
 
     // Initialize the core. Returns true on success.
     bool Init(MDnsSocketFactory* socket_factory);
 
     // Send a query with a specific rrtype and name. Returns true on success.
     bool SendQuery(uint16 rrtype, std::string name);
 
     // Add/remove a listener to the list of listeners.
     void AddListener(MDnsListenerImpl* listener);
     void RemoveListener(MDnsListenerImpl* listener);
 
     // Query the cache for records of a specific type and name.
     void QueryCache(uint16 rrtype, const std::string& name,
                     std::vector<const RecordParsed*>* records) const;
 
     // Parse the response and alert relevant listeners.
     void HandlePacket(DnsResponse* response, int bytes_read) override;
 
     void OnConnectionError(int error) override;
 
    private:
+    FRIEND_TEST_ALL_PREFIXES(MDnsTest, CacheCleanupWithShortTTL);
+
     typedef std::pair<std::string, uint16> ListenerKey;
     typedef std::map<ListenerKey, ObserverList<MDnsListenerImpl>* >
     ListenerMap;
 
     // Alert listeners of an update to the cache.
     void AlertListeners(MDnsCache::UpdateType update_type,
                         const ListenerKey& key, const RecordParsed* record);
 
     // Schedule a cache cleanup to a specific time, cancelling other cleanups.
     void ScheduleCleanup(base::Time cleanup);
 
     // Clean up the cache and schedule a new cleanup.
     void DoCleanup();
 
     // Callback for when a record is removed from the cache.
     void OnRecordRemoved(const RecordParsed* record);
 
     void NotifyNsecRecord(const RecordParsed* record);
 
     // Delete and erase the observer list for |key|. Only deletes the observer
     // list if is empty.
     void CleanupObserverList(const ListenerKey& key);
 
     ListenerMap listeners_;
 
     MDnsCache cache_;
 
-    base::CancelableClosure cleanup_callback_;
+    base::Clock* clock_;
+    base::Timer* cleanup_timer_;
     base::Time scheduled_cleanup_;
 
     scoped_ptr<MDnsConnection> connection_;
 
     DISALLOW_COPY_AND_ASSIGN(Core);
   };
 
   MDnsClientImpl();
   ~MDnsClientImpl() override;
 
   // MDnsClient implementation:
   scoped_ptr<MDnsListener> CreateListener(
       uint16 rrtype,
       const std::string& name,
       MDnsListener::Delegate* delegate) override;
 
   scoped_ptr<MDnsTransaction> CreateTransaction(
       uint16 rrtype,
       const std::string& name,
       int flags,
       const MDnsTransaction::ResultCallback& callback) override;
 
   bool StartListening(MDnsSocketFactory* socket_factory) override;
   void StopListening() override;
   bool IsListening() const override;
 
   Core* core() { return core_.get(); }
 
  private:
+  FRIEND_TEST_ALL_PREFIXES(MDnsTest, CacheCleanupWithShortTTL);
+
+  // Test constructor, takes a mock clock and mock timer.
+  MDnsClientImpl(scoped_ptr<base::Clock> clock,
+                 scoped_ptr<base::Timer> cleanup_timer);
+
   scoped_ptr<Core> core_;
+  scoped_ptr<base::Clock> clock_;
+  scoped_ptr<base::Timer> cleanup_timer_;
 
   DISALLOW_COPY_AND_ASSIGN(MDnsClientImpl);
 };
 
 class MDnsListenerImpl : public MDnsListener,
                          public base::SupportsWeakPtr<MDnsListenerImpl> {
  public:
   MDnsListenerImpl(uint16 rrtype,
                    const std::string& name,
+                   base::Clock* clock,
                    MDnsListener::Delegate* delegate,
                    MDnsClientImpl* client);
 
   ~MDnsListenerImpl() override;
 
   // MDnsListener implementation:
   bool Start() override;
 
   // Actively refresh any received records.
   void SetActiveRefresh(bool active_refresh) override;
@@ skipping 10 matching lines @@
 
   // Alert the delegate of the existence of an Nsec record.
   void AlertNsecRecord();
 
  private:
   void ScheduleNextRefresh();
   void DoRefresh();
 
   uint16 rrtype_;
   std::string name_;
+  base::Clock* clock_;
   MDnsClientImpl* client_;
   MDnsListener::Delegate* delegate_;
 
   base::Time last_update_;
   uint32 ttl_;
   bool started_;
   bool active_refresh_;
 
   base::CancelableClosure next_refresh_;
   DISALLOW_COPY_AND_ASSIGN(MDnsListenerImpl);
@@ skipping 57 matching lines @@
   MDnsClientImpl* client_;
 
   bool started_;
   int flags_;
 
   DISALLOW_COPY_AND_ASSIGN(MDnsTransactionImpl);
 };
 
 }  // namespace net
 #endif  // NET_DNS_MDNS_CLIENT_IMPL_H_