Clean up nsLocation logic.

1 files changed, 18 insertions, 6 deletions

diff --git a/dom/base/nsLocation.cpp b/dom/base/nsLocation.cpp
index 20acf293d..2789934ef 100644
--- a/dom/base/nsLocation.cpp
+++ b/dom/base/nsLocation.cpp
@@ -720,15 +720,27 @@ nsLocation::SetProtocol(const nsAString& aProtocol)
 
   nsCOMPtr<nsIURI> uri;
   nsresult rv = GetWritableURI(getter_AddRefs(uri));
+  if (NS_FAILED(rv) || !uri) {
+    return rv;
+  }
 
-  if (uri) {
-    rv = uri->SetScheme(NS_ConvertUTF16toUTF8(aProtocol));
-    if (NS_SUCCEEDED(rv)) {
-      SetURI(uri);
-    }
+  rv = uri->SetScheme(NS_ConvertUTF16toUTF8(aProtocol));
+  if (NS_FAILED(rv)) {
+    return rv;
   }
 
-  return rv;
+  nsAutoCString newSpec;
+  rv = uri->GetSpec(newSpec);
+  if (NS_FAILED(rv)) {
+    return rv;
+  }
+  // We may want a new URI class for the new URI, so recreate it:
+  rv = NS_NewURI(getter_AddRefs(uri), newSpec);
+  if (NS_FAILED(rv)) {
+    return rv;
+  }
+
+  return SetURI(uri);
 }
 
 NS_IMETHODIMP