cli http tests: allow response header capture

checks/http.go

@@ -76,7 +76,16 @@ func runHTTPRequest(
 		req.SetBasicAuth(requestStep.Request.BasicAuth.Username, requestStep.Request.BasicAuth.Password)
 	}
 
-	resp, err := client.Do(req)
+	requestClient := client
+	if requestStep.Request.FollowRedirects != nil && !*requestStep.Request.FollowRedirects {
+		clientCopy := *client
+		clientCopy.CheckRedirect = func(req *http.Request, via []*http.Request) error {
+			return http.ErrUseLastResponse
+		}
+		requestClient = &clientCopy
+	}
+
+	resp, err := requestClient.Do(req)
 	if err != nil {
 		errString := fmt.Sprintf("Failed to fetch: %s", err.Error())
 		result = api.HTTPRequestResult{Err: errString}
@@ -100,7 +109,12 @@ func runHTTPRequest(
 		trailers[k] = strings.Join(v, ",")
 	}
 
-	parseVariables(body, requestStep.ResponseVariables, variables)
+	if err := parseVariables(body, requestStep.ResponseVariables, variables); err != nil {
+		return api.HTTPRequestResult{Err: fmt.Sprintf("Failed to parse response variable: %s", err)}
+	}
+	if err := parseHeaderVariables(headers, requestStep.ResponseHeaderVariables, variables); err != nil {
+		return api.HTTPRequestResult{Err: fmt.Sprintf("Failed to parse response header variable: %s", err)}
+	}
 
 	result = api.HTTPRequestResult{
 		StatusCode:       resp.StatusCode,
@@ -186,15 +200,57 @@ func truncateAndStringifyBody(body []byte) string {
 
 func parseVariables(body []byte, vardefs []api.HTTPRequestResponseVariable, variables map[string]string) error {
 	for _, vardef := range vardefs {
-		val, err := valFromJqPath(vardef.Path, string(body))
+		vals, err := valsFromJqPath(vardef.Path, string(body))
 		if err != nil {
 			return err
 		}
-		variables[vardef.Name] = fmt.Sprintf("%v", val)
+		if len(vals) != 1 || vals[0] == nil {
+			continue
+		}
+		variables[vardef.Name] = fmt.Sprintf("%v", vals[0])
 	}
 	return nil
 }
 
+func parseHeaderVariables(headers map[string]string, vardefs []api.HTTPRequestResponseHeaderVariable, variables map[string]string) error {
+	for _, vardef := range vardefs {
+		headerValue, ok := findHeaderValue(headers, vardef.Header)
+		if !ok || headerValue == "" {
+			continue
+		}
+
+		value := headerValue
+		if vardef.Regex != "" {
+			re, err := regexp.Compile(vardef.Regex)
+			if err != nil {
+				return err
+			}
+			if re.NumSubexp() != 1 {
+				return fmt.Errorf("regex for header variable %q must have exactly one capture group", vardef.Name)
+			}
+
+			matches := re.FindStringSubmatch(headerValue)
+			if len(matches) != 2 || matches[1] == "" {
+				continue
+			}
+			value = matches[1]
+		}
+
+		variables[vardef.Name] = value
+	}
+
+	return nil
+}
+
+func findHeaderValue(headers map[string]string, key string) (string, bool) {
+	for actualKey, value := range headers {
+		if strings.EqualFold(actualKey, key) {
+			return value, true
+		}
+	}
+	return "", false
+}
+
 func InterpolateVariables(template string, vars map[string]string) string {
 	r := regexp.MustCompile(`\$\{([^}]+)\}`)
 	return r.ReplaceAllStringFunc(template, func(m string) string {

checks/http_test.go

@@ -126,3 +126,103 @@ func TestTruncateAndStringifyBodyCapsBinaryBody(t *testing.T) {
 		t.Fatalf("len(truncateAndStringifyBody(binary)) = %d, want %d", len(got), 16*1024)
 	}
 }
+
+func TestRunHTTPRequestCapturesResponseHeaderVariableAndDoesNotFollowRedirect(t *testing.T) {
+	followRedirects := false
+
+	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		if r.URL.Path != "/login" {
+			t.Errorf("path = %q, want /login", r.URL.Path)
+			return
+		}
+
+		w.Header().Set("Set-Cookie", "session_id=abc123; Path=/; HttpOnly")
+		w.Header().Set("Location", "/account")
+		w.WriteHeader(http.StatusFound)
+		_, _ = w.Write([]byte("Found. Redirecting to /account"))
+	}))
+	defer server.Close()
+
+	variables := map[string]string{}
+	requestStep := api.CLIStepHTTPRequest{
+		ResponseHeaderVariables: []api.HTTPRequestResponseHeaderVariable{{
+			Name:   "sessionID",
+			Header: "Set-Cookie",
+			Regex:  "session_id=([^;]+)",
+		}},
+		Request: api.HTTPRequest{
+			Method:          http.MethodPost,
+			FullURL:         api.BaseURLPlaceholder + "/login",
+			FollowRedirects: &followRedirects,
+			BodyForm: map[string]string{
+				"email":    "pacifica@example.com",
+				"password": "password123",
+				"returnTo": "/account",
+			},
+		},
+	}
+
+	result := runHTTPRequest(server.Client(), server.URL, variables, requestStep)
+	if result.Err != "" {
+		t.Fatalf("unexpected request error: %s", result.Err)
+	}
+	if result.StatusCode != http.StatusFound {
+		t.Fatalf("StatusCode = %d, want %d", result.StatusCode, http.StatusFound)
+	}
+	if result.ResponseHeaders["Set-Cookie"] == "" {
+		t.Fatalf("expected Set-Cookie response header")
+	}
+	if result.Variables["sessionID"] != "abc123" {
+		t.Fatalf("captured sessionID = %q, want abc123", result.Variables["sessionID"])
+	}
+}
+
+func TestParseVariablesLeavesMissingValuesUnset(t *testing.T) {
+	variables := map[string]string{}
+	err := parseVariables(
+		[]byte(`{"token":"abc123","missing":null}`),
+		[]api.HTTPRequestResponseVariable{
+			{Name: "token", Path: ".token"},
+			{Name: "missing", Path: ".missing"},
+			{Name: "notFound", Path: ".not_found"},
+		},
+		variables,
+	)
+	if err != nil {
+		t.Fatalf("unexpected parseVariables error: %v", err)
+	}
+	if variables["token"] != "abc123" {
+		t.Fatalf("token = %q, want abc123", variables["token"])
+	}
+	if _, ok := variables["missing"]; ok {
+		t.Fatalf("expected null variable to remain unset")
+	}
+	if _, ok := variables["notFound"]; ok {
+		t.Fatalf("expected missing variable to remain unset")
+	}
+}
+
+func TestParseHeaderVariablesLeavesMissingValuesUnset(t *testing.T) {
+	variables := map[string]string{}
+	err := parseHeaderVariables(
+		map[string]string{"Set-Cookie": "session_id=abc123; Path=/; HttpOnly"},
+		[]api.HTTPRequestResponseHeaderVariable{
+			{Name: "sessionID", Header: "Set-Cookie", Regex: "session_id=([^;]+)"},
+			{Name: "missingHeader", Header: "X-Missing"},
+			{Name: "missingMatch", Header: "Set-Cookie", Regex: "missing=([^;]+)"},
+		},
+		variables,
+	)
+	if err != nil {
+		t.Fatalf("unexpected parseHeaderVariables error: %v", err)
+	}
+	if variables["sessionID"] != "abc123" {
+		t.Fatalf("sessionID = %q, want abc123", variables["sessionID"])
+	}
+	if _, ok := variables["missingHeader"]; ok {
+		t.Fatalf("expected missing header variable to remain unset")
+	}
+	if _, ok := variables["missingMatch"]; ok {
+		t.Fatalf("expected non-matching header variable to remain unset")
+	}
+}

client/lessons.go

@@ -72,10 +72,11 @@ const (
 )
 
 type CLIStepHTTPRequest struct {
-	ResponseVariables []HTTPRequestResponseVariable
-	Tests             []HTTPRequestTest
-	Request           HTTPRequest
-	SleepAfterMs      *int
+	ResponseVariables       []HTTPRequestResponseVariable
+	ResponseHeaderVariables []HTTPRequestResponseHeaderVariable
+	Tests                   []HTTPRequestTest
+	Request                 HTTPRequest
+	SleepAfterMs            *int
 }
 
 type Sleepable interface {
@@ -93,11 +94,12 @@ func (h *CLIStepHTTPRequest) GetSleepAfterMs() *int {
 const BaseURLPlaceholder = "${baseURL}"
 
 type HTTPRequest struct {
-	Method   string
-	FullURL  string
-	Headers  map[string]string
-	BodyJSON map[string]any
-	BodyForm map[string]string
+	Method          string
+	FullURL         string
+	Headers         map[string]string
+	BodyJSON        map[string]any
+	BodyForm        map[string]string
+	FollowRedirects *bool
 
 	BasicAuth *HTTPBasicAuth
 }
@@ -112,6 +114,12 @@ type HTTPRequestResponseVariable struct {
 	Path string
 }
 
+type HTTPRequestResponseHeaderVariable struct {
+	Name   string
+	Header string
+	Regex  string
+}
+
 // HTTPRequestTest should have only one field set
 type HTTPRequestTest struct {
 	StatusCode       *int

render/variables.go

@@ -52,6 +52,17 @@ func savedVariablesForHTTPResult(result api.HTTPRequestResult) []variableEntry {
 			description: "JSON Body " + responseVariable.Path,
 		})
 	}
+	for _, responseHeaderVariable := range result.Request.ResponseHeaderVariables {
+		value := result.Variables[responseHeaderVariable.Name]
+		if value == "" {
+			continue
+		}
+		entries = append(entries, variableEntry{
+			name:        responseHeaderVariable.Name,
+			value:       value,
+			description: responseHeaderVariableDescription(responseHeaderVariable),
+		})
+	}
 	return entries
 }
 
@@ -66,9 +77,25 @@ func missingSaveVariablesForHTTPResult(result api.HTTPRequestResult) []variableE
 			description: "JSON Body " + responseVariable.Path,
 		})
 	}
+	for _, responseHeaderVariable := range result.Request.ResponseHeaderVariables {
+		if result.Variables[responseHeaderVariable.Name] != "" {
+			continue
+		}
+		entries = append(entries, variableEntry{
+			name:        responseHeaderVariable.Name,
+			description: responseHeaderVariableDescription(responseHeaderVariable),
+		})
+	}
 	return entries
 }
 
+func responseHeaderVariableDescription(v api.HTTPRequestResponseHeaderVariable) string {
+	if v.Regex == "" {
+		return "Response Header " + v.Header
+	}
+	return fmt.Sprintf("Response Header %s matching %s", v.Header, v.Regex)
+}
+
 func availableVariablesForHTTPResult(result api.HTTPRequestResult) (entries []variableEntry, expectsVariables bool) {
 	seen := map[string]bool{}
 

render/variables_test.go

@@ -12,12 +12,17 @@ func TestHTTPVariableSections(t *testing.T) {
 		Variables: map[string]string{
 			"authToken": "token-123",
 			"shortCode": "abc123",
+			"sessionID": "session-123",
 		},
 		Request: api.CLIStepHTTPRequest{
 			ResponseVariables: []api.HTTPRequestResponseVariable{
 				{Name: "shortCode", Path: ".short_code"},
 				{Name: "missingCode", Path: ".missing_code"},
 			},
+			ResponseHeaderVariables: []api.HTTPRequestResponseHeaderVariable{
+				{Name: "sessionID", Header: "Set-Cookie", Regex: "session_id=([^;]+)"},
+				{Name: "missingSessionID", Header: "Set-Cookie", Regex: "missing=([^;]+)"},
+			},
 			Request: api.HTTPRequest{
 				FullURL: "${baseURL}/api/links/${shortCode}",
 				Headers: map[string]string{
@@ -37,9 +42,11 @@ func TestHTTPVariableSections(t *testing.T) {
 
 	wantContains := []string{
 		"Variables Saved:",
+		"sessionID: session-123 (Response Header Set-Cookie matching session_id=([^;]+))",
 		"shortCode: abc123 (JSON Body .short_code)",
 		"Variables Missing:",
 		"missingCode: [not found] (JSON Body .missing_code)",
+		"missingSessionID: [not found] (Response Header Set-Cookie matching missing=([^;]+))",
 		"Variables Available:",
 		"authToken: token-123 (Request Header \"Authorization\")",
 		"shortCode: abc123 (Request URL)",

render/view.go

@@ -151,9 +151,12 @@ func (m rootModel) View() string {
 	}
 
 	if m.result == api.VerificationResultSlugSuccess && m.isSubmit {
-		str.WriteString("\n\n" + green.Render("All tests passed! 🎉") + "\n")
+		str.WriteByte('\n')
+		str.WriteByte('\n')
+		str.WriteString(green.Render("All tests passed! 🎉"))
+		str.WriteByte('\n')
 		if m.xpReward >= 0 {
-			str.WriteString("\n")
+			str.WriteByte('\n')
 			str.WriteString(green.Bold(true).Render(fmt.Sprintf("Gained +%d XP", m.xpReward)))
 			str.WriteByte('\n')
 			for _, item := range m.xpBreakdown {
@@ -174,26 +177,36 @@ func (m rootModel) View() string {
 				str.WriteByte('\n')
 			}
 		}
-		str.WriteString("\n" + green.Render("Return to your browser to continue with the next lesson.") + "\n\n")
+		str.WriteByte('\n')
+		str.WriteString(green.Render("Return to your browser to continue with the next lesson."))
+		str.WriteByte('\n')
+		str.WriteByte('\n')
 	} else if m.result == api.VerificationResultSlugNoop {
 		str.WriteString("\n\nTests failed! ❌")
 		fmt.Fprintf(&str, "\n\nFailed Step: %v", m.failure.FailedStepIndex+1)
-		str.WriteString("\nError: " + m.failure.ErrorMessage + "\n")
-		str.WriteString("\n" + white.Render(safeStepIcon) + " This was a safe step.\n")
+		str.WriteString("\nError: ")
+		str.WriteString(m.failure.ErrorMessage)
+		str.WriteByte('\n')
+		str.WriteByte('\n')
+		str.WriteString(white.Render(safeStepIcon))
+		str.WriteString(" This was a safe step.\n")
 		str.WriteString("You haven't passed, but you also haven't lost armor or Sharpshooter progress.\n\n")
 	} else if m.result == api.VerificationResultSlugFailure {
-		str.WriteString("\n\n" + red.Render("Tests failed! ❌"))
+		str.WriteByte('\n')
+		str.WriteByte('\n')
+		str.WriteString(red.Render("Tests failed! ❌"))
 		if m.failure != nil {
 			if m.failure.FailedStepIndex >= 0 && m.failure.FailedStepIndex < len(m.steps) {
-				fmt.Fprintf(&str, "%s", red.Render(fmt.Sprintf("\n\nFailed Command: %s", m.steps[m.failure.FailedStepIndex].step)))
+				str.WriteString(red.Render(fmt.Sprintf("\n\nFailed Command: %s", m.steps[m.failure.FailedStepIndex].step)))
 			}
-			fmt.Fprintf(&str, "%s", red.Render(fmt.Sprintf("\n\nFailed Step: %v", m.failure.FailedStepIndex+1)))
-			fmt.Fprintf(&str, "%s", red.Render("\nError: "+m.failure.ErrorMessage))
+			str.WriteString(red.Render(fmt.Sprintf("\n\nFailed Step: %v", m.failure.FailedStepIndex+1)))
+			str.WriteString(red.Render(fmt.Sprintf("\nError: %s", m.failure.ErrorMessage)))
 		} else {
-			fmt.Fprintf(&str, "%s", red.Render("\n\nFailed Step: unknown"))
-			fmt.Fprintf(&str, "%s", red.Render("\nError: unknown"))
+			str.WriteString(red.Render("\n\nFailed Step: unknown"))
+			str.WriteString(red.Render("\nError: unknown"))
 		}
-		str.WriteString("\n\n")
+		str.WriteByte('\n')
+		str.WriteByte('\n')
 		currentDate := time.Now().Format("2006-01-02")
 		if strings.HasSuffix(currentDate, "04-01") {
 			str.WriteString(magenta.Render(fmt.Sprintf("This incident has been reported to your system administrator. [%s]\n", currentDate)))