nezha_domains/cmd/dashboard/controller/force_update_ownership_test.go

package controller

import (
	"bytes"
	"encoding/json"
	"net/http"
	"net/http/httptest"
	"testing"

	"github.com/gin-gonic/gin"
	"github.com/stretchr/testify/assert"
	"gorm.io/driver/sqlite"
	"gorm.io/gorm"

	"github.com/nezhahq/nezha/model"
	"github.com/nezhahq/nezha/pkg/i18n"
	pb "github.com/nezhahq/nezha/proto"
	"github.com/nezhahq/nezha/service/singleton"
)

// fakeTaskStream is the minimum stub of pb.NezhaService_RequestTaskServer
// required to make a server look "online" to forceUpdateServer. Only Send is
// called; we capture its argument so the test can verify the upgrade task is
// NOT dispatched for foreign IDs.
type fakeTaskStream struct {
	pb.NezhaService_RequestTaskServer
	sentTasks []*pb.Task
}

func (f *fakeTaskStream) Send(t *pb.Task) error {
	f.sentTasks = append(f.sentTasks, t)
	return nil
}

// setupServerOwnershipFixture seeds an in-memory DB with alice's server
// (UserID=100, ID=1). The returned stream is wired in so the server is
// "online" — i.e. exercises the path that previously returned permission
// denied for foreign callers (the actual leak channel).
func setupServerOwnershipFixture(t *testing.T) (stream *fakeTaskStream, reset func()) {
	t.Helper()
	if singleton.Localizer == nil {
		singleton.Localizer = i18n.NewLocalizer("en_US", "nezha", "translations", i18n.Translations)
	}
	originalDB := singleton.DB
	originalShared := singleton.ServerShared

	db, err := gorm.Open(sqlite.Open(":memory:"), &gorm.Config{})
	assert.NoError(t, err)
	assert.NoError(t, db.AutoMigrate(&model.Server{}))
	assert.NoError(t, db.Create(&model.Server{
		Common: model.Common{ID: 1, UserID: 100},
		Name:   "alice-online",
	}).Error)
	singleton.DB = db
	singleton.ServerShared = singleton.NewServerClass()

	alice, _ := singleton.ServerShared.Get(1)
	stream = &fakeTaskStream{}
	alice.SetTaskStream(stream)

	return stream, func() {
		singleton.DB = originalDB
		singleton.ServerShared = originalShared
	}
}

func runForceUpdate(t *testing.T, callerID uint64, ids []uint64) []byte {
	t.Helper()
	r := gin.New()
	r.Use(func(c *gin.Context) {
		setAuthUser(c, callerID, model.RoleMember)
		c.Next()
	})
	r.POST("/force-update/server", commonHandler(forceUpdateServer))

	body, _ := json.Marshal(ids)
	w := httptest.NewRecorder()
	req := httptest.NewRequest(http.MethodPost, "/force-update/server", bytes.NewReader(body))
	req.Header.Set("Content-Type", "application/json")
	r.ServeHTTP(w, req)
	return w.Body.Bytes()
}

type forceUpdateBody struct {
	Success bool   `json:"success"`
	Error   string `json:"error"`
	Data    struct {
		Offline []uint64 `json:"offline"`
		Success []uint64 `json:"success"`
		Failure []uint64 `json:"failure"`
	} `json:"data"`
}

func decodeForceUpdate(t *testing.T, body []byte) forceUpdateBody {
	t.Helper()
	var resp forceUpdateBody
	assert.NoError(t, json.Unmarshal(body, &resp))
	return resp
}

// Core regression: bob submitting alice's online server ID must NOT produce
// a distinct response from bob submitting an unknown ID. The original code
// returned "permission denied" for the former and a structured success/Offline
// response for the latter — that delta is the enumeration oracle for server
// IDs and online state.
func TestForceUpdateServerOnlineForeignIDIndistinguishableFromUnknown(t *testing.T) {
	gin.SetMode(gin.TestMode)
	_, reset := setupServerOwnershipFixture(t)
	defer reset()

	const bobID = uint64(200)
	foreignResp := decodeForceUpdate(t, runForceUpdate(t, bobID, []uint64{1}))    // alice's online
	unknownResp := decodeForceUpdate(t, runForceUpdate(t, bobID, []uint64{9999})) // does not exist

	assert.Equal(t, foreignResp.Success, unknownResp.Success,
		"top-level success flag must not differ between foreign-online and unknown IDs")
	assert.Equal(t, foreignResp.Error, unknownResp.Error,
		"error string must not differ — distinct error reveals existence/state of foreign servers")
	assert.Equal(t, foreignResp.Data.Success, unknownResp.Data.Success)
	assert.Equal(t, foreignResp.Data.Failure, unknownResp.Data.Failure)
}

// Submitting a foreign online server must NOT actually trigger the upgrade
// task on it — that would be a write primitive on someone else's machine.
func TestForceUpdateServerForeignOnlineDoesNotDispatchUpgrade(t *testing.T) {
	gin.SetMode(gin.TestMode)
	stream, reset := setupServerOwnershipFixture(t)
	defer reset()

	_ = runForceUpdate(t, 200, []uint64{1}) // bob hits alice's online server
	assert.Empty(t, stream.sentTasks,
		"foreign server must not receive the upgrade task even when online")
}

// Sanity: owner submitting their own online server must still get the upgrade
// dispatched and a structured success response — the hardening must not
// regress the legitimate case.
func TestForceUpdateServerOwnerOnlineStillDispatches(t *testing.T) {
	gin.SetMode(gin.TestMode)
	stream, reset := setupServerOwnershipFixture(t)
	defer reset()

	resp := decodeForceUpdate(t, runForceUpdate(t, 100, []uint64{1})) // alice on her own server
	assert.True(t, resp.Success)
	assert.Equal(t, []uint64{1}, resp.Data.Success)
	assert.Empty(t, resp.Data.Offline)
	assert.Empty(t, resp.Data.Failure)
	assert.Len(t, stream.sentTasks, 1, "owner's own server must receive the upgrade task exactly once")
}