Fix ChromeDriver flakiness with the C# webdriver bindings.

chrome/test/webdriver/webdriver_dispatch.cc

 // Copyright (c) 2012 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.
 
 #include "chrome/test/webdriver/webdriver_dispatch.h"
 
 #include <sstream>
 #include <string>
 #include <vector>
 
+#include "base/command_line.h"
 #include "base/format_macros.h"
 #include "base/json/json_reader.h"
 #include "base/logging.h"
 #include "base/memory/scoped_ptr.h"
 #include "base/message_loop_proxy.h"
 #include "base/string_split.h"
 #include "base/string_util.h"
 #include "base/stringprintf.h"
 #include "base/synchronization/waitable_event.h"
 #include "base/sys_info.h"
 #include "base/threading/platform_thread.h"
 #include "base/threading/thread.h"
 #include "chrome/common/chrome_version_info.h"
 #include "chrome/test/webdriver/commands/command.h"
 #include "chrome/test/webdriver/http_response.h"
 #include "chrome/test/webdriver/webdriver_logging.h"
 #include "chrome/test/webdriver/webdriver_session_manager.h"
+#include "chrome/test/webdriver/webdriver_switches.h"
 #include "chrome/test/webdriver/webdriver_util.h"
 
 namespace webdriver {
 
 namespace {
 
 // Maximum safe size of HTTP response message. Any larger than this,
 // the message may not be transferred at all.
 const size_t kMaxHttpMessageSize = 1024 * 1024 * 16;  // 16MB
 
-bool ForbidsMessageBody(const std::string& request_method,
-                        const HttpResponse& response) {
-  return request_method == "HEAD" ||
-         response.status() == HttpResponse::kNoContent ||
-         response.status() == HttpResponse::kNotModified ||
-         (response.status() >= 100 && response.status() < 200);
-}
-
 void ReadRequestBody(const struct mg_request_info* const request_info,
                      struct mg_connection* const connection,
                      std::string* request_body) {
   int content_length = 0;
   // 64 maximum header count hard-coded in mongoose.h
   for (int header_index = 0; header_index < 64; ++header_index) {
     if (request_info->http_headers[header_index].name == NULL) {
       break;
     }
     if (strcmp(request_info->http_headers[header_index].name,
                "Content-Length") == 0) {
       content_length = atoi(request_info->http_headers[header_index].value);
       break;
     }
   }
   if (content_length > 0) {
     request_body->resize(content_length);
     int bytes_read = 0;
     while (bytes_read < content_length) {
       bytes_read += mg_read(connection,
                             &(*request_body)[bytes_read],
                             content_length - bytes_read);
     }
   }
 }
 
+void WriteHttpResponse(struct mg_connection* connection,
+                       const HttpResponse& response) {
+  HttpResponse modified_response(response);
+  if (!CommandLine::ForCurrentProcess()->HasSwitch(kEnableKeepAlive))
+    modified_response.AddHeader("connection", "close");
+  std::string data;
+  modified_response.GetData(&data);
+  mg_write(connection, data.data(), data.length());
+}
+
 void DispatchCommand(Command* const command,
                      const std::string& method,
                      Response* response) {
   if (!command->Init(response))
     return;
 
   if (method == "POST") {
     command->ExecutePost(response);
   } else if (method == "GET") {
     command->ExecuteGet(response);
   } else if (method == "DELETE") {
     command->ExecuteDelete(response);
   } else {
     NOTREACHED();
   }
   command->Finish(response);
 }
 
 void SendOkWithBody(struct mg_connection* connection,
                     const std::string& content) {
-  const char* response_fmt = "HTTP/1.1 200 OK\r\n"
-                             "Content-Length:%d\r\n\r\n"
-                             "%s";
-  std::string response = base::StringPrintf(
-      response_fmt, content.length(), content.c_str());
-  mg_write(connection, response.data(), response.length());
+  HttpResponse response;
+  response.set_body(content);
+  WriteHttpResponse(connection, response);
 }
 
 void Shutdown(struct mg_connection* connection,
               const struct mg_request_info* request_info,
               void* user_data) {
   base::WaitableEvent* shutdown_event =
       reinterpret_cast<base::WaitableEvent*>(user_data);
-  mg_printf(connection, "HTTP/1.1 200 OK\r\n\r\n\r\n");
+  WriteHttpResponse(connection, HttpResponse());
   shutdown_event->Signal();
 }
 
 void SendStatus(struct mg_connection* connection,
                 const struct mg_request_info* request_info,
                 void* user_data) {
   chrome::VersionInfo version_info;
   DictionaryValue* build_info = new DictionaryValue;
   build_info->SetString("time",
                         base::StringPrintf("%s %s PST", __DATE__, __TIME__));
@@ skipping 38 matching lines @@
 
 void SimulateHang(struct mg_connection* connection,
                   const struct mg_request_info* request_info,
                   void* user_data) {
   base::PlatformThread::Sleep(base::TimeDelta::FromMinutes(5));
 }
 
 void SendNoContentResponse(struct mg_connection* connection,
                            const struct mg_request_info* request_info,
                            void* user_data) {
-  std::string response = "HTTP/1.1 204 No Content\r\n"
-                         "Content-Length:0\r\n"
-                         "\r\n";
-  mg_write(connection, response.data(), response.length());
+  WriteHttpResponse(connection, HttpResponse(HttpResponse::kNoContent));
 }
 
 void SendForbidden(struct mg_connection* connection,
                    const struct mg_request_info* request_info,
                    void* user_data) {
-  mg_printf(connection, "HTTP/1.1 403 Forbidden\r\n\r\n");
+  WriteHttpResponse(connection, HttpResponse(HttpResponse::kForbidden));
 }
 
 void SendNotImplementedError(struct mg_connection* connection,
                              const struct mg_request_info* request_info,
                              void* user_data) {
   // Send a well-formed WebDriver JSON error response to ensure clients
   // handle it correctly.
   std::string body = base::StringPrintf(
       "{\"status\":%d,\"value\":{\"message\":"
       "\"Command has not been implemented yet: %s %s\"}}",
       kUnknownCommand, request_info->request_method, request_info->uri);
 
-  std::string header = base::StringPrintf(
-      "HTTP/1.1 501 Not Implemented\r\n"
-      "Content-Type:application/json\r\n"
-      "Content-Length:%" PRIuS "\r\n"
-      "\r\n", body.length());
-
-  mg_write(connection, header.data(), header.length());
-  mg_write(connection, body.data(), body.length());
+  HttpResponse response(HttpResponse::kNotImplemented);
+  response.AddHeader("Content-Type", "application/json");
+  response.set_body(body);
+  WriteHttpResponse(connection, response);
 }
 
 }  // namespace
 
 namespace internal {
 
 void PrepareHttpResponse(const Response& command_response,
                          HttpResponse* const http_response) {
   ErrorCode status = command_response.GetStatus();
   switch (status) {
     case kSuccess:
       http_response->set_status(HttpResponse::kOk);
       break;
 
     // TODO(jleyba): kSeeOther, kBadRequest, kSessionNotFound,
     // and kMethodNotAllowed should be detected before creating
     // a command_response, and should thus not need conversion.
     case kSeeOther: {
       const Value* const value = command_response.GetValue();
       std::string location;
       if (!value->GetAsString(&location)) {
         // This should never happen.
         http_response->set_status(HttpResponse::kInternalServerError);
-        http_response->SetBody("Unable to set 'Location' header: response "
+        http_response->set_body("Unable to set 'Location' header: response "
                                "value is not a string: " +
                                command_response.ToJSON());
         return;
       }
       http_response->AddHeader("Location", location);
       http_response->set_status(HttpResponse::kSeeOther);
       break;
     }
 
     case kBadRequest:
     case kSessionNotFound:
       http_response->set_status(status);
       break;
 
     case kMethodNotAllowed: {
       const Value* const value = command_response.GetValue();
       if (!value->IsType(Value::TYPE_LIST)) {
         // This should never happen.
         http_response->set_status(HttpResponse::kInternalServerError);
-        http_response->SetBody(
+        http_response->set_body(
             "Unable to set 'Allow' header: response value was "
             "not a list of strings: " + command_response.ToJSON());
         return;
       }
 
       const ListValue* const list_value =
           static_cast<const ListValue* const>(value);
       std::vector<std::string> allowed_methods;
       for (size_t i = 0; i < list_value->GetSize(); ++i) {
         std::string method;
         if (list_value->GetString(i, &method)) {
           allowed_methods.push_back(method);
         } else {
           // This should never happen.
           http_response->set_status(HttpResponse::kInternalServerError);
-          http_response->SetBody(
+          http_response->set_body(
               "Unable to set 'Allow' header: response value was "
               "not a list of strings: " + command_response.ToJSON());
           return;
         }
       }
       http_response->AddHeader("Allow", JoinString(allowed_methods, ','));
       http_response->set_status(HttpResponse::kMethodNotAllowed);
       break;
     }
 
     // All other errors should be treated as generic 500s. The client
     // will be responsible for inspecting the message body for details.
     case kInternalServerError:
     default:
       http_response->set_status(HttpResponse::kInternalServerError);
       break;
   }
 
   http_response->SetMimeType("application/json; charset=utf-8");
-  http_response->SetBody(command_response.ToJSON());
+  http_response->set_body(command_response.ToJSON());
 }
 
 void SendResponse(struct mg_connection* const connection,
                   const std::string& request_method,
                   const Response& response) {
   HttpResponse http_response;
   PrepareHttpResponse(response, &http_response);
-
-  std::string message_header = base::StringPrintf("HTTP/1.1 %d %s\r\n",
-      http_response.status(), http_response.GetReasonPhrase().c_str());
-
-  typedef HttpResponse::HeaderMap::const_iterator HeaderIter;
-  for (HeaderIter header = http_response.headers()->begin();
-       header != http_response.headers()->end();
-       ++header) {
-    message_header.append(base::StringPrintf("%s:%s\r\n",
-        header->first.c_str(), header->second.c_str()));
-  }
-  message_header.append("\r\n");
-
-  mg_write(connection, message_header.data(), message_header.length());
-  if (!ForbidsMessageBody(request_method, http_response))
-    mg_write(connection, http_response.data(), http_response.length());
+  WriteHttpResponse(connection, http_response);
 }
 
 bool ParseRequestInfo(const struct mg_request_info* const request_info,
                       struct mg_connection* const connection,
                       std::string* method,
                       std::vector<std::string>* path_segments,
                       DictionaryValue** parameters,
                       Response* const response) {
   *method = request_info->request_method;
   if (*method == "HEAD")
@@ skipping 111 matching lines @@
        ++callback) {
     if (MatchPattern(request_info->uri, callback->uri_regex_)) {
       callback->func_(connection, request_info, callback->user_data_);
       return true;
     }
   }
   return false;
 }
 
 }  // namespace webdriver