mirror of
https://github.com/Buriburizaem0n/nezha_domains.git
synced 2026-09-19 09:40:12 +00:00
feat(agentcompat): expose dashboard capability routes
Co-authored-by: naiba/CloudCode <hi+cloudcode@nai.ba>
This commit is contained in:
naiba
naiba/CloudCode
@@ -0,0 +1,245 @@
|
||||
//go:build agentcompat
|
||||
|
||||
package rpc
|
||||
|
||||
import (
|
||||
"context"
|
||||
"io"
|
||||
"net/http"
|
||||
"net/url"
|
||||
"sync"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/goccy/go-json"
|
||||
"github.com/nezhahq/nezha/model"
|
||||
"github.com/nezhahq/nezha/pkg/agentcompatcontract"
|
||||
"github.com/nezhahq/nezha/proto"
|
||||
serviceRPC "github.com/nezhahq/nezha/service/rpc"
|
||||
"github.com/nezhahq/nezha/service/singleton"
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
func TestServeNATAgentCompatReleasesCapabilityBeforeServerDispatch(t *testing.T) {
|
||||
fixture := newServeNATFixture(t)
|
||||
capability, access := registerNATCapabilityForTest(t, fixture.handler, fixture.server.ID, 101)
|
||||
originalServers := singleton.ServerShared
|
||||
singleton.ServerShared = singleton.NewServerClass()
|
||||
t.Cleanup(func() { singleton.ServerShared = originalServers })
|
||||
request := capabilityRequest(capability)
|
||||
|
||||
ServeNAT(&serveNATResponseWriter{}, request, &model.NAT{Common: model.Common{ID: 101}, ServerID: fixture.server.ID, Host: "target.example"})
|
||||
|
||||
assertNATCapabilitySlotReleased(t, fixture.handler, access)
|
||||
require.Equal(t, 0, fixture.handler.StreamCount())
|
||||
require.Empty(t, fixture.taskStream.sent)
|
||||
}
|
||||
|
||||
func TestServeNATAgentCompatTaskStreamUnavailableReleasesCapabilityBeforeActivity(t *testing.T) {
|
||||
fixture := newServeNATFixture(t)
|
||||
capability, access := registerNATCapabilityForTest(t, fixture.handler, fixture.server.ID, 105)
|
||||
fixture.server.SetTaskStream(nil)
|
||||
request := capabilityRequest(capability)
|
||||
|
||||
ServeNAT(&serveNATResponseWriter{}, request, &model.NAT{Common: model.Common{ID: 105}, ServerID: fixture.server.ID, Host: "target.example"})
|
||||
|
||||
assertNATCapabilitySlotReleased(t, fixture.handler, access)
|
||||
require.Equal(t, 0, fixture.handler.StreamCount())
|
||||
require.Empty(t, fixture.taskStream.sent)
|
||||
if request.Header.Get(agentcompatcontract.IOStreamCapabilityHeader) != "" {
|
||||
t.Fatal("capability header remained after task stream unavailable")
|
||||
}
|
||||
}
|
||||
|
||||
func TestServeNATAgentCompatCreateQuotaFailurePreservesExistingStreams(t *testing.T) {
|
||||
fixture := newServeNATFixture(t)
|
||||
capability, access := registerNATCapabilityForTest(t, fixture.handler, fixture.server.ID, 106)
|
||||
const existingStreamCount = 40
|
||||
for index := 0; index < existingStreamCount; index++ {
|
||||
require.NoError(t, fixture.handler.CreateStreamWithPurpose("quota-existing-"+string(rune(index)), 0, fixture.server.ID, serviceRPC.PurposeNAT))
|
||||
}
|
||||
t.Cleanup(func() {
|
||||
for index := 0; index < existingStreamCount; index++ {
|
||||
_ = fixture.handler.CloseStream("quota-existing-" + string(rune(index)))
|
||||
}
|
||||
})
|
||||
request := capabilityRequest(capability)
|
||||
|
||||
ServeNAT(&serveNATResponseWriter{}, request, &model.NAT{Common: model.Common{ID: 106}, ServerID: fixture.server.ID, Host: "target.example"})
|
||||
|
||||
assertNATCapabilitySlotReleased(t, fixture.handler, access)
|
||||
require.Equal(t, existingStreamCount, fixture.handler.StreamCount())
|
||||
require.Empty(t, fixture.taskStream.sent)
|
||||
}
|
||||
|
||||
func TestServeNATAgentCompatPublishCancellationClosesRegistryOwnedEndpoints(t *testing.T) {
|
||||
fixture := newServeNATFixture(t)
|
||||
capability, access := registerNATCapabilityForTest(t, fixture.handler, fixture.server.ID, 107)
|
||||
connection := newServeNATConn()
|
||||
agent := &serveNATAgent{readErr: io.EOF, writeDone: make(chan struct{})}
|
||||
observerEntered := make(chan struct{})
|
||||
observerRelease := make(chan struct{})
|
||||
var observerOnce sync.Once
|
||||
fixture.handler.SetAgentCompatCapabilityPublishObserverForTest(func() {
|
||||
observerOnce.Do(func() { close(observerEntered) })
|
||||
<-observerRelease
|
||||
})
|
||||
t.Cleanup(func() { fixture.handler.SetAgentCompatCapabilityPublishObserverForTest(nil) })
|
||||
fixture.taskStream.onSend = func(task *proto.Task) error {
|
||||
var nat model.TaskNAT
|
||||
require.NoError(t, json.Unmarshal([]byte(task.Data), &nat))
|
||||
return fixture.handler.AgentConnected(nat.StreamID, agent)
|
||||
}
|
||||
request := capabilityRequest(capability)
|
||||
writer := &serveNATResponseWriter{conn: connection}
|
||||
serveDone := make(chan struct{})
|
||||
go func() {
|
||||
ServeNAT(writer, request, &model.NAT{Common: model.Common{ID: 107}, ServerID: fixture.server.ID, Host: "target.example"})
|
||||
close(serveDone)
|
||||
}()
|
||||
select {
|
||||
case <-observerEntered:
|
||||
case <-serveDone:
|
||||
t.Fatal("ServeNAT returned before publish pause")
|
||||
}
|
||||
fixture.handler.CreateStreamWithPurpose("publish-replacement", 0, fixture.server.ID, serviceRPC.PurposeNAT)
|
||||
require.NoError(t, fixture.handler.CancelAgentCompatIOStreamCapability(access))
|
||||
close(observerRelease)
|
||||
completionContext, cancel := context.WithTimeout(context.Background(), time.Second)
|
||||
defer cancel()
|
||||
select {
|
||||
case <-serveDone:
|
||||
case <-completionContext.Done():
|
||||
t.Fatal("ServeNAT did not finish after publish cancellation")
|
||||
}
|
||||
assertCapabilityReleased(t, fixture.handler, access)
|
||||
require.Equal(t, int32(1), connection.closeCount.Load())
|
||||
require.Equal(t, int32(1), agent.closeCount.Load())
|
||||
require.Equal(t, 1, fixture.handler.StreamCount())
|
||||
select {
|
||||
case <-agent.writeDone:
|
||||
t.Fatal("agent endpoint was used after publish cancellation")
|
||||
default:
|
||||
}
|
||||
}
|
||||
|
||||
func TestServeNATAgentCompatSendTaskFailureReleasesStreamAndCapabilityBeforeWrapper(t *testing.T) {
|
||||
fixture := newServeNATFixture(t)
|
||||
capability, access := registerNATCapabilityForTest(t, fixture.handler, fixture.server.ID, 102)
|
||||
fixture.taskStream.sendErr = io.ErrClosedPipe
|
||||
connection := newServeNATConn()
|
||||
request := capabilityRequest(capability)
|
||||
request.Body = &serveNATBody{}
|
||||
writer := &serveNATResponseWriter{conn: connection}
|
||||
|
||||
ServeNAT(writer, request, &model.NAT{Common: model.Common{ID: 102}, ServerID: fixture.server.ID, Host: "target.example"})
|
||||
|
||||
assertCapabilityReleased(t, fixture.handler, access)
|
||||
require.Equal(t, 0, fixture.handler.StreamCount())
|
||||
require.Equal(t, int32(0), connection.closeCount.Load())
|
||||
}
|
||||
|
||||
func TestServeNATAgentCompatWrapperFailureReleasesCapabilityWithoutSecondResponse(t *testing.T) {
|
||||
fixture := newServeNATFixture(t)
|
||||
capability, access := registerNATCapabilityForTest(t, fixture.handler, fixture.server.ID, 103)
|
||||
request := capabilityRequest(capability)
|
||||
writer := &nonHijackingResponseWriter{}
|
||||
|
||||
ServeNAT(writer, request, &model.NAT{Common: model.Common{ID: 103}, ServerID: fixture.server.ID, Host: "target.example"})
|
||||
|
||||
assertCapabilityReleased(t, fixture.handler, access)
|
||||
require.Equal(t, 0, fixture.handler.StreamCount())
|
||||
require.Equal(t, 0, writer.writes)
|
||||
}
|
||||
|
||||
func TestServeNATAgentCompatStartFailureCancelsPublishedCapabilityAndClosesEndpoints(t *testing.T) {
|
||||
fixture := newServeNATFixture(t)
|
||||
capability, access := registerNATCapabilityForTest(t, fixture.handler, fixture.server.ID, 104)
|
||||
connection := newServeNATConn()
|
||||
agent := &serveNATAgent{readErr: io.EOF, writeDone: make(chan struct{})}
|
||||
fixture.taskStream.onSend = func(task *proto.Task) error {
|
||||
var nat model.TaskNAT
|
||||
if err := json.Unmarshal([]byte(task.Data), &nat); err != nil {
|
||||
return err
|
||||
}
|
||||
return fixture.handler.AgentConnected(nat.StreamID, agent)
|
||||
}
|
||||
request := capabilityRequest(capability)
|
||||
writer := &serveNATResponseWriter{conn: connection}
|
||||
|
||||
ServeNAT(writer, request, &model.NAT{Common: model.Common{ID: 104}, ServerID: fixture.server.ID, Host: "target.example"})
|
||||
|
||||
assertCapabilityReleased(t, fixture.handler, access)
|
||||
require.Equal(t, 0, fixture.handler.StreamCount())
|
||||
require.Equal(t, int32(1), connection.closeCount.Load())
|
||||
require.Equal(t, int32(1), agent.closeCount.Load())
|
||||
}
|
||||
|
||||
func registerNATCapabilityForTest(t *testing.T, handler *serviceRPC.NezhaHandler, serverID, resourceID uint64) (string, serviceRPC.AgentCompatCapabilityAccess) {
|
||||
t.Helper()
|
||||
capability, err := handler.RegisterAgentCompatIOStreamCapability(context.Background(), serviceRPC.AgentCompatCapabilityRegistration{
|
||||
Owner: serviceRPC.AgentCompatCapabilityOwner{PATID: resourceID, UserID: resourceID + 1}, Purpose: serviceRPC.AgentCompatCapabilityNAT,
|
||||
TargetServerID: serverID, ResourceID: resourceID, ServerAccessAllowed: true,
|
||||
})
|
||||
require.NoError(t, err)
|
||||
parsed, err := serviceRPC.ParseAgentCompatIOStreamCapability(capability.String())
|
||||
require.NoError(t, err)
|
||||
return capability.String(), serviceRPC.AgentCompatCapabilityAccess{
|
||||
Capability: parsed, Owner: serviceRPC.AgentCompatCapabilityOwner{PATID: resourceID, UserID: resourceID + 1},
|
||||
Purpose: serviceRPC.AgentCompatCapabilityNAT, TargetServerID: serverID, ResourceID: resourceID, ServerAccessAllowed: true,
|
||||
}
|
||||
}
|
||||
|
||||
func capabilityRequest(value string) *http.Request {
|
||||
request := &http.Request{Method: http.MethodPost, URL: &url.URL{Scheme: "http", Host: "example.test", Path: "/nat"}, Header: make(http.Header), Body: &serveNATBody{}}
|
||||
request.Header.Set(agentcompatcontract.IOStreamCapabilityHeader, value)
|
||||
return request
|
||||
}
|
||||
|
||||
func assertCapabilityReleased(t *testing.T, handler *serviceRPC.NezhaHandler, access serviceRPC.AgentCompatCapabilityAccess) {
|
||||
t.Helper()
|
||||
_, _, err := handler.ConsumeAgentCompatNATCapabilityForProfile(access.Capability.String(), access.TargetServerID, access.ResourceID)
|
||||
require.ErrorIs(t, err, serviceRPC.ErrAgentCompatCapabilityHidden)
|
||||
}
|
||||
|
||||
func assertNATCapabilitySlotReleased(t *testing.T, handler *serviceRPC.NezhaHandler, access serviceRPC.AgentCompatCapabilityAccess) {
|
||||
t.Helper()
|
||||
assertCapabilityReleased(t, handler, access)
|
||||
type activeCapability struct {
|
||||
capability serviceRPC.AgentCompatIOStreamCapability
|
||||
resourceID uint64
|
||||
}
|
||||
active := make([]activeCapability, 0, 16)
|
||||
for index := uint64(0); index < 16; index++ {
|
||||
capability, err := handler.RegisterAgentCompatIOStreamCapability(context.Background(), serviceRPC.AgentCompatCapabilityRegistration{
|
||||
Owner: access.Owner, Purpose: serviceRPC.AgentCompatCapabilityNAT, TargetServerID: access.TargetServerID,
|
||||
ResourceID: access.ResourceID + index + 1000, ServerAccessAllowed: true,
|
||||
})
|
||||
require.NoError(t, err)
|
||||
active = append(active, activeCapability{capability: capability, resourceID: access.ResourceID + index + 1000})
|
||||
}
|
||||
_, err := handler.RegisterAgentCompatIOStreamCapability(context.Background(), serviceRPC.AgentCompatCapabilityRegistration{
|
||||
Owner: access.Owner, Purpose: serviceRPC.AgentCompatCapabilityNAT, TargetServerID: access.TargetServerID,
|
||||
ResourceID: access.ResourceID + 2000, ServerAccessAllowed: true,
|
||||
})
|
||||
require.ErrorIs(t, err, serviceRPC.ErrAgentCompatCapabilityUnavailable)
|
||||
for _, item := range active {
|
||||
parsed, parseErr := serviceRPC.ParseAgentCompatIOStreamCapability(item.capability.String())
|
||||
require.NoError(t, parseErr)
|
||||
require.NoError(t, handler.CancelAgentCompatIOStreamCapability(serviceRPC.AgentCompatCapabilityAccess{
|
||||
Capability: parsed, Owner: access.Owner, Purpose: serviceRPC.AgentCompatCapabilityNAT,
|
||||
TargetServerID: access.TargetServerID, ResourceID: item.resourceID, ServerAccessAllowed: true,
|
||||
}))
|
||||
}
|
||||
}
|
||||
|
||||
type nonHijackingResponseWriter struct {
|
||||
writes int
|
||||
}
|
||||
|
||||
func (writer *nonHijackingResponseWriter) Header() http.Header { return make(http.Header) }
|
||||
func (writer *nonHijackingResponseWriter) Write(data []byte) (int, error) {
|
||||
writer.writes += len(data)
|
||||
return len(data), nil
|
||||
}
|
||||
func (*nonHijackingResponseWriter) WriteHeader(int) {}
|
||||
Reference in New Issue
Block a user