From f6d0de19ee65300873e8a3e2091ab3d9df81d954 Mon Sep 17 00:00:00 2001
From: "maruel@chromium.org"
 <maruel@chromium.org@0039d316-1c4b-4281-b951-d872f2087c98>
Date: Fri, 19 Nov 2010 20:42:41 +0000
Subject: [PATCH] Apply directly http://codereview.appspot.com/3227041 to add
 support for HTTP 301

We need this now for an imminent rietveld instance move.

TEST=worked with manual testing
BUG=none

Review URL: http://codereview.chromium.org/5138007

git-svn-id: svn://svn.chromium.org/chrome/trunk/tools/depot_tools@66810 0039d316-1c4b-4281-b951-d872f2087c98
---
 third_party/upload.py | 51 ++++++++++++++++++++++++++++---------------
 1 file changed, 34 insertions(+), 17 deletions(-)

diff --git a/third_party/upload.py b/third_party/upload.py
index dac899708..8dc48f24a 100644
--- a/third_party/upload.py
+++ b/third_party/upload.py
@@ -261,10 +261,12 @@ class AbstractRpcServer(object):
       else:
         raise
 
-  def _GetAuthCookie(self, auth_token):
+  def _GetAuthCookie(self, host, auth_token):
     """Fetches authentication cookies for an authentication token.
 
     Args:
+      host: The host to get a cookie against. Because of 301, it may be a
+            different host than self.host.
       auth_token: The authentication token returned by ClientLogin.
 
     Raises:
@@ -273,21 +275,32 @@ class AbstractRpcServer(object):
     # This is a dummy value to allow us to identify when we're successful.
     continue_location = "http://localhost/"
     args = {"continue": continue_location, "auth": auth_token}
-    req = self._CreateRequest("%s/_ah/login?%s" %
-                              (self.host, urllib.urlencode(args)))
-    try:
-      response = self.opener.open(req)
-    except urllib2.HTTPError, e:
-      response = e
+    tries = 0
+    url = "%s/_ah/login?%s" % (host, urllib.urlencode(args))
+    while tries < 3:
+      req = self._CreateRequest(url)
+      try:
+        response = self.opener.open(req)
+      except urllib2.HTTPError, e:
+        response = e
+        if e.code == 301:
+          # Handle permanent redirect manually.
+          url = e.info()["location"]
+          continue
+      break
     if (response.code != 302 or
         response.info()["location"] != continue_location):
       raise urllib2.HTTPError(req.get_full_url(), response.code, response.msg,
                               response.headers, response.fp)
     self.authenticated = True
 
-  def _Authenticate(self):
+  def _Authenticate(self, host):
     """Authenticates the user.
 
+    Args:
+      host: The host to get a cookie against. Because of 301, it may be a
+            different host than self.host.
+
     The authentication process works as follows:
      1) We get a username and password from the user
      2) We use ClientLogin to obtain an AUTH token for the user
@@ -336,7 +349,7 @@ class AbstractRpcServer(object):
           print >>sys.stderr, "The service is not available; try again later."
           break
         raise
-      self._GetAuthCookie(auth_token)
+      self._GetAuthCookie(host, auth_token)
       return
 
   def Send(self, request_path, payload=None,
@@ -363,18 +376,18 @@ class AbstractRpcServer(object):
     # TODO: Don't require authentication.  Let the server say
     # whether it is necessary.
     if not self.authenticated:
-      self._Authenticate()
+      self._Authenticate(self.host)
 
     old_timeout = socket.getdefaulttimeout()
     socket.setdefaulttimeout(timeout)
     try:
       tries = 0
+      args = dict(kwargs)
+      url = "%s%s" % (self.host, request_path)
+      if args:
+        url += "?" + urllib.urlencode(args)
       while True:
         tries += 1
-        args = dict(kwargs)
-        url = "%s%s" % (self.host, request_path)
-        if args:
-          url += "?" + urllib.urlencode(args)
         req = self._CreateRequest(url=url, data=payload)
         req.add_header("Content-Type", content_type)
         if extra_headers:
@@ -389,10 +402,14 @@ class AbstractRpcServer(object):
           if tries > 3:
             raise
           elif e.code == 401 or e.code == 302:
-            self._Authenticate()
+            url_loc = urlparse.urlparse(url)
+            self._Authenticate('%s://%s' % (url_loc.scheme, url_loc.netloc))
 ##           elif e.code >= 500 and e.code < 600:
 ##             # Server Error - try again.
 ##             continue
+          elif e.code == 301:
+            # Handle permanent redirect manually.
+            url = e.info()["location"]
           else:
             raise
     finally:
@@ -402,9 +419,9 @@ class AbstractRpcServer(object):
 class HttpRpcServer(AbstractRpcServer):
   """Provides a simplified RPC-style interface for HTTP requests."""
 
-  def _Authenticate(self):
+  def _Authenticate(self, *args):
     """Save the cookie jar after authentication."""
-    super(HttpRpcServer, self)._Authenticate()
+    super(HttpRpcServer, self)._Authenticate(*args)
     if self.save_cookies:
       StatusUpdate("Saving authentication cookies to %s" % self.cookie_file)
       self.cookie_jar.save()