Fix automatic login when boot.dev redirects to www

cmd/login.go

@@ -12,6 +12,8 @@ import (
 	"net/url"
 	"os"
 	"regexp"
+	"slices"
+	"strings"
 	"sync"
 	"time"
 
@@ -101,15 +103,16 @@ var loginCmd = &cobra.Command{
 
 const maxLoginCodeBytes = 4 * 1024
 
-func cors(allowedOrigin string, next http.Handler) http.Handler {
+func cors(allowedOrigins []string, next http.Handler) http.Handler {
 	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
-		if allowedOrigin == "" || r.Header.Get("Origin") != allowedOrigin {
+		origin := r.Header.Get("Origin")
+		if origin == "" || !slices.Contains(allowedOrigins, origin) {
 			http.Error(w, "origin not allowed", http.StatusForbidden)
 			return
 		}
 
 		w.Header().Set("Access-Control-Allow-Methods", "GET, POST, OPTIONS")
-		w.Header().Set("Access-Control-Allow-Origin", allowedOrigin)
+		w.Header().Set("Access-Control-Allow-Origin", origin)
 		w.Header().Set("Access-Control-Allow-Headers", "Content-Type")
 		w.Header().Set("Vary", "Origin")
 		if r.Method == http.MethodOptions {
@@ -120,17 +123,26 @@ func cors(allowedOrigin string, next http.Handler) http.Handler {
 	})
 }
 
-func frontendOrigin(frontendURL string) string {
+// The site redirects between the apex and www hosts, so the browser's Origin
+// may not match the configured frontend_url exactly.
+func frontendOrigins(frontendURL string) []string {
 	parsed, err := url.Parse(frontendURL)
 	if err != nil || parsed.Scheme == "" || parsed.Host == "" {
-		return ""
+		return nil
+	}
+	alternateHost := "www." + parsed.Host
+	if host, ok := strings.CutPrefix(parsed.Host, "www."); ok {
+		alternateHost = host
+	}
+	return []string{
+		parsed.Scheme + "://" + parsed.Host,
+		parsed.Scheme + "://" + alternateHost,
 	}
-	return parsed.Scheme + "://" + parsed.Host
 }
 
 func newLoginHTTPHandler(inputChan chan<- string, frontendURL string) http.Handler {
 	mux := http.NewServeMux()
-	allowedOrigin := frontendOrigin(frontendURL)
+	allowedOrigins := frontendOrigins(frontendURL)
 
 	handleSubmit := func(w http.ResponseWriter, r *http.Request) {
 		if r.Method != http.MethodPost {
@@ -170,8 +182,8 @@ func newLoginHTTPHandler(inputChan chan<- string, frontendURL string) http.Handl
 		http.Redirect(w, r, loginURL, http.StatusSeeOther)
 	}
 
-	mux.Handle("/submit", cors(allowedOrigin, http.HandlerFunc(handleSubmit)))
-	mux.Handle("/health", cors(allowedOrigin, http.HandlerFunc(handleHealth)))
+	mux.Handle("/submit", cors(allowedOrigins, http.HandlerFunc(handleSubmit)))
+	mux.Handle("/health", cors(allowedOrigins, http.HandlerFunc(handleHealth)))
 	mux.Handle("/{$}", http.HandlerFunc(handleRedirect))
 	return mux
 }

cmd/login_test.go

@@ -9,33 +9,46 @@ import (
 
 const testFrontendURL = "https://boot.dev"
 
-func TestLoginHTTPHandlerAcceptsConfiguredOrigin(t *testing.T) {
-	inputChan := make(chan string, 1)
-	handler := newLoginHTTPHandler(inputChan, testFrontendURL)
-	request := httptest.NewRequest(http.MethodPost, "/submit", strings.NewReader("valid-code"))
-	request.Header.Set("Origin", testFrontendURL)
-	response := httptest.NewRecorder()
+func TestLoginHTTPHandlerAcceptsFrontendOrigins(t *testing.T) {
+	tests := []struct {
+		frontendURL string
+		origin      string
+	}{
+		{"https://boot.dev", "https://boot.dev"},
+		{"https://boot.dev", "https://www.boot.dev"},
+		{"https://www.boot.dev", "https://boot.dev"},
+		{"http://localhost:3000", "http://localhost:3000"},
+	}
+	for _, tt := range tests {
+		t.Run(tt.frontendURL+" "+tt.origin, func(t *testing.T) {
+			inputChan := make(chan string, 1)
+			handler := newLoginHTTPHandler(inputChan, tt.frontendURL)
+			request := httptest.NewRequest(http.MethodPost, "/submit", strings.NewReader("valid-code"))
+			request.Header.Set("Origin", tt.origin)
+			response := httptest.NewRecorder()
 
-	handler.ServeHTTP(response, request)
+			handler.ServeHTTP(response, request)
 
-	if response.Code != http.StatusNoContent {
-		t.Fatalf("status = %d, want %d", response.Code, http.StatusNoContent)
-	}
-	select {
-	case code := <-inputChan:
-		if code != "valid-code" {
-			t.Fatalf("login code = %q, want valid-code", code)
-		}
-	default:
-		t.Fatal("login code was not delivered")
-	}
-	if origin := response.Header().Get("Access-Control-Allow-Origin"); origin != testFrontendURL {
-		t.Fatalf("allowed origin = %q, want %q", origin, testFrontendURL)
+			if response.Code != http.StatusNoContent {
+				t.Fatalf("status = %d, want %d", response.Code, http.StatusNoContent)
+			}
+			select {
+			case code := <-inputChan:
+				if code != "valid-code" {
+					t.Fatalf("login code = %q, want valid-code", code)
+				}
+			default:
+				t.Fatal("login code was not delivered")
+			}
+			if origin := response.Header().Get("Access-Control-Allow-Origin"); origin != tt.origin {
+				t.Fatalf("allowed origin = %q, want %q", origin, tt.origin)
+			}
+		})
 	}
 }
 
 func TestLoginHTTPHandlerRejectsUntrustedOrigin(t *testing.T) {
-	for _, origin := range []string{"https://example.com", ""} {
+	for _, origin := range []string{"https://example.com", "https://evil.boot.dev", "http://www.boot.dev", ""} {
 		t.Run(origin, func(t *testing.T) {
 			inputChan := make(chan string, 1)
 			request := httptest.NewRequest(http.MethodPost, "/submit", strings.NewReader("attacker-code"))