oauth: add HTMLSuccess + HTMLError fields for branded login pages

The htmlSuccess / htmlError templates are good defaults but shipped as
unexported package vars, so callers can't brand the post-login landing
page short of forking. Adds two optional fields on
AuthorizationCodeTokenSource that override the defaults when set, and
plumbs them through to authHandler via two matching unexported fields.

When empty, the new fields are no-ops and the existing constants are
served, so existing callers see no behaviour change.

The custom error HTML supports the same $ERROR / $DETAILS substitution
the built-in errorHTML does, so callers can reuse the documented format
without having to learn two templating systems.

Tests: five new cases in oauth/authcode_test.go covering default + custom
success, default + custom error, and field independence (overriding
success alone leaves the error path unaffected).

Author: davidlwg

oauth/authcode.go

@@ -177,24 +177,37 @@ func getInput(input chan string) {
 }
 
 // authHandler is an HTTP handler that takes a channel and sends the `code`
-// query param when it gets a request.
+// query param when it gets a request. The successHTML and errorHTML fields
+// allow callers to override the built-in HTML pages; an empty string falls
+// back to the package-level htmlSuccess / htmlError defaults.
 type authHandler struct {
-	c chan string
+	c           chan string
+	successHTML string
+	errorHTML   string
 }
 
 func (h authHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
 	w.Header().Set("Content-Type", "text/html")
 
+	successBody := h.successHTML
+	if successBody == "" {
+		successBody = htmlSuccess
+	}
+	errorBody := h.errorHTML
+	if errorBody == "" {
+		errorBody = htmlError
+	}
+
 	if err := r.URL.Query().Get("error"); err != "" {
 		details := r.URL.Query().Get("error_description")
-		rendered := strings.Replace(strings.Replace(htmlError, "$ERROR", err, 1), "$DETAILS", details, 1)
+		rendered := strings.Replace(strings.Replace(errorBody, "$ERROR", err, 1), "$DETAILS", details, 1)
 		w.Write([]byte(rendered))
 		h.c <- ""
 		return
 	}
 
 	h.c <- r.URL.Query().Get("code")
-	w.Write([]byte(htmlSuccess))
+	w.Write([]byte(successBody))
 }
 
 // AuthorizationCodeTokenSource with PKCE as described in:
@@ -212,6 +225,8 @@ type AuthorizationCodeTokenSource struct {
 	RedirectURL    string
 	EndpointParams *url.Values
 	Scopes         []string
+	HTMLSuccess    string
+	HTMLError      string
 }
 
 func (ac *AuthorizationCodeTokenSource) getRedirectUrl() string {
@@ -260,7 +275,9 @@ func (ac *AuthorizationCodeTokenSource) Token() (*oauth2.Token, error) {
 	// Run server before opening the user's browser so we are ready for any redirect.
 	codeChan := make(chan string)
 	handler := authHandler{
-		c: codeChan,
+		c:           codeChan,
+		successHTML: ac.HTMLSuccess,
+		errorHTML:   ac.HTMLError,
 	}
 
 	// strip protocol prefix from configured redirect url for local webserver

oauth/authcode_test.go

@@ -1,10 +1,14 @@
 package oauth
 
 import (
+	"io"
+	"net/http"
+	"net/http/httptest"
 	"strings"
 	"testing"
 
 	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
 )
 
 func TestEncodeUrlWindowsSuccess(t *testing.T) {
@@ -18,3 +22,77 @@ func TestEncodeUrlWindowsSuccess(t *testing.T) {
 	assert.False(t, strings.HasPrefix(r, "^&"))
 	assert.False(t, strings.HasSuffix(r, "^&"))
 }
+
+// runAuthHandler drives the authHandler with the given query string and
+// returns the response body, allowing assertions on what HTML was served.
+func runAuthHandler(t *testing.T, h authHandler, query string) string {
+	t.Helper()
+
+	rec := httptest.NewRecorder()
+	req := httptest.NewRequest(http.MethodGet, "http://localhost/callback?"+query, nil)
+
+	// authHandler sends on h.c synchronously; drain in a goroutine so the
+	// handler can complete without blocking on an unbuffered channel.
+	done := make(chan struct{})
+	go func() {
+		<-h.c
+		close(done)
+	}()
+
+	h.ServeHTTP(rec, req)
+	<-done
+
+	body, err := io.ReadAll(rec.Body)
+	require.NoError(t, err)
+
+	return string(body)
+}
+
+func TestAuthHandlerServesDefaultSuccessHTML(t *testing.T) {
+	h := authHandler{c: make(chan string)}
+
+	body := runAuthHandler(t, h, "code=abc123")
+
+	assert.Contains(t, body, "Login Successful!", "default success HTML should be served when successHTML is unset")
+}
+
+func TestAuthHandlerServesCustomSuccessHTML(t *testing.T) {
+	const custom = `<html><body><h1>Welcome to my-tool</h1></body></html>`
+	h := authHandler{c: make(chan string), successHTML: custom}
+
+	body := runAuthHandler(t, h, "code=abc123")
+
+	assert.Equal(t, custom, body, "custom successHTML should be served verbatim")
+	assert.NotContains(t, body, "Login Successful!", "default content should not leak through")
+}
+
+func TestAuthHandlerServesDefaultErrorHTML(t *testing.T) {
+	h := authHandler{c: make(chan string)}
+
+	body := runAuthHandler(t, h, "error=access_denied&error_description=user+denied")
+
+	assert.Contains(t, body, "access_denied", "$ERROR should be substituted in the default error HTML")
+	assert.Contains(t, body, "user denied", "$DETAILS should be substituted in the default error HTML")
+}
+
+func TestAuthHandlerServesCustomErrorHTML(t *testing.T) {
+	const custom = `<html><body><h1>Login failed: $ERROR</h1><p>$DETAILS</p></body></html>`
+	h := authHandler{c: make(chan string), errorHTML: custom}
+
+	body := runAuthHandler(t, h, "error=invalid_grant&error_description=code+expired")
+
+	assert.Equal(t, `<html><body><h1>Login failed: invalid_grant</h1><p>code expired</p></body></html>`, body,
+		"custom errorHTML should be served with $ERROR and $DETAILS substituted")
+}
+
+func TestAuthHandlerCustomSuccessDoesNotAffectErrorPath(t *testing.T) {
+	// A caller that only overrides successHTML should still see the default
+	// error HTML when an error comes through. Each field is independently
+	// optional.
+	h := authHandler{c: make(chan string), successHTML: "<html>custom success</html>"}
+
+	body := runAuthHandler(t, h, "error=server_error")
+
+	assert.Contains(t, body, "server_error", "default error HTML should still substitute $ERROR")
+	assert.NotContains(t, body, "custom success")
+}