From efae5f6bf1166324e56cc03115679922b883ee83 Mon Sep 17 00:00:00 2001
From: Luke Barnard <lukeb@openmarket.com>
Date: Fri, 3 Feb 2017 11:48:24 +0000
Subject: [PATCH] Use localStorage team token with fall-back on query parameter

---
 src/components/structures/LoggedInView.js | 4 +++-
 src/components/structures/MatrixChat.js   | 8 +++++++-
 2 files changed, 10 insertions(+), 2 deletions(-)

diff --git a/src/components/structures/LoggedInView.js b/src/components/structures/LoggedInView.js
index 7dd37683b7..ba63794f60 100644
--- a/src/components/structures/LoggedInView.js
+++ b/src/components/structures/LoggedInView.js
@@ -42,6 +42,8 @@ export default React.createClass({
         onRoomCreated: React.PropTypes.func,
         onUserSettingsClose: React.PropTypes.func,
 
+        teamToken: React.PropTypes.string,
+
         // and lots and lots of other stuff.
     },
 
@@ -197,7 +199,7 @@ export default React.createClass({
                 page_element = <HomePage
                     collapsedRhs={this.props.collapse_rhs}
                     teamServerUrl={this.props.config.teamServerConfig.teamServerURL}
-                    teamToken={window.localStorage.getItem("mx_team_token")}
+                    teamToken={this.props.teamToken}
                 />
                 if (!this.props.collapse_rhs) right_panel = <RightPanel opacity={this.props.sideOpacity}/>
                 break;
diff --git a/src/components/structures/MatrixChat.js b/src/components/structures/MatrixChat.js
index 3cc34ba69c..b986cd88f0 100644
--- a/src/components/structures/MatrixChat.js
+++ b/src/components/structures/MatrixChat.js
@@ -190,6 +190,11 @@ module.exports = React.createClass({
         if (this.props.config.sync_timeline_limit) {
             MatrixClientPeg.opts.initialSyncLimit = this.props.config.sync_timeline_limit;
         }
+
+        // Use the locally-stored team token first, then as a fall-back, check to see if
+        // a referral link was used, which will contain a query parameter `team_token`.
+        this._teamToken = window.localStorage.getItem('mx_team_token') ||
+            startingFragmentQueryParams.team_token;
     },
 
     componentDidMount: function() {
@@ -694,7 +699,7 @@ module.exports = React.createClass({
                         )[0].roomId;
                         self.setState({ready: true, currentRoomId: firstRoom, page_type: PageTypes.RoomView});
                     } else {
-                        if (window.localStorage.getItem('mx_team_token')) {
+                        if (this._teamToken) {
                             self.setState({ready: true, page_type: PageTypes.HomePage});
                         }
                         else {
@@ -1051,6 +1056,7 @@ module.exports = React.createClass({
                     onRoomIdResolved={this.onRoomIdResolved}
                     onRoomCreated={this.onRoomCreated}
                     onUserSettingsClose={this.onUserSettingsClose}
+                    teamToken={this._teamToken}
                     {...this.props}
                     {...this.state}
                 />