From 869fd8c6bca3363c1165615058080a8d2a011566 Mon Sep 17 00:00:00 2001
From: Dan Sosedoff <dan.sosedoff@gmail.com>
Date: Thu, 21 Sep 2017 01:21:26 -0500
Subject: [PATCH] Refactor the third-party connect backend functionality

---
 pkg/api/api.go         | 44 ++++++++-----------------
 pkg/api/backend.go     | 75 ++++++++++++++++++++++++++++++++++++++++++
 pkg/command/options.go |  1 +
 3 files changed, 90 insertions(+), 30 deletions(-)
 create mode 100644 pkg/api/backend.go

diff --git a/pkg/api/api.go b/pkg/api/api.go
index c40ab4e..ab0f435 100644
--- a/pkg/api/api.go
+++ b/pkg/api/api.go
@@ -2,11 +2,8 @@ package api
 
 import (
 	"encoding/base64"
-	"encoding/json"
 	"errors"
 	"fmt"
-	"io/ioutil"
-	"net/http"
 	neturl "net/url"
 	"strings"
 	"time"
@@ -74,27 +71,21 @@ func GetSessions(c *gin.Context) {
 }
 
 func ConnectWithBackend(c *gin.Context) {
-	resp, err := http.PostForm(command.Opts.ConnectBackend, neturl.Values{
-		"resource": {c.Param("resource")},
-		"token":    {command.Opts.ConnectToken},
-	})
+	// Setup a new backend client
+	backend := Backend{
+		Endpoint:    command.Opts.ConnectBackend,
+		Token:       command.Opts.ConnectToken,
+		PassHeaders: command.Opts.ConnectHeaders,
+	}
+
+	// Fetch connection credentials
+	cred, err := backend.FetchCredential(c.Param("resource"), c)
 	if err != nil {
-		c.JSON(400, err)
-		return
-	}
-	defer resp.Body.Close()
-
-	if resp.StatusCode != 200 {
-		c.JSON(400, Error{"Unable to fetch connection settings"})
-		return
-	}
-
-	data, err := ioutil.ReadAll(resp.Body)
-	if err != nil {
-		c.JSON(400, err)
+		c.JSON(400, Error{err.Error()})
 		return
 	}
 
+	// Make the new session
 	sessionId, err := securerandom.Uuid()
 	if err != nil {
 		c.JSON(400, Error{err.Error()})
@@ -102,22 +93,15 @@ func ConnectWithBackend(c *gin.Context) {
 	}
 	c.Request.Header.Add("x-session-id", sessionId)
 
-	config := struct {
-		DatabaseUrl string `json:"database_url"`
-	}{}
-
-	if err := json.Unmarshal(data, &config); err != nil {
-		c.JSON(400, Error{err.Error()})
-		return
-	}
-
-	cl, err := client.NewFromUrl(config.DatabaseUrl, nil)
+	// Connect to the database
+	cl, err := client.NewFromUrl(cred.DatabaseUrl, nil)
 	if err != nil {
 		c.JSON(400, Error{err.Error()})
 		return
 	}
 	cl.External = true
 
+	// Finalize session seetup
 	_, err = cl.Info()
 	if err == nil {
 		err = setClient(c, cl)
diff --git a/pkg/api/backend.go b/pkg/api/backend.go
new file mode 100644
index 0000000..90baf79
--- /dev/null
+++ b/pkg/api/backend.go
@@ -0,0 +1,75 @@
+package api
+
+import (
+	"bytes"
+	"encoding/json"
+	"fmt"
+	"io/ioutil"
+	"log"
+	"net/http"
+	"strings"
+
+	"github.com/gin-gonic/gin"
+)
+
+type Backend struct {
+	Endpoint    string
+	Token       string
+	PassHeaders string
+}
+
+type BackendRequest struct {
+	Resource string            `json:"resource"`
+	Token    string            `json:"token"`
+	Headers  map[string]string `json:"headers"`
+}
+
+type BackendCredential struct {
+	DatabaseUrl string `json:"database_url"`
+}
+
+func (be Backend) FetchCredential(resource string, c *gin.Context) (*BackendCredential, error) {
+	request := BackendRequest{
+		Resource: resource,
+		Token:    be.Token,
+		Headers:  map[string]string{},
+	}
+
+	for _, name := range strings.Split(be.PassHeaders, ",") {
+		request.Headers[strings.ToLower(name)] = c.Request.Header.Get(name)
+	}
+
+	body, err := json.Marshal(request)
+	if err != nil {
+		return nil, err
+	}
+
+	resp, err := http.Post(be.Endpoint, "application/json", bytes.NewReader(body))
+	if err != nil {
+		// Any connection-related issues will show up in the server log
+		log.Println("Unable to fetch backend credential:", err)
+
+		// We dont want to expose the url of the backend here, so reply with generic error
+		return nil, fmt.Errorf("Unable to connect to the auth backend")
+	}
+	defer resp.Body.Close()
+
+	if resp.StatusCode != 200 {
+		return nil, fmt.Errorf("Got HTTP error %v from backend", resp.StatusCode)
+	}
+
+	respBody, err := ioutil.ReadAll(resp.Body)
+	if err != nil {
+		return nil, err
+	}
+
+	cred := &BackendCredential{}
+	if err := json.Unmarshal(respBody, cred); err != nil {
+		return nil, err
+	}
+	if cred.DatabaseUrl == "" {
+		return nil, fmt.Errorf("Database url was not provided")
+	}
+
+	return cred, nil
+}
diff --git a/pkg/command/options.go b/pkg/command/options.go
index b7468ab..d9f1c1c 100644
--- a/pkg/command/options.go
+++ b/pkg/command/options.go
@@ -32,6 +32,7 @@ type Options struct {
 	DisablePrettyJson bool   `long:"no-pretty-json" description:"Disable JSON formatting feature for result export" default:"false"`
 	ConnectBackend    string `long:"connect-backend" description:"Enable database authentication through a third party backend"`
 	ConnectToken      string `long:"connect-token" description:"Authentication token for the third-party connect backend"`
+	ConnectHeaders    string `long:"connect-headers" description:"List of headers to pass to the connect backend"`
 }
 
 var Opts Options