NixMsg/internal/admin/admin_test.go

package admin_test

import (
	"bytes"
	"context"
	"encoding/json"
	"net/http"
	"net/http/cookiejar"
	"net/http/httptest"
	"net/url"
	"path/filepath"
	"strings"
	"testing"

	"git.asio.asia/nixevol/NixMsg/internal/admin"
	"git.asio.asia/nixevol/NixMsg/internal/auth"
	"git.asio.asia/nixevol/NixMsg/internal/store"
)

const testPassword = "admin-password-ok"

type envelope struct {
	OK    bool            `json:"ok"`
	Data  json.RawMessage `json:"data"`
	Error *struct {
		Code    string `json:"code"`
		Message string `json:"message"`
	} `json:"error"`
}

func setup(t *testing.T) (*admin.Handler, *httptest.Server, *http.Client, auth.HashPool) {
	t.Helper()
	dir := t.TempDir()
	db, err := store.Open(filepath.Join(dir, "data"), "FULL")
	if err != nil {
		t.Fatal(err)
	}
	t.Cleanup(func() { _ = db.Close() })

	hash := auth.NewStubHashPool()
	if seedErr := admin.SeedAdminPassword(context.Background(), db, hash, testPassword); seedErr != nil {
		t.Fatal(seedErr)
	}
	locks := admin.NewMemoryLoginLocks()
	h := admin.New(admin.Deps{
		DB:     db,
		Hash:   hash,
		Tokens: admin.NewRandomAPITokens(),
		Locks:  locks,
	})
	srv := httptest.NewServer(h)
	t.Cleanup(srv.Close)

	jar, err := cookiejar.New(nil)
	if err != nil {
		t.Fatal(err)
	}
	client := &http.Client{Jar: jar}
	return h, srv, client, hash
}

func decodeEnv(t *testing.T, res *http.Response) envelope {
	t.Helper()
	defer func() { _ = res.Body.Close() }()
	var env envelope
	if err := json.NewDecoder(res.Body).Decode(&env); err != nil {
		t.Fatal(err)
	}
	return env
}

func postJSON(t *testing.T, client *http.Client, url, body string, headers map[string]string) *http.Response {
	t.Helper()
	req, err := http.NewRequest(http.MethodPost, url, strings.NewReader(body))
	if err != nil {
		t.Fatal(err)
	}
	req.Header.Set("Content-Type", "application/json")
	for k, v := range headers {
		req.Header.Set(k, v)
	}
	res, err := client.Do(req)
	if err != nil {
		t.Fatal(err)
	}
	return res
}

func doReq(t *testing.T, client *http.Client, method, rawURL, body string, headers map[string]string) *http.Response {
	t.Helper()
	req, err := http.NewRequest(method, rawURL, strings.NewReader(body))
	if err != nil {
		t.Fatal(err)
	}
	if body != "" {
		req.Header.Set("Content-Type", "application/json")
	}
	for k, v := range headers {
		req.Header.Set(k, v)
	}
	res, err := client.Do(req)
	if err != nil {
		t.Fatal(err)
	}
	return res
}

func login(t *testing.T, client *http.Client, base string) {
	t.Helper()
	res := postJSON(t, client, base+"/api/admin/login",
		`{"username":"admin","password":"`+testPassword+`"}`, nil)
	env := decodeEnv(t, res)
	if res.StatusCode != 200 || !env.OK {
		t.Fatalf("login: status=%d env=%+v", res.StatusCode, env)
	}
}

func TestLoginMePasswordLogout(t *testing.T) {
	_, srv, client, _ := setup(t)
	base := srv.URL

	login(t, client, base)

	res := doReq(t, client, http.MethodGet, base+"/api/admin/me", "", nil)
	env := decodeEnv(t, res)
	if res.StatusCode != 200 || !env.OK {
		t.Fatalf("me: %d %+v", res.StatusCode, env)
	}
	var me map[string]any
	_ = json.Unmarshal(env.Data, &me)
	if me["username"] != "admin" || me["auth"] != "cookie" {
		t.Fatalf("me data=%v", me)
	}

	res = postJSON(t, client, base+"/api/admin/password",
		`{"old_password":"`+testPassword+`","new_password":"new-password-12"}`,
		map[string]string{"X-Nixmsg-Request": "1"})
	env = decodeEnv(t, res)
	if res.StatusCode != 200 || !env.OK {
		t.Fatalf("password: %d %+v", res.StatusCode, env)
	}

	res = postJSON(t, client, base+"/api/admin/logout", `{}`,
		map[string]string{"X-Nixmsg-Request": "1"})
	env = decodeEnv(t, res)
	if res.StatusCode != 200 || !env.OK {
		t.Fatalf("logout: %d %+v", res.StatusCode, env)
	}

	res = doReq(t, client, http.MethodGet, base+"/api/admin/me", "", nil)
	env = decodeEnv(t, res)
	if res.StatusCode != 401 {
		t.Fatalf("after logout want 401 got %d", res.StatusCode)
	}

	// 用新密码再登录
	res = postJSON(t, client, base+"/api/admin/login",
		`{"username":"admin","password":"new-password-12"}`, nil)
	env = decodeEnv(t, res)
	if res.StatusCode != 200 || !env.OK {
		t.Fatalf("relogin: %d %+v", res.StatusCode, env)
	}
}

func TestCSRFRequiredForCookieMutating(t *testing.T) {
	_, srv, client, _ := setup(t)
	base := srv.URL
	login(t, client, base)

	res := postJSON(t, client, base+"/api/admin/password",
		`{"old_password":"`+testPassword+`","new_password":"new-password-12"}`,
		nil) // 无 CSRF 头
	env := decodeEnv(t, res)
	if res.StatusCode != 403 || env.Error == nil || env.Error.Code != "forbidden" {
		t.Fatalf("want 403 forbidden, got %d %+v", res.StatusCode, env)
	}
}

func TestAPITokenAuthAndRestrictions(t *testing.T) {
	_, srv, client, _ := setup(t)
	base := srv.URL
	login(t, client, base)

	res := postJSON(t, client, base+"/api/admin/tokens",
		`{"name":"ops"}`,
		map[string]string{"X-Nixmsg-Request": "1"})
	env := decodeEnv(t, res)
	if res.StatusCode != 200 || !env.OK {
		t.Fatalf("create token: %d %+v", res.StatusCode, env)
	}
	var created struct {
		ID    string `json:"id"`
		Token string `json:"token"`
		Name  string `json:"name"`
	}
	if err := json.Unmarshal(env.Data, &created); err != nil {
		t.Fatal(err)
	}
	if !strings.HasPrefix(created.Token, "nxm_") {
		t.Fatalf("token prefix: %q", created.Token)
	}

	tokClient := &http.Client{}
	hdr := map[string]string{"Authorization": "Bearer " + created.Token}

	res = doReq(t, tokClient, http.MethodGet, base+"/api/admin/me", "", hdr)
	env = decodeEnv(t, res)
	if res.StatusCode != 200 || !env.OK {
		t.Fatalf("token me: %d %+v", res.StatusCode, env)
	}
	var me map[string]any
	_ = json.Unmarshal(env.Data, &me)
	if me["auth"] != "token" {
		t.Fatalf("auth=%v", me["auth"])
	}

	// 普通管理接口鉴权通过（业务 501）
	res = doReq(t, tokClient, http.MethodGet, base+"/api/admin/overview", "", hdr)
	env = decodeEnv(t, res)
	if res.StatusCode != http.StatusNotImplemented {
		t.Fatalf("overview want 501 got %d %+v", res.StatusCode, env)
	}

	// 禁止 password / tokens
	res = postJSON(t, tokClient, base+"/api/admin/password",
		`{"old_password":"x","new_password":"new-password-12"}`, hdr)
	env = decodeEnv(t, res)
	if res.StatusCode != 403 {
		t.Fatalf("token password want 403 got %d", res.StatusCode)
	}

	res = doReq(t, tokClient, http.MethodGet, base+"/api/admin/tokens", "", hdr)
	env = decodeEnv(t, res)
	if res.StatusCode != 403 {
		t.Fatalf("token list want 403 got %d", res.StatusCode)
	}

	// 停用后立即失效
	res = doReq(t, client, http.MethodPatch, base+"/api/admin/tokens/"+created.ID,
		`{"enabled":false}`,
		map[string]string{"X-Nixmsg-Request": "1", "Content-Type": "application/json"})
	env = decodeEnv(t, res)
	if res.StatusCode != 200 || !env.OK {
		t.Fatalf("disable: %d %+v", res.StatusCode, env)
	}

	res = doReq(t, tokClient, http.MethodGet, base+"/api/admin/me", "", hdr)
	env = decodeEnv(t, res)
	if res.StatusCode != 401 {
		t.Fatalf("disabled token want 401 got %d %+v", res.StatusCode, env)
	}
}

func TestLoginLock(t *testing.T) {
	_, srv, _, _ := setup(t)
	base := srv.URL

	for i := 0; i < 9; i++ {
		client := &http.Client{}
		res := postJSON(t, client, base+"/api/admin/login",
			`{"username":"admin","password":"wrong-password!!"}`, nil)
		env := decodeEnv(t, res)
		if res.StatusCode != 401 {
			t.Fatalf("fail %d: want 401 got %d %+v", i, res.StatusCode, env)
		}
	}
	client := &http.Client{}
	res := postJSON(t, client, base+"/api/admin/login",
		`{"username":"admin","password":"wrong-password!!"}`, nil)
	env := decodeEnv(t, res)
	if res.StatusCode != 429 || env.Error == nil || env.Error.Code != "rate_limited" {
		t.Fatalf("want 429 rate_limited got %d %+v", res.StatusCode, env)
	}

	res = postJSON(t, client, base+"/api/admin/login",
		`{"username":"admin","password":"`+testPassword+`"}`, nil)
	env = decodeEnv(t, res)
	if res.StatusCode != 429 {
		t.Fatalf("locked correct login want 429 got %d", res.StatusCode)
	}
}

func TestBadAPITokenCountsTowardLock(t *testing.T) {
	_, srv, _, _ := setup(t)
	base := srv.URL
	tokClient := &http.Client{}
	hdr := map[string]string{"Authorization": "Bearer nxm_" + strings.Repeat("a", 43)}

	for i := 0; i < 9; i++ {
		res := doReq(t, tokClient, http.MethodGet, base+"/api/admin/me", "", hdr)
		env := decodeEnv(t, res)
		if res.StatusCode != 401 {
			t.Fatalf("bad token %d: want 401 got %d %+v", i, res.StatusCode, env)
		}
	}
	res := doReq(t, tokClient, http.MethodGet, base+"/api/admin/me", "", hdr)
	env := decodeEnv(t, res)
	if res.StatusCode != 429 {
		t.Fatalf("want lock 429 got %d %+v", res.StatusCode, env)
	}
}

func TestCookieSetAttributes(t *testing.T) {
	_, srv, client, _ := setup(t)
	base := srv.URL
	res := postJSON(t, client, base+"/api/admin/login",
		`{"username":"admin","password":"`+testPassword+`"}`, nil)
	_ = decodeEnv(t, res)

	u, _ := url.Parse(base)
	cookies := client.Jar.Cookies(u)
	found := false
	for _, c := range cookies {
		if c.Name == "nixmsg_admin" {
			found = true
			if c.Value == "" {
				t.Fatal("empty cookie")
			}
		}
	}
	if !found {
		t.Fatal("cookie not set")
	}
}

func TestMountableHandler(t *testing.T) {
	dir := t.TempDir()
	db, err := store.Open(filepath.Join(dir, "data"), "FULL")
	if err != nil {
		t.Fatal(err)
	}
	t.Cleanup(func() { _ = db.Close() })
	hash := auth.NewStubHashPool()
	_ = admin.SeedAdminPassword(context.Background(), db, hash, testPassword)
	h := admin.New(admin.Deps{
		DB:     db,
		Hash:   hash,
		Tokens: admin.NewRandomAPITokens(),
		Locks:  admin.NewMemoryLoginLocks(),
	})
	mux := http.NewServeMux()
	mux.Handle("/api/admin/", h)
	srv := httptest.NewServer(mux)
	t.Cleanup(srv.Close)

	body := `{"username":"admin","password":"` + testPassword + `"}`
	res, err := http.Post(srv.URL+"/api/admin/login", "application/json", bytes.NewBufferString(body))
	if err != nil {
		t.Fatal(err)
	}
	env := decodeEnv(t, res)
	if !env.OK {
		t.Fatalf("mount login failed: %+v", env)
	}
}