package main

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

func TestRemoteImageMiddlewareProxiesPNG(t *testing.T) {
	upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		w.Header().Set("Content-Type", "image/png")
		_, _ = w.Write([]byte("\x89PNG\r\n"))
	}))
	defer upstream.Close()

	h := remoteImageMiddleware(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		t.Fatal("next should not run for /__ccimg")
	}))
	req := httptest.NewRequest(http.MethodGet, "/__ccimg?u="+upstream.URL+"/a.png", nil)
	rec := httptest.NewRecorder()
	h.ServeHTTP(rec, req)
	if rec.Code != 200 {
		t.Fatalf("status %d: %s", rec.Code, rec.Body.String())
	}
	if !strings.HasPrefix(rec.Header().Get("Content-Type"), "image/png") {
		t.Fatalf("content-type %q", rec.Header().Get("Content-Type"))
	}
	body, _ := io.ReadAll(rec.Body)
	if len(body) < 4 || string(body[:4]) != "\x89PNG" {
		t.Fatalf("body not proxied: %q", body)
	}
}

func TestRemoteImageMiddlewareRejectsBadURL(t *testing.T) {
	h := remoteImageMiddleware(http.NotFoundHandler())
	req := httptest.NewRequest(http.MethodGet, "/__ccimg?u=file:///etc/passwd", nil)
	rec := httptest.NewRecorder()
	h.ServeHTTP(rec, req)
	if rec.Code != http.StatusBadRequest {
		t.Fatalf("want 400, got %d", rec.Code)
	}
}