package main

import (
	"archive/tar"
	"bytes"
	"encoding/json"
	"mime/multipart"
	"net/http"
	"net/http/httptest"
	"net/url"
	"os"
	"path/filepath"
	"strings"
	"testing"
)

type entry struct {
	name string
	body string
	typ  byte
	link string
}

func makeBundle(t *testing.T, entries ...entry) []byte {
	t.Helper()
	var buf bytes.Buffer
	tw := tar.NewWriter(&buf)
	for _, e := range entries {
		typ := e.typ
		if typ == 0 {
			typ = tar.TypeReg
		}
		hdr := &tar.Header{Name: e.name, Typeflag: typ, Mode: 0o644, Size: int64(len(e.body)), Linkname: e.link}
		if typ != tar.TypeReg {
			hdr.Size = 0
		}
		if err := tw.WriteHeader(hdr); err != nil {
			t.Fatal(err)
		}
		if typ == tar.TypeReg {
			tw.Write([]byte(e.body))
		}
	}
	tw.Close()
	return buf.Bytes()
}

const hash32 = "00fs4zz9c0jjcvs62wg39b91dydz7zqb"

func goodBundle(t *testing.T) []byte {
	return makeBundle(t,
		entry{name: "latest", body: `{"release":"/nix/store/x","version":"0.3.0","serial":2000,"schema":0}`},
		entry{name: "cache/", typ: tar.TypeDir},
		entry{name: "cache/log/", typ: tar.TypeDir},
		entry{name: "cache/nix-cache-info", body: "StoreDir: /nix/store\n"},
		entry{name: "cache/" + hash32 + ".narinfo", body: "StorePath: ..."},
		entry{name: "cache/nar/" + strings.Repeat("a1", 26) + ".nar.zst", body: "nar"},
	)
}

func TestUnbundle(t *testing.T) {
	dir := filepath.Join(t.TempDir(), "staged")
	os.MkdirAll(dir, 0o700)
	os.WriteFile(filepath.Join(dir, "old"), []byte("x"), 0o600)
	if err := unbundle(bytes.NewReader(goodBundle(t)), dir); err != nil {
		t.Fatal(err)
	}
	for _, f := range []string{"latest", "cache/nix-cache-info", "cache/" + hash32 + ".narinfo"} {
		if _, err := os.Stat(filepath.Join(dir, f)); err != nil {
			t.Fatalf("%s: %v", f, err)
		}
	}
	if _, err := os.Stat(filepath.Join(dir, "old")); err == nil {
		t.Fatal("the previous staging was kept")
	}
}

func TestUnbundleRefusesHostileBundles(t *testing.T) {
	good := []entry{
		{name: "latest", body: "{}"},
		{name: "cache/nix-cache-info", body: "x"},
	}
	for name, extra := range map[string]entry{
		"traversal":   {name: "../../etc/passwd", body: "x"},
		"absolute":    {name: "/etc/passwd", body: "x"},
		"nested":      {name: "cache/nar/../../../x", body: "x"},
		"other file":  {name: "cache/evil.sh", body: "x"},
		"symlink":     {name: "cache/" + hash32 + ".narinfo", typ: tar.TypeSymlink, link: "/etc/passwd"},
		"hardlink":    {name: "cache/" + hash32 + ".narinfo", typ: tar.TypeLink, link: "/etc/passwd"},
		"directory":   {name: "etc/", typ: tar.TypeDir},
		"duplicate":   {name: "latest", body: "{}"},
		"huge latest": {name: "latest", body: strings.Repeat("x", maxLatest+1)},
	} {
		dir := filepath.Join(t.TempDir(), "staged")
		entries := append([]entry{}, good...)
		if name == "huge latest" {
			entries = good[1:] // so the size, not a duplicate, is what is refused
		}
		err := unbundle(bytes.NewReader(makeBundle(t, append(entries, extra)...)), dir)
		if err == nil {
			t.Errorf("%s: accepted", name)
		}
		if _, statErr := os.Stat(dir); statErr == nil {
			t.Errorf("%s: left %s behind", name, dir)
		}
		if _, statErr := os.Stat(dir + ".tmp"); statErr == nil {
			t.Errorf("%s: left a temporary directory", name)
		}
	}
	// Incomplete or not a tar at all.
	if err := unbundle(bytes.NewReader(makeBundle(t, entry{name: "latest", body: "{}"})), filepath.Join(t.TempDir(), "s")); err == nil {
		t.Error("a bundle without a cache was accepted")
	}
	if err := unbundle(strings.NewReader("this is not a tar file"), filepath.Join(t.TempDir(), "s")); err == nil {
		t.Error("garbage was accepted")
	}
}

func TestParseUpdateURL(t *testing.T) {
	for in, want := range map[string]string{
		"https://updates.example.com/router/": "https://updates.example.com/router",
		" http://192.168.1.5:8000 ":           "http://192.168.1.5:8000",
	} {
		if got, err := parseUpdateURL(in); err != nil || got != want {
			t.Errorf("%q: %q %v", in, got, err)
		}
	}
	for _, in := range []string{"", "ftp://x/y", "file:///etc", "https://user:pw@host/x", "https://host/x?y=1", "updates.example.com", "https://"} {
		if _, err := parseUpdateURL(in); err == nil {
			t.Errorf("%q accepted", in)
		}
	}
}

func TestCheckUpdate(t *testing.T) {
	a, sys := newTestApp(t)
	_ = a
	sys.ReleaseFile = filepath.Join(t.TempDir(), "router-release")
	os.WriteFile(sys.ReleaseFile, []byte(`{"version":"0.2.0","serial":1000,"schema":0}`), 0o600)
	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		switch r.URL.Path {
		case "/new/latest":
			w.Write([]byte(`{"release":"/nix/store/x","version":"0.3.0","serial":2000,"schema":1}`))
		case "/same/latest":
			w.Write([]byte(`{"version":"0.2.0","serial":1000}`))
		case "/junk/latest":
			w.Write([]byte(`hello`))
		default:
			http.NotFound(w, r)
		}
	}))
	defer srv.Close()
	if o, err := sys.checkUpdate(srv.URL + "/new"); err != nil || !o.Newer || o.Version != "0.3.0" {
		t.Fatalf("new: %+v %v", o, err)
	}
	if o, err := sys.checkUpdate(srv.URL + "/same"); err != nil || o.Newer {
		t.Fatalf("same: %+v %v", o, err)
	}
	for _, p := range []string{"/junk", "/missing"} {
		if _, err := sys.checkUpdate(srv.URL + p); err == nil {
			t.Fatalf("%s: no error", p)
		}
	}
}

func upload(t *testing.T, h http.Handler, token string, cookie *http.Cookie, bundle []byte, force bool) *httptest.ResponseRecorder {
	t.Helper()
	var body bytes.Buffer
	mw := multipart.NewWriter(&body)
	if force {
		mw.WriteField("force", "on")
	}
	fw, _ := mw.CreateFormFile("bundle", "router-update.bundle")
	fw.Write(bundle)
	mw.Close()
	req := httptest.NewRequest("POST", "/update/upload?csrf="+url.QueryEscape(token), &body)
	req.Header.Set("Content-Type", mw.FormDataContentType())
	if cookie != nil {
		req.AddCookie(cookie)
	}
	rec := httptest.NewRecorder()
	h.ServeHTTP(rec, req)
	return rec
}

func TestUpdateUploadAndInstall(t *testing.T) {
	a, sys := newTestApp(t)
	var started []string
	startUnit = func(unit string) error { started = append(started, unit); return nil }
	orig := startUnit
	defer func() { startUnit = orig }()
	h := a.routes()
	cookie := login(t, h)
	token := csrfFrom(t, do(h, "GET", "/update", nil, cookie).Body.String())

	if rec := upload(t, h, "wrong", cookie, goodBundle(t), false); rec.Code != http.StatusForbidden {
		t.Fatalf("upload with a wrong token: %d", rec.Code)
	}
	if rec := upload(t, h, token, nil, goodBundle(t), false); rec.Code != http.StatusSeeOther || !strings.HasPrefix(rec.Header().Get("Location"), "/login") {
		t.Fatalf("upload without a session: %d %s", rec.Code, rec.Header().Get("Location"))
	}
	if len(started) != 0 {
		t.Fatal("the update unit was started by a rejected upload")
	}

	rec := upload(t, h, token, cookie, makeBundle(t, entry{name: "evil", body: "x"}), false)
	if loc := rec.Header().Get("Location"); !strings.Contains(loc, "Rejected") || len(started) != 0 {
		t.Fatalf("hostile bundle: %q, started %v", loc, started)
	}

	rec = upload(t, h, token, cookie, goodBundle(t), true)
	if loc := rec.Header().Get("Location"); !strings.Contains(loc, "Bundle+received") {
		t.Fatalf("upload: %d %q", rec.Code, loc)
	}
	if len(started) != 1 || started[0] != "router-update.service" {
		t.Fatalf("started %v", started)
	}
	var req struct {
		Mode  string `json:"mode"`
		Force bool   `json:"force"`
	}
	data, _ := os.ReadFile(filepath.Join(sys.updateDir(), "request.json"))
	if err := json.Unmarshal(data, &req); err != nil || req.Mode != "staged" || !req.Force {
		t.Fatalf("request: %s", data)
	}
	if _, err := os.Stat(filepath.Join(sys.updateDir(), "staged", "latest")); err != nil {
		t.Fatalf("not staged: %v", err)
	}

	// Pulling: the saved server, through install.
	if rec := do(h, "POST", "/update/url", url.Values{"csrf": {token}, "url": {"ftp://nope"}}, cookie); !strings.Contains(rec.Header().Get("Location"), "Not+saved") {
		t.Fatalf("bad url: %q", rec.Header().Get("Location"))
	}
	do(h, "POST", "/update/url", url.Values{"csrf": {token}, "url": {"http://192.168.10.60:8000/"}}, cookie)
	if sys.updateURL() != "http://192.168.10.60:8000" {
		t.Fatalf("url: %q", sys.updateURL())
	}
	do(h, "POST", "/update/install", url.Values{"csrf": {token}}, cookie)
	data, _ = os.ReadFile(filepath.Join(sys.updateDir(), "request.json"))
	if !strings.Contains(string(data), `"mode":"url"`) || !strings.Contains(string(data), "192.168.10.60:8000") {
		t.Fatalf("request: %s", data)
	}
	if page := do(h, "GET", "/update", nil, cookie).Body.String(); !strings.Contains(page, "192.168.10.60:8000") {
		t.Fatalf("the page lacks the saved server:\n%s", page)
	}
}