From 0eb5379a452b7bb0ded3960ee750050fd65040b5 Mon Sep 17 00:00:00 2001
From: Clement Michaud <clement.michaud34@gmail.com>
Date: Wed, 15 Mar 2017 23:47:59 +0100
Subject: [PATCH] Handle redirection after registration either with U2F or TOTP

---
 src/lib/identity_check.js           |  4 ++++
 src/public_html/js/login.js         | 20 ++++++++++++++++++--
 src/public_html/js/totp-register.js | 14 ++++++++++++++
 src/public_html/js/u2f-register.js  |  2 +-
 src/views/totp-register.ejs         |  2 +-
 5 files changed, 38 insertions(+), 4 deletions(-)

diff --git a/src/lib/identity_check.js b/src/lib/identity_check.js
index 1b37c28b1..bef355343 100644
--- a/src/lib/identity_check.js
+++ b/src/lib/identity_check.js
@@ -109,8 +109,12 @@ function identity_check_post(endpoint, icheck_interface) {
       throw new exceptions.AccessDeniedError();
     })
     .then(function(token) {
+      var redirect_url = objectPath.get(req, 'body.redirect');
       var original_url = util.format('https://%s%s', req.headers.host, req.headers['x-original-uri']);
       var link_url = util.format('%s?identity_token=%s', original_url, token); 
+      if(redirect_url) {
+        link_url = util.format('%s&redirect=%s', link_url, redirect_url); 
+      }
 
       logger.info('POST identity_check: notify to %s', identity.userid);
       return notifier.notify(identity, icheck_interface.email_subject, link_url);
diff --git a/src/public_html/js/login.js b/src/public_html/js/login.js
index 6fa227f86..5da20d41d 100644
--- a/src/public_html/js/login.js
+++ b/src/public_html/js/login.js
@@ -3,6 +3,11 @@
 params={};
 location.search.replace(/[?&]+([^=&]+)=([^&]*)/gi,function(s,k,v){params[k]=v});
 
+function get_redirect_param() {
+  if('redirect' in params)
+    return params['redirect'];
+  return;
+}
 
 function setupEnterKeypressListener(filter, fn) {
   $(filter).on('keydown', 'input', function (e) {
@@ -49,7 +54,12 @@ function onTotpSignButtonClicked() {
 function onTotpRegisterButtonClicked() {
   $.ajax({
     type: 'POST',
-    url: '/authentication/totp-register'
+    url: '/authentication/totp-register',
+    data: JSON.stringify({
+      redirect: get_redirect_param()
+    }),
+    contentType: 'application/json',
+    dataType: 'json',
   })
   .done(function(data) {
     $.notify('An email has been sent to your email address', 'info');
@@ -82,7 +92,12 @@ function onU2fRegistrationButtonClicked() {
 function askForU2fRegistration(fn) {
   $.ajax({
     type: 'POST',
-    url: '/authentication/u2f-register'
+    url: '/authentication/u2f-register',
+    data: JSON.stringify({
+      redirect: get_redirect_param()
+    }),
+    contentType: 'application/json',
+    dataType: 'json',
   })
   .done(function(data) {
     fn(undefined, data);
@@ -158,6 +173,7 @@ function validateFirstFactor(username, password, fn) {
   });
 }
 
+
 function redirect() {
   var redirect_uri = '/';
   if('redirect' in params) {
diff --git a/src/public_html/js/totp-register.js b/src/public_html/js/totp-register.js
index e4f4b7eb0..51a0c7686 100644
--- a/src/public_html/js/totp-register.js
+++ b/src/public_html/js/totp-register.js
@@ -1,5 +1,8 @@
 (function() {
 
+params={};
+location.search.replace(/[?&]+([^=&]+)=([^&]*)/gi,function(s,k,v){params[k]=v});
+
 function generateSecret(fn) {
   $.ajax({
     type: 'POST',
@@ -22,7 +25,18 @@ function onSecretGenerated(err, secret) {
   $("#secret").text(secret.base32);
 }
 
+function redirect() {
+  var redirect_uri = '/authentication/login';
+  if('redirect' in params) {
+    redirect_uri = params['redirect'];
+  }
+  window.location.replace(redirect_uri);
+}
+
 $(document).ready(function() {
   generateSecret(onSecretGenerated);
+  $('#login-button').on('click', function() {
+    redirect();
+  });
 });
 })();
diff --git a/src/public_html/js/u2f-register.js b/src/public_html/js/u2f-register.js
index 619e6ad68..53614bf65 100644
--- a/src/public_html/js/u2f-register.js
+++ b/src/public_html/js/u2f-register.js
@@ -39,7 +39,7 @@ function startRegister(fn, timeout) {
 }
 
 function redirect() {
-  var redirect_uri = '/';
+  var redirect_uri = '/authentication/login';
   if('redirect' in params) {
     redirect_uri = params['redirect'];
   }
diff --git a/src/views/totp-register.ejs b/src/views/totp-register.ejs
index 2652459e8..03f067637 100644
--- a/src/views/totp-register.ejs
+++ b/src/views/totp-register.ejs
@@ -9,7 +9,7 @@
       <p>Insert your secret in Google Authenticator</p>
       <p id="secret"></p>
       <div id="qrcode"></div>
-      <p><a href="/authentication/login">Login</a></p>
+      <p><a href="#" id="login-button">Login</a></p>
     </div>
   </body>