Template Fixes

devalin · Dec 23, 2021 · 1c7cb98 · 1c7cb98
1 parent fb634cd
commit 1c7cb98
Show file tree

Hide file tree

Showing 4 changed files with 91 additions and 37 deletions.
diff --git a/api.go b/api.go
@@ -1,10 +1,12 @@
 package headscale
 
 import (
+	"bytes"
 	"encoding/binary"
 	"encoding/json"
 	"errors"
 	"fmt"
+	"html/template"
 	"io"
 	"net/http"
 	"strings"
@@ -38,18 +40,12 @@ func (h *Headscale) KeyHandler(ctx *gin.Context) {
 	)
 }
 
-// RegisterWebAPI shows a simple message in the browser to point to the CLI
-// Listens in /register.
-func (h *Headscale) RegisterWebAPI(ctx *gin.Context) {
-	machineKeyStr := ctx.Query("key")
-	if machineKeyStr == "" {
-		ctx.String(http.StatusBadRequest, "Wrong params")
-
-		return
-	}
+type registerWebAPITemplateConfig struct {
+	Key string
+}
 
-	ctx.Data(http.StatusOK, "text/html; charset=utf-8", []byte(fmt.Sprintf(`
-	<html>
+var registerWebAPITemplate = template.Must(
+	template.New("registerweb").Parse(`<html>
 	<body>
 	<h1>headscale</h1>
 	<p>
@@ -58,14 +54,40 @@ func (h *Headscale) RegisterWebAPI(ctx *gin.Context) {
 
 	<p>
 		<code>
-			<b>headscale -n NAMESPACE nodes register --key %s</b>
+			<b>headscale -n NAMESPACE nodes register --key {{.Key}}</b>
 		</code>
 	</p>
 
 	</body>
-	</html>
+	</html>`),
+)
+
+// RegisterWebAPI shows a simple message in the browser to point to the CLI
+// Listens in /register.
+func (h *Headscale) RegisterWebAPI(ctx *gin.Context) {
+	machineKeyStr := ctx.Query("key")
+	if machineKeyStr == "" {
+		ctx.String(http.StatusBadRequest, "Wrong params")
+
+		return
+	}
+
+	var content bytes.Buffer
+	if err := registerWebAPITemplate.Execute(&content, registerWebAPITemplateConfig{
+		Key: machineKeyStr,
+	}); err != nil {
+		log.Error().
+			Str("func", "RegisterWebAPI").
+			Err(err).
+			Msg("Could not render register web API template")
+		ctx.Data(
+			http.StatusInternalServerError,
+			"text/html; charset=utf-8",
+			[]byte("Could not render register web API template"),
+		)
+	}
 
-	`, machineKeyStr)))
+	ctx.Data(http.StatusOK, "text/html; charset=utf-8", content.Bytes())
 }
 
 // RegistrationHandler handles the actual registration process of a machine

diff --git a/apple_mobileconfig.go b/apple_mobileconfig.go
@@ -2,8 +2,8 @@ package headscale
 
 import (
 	"bytes"
+	"html/template"
 	"net/http"
-	"text/template"
 
 	"github.com/gin-gonic/gin"
 	"github.com/gofrs/uuid"

diff --git a/oidc.go b/oidc.go
@@ -1,11 +1,13 @@
 package headscale
 
 import (
+	"bytes"
 	"context"
 	"crypto/rand"
 	"encoding/hex"
 	"errors"
 	"fmt"
+	"html/template"
 	"net/http"
 	"regexp"
 	"strings"
@@ -108,6 +110,22 @@ func (h *Headscale) RegisterOIDC(ctx *gin.Context) {
 	ctx.Redirect(http.StatusFound, authURL)
 }
 
+type oidcCallbackTemplateConfig struct {
+	User string
+	Verb string
+}
+
+var oidcCallbackTemplate = template.Must(
+	template.New("oidccallback").Parse(`<html>
+	<body>
+	<h1>headscale</h1>
+	<p>
+			{{.Verb}} as {{.User}}, you can now close this window.
+	</p>
+	</body>
+	</html>`),
+)
+
 // OIDCCallback handles the callback from the OIDC endpoint
 // Retrieves the mkey from the state cache and adds the machine to the users email namespace
 // TODO: A confirmation page for new machines should be added to avoid phishing vulnerabilities
@@ -239,17 +257,24 @@ func (h *Headscale) OIDCCallback(ctx *gin.Context) {
 
 		h.RefreshMachine(machine, requestedTime)
 
-		ctx.Data(http.StatusOK, "text/html; charset=utf-8", []byte(fmt.Sprintf(`
-<html>
-<body>
-<h1>headscale</h1>
-<p>
-    Reuthenticated as %s, you can now close this window.
-</p>
-</body>
-</html>
+		var content bytes.Buffer
+		if err := oidcCallbackTemplate.Execute(&content, oidcCallbackTemplateConfig{
+			User: claims.Email,
+			Verb: "Reauthenticated",
+		}); err != nil {
+			log.Error().
+				Str("func", "OIDCCallback").
+				Str("type", "reauthenticate").
+				Err(err).
+				Msg("Could not render OIDC callback template")
+			ctx.Data(
+				http.StatusInternalServerError,
+				"text/html; charset=utf-8",
+				[]byte("Could not render OIDC callback template"),
+			)
+		}
 
-`, claims.Email)))
+		ctx.Data(http.StatusOK, "text/html; charset=utf-8", content.Bytes())
 
 		return
 	}
@@ -314,17 +339,24 @@ func (h *Headscale) OIDCCallback(ctx *gin.Context) {
 			h.db.Save(&machine)
 		}
 
-		ctx.Data(http.StatusOK, "text/html; charset=utf-8", []byte(fmt.Sprintf(`
-<html>
-<body>
-<h1>headscale</h1>
-<p>
-    Authenticated as %s, you can now close this window.
-</p>
-</body>
-</html>
-
-`, claims.Email)))
+		var content bytes.Buffer
+		if err := oidcCallbackTemplate.Execute(&content, oidcCallbackTemplateConfig{
+			User: claims.Email,
+			Verb: "Authenticated",
+		}); err != nil {
+			log.Error().
+				Str("func", "OIDCCallback").
+				Str("type", "authenticate").
+				Err(err).
+				Msg("Could not render OIDC callback template")
+			ctx.Data(
+				http.StatusInternalServerError,
+				"text/html; charset=utf-8",
+				[]byte("Could not render OIDC callback template"),
+			)
+		}
+
+		ctx.Data(http.StatusOK, "text/html; charset=utf-8", content.Bytes())
 
 		return
 	}

diff --git a/swagger.go b/swagger.go
@@ -3,8 +3,8 @@ package headscale
 import (
 	"bytes"
 	_ "embed"
+	"html/template"
 	"net/http"
-	"text/template"
 
 	"github.com/gin-gonic/gin"
 	"github.com/rs/zerolog/log"