Enable sandbox in local discovery utility process.

chrome/utility/local_discovery/service_discovery_message_handler.cc

Index: chrome/utility/local_discovery/service_discovery_message_handler.cc
diff --git a/chrome/utility/local_discovery/service_discovery_message_handler.cc b/chrome/utility/local_discovery/service_discovery_message_handler.cc
index e401226d545909b99eff8bbb9a6d1ab16be2e0d3..eb2ccab605915a14dab31d9b42d165ef9418f694 100644
--- a/chrome/utility/local_discovery/service_discovery_message_handler.cc
+++ b/chrome/utility/local_discovery/service_discovery_message_handler.cc
@@ -4,10 +4,118 @@
 
 #include "chrome/utility/local_discovery/service_discovery_message_handler.h"
 
+#include "base/command_line.h"
 #include "chrome/common/local_discovery/local_discovery_messages.h"
 #include "chrome/utility/local_discovery/service_discovery_client_impl.h"
+#include "content/public/common/content_switches.h"
 #include "content/public/utility/utility_thread.h"
 
+#if defined(OS_WIN)
+
+#include "base/lazy_instance.h"
+#include "net/base/winsock_init.h"
+#include "net/base/winsock_util.h"
+
+#endif  // OS_WIN
+
+namespace {
+
+bool NeedsSockets() {
+  return !CommandLine::ForCurrentProcess()->HasSwitch(switches::kNoSandbox) &&
+         CommandLine::ForCurrentProcess()->HasSwitch(
+             switches::kUtilityProcessEnableMDns);
+}
+
+#if defined(OS_WIN)
+
+class SocketFactory : public net::PlatformSocketFactory {
+ public:
+  SocketFactory()
+      : socket_v4_(NULL),
+        socket_v6_(NULL) {
+    net::EnsureWinsockInit();
+    socket_v4_ = WSASocket(AF_INET, SOCK_DGRAM, IPPROTO_UDP, NULL, 0,
+                           WSA_FLAG_OVERLAPPED);
+    socket_v6_ = WSASocket(AF_INET6, SOCK_DGRAM, IPPROTO_UDP, NULL, 0,
+                           WSA_FLAG_OVERLAPPED);
+  }
+
+  void Reset() {
+    if (socket_v4_ != INVALID_SOCKET) {
+      closesocket(socket_v4_);
+      socket_v4_ = INVALID_SOCKET;
+    }
+    if (socket_v6_ != INVALID_SOCKET) {
+      closesocket(socket_v6_);
+      socket_v6_ = INVALID_SOCKET;
+    }
+  }
+
+  virtual ~SocketFactory() {
+    Reset();
+  }
+
+  virtual SOCKET CreateSocket(int family, int type, int protocol) OVERRIDE {
+    SOCKET result = INVALID_SOCKET;
+    if (type != SOCK_DGRAM && protocol != IPPROTO_UDP) {
+      NOTREACHED();
+    } else if (family == AF_INET) {
+      std::swap(result, socket_v4_);
+    } else if (family == AF_INET6) {
+      std::swap(result, socket_v6_);
+    }
+    return result;
+  }
+
+  SOCKET socket_v4_;
+  SOCKET socket_v6_;
+
+  DISALLOW_COPY_AND_ASSIGN(SocketFactory);
+};
+
+base::LazyInstance<SocketFactory>
+    g_local_discovery_socket_factory = LAZY_INSTANCE_INITIALIZER;
+
+class ScopedSocketFactorySetter {
+ public:
+  ScopedSocketFactorySetter() {
+    if (NeedsSockets()) {
+      net::PlatformSocketFactory::SetInstance(
+          &g_local_discovery_socket_factory.Get());
+    }
+  }
+
+  ~ScopedSocketFactorySetter() {
+    if (NeedsSockets()) {
+      net::PlatformSocketFactory::SetInstance(NULL);
+      g_local_discovery_socket_factory.Get().Reset();
+    }
+  }
+
+  static void Initialize() {
+    if (NeedsSockets()) {
+      g_local_discovery_socket_factory.Get();
+    }
+  }
+
+ private:
+  DISALLOW_COPY_AND_ASSIGN(ScopedSocketFactorySetter);
+};
+
+#elif  // OS_WIN
+
+class ScopedSocketFactorySetter {
+  static void Initialize() {
+    // TODO(vitalybuka) : implement socket access from sandbox for other
+    // platforms.
+    DCHECK(!NeedsSockets());
+  }
+};
+
+#endif  // OS_WIN
+
+}  // namespace
+
 namespace local_discovery {
 
 ServiceDiscoveryMessageHandler::ServiceDiscoveryMessageHandler() {
@@ -16,13 +124,28 @@ ServiceDiscoveryMessageHandler::ServiceDiscoveryMessageHandler() {
 ServiceDiscoveryMessageHandler::~ServiceDiscoveryMessageHandler() {
 }
 
-void ServiceDiscoveryMessageHandler::Initialize() {
-  if (!service_discovery_client_) {
-    mdns_client_ = net::MDnsClient::CreateDefault();
-    mdns_client_->StartListening();
-    service_discovery_client_.reset(
-        new local_discovery::ServiceDiscoveryClientImpl(mdns_client_.get()));
+void ServiceDiscoveryMessageHandler::PreSandboxStartup() {
+  ScopedSocketFactorySetter::Initialize();
+}
+
+bool ServiceDiscoveryMessageHandler::Initialize() {
+  if (service_discovery_client_)
+    return true;
+
+  if (mdns_client_)  // We tried but failed before.
+    return false;
+
+  mdns_client_ = net::MDnsClient::CreateDefault();
+  {
+    // Temporarily redirect network code to use pre-created sockets.
+    ScopedSocketFactorySetter socket_factory_setter;
+    if (!mdns_client_->StartListening())
+      return false;
   }
+
+  service_discovery_client_.reset(
+      new local_discovery::ServiceDiscoveryClientImpl(mdns_client_.get()));
+  return true;
 }
 
 bool ServiceDiscoveryMessageHandler::OnMessageReceived(
@@ -42,7 +165,8 @@ bool ServiceDiscoveryMessageHandler::OnMessageReceived(
 void ServiceDiscoveryMessageHandler::OnStartWatcher(
     uint64 id,
     const std::string& service_type) {
-  Initialize();
+  if (!Initialize())
+    return;
   DCHECK(!ContainsKey(service_watchers_, id));
   scoped_ptr<ServiceWatcher> watcher(
       service_discovery_client_->CreateServiceWatcher(
@@ -54,12 +178,16 @@ void ServiceDiscoveryMessageHandler::OnStartWatcher(
 }
 
 void ServiceDiscoveryMessageHandler::OnDiscoverServices(uint64 id,
-                                                       bool force_update) {
+                                                        bool force_update) {
+  if (!service_discovery_client_)
+    return;
   DCHECK(ContainsKey(service_watchers_, id));
   service_watchers_[id]->DiscoverNewServices(force_update);
 }
 
 void ServiceDiscoveryMessageHandler::OnDestroyWatcher(uint64 id) {
+  if (!service_discovery_client_)
+    return;
   DCHECK(ContainsKey(service_watchers_, id));
   service_watchers_.erase(id);
 }
@@ -67,7 +195,8 @@ void ServiceDiscoveryMessageHandler::OnDestroyWatcher(uint64 id) {
 void ServiceDiscoveryMessageHandler::OnResolveService(
     uint64 id,
     const std::string& service_name) {
-  Initialize();
+  if (!Initialize())
+    return;
   DCHECK(!ContainsKey(service_resolvers_, id));
   scoped_ptr<ServiceResolver> resolver(
       service_discovery_client_->CreateServiceResolver(
@@ -79,6 +208,8 @@ void ServiceDiscoveryMessageHandler::OnResolveService(
 }
 
 void ServiceDiscoveryMessageHandler::OnDestroyResolver(uint64 id) {
+  if (!service_discovery_client_)
+    return;
   DCHECK(ContainsKey(service_resolvers_, id));
   service_resolvers_.erase(id);
 }
@@ -87,6 +218,7 @@ void ServiceDiscoveryMessageHandler::OnServiceUpdated(
     uint64 id,
     ServiceWatcher::UpdateType update,
     const std::string& name) {
+  DCHECK(service_discovery_client_);
   content::UtilityThread::Get()->Send(
       new LocalDiscoveryHostMsg_WatcherCallback(id, update, name));
 }
@@ -95,6 +227,7 @@ void ServiceDiscoveryMessageHandler::OnServiceResolved(
     uint64 id,
     ServiceResolver::RequestStatus status,
     const ServiceDescription& description) {
+  DCHECK(service_discovery_client_);
   content::UtilityThread::Get()->Send(
       new LocalDiscoveryHostMsg_ResolverCallback(id, status, description));
 }