wehub-resource-sync
100 lines
4.2 KiB
Go
100 lines
4.2 KiB
Go
package distributed_test
|
|
|
|
import (
|
|
"time"
|
|
|
|
"github.com/mudler/LocalAI/core/services/messaging"
|
|
"github.com/mudler/LocalAI/pkg/natsauth"
|
|
|
|
. "github.com/onsi/ginkgo/v2"
|
|
. "github.com/onsi/gomega"
|
|
)
|
|
|
|
var _ = Describe("NATS JWT Auth", Label("Distributed", "NatsJWT"), func() {
|
|
var infra *JWTTestInfra
|
|
|
|
BeforeEach(func() {
|
|
infra = SetupJWTInfra()
|
|
})
|
|
|
|
It("connects with a minted backend worker JWT and publishes on allowed subjects", func() {
|
|
// Backend workers may publish under nodes.<id>.files.> (see pkg/natsauth permissions).
|
|
subject := nodeSubjectPrefix(infra.NodeID) + ".files.in"
|
|
Expect(infra.NC.Publish(subject, map[string]string{"path": "/tmp/model"})).To(Succeed())
|
|
Expect(infra.NC.Conn().FlushTimeout(2 * time.Second)).To(Succeed())
|
|
Expect(infra.NC.Conn().IsConnected()).To(BeTrue())
|
|
})
|
|
|
|
It("allows backend subscribe on the node prefix", func() {
|
|
wild := nodeSubjectPrefix(infra.NodeID) + ".>"
|
|
sub, err := infra.NC.Subscribe(wild, func(_ []byte) {})
|
|
Expect(err).ToNot(HaveOccurred())
|
|
defer func() { _ = sub.Unsubscribe() }()
|
|
Expect(infra.NC.Conn().FlushTimeout(2 * time.Second)).To(Succeed())
|
|
Expect(infra.NC.Conn().IsConnected()).To(BeTrue())
|
|
})
|
|
|
|
It("rejects anonymous publish on the JWT-enabled server", func() {
|
|
anon, err := messaging.New(infra.NatsURL)
|
|
Expect(err).ToNot(HaveOccurred())
|
|
defer anon.Close()
|
|
|
|
err = anon.Publish("nodes.any.files.x", map[string]string{"x": "1"})
|
|
Expect(err).ToNot(HaveOccurred())
|
|
Expect(anon.Conn().FlushTimeout(2 * time.Second)).To(HaveOccurred())
|
|
})
|
|
|
|
It("denies backend publish to another node's subjects", func() {
|
|
other := nodeSubjectPrefix("other-node-id") + ".files.stage"
|
|
Expect(infra.NC.Publish(other, map[string]string{"stage": "nope"})).To(Succeed())
|
|
Eventually(func() error {
|
|
_ = infra.NC.Conn().FlushTimeout(500 * time.Millisecond)
|
|
return infra.NC.Conn().LastError()
|
|
}, "3s", "50ms").Should(HaveOccurred())
|
|
})
|
|
|
|
It("mints agent JWT without backend.install in claims", func() {
|
|
cfg := natsauth.Config{AccountSeed: infra.AccountSeed}
|
|
token, _, err := cfg.MintWorkerJWT("agent-node-1", "agent")
|
|
Expect(err).ToNot(HaveOccurred())
|
|
|
|
claims, err := natsauth.DecodeUserClaims(token)
|
|
Expect(err).ToNot(HaveOccurred())
|
|
Expect(claims.Permissions.Sub.Allow).To(ContainElement("agent.execute"))
|
|
for _, subj := range claims.Permissions.Sub.Allow {
|
|
Expect(subj).NotTo(ContainSubstring("backend.install"))
|
|
}
|
|
})
|
|
|
|
// Regression guard for the silent permission gaps: decoding the JWT claims
|
|
// (above) only proves the agent JWT is *restrictive*, not that it is
|
|
// *sufficient*. Stand a real agent connection up against the enforcing
|
|
// server and exercise every subscription core/cli/agent_worker.go actually
|
|
// makes — a denied SUB now surfaces synchronously via confirmSubscription,
|
|
// so a missing allow rule fails this test instead of silently dropping
|
|
// backend.stop / MCP-CI deliveries at runtime.
|
|
It("lets an agent-minted JWT establish all the subscriptions the agent worker uses", func() {
|
|
const nodeID = "agent-node-subs"
|
|
cfg := natsauth.Config{AccountSeed: infra.AccountSeed, WorkerJWTTTL: time.Hour}
|
|
token, seed, err := cfg.MintWorkerJWT(nodeID, "agent")
|
|
Expect(err).ToNot(HaveOccurred())
|
|
|
|
nc, err := messaging.New(infra.NatsURL, messaging.WithUserJWT(token, seed))
|
|
Expect(err).ToNot(HaveOccurred())
|
|
DeferCleanup(nc.Close)
|
|
|
|
// Mirror core/cli/agent_worker.go exactly.
|
|
_, err = nc.QueueSubscribeReply(messaging.SubjectMCPToolExecute, messaging.QueueAgentWorkers, func([]byte, func([]byte)) {})
|
|
Expect(err).ToNot(HaveOccurred(), "agent JWT must allow %s", messaging.SubjectMCPToolExecute)
|
|
|
|
_, err = nc.QueueSubscribeReply(messaging.SubjectMCPDiscovery, messaging.QueueAgentWorkers, func([]byte, func([]byte)) {})
|
|
Expect(err).ToNot(HaveOccurred(), "agent JWT must allow %s", messaging.SubjectMCPDiscovery)
|
|
|
|
_, err = nc.QueueSubscribe(messaging.SubjectMCPCIJobsNew, messaging.QueueWorkers, func([]byte) {})
|
|
Expect(err).ToNot(HaveOccurred(), "agent JWT must allow %s (MCP CI jobs)", messaging.SubjectMCPCIJobsNew)
|
|
|
|
_, err = nc.Subscribe(messaging.SubjectNodeBackendStop(nodeID), func([]byte) {})
|
|
Expect(err).ToNot(HaveOccurred(), "agent JWT must allow %s (MCP session cleanup)", messaging.SubjectNodeBackendStop(nodeID))
|
|
})
|
|
})
|