package middleware

import (
	"net/http"
	"net/http/httptest"
	"strings"
	"testing"

	"github.com/gorilla/websocket"
)

func TestProtectRefreshedTokenResponsesPreservesWebSocketHijacking(t *testing.T) {
	upgradeResult := make(chan error, 1)
	upgrader := websocket.Upgrader{CheckOrigin: func(*http.Request) bool { return true }}
	handler := ProtectRefreshedTokenResponses(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		setRefreshedToken(r.Context(), "must-not-send")
		conn, err := upgrader.Upgrade(w, r, nil)
		if err != nil {
			upgradeResult <- err
			return
		}
		upgradeResult <- conn.Close()
	}))
	server := httptest.NewServer(handler)
	defer server.Close()

	wsURL := "ws" + strings.TrimPrefix(server.URL, "http")
	conn, response, err := websocket.DefaultDialer.Dial(wsURL, nil)
	if err != nil {
		t.Fatalf("WebSocket dial failed: %v", err)
	}
	defer conn.Close()
	if response.Body != nil {
		defer response.Body.Close()
	}
	if response.StatusCode != http.StatusSwitchingProtocols {
		t.Fatalf("status = %d, want %d", response.StatusCode, http.StatusSwitchingProtocols)
	}
	if got := response.Header.Get(HeaderRefreshedToken); got != "" {
		t.Fatalf("%s = %q, want empty on hijacked response", HeaderRefreshedToken, got)
	}
	if err := <-upgradeResult; err != nil {
		t.Fatalf("WebSocket handler failed: %v", err)
	}
}

func TestProtectRefreshedTokenResponsesProtectsFlushCommit(t *testing.T) {
	recorder := httptest.NewRecorder()
	handler := ProtectRefreshedTokenResponses(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		setRefreshedToken(r.Context(), "renewed-token")
		w.Header()["x-refreshed-token"] = []string{"spoofed-token"}
		w.Header()["cache-control"] = []string{"public, max-age=3600"}
		w.(http.Flusher).Flush()
	}))

	handler.ServeHTTP(recorder, httptest.NewRequest(http.MethodGet, "/", nil))

	if !recorder.Flushed {
		t.Fatal("response was not flushed")
	}
	if got := headerValuesEqualFold(recorder.Result().Header, HeaderRefreshedToken); len(got) != 1 || got[0] != "renewed-token" {
		t.Fatalf("%s values = %q, want [renewed-token]", HeaderRefreshedToken, got)
	}
	if got := headerValuesEqualFold(recorder.Result().Header, headerCacheControl); len(got) != 1 || got[0] != cacheControlNoStore {
		t.Fatalf("%s values = %q, want [%s]", headerCacheControl, got, cacheControlNoStore)
	}
}