Refactor config and improve security

download.go

@@ -4,11 +4,9 @@ import (
 	"encoding/hex"
 	"encoding/json"
 	"github.com/patrickmn/go-cache"
-	"github.com/spf13/viper"
 	"golang.org/x/oauth2"
 	"log"
 	"math/rand"
-	"net"
 	"net/http"
 	"strings"
 	"time"
@@ -30,7 +28,12 @@ func handleRdpDownload(w http.ResponseWriter, r *http.Request) {
 		return
 	}
 
-	host := strings.Replace(viper.GetString("hostTemplate"), "%%", data.(string), 1)
+	var host = conf.Server.HostTemplate
+	for k, v := range data.(map[string]interface{}) {
+		if val, ok := v.(string); ok == true {
+			host = strings.Replace(host, "{{ " + k + " }}", val, 1)
+		}
+	}
 
 	// authenticated
 	seed := make([]byte, 16)
@@ -41,7 +44,7 @@ func handleRdpDownload(w http.ResponseWriter, r *http.Request) {
 	w.Header().Set("Content-Type", "application/x-rdp")
 	http.ServeContent(w, r, fn, time.Now(), strings.NewReader(
 		"full address:s:" + host + "\r\n"+
-			"gatewayhostname:s:" + net.JoinHostPort(conf.Server.GatewayAddress, string(conf.Server.Port)) +"\r\n"+
+			"gatewayhostname:s:" + conf.Server.GatewayAddress +"\r\n"+
 			"gatewaycredentialssource:i:5\r\n"+
 			"gatewayusagemethod:i:1\r\n"+
 			"gatewayaccesstoken:s:" + cookie.Value + "\r\n"))
@@ -99,7 +102,7 @@ func handleCallback(w http.ResponseWriter, r *http.Request) {
 	}
 
 	// TODO: make dynamic
-	tokens.Set(token, data["preferred_username"].(string), cache.DefaultExpiration)
+	tokens.Set(token, data, cache.DefaultExpiration)
 
 	http.SetCookie(w, &cookie)
 	http.Redirect(w, r, "/connect", http.StatusFound)

main.go

@@ -59,13 +59,13 @@ func main() {
 		log.Fatalf("Cannot get oidc provider: %s", err)
 	}
 	oidcConfig := &oidc.Config{
-		ClientID: viper.GetString("clientId"),
+		ClientID: conf.OpenId.ClientId,
 	}
 	verifier = provider.Verifier(oidcConfig)
 
 	oauthConfig = oauth2.Config{
-		ClientID: viper.GetString("clientId"),
-		ClientSecret: viper.GetString("clientSecret"),
+		ClientID: conf.OpenId.ClientId,
+		ClientSecret: conf.OpenId.ClientSecret,
 		RedirectURL: "https://" + conf.Server.GatewayAddress + "/callback",
 		Endpoint: provider.Endpoint(),
 		Scopes: []string{oidc.ScopeOpenID, "profile", "email"},

rdg.go

@@ -225,7 +225,12 @@ func handleWebsocketProtocol(conn *websocket.Conn) {
 				log.Printf("Invalid PAA cookie: %s from %s", cookie, conn.RemoteAddr())
 				return
 			}
-			host = strings.Replace(conf.Server.HostTemplate, "%%", data.(string), 1)
+			host = conf.Server.HostTemplate
+			for k, v := range data.(map[string]interface{}) {
+				if val, ok := v.(string); ok == true {
+					host = strings.Replace(host, "{{ " + k + " }}", val, 1)
+				}
+			}
 			msg := createTunnelResponse()
 			log.Printf("Create tunnel response: %x", msg)
 			conn.WriteMessage(mt, msg)