Full hack mode: openid connect, host template and short term tokens

download.go

@@ -2,7 +2,11 @@ package main
 
 import (
 	"encoding/hex"
+	"encoding/json"
 	"github.com/patrickmn/go-cache"
+	"github.com/spf13/viper"
+	"golang.org/x/oauth2"
+	"log"
 	"math/rand"
 	"net/http"
 	"strings"
@@ -10,21 +14,93 @@ import (
 )
 
 func handleRdpDownload(w http.ResponseWriter, r *http.Request) {
+	cookie, err := r.Cookie("RDPGWSESSIONV1")
+	if err != nil {
+		http.Redirect(w, r, oauthConfig.AuthCodeURL(state), http.StatusFound)
+		return
+	}
+
+	data, found := tokens.Get(cookie.Value)
+	if found == false {
+		log.Printf("Found expired or non existent session: %s", cookie.Value)
+		http.Redirect(w, r, oauthConfig.AuthCodeURL(state), http.StatusFound)
+		return
+	}
+
+	// authenticated
 	seed := make([]byte, 16)
 	rand.Read(seed)
 	fn := hex.EncodeToString(seed) + ".rdp"
 
-	rand.Read(seed)
-	token := hex.EncodeToString(seed)
-
-	tokens.Set(token, token, cache.DefaultExpiration)
-
 	w.Header().Set("Content-Disposition", "attachment; filename="+fn)
 	w.Header().Set("Content-Type", "application/x-rdp")
 	http.ServeContent(w, r, fn, time.Now(), strings.NewReader(
-		"full address:s:localhost\r\n"+
-			"gatewayhostname:s:localhost\r\n"+
+		"full address:s:" + host + "\r\n"+
+			"gatewayhostname:s:" + gateway +"\r\n"+
 			"gatewaycredentialssource:i:5\r\n"+
 			"gatewayusagemethod:i:1\r\n"+
-			"gatewayaccesstoken:s:" + token + "\r\n"))
+			"gatewayaccesstoken:s:" + cookie.Value + "\r\n"))
 }
+
+func handleCallback(w http.ResponseWriter, r *http.Request) {
+	if r.URL.Query().Get("state") != state {
+		http.Error(w, "state did not match", http.StatusBadRequest)
+		return
+	}
+
+	oauthToken, err := oauthConfig.Exchange(ctx, r.URL.Query().Get("code"))
+	if err != nil {
+		http.Error(w, "Failed to exchange token: "+err.Error(), http.StatusInternalServerError)
+		return
+	}
+
+	rawIDToken, ok := oauthToken.Extra("id_token").(string)
+	if !ok {
+		http.Error(w, "No id_token field in oauth2 token.", http.StatusInternalServerError)
+		return
+	}
+	idToken, err := verifier.Verify(ctx, rawIDToken)
+	if err != nil {
+		http.Error(w, "Failed to verify ID Token: "+err.Error(), http.StatusInternalServerError)
+		return
+	}
+
+	resp := struct {
+		OAuth2Token   *oauth2.Token
+		IDTokenClaims *json.RawMessage // ID Token payload is just JSON.
+	}{oauthToken, new(json.RawMessage)}
+
+	if err := idToken.Claims(&resp.IDTokenClaims); err != nil {
+		http.Error(w, err.Error(), http.StatusInternalServerError)
+		return
+	}
+
+	q, err := json.MarshalIndent(resp, "", "    ")
+	if err != nil {
+		http.Error(w, err.Error(), http.StatusInternalServerError)
+		return
+	}
+
+	var data map[string]interface{}
+	if err := json.Unmarshal(*resp.IDTokenClaims, &data); err != nil {
+		http.Error(w, err.Error(), http.StatusInternalServerError)
+		return
+	}
+
+	seed := make([]byte, 16)
+	rand.Read(seed)
+	token := hex.EncodeToString(seed)
+
+	cookie := http.Cookie{
+		Name: "RDPGWSESSIONV1",
+		Value: token,
+		Path: "/",
+		Secure: true,
+		HttpOnly: true,
+	}
+
+	tokens.Set(token, data[claim].(string), cache.DefaultExpiration)
+
+	http.SetCookie(w, &cookie)
+	http.Redirect(w, r, "/connect", http.StatusFound)
+}