From c8a89eca0800ca64da09ab09f758d78b98bc974f Mon Sep 17 00:00:00 2001
From: leyshon <leyshon@gmail.com>
Date: Thu, 29 Aug 2019 14:32:01 +0100
Subject: [PATCH] Adding the IDToken to the session for the Azure Provider.

---
 providers/azure.go      | 65 +++++++++++++++++++++++++++++++++++++++++
 providers/azure_test.go | 18 +++++++++++-
 2 files changed, 82 insertions(+), 1 deletion(-)

diff --git a/providers/azure.go b/providers/azure.go
index 653090b..6729e4f 100644
--- a/providers/azure.go
+++ b/providers/azure.go
@@ -5,6 +5,10 @@ import (
 	"fmt"
 	"net/http"
 	"net/url"
+	"encoding/json"
+	"io/ioutil"
+	"bytes"
+	"time"
 
 	"github.com/bitly/go-simplejson"
 	"github.com/pusher/oauth2_proxy/pkg/apis/sessions"
@@ -65,6 +69,67 @@ func (p *AzureProvider) Configure(tenant string) {
 	}
 }
 
+func (p *AzureProvider) Redeem(redirectURL, code string) (s *sessions.SessionState, err error) {
+	if code == "" {
+		err = errors.New("missing code")
+		return
+	}
+
+	params := url.Values{}
+	params.Add("redirect_uri", redirectURL)
+	params.Add("client_id", p.ClientID)
+	params.Add("client_secret", p.ClientSecret)
+	params.Add("code", code)
+	params.Add("grant_type", "authorization_code")
+	if p.ProtectedResource != nil && p.ProtectedResource.String() != "" {
+		params.Add("resource", p.ProtectedResource.String())
+	}
+
+	var req *http.Request
+	req, err = http.NewRequest("POST", p.RedeemURL.String(), bytes.NewBufferString(params.Encode()))
+	if err != nil {
+		return
+	}
+	req.Header.Set("Content-Type", "application/x-www-form-urlencoded")
+
+	var resp *http.Response
+	resp, err = http.DefaultClient.Do(req)
+	if err != nil {
+		return nil, err
+	}
+	var body []byte
+	body, err = ioutil.ReadAll(resp.Body)
+	resp.Body.Close()
+	if err != nil {
+		return
+	}
+
+	if resp.StatusCode != 200 {
+		err = fmt.Errorf("got %d from %q %s", resp.StatusCode, p.RedeemURL.String(), body)
+		return
+	}
+
+	var jsonResponse struct {
+		AccessToken  string `json:"access_token"`
+		RefreshToken string `json:"refresh_token"`
+		ExpiresOn    int64  `json:"expires_on,string"`
+		IDToken      string `json:"id_token"`
+	}
+	err = json.Unmarshal(body, &jsonResponse)
+	if err != nil {
+		return
+	}
+	
+	s = &sessions.SessionState{
+		AccessToken:  jsonResponse.AccessToken,
+		IDToken:      jsonResponse.IDToken,
+		CreatedAt:    time.Now(),
+		ExpiresOn:    time.Unix(jsonResponse.ExpiresOn, 0),
+		RefreshToken: jsonResponse.RefreshToken,
+	}
+	return
+}
+
 func getAzureHeader(accessToken string) http.Header {
 	header := make(http.Header)
 	header.Set("Authorization", fmt.Sprintf("Bearer %s", accessToken))
diff --git a/providers/azure_test.go b/providers/azure_test.go
index 8d34bdc..8ec5978 100644
--- a/providers/azure_test.go
+++ b/providers/azure_test.go
@@ -20,6 +20,7 @@ func testAzureProvider(hostname string) *AzureProvider {
 			ValidateURL:       &url.URL{},
 			ProtectedResource: &url.URL{},
 			Scope:             ""})
+
 	if hostname != "" {
 		updateURL(p.Data().LoginURL, hostname)
 		updateURL(p.Data().RedeemURL, hostname)
@@ -111,8 +112,11 @@ func testAzureBackend(payload string) *httptest.Server {
 
 	return httptest.NewServer(http.HandlerFunc(
 		func(w http.ResponseWriter, r *http.Request) {
-			if r.URL.Path != path || r.URL.RawQuery != query {
+			if (r.URL.Path != path || r.URL.RawQuery != query) && r.Method != "POST" {
 				w.WriteHeader(404)
+			} else if r.Method == "POST" && r.Body != nil {
+				w.WriteHeader(200)
+				w.Write([]byte(payload))
 			} else if r.Header.Get("Authorization") != "Bearer imaginary_access_token" {
 				w.WriteHeader(403)
 			} else {
@@ -199,3 +203,15 @@ func TestAzureProviderGetEmailAddressIncorrectOtherMails(t *testing.T) {
 	assert.Equal(t, "type assertion to string failed", err.Error())
 	assert.Equal(t, "", email)
 }
+
+func TestAzureProviderRedeemReturnsIdToken(t *testing.T) {
+	b := testAzureBackend(`{ "id_token": "testtoken1234" }`)
+	defer b.Close()
+
+	bURL, _ := url.Parse(b.URL)
+	p := testAzureProvider(bURL.Host)
+	p.Data().RedeemURL.Path = "/common/oauth2/token"
+	s, err := p.Redeem("https://localhost", "1234")
+	assert.Equal(t, nil, err)
+	assert.Equal(t, "testtoken1234", s.IDToken)
+}
\ No newline at end of file