Fix redirect after choosing source

ui/src/CheckDataNodes.js

@@ -17,6 +17,9 @@ const CheckDataNodes = React.createClass({
     router: PropTypes.shape({
       push: PropTypes.func.isRequired,
     }).isRequired,
+    location: PropTypes.shape({
+      pathname: PropTypes.string.isRequired,
+    }).isRequired,
   },
 
   contextTypes: {
@@ -41,7 +44,8 @@ const CheckDataNodes = React.createClass({
       const source = sources.find((s) => s.id === sourceID);
 
       if (!source) { // would be great to check source.status or similar here
-        return this.props.router.push(`/?error="bad id: ${sourceID}"`);
+        const {router, location} = this.props;
+        return router.push(`/?redirectPath=${location.pathname}`);
       }
 
       this.setState({

ui/src/select_source/containers/SelectSourcePage.js

@@ -10,7 +10,7 @@ export const SelectSourcePage = React.createClass({
     }).isRequired,
     location: PropTypes.shape({
       query: PropTypes.shape({
-        error: PropTypes.string,
+        redirectPath: PropTypes.string,
       }).isRequired,
     }).isRequired,
   },
@@ -32,7 +32,7 @@ export const SelectSourcePage = React.createClass({
   handleSelectSource(e) {
     e.preventDefault();
     const source = this.state.sources.find((s) => s.name === this.selectedSource.value);
-    this.props.router.push(`/sources/${source.id}/hosts`);
+    this.redirectToApp(source);
   },
 
   handleNewSource(e) {
@@ -44,12 +44,22 @@ export const SelectSourcePage = React.createClass({
       password: this.sourcePassword.value,
     };
     createSource(source).then(() => {
-      // this.props.router.push(`/sources/${source.id}/hosts`);
+      // this.redirectToApp(sourceFromServer)
     });
   },
 
+  redirectToApp(source) {
+    const {redirectPath} = this.props.location.query;
+    if (!redirectPath) {
+      this.props.router.push(`/sources/${source.id}/hosts`);
+    }
+
+    const fixedPath = redirectPath.replace(/\/sources\/[^/]*/, `/sources/${source.id}`);
+    return this.props.router.push(fixedPath);
+  },
+
   render() {
-    const error = !!this.props.location.query.error;
+    const error = !!this.props.location.query.redirectPath;
     return (
       <div id="select-source-page">
         <div className="container">