mirror of
https://github.com/gravitational/teleport.git
synced 2026-10-12 02:15:01 +00:00
Add Classifier backend storage, cache, gRPC service, and tctl support (#67761)
* Add Classifier backend storage, cache, gRPC service, and tctl support * Fix formatting in lib/services/summarizer_test.go * Replace Cache.AllClassifiers with RangeClassifiers using genericLister * Validate that the filter is a boolean expression * Code review fixes: sync test, range fallback
This commit is contained in:
@@ -204,6 +204,10 @@ func EventToGRPC(in types.Event) (*proto.Event, error) {
|
||||
out.Resource = &proto.Event_InferencePolicy{
|
||||
InferencePolicy: r.UnwrapT(),
|
||||
}
|
||||
case types.Resource153UnwrapperT[*summaryv1.Classifier]:
|
||||
out.Resource = &proto.Event_Classifier{
|
||||
Classifier: r.UnwrapT(),
|
||||
}
|
||||
case types.Resource153UnwrapperT[*summaryv1.RetrievalModel]:
|
||||
out.Resource = &proto.Event_RetrievalModel{
|
||||
RetrievalModel: r.UnwrapT(),
|
||||
@@ -743,6 +747,9 @@ func EventFromGRPC(in *proto.Event) (*types.Event, error) {
|
||||
} else if r := in.GetInferencePolicy(); r != nil {
|
||||
out.Resource = types.ProtoResource153ToLegacy(r)
|
||||
return &out, nil
|
||||
} else if r := in.GetClassifier(); r != nil {
|
||||
out.Resource = types.ProtoResource153ToLegacy(r)
|
||||
return &out, nil
|
||||
} else if r := in.GetRetrievalModel(); r != nil {
|
||||
out.Resource = types.Resource153ToLegacy(r)
|
||||
return &out, nil
|
||||
|
||||
@@ -161,6 +161,47 @@ func (c *Client) DeleteInferencePolicy(ctx context.Context, name string) error {
|
||||
return trace.Wrap(err)
|
||||
}
|
||||
|
||||
// GetClassifier retrieves an existing Classifier by name.
|
||||
func (c *Client) GetClassifier(ctx context.Context, name string) (*summarizerv1.Classifier, error) {
|
||||
resp, err := c.grpcClient.GetClassifier(ctx, &summarizerv1.GetClassifierRequest{
|
||||
Name: name,
|
||||
})
|
||||
if err != nil {
|
||||
return nil, trace.Wrap(err)
|
||||
}
|
||||
return resp.Classifier, nil
|
||||
}
|
||||
|
||||
// CreateClassifier creates a new Classifier.
|
||||
func (c *Client) CreateClassifier(ctx context.Context, classifier *summarizerv1.Classifier) (*summarizerv1.Classifier, error) {
|
||||
resp, err := c.grpcClient.CreateClassifier(ctx, &summarizerv1.CreateClassifierRequest{
|
||||
Classifier: classifier,
|
||||
})
|
||||
if err != nil {
|
||||
return nil, trace.Wrap(err)
|
||||
}
|
||||
return resp.Classifier, nil
|
||||
}
|
||||
|
||||
// UpsertClassifier creates a new Classifier or updates an existing one.
|
||||
func (c *Client) UpsertClassifier(ctx context.Context, classifier *summarizerv1.Classifier) (*summarizerv1.Classifier, error) {
|
||||
resp, err := c.grpcClient.UpsertClassifier(ctx, &summarizerv1.UpsertClassifierRequest{
|
||||
Classifier: classifier,
|
||||
})
|
||||
if err != nil {
|
||||
return nil, trace.Wrap(err)
|
||||
}
|
||||
return resp.Classifier, nil
|
||||
}
|
||||
|
||||
// DeleteClassifier deletes an existing Classifier by name.
|
||||
func (c *Client) DeleteClassifier(ctx context.Context, name string) error {
|
||||
_, err := c.grpcClient.DeleteClassifier(ctx, &summarizerv1.DeleteClassifierRequest{
|
||||
Name: name,
|
||||
})
|
||||
return trace.Wrap(err)
|
||||
}
|
||||
|
||||
// GetRetrievalModel retrieves the existing RetrievalModel.
|
||||
func (c *Client) GetRetrievalModel(ctx context.Context) (*summarizerv1.RetrievalModel, error) {
|
||||
resp, err := c.grpcClient.GetRetrievalModel(ctx, &summarizerv1.GetRetrievalModelRequest{})
|
||||
|
||||
@@ -1177,6 +1177,18 @@
|
||||
"delete"
|
||||
]
|
||||
},
|
||||
{
|
||||
"resources": [
|
||||
"classifier"
|
||||
],
|
||||
"verbs": [
|
||||
"list",
|
||||
"create",
|
||||
"read",
|
||||
"update",
|
||||
"delete"
|
||||
]
|
||||
},
|
||||
{
|
||||
"resources": [
|
||||
"retrieval_model"
|
||||
|
||||
@@ -204,6 +204,62 @@ func (s *UnimplementedService) ListInferencePolicies(
|
||||
return nil, requireEnterprise()
|
||||
}
|
||||
|
||||
// CRUD operations for classifiers
|
||||
|
||||
// CreateClassifier is supposed to create a new classifier, but returns an
|
||||
// error indicating that this feature is only available in the enterprise
|
||||
// version of Teleport.
|
||||
func (s *UnimplementedService) CreateClassifier(
|
||||
ctx context.Context, req *summarizerv1pb.CreateClassifierRequest,
|
||||
) (*summarizerv1pb.CreateClassifierResponse, error) {
|
||||
return nil, requireEnterprise()
|
||||
}
|
||||
|
||||
// GetClassifier is supposed to retrieve a classifier by name, but returns an
|
||||
// error indicating that this feature is only available in the enterprise
|
||||
// version of Teleport.
|
||||
func (s *UnimplementedService) GetClassifier(
|
||||
ctx context.Context, req *summarizerv1pb.GetClassifierRequest,
|
||||
) (*summarizerv1pb.GetClassifierResponse, error) {
|
||||
return nil, requireEnterprise()
|
||||
}
|
||||
|
||||
// UpdateClassifier is supposed to update an existing classifier, but returns
|
||||
// an error indicating that this feature is only available in the enterprise
|
||||
// version of Teleport.
|
||||
func (s *UnimplementedService) UpdateClassifier(
|
||||
ctx context.Context, req *summarizerv1pb.UpdateClassifierRequest,
|
||||
) (*summarizerv1pb.UpdateClassifierResponse, error) {
|
||||
return nil, requireEnterprise()
|
||||
}
|
||||
|
||||
// UpsertClassifier is supposed to create or update a classifier, but returns
|
||||
// an error indicating that this feature is only available in the enterprise
|
||||
// version of Teleport.
|
||||
func (s *UnimplementedService) UpsertClassifier(
|
||||
ctx context.Context, req *summarizerv1pb.UpsertClassifierRequest,
|
||||
) (*summarizerv1pb.UpsertClassifierResponse, error) {
|
||||
return nil, requireEnterprise()
|
||||
}
|
||||
|
||||
// DeleteClassifier is supposed to delete a classifier by name, but returns an
|
||||
// error indicating that this feature is only available in the enterprise
|
||||
// version of Teleport.
|
||||
func (s *UnimplementedService) DeleteClassifier(
|
||||
ctx context.Context, req *summarizerv1pb.DeleteClassifierRequest,
|
||||
) (*summarizerv1pb.DeleteClassifierResponse, error) {
|
||||
return nil, requireEnterprise()
|
||||
}
|
||||
|
||||
// ListClassifiers is supposed to list classifiers with pagination support,
|
||||
// but returns an error indicating that this feature is only available in the
|
||||
// enterprise version of Teleport.
|
||||
func (s *UnimplementedService) ListClassifiers(
|
||||
ctx context.Context, req *summarizerv1pb.ListClassifiersRequest,
|
||||
) (*summarizerv1pb.ListClassifiersResponse, error) {
|
||||
return nil, requireEnterprise()
|
||||
}
|
||||
|
||||
// IsEnabled is supposed to tell if the summarizer is enabled by the license
|
||||
// and configured, but it always returns false. In contrast to other methods
|
||||
// here, it doesn't return an error, since not having an enterprise license is
|
||||
|
||||
Vendored
+1
@@ -224,6 +224,7 @@ func ForAuth(cfg Config) Config {
|
||||
{Kind: types.KindInferenceModel},
|
||||
{Kind: types.KindInferencePolicy},
|
||||
{Kind: types.KindInferenceSecret},
|
||||
{Kind: types.KindClassifier},
|
||||
{Kind: types.KindRetrievalModel},
|
||||
{Kind: types.KindValidatedMFAChallenge},
|
||||
}
|
||||
|
||||
Vendored
+3
@@ -2013,6 +2013,7 @@ func TestCacheWatchKindExistsInEvents(t *testing.T) {
|
||||
types.KindInferenceModel: types.Resource153ToLegacy(new(summaryv1.InferenceModel)),
|
||||
types.KindInferenceSecret: types.Resource153ToLegacy(new(summaryv1.InferenceSecret)),
|
||||
types.KindInferencePolicy: types.Resource153ToLegacy(new(summaryv1.InferencePolicy)),
|
||||
types.KindClassifier: types.Resource153ToLegacy(new(summaryv1.Classifier)),
|
||||
types.KindRetrievalModel: types.Resource153ToLegacy(new(summaryv1.RetrievalModel)),
|
||||
types.KindCertAuthorityOverride: types.Resource153ToLegacy(&subcav1.CertAuthorityOverride{}),
|
||||
types.KindValidatedMFAChallenge: types.Resource153ToLegacy(new(mfav2.ValidatedMFAChallenge)),
|
||||
@@ -2097,6 +2098,8 @@ func TestCacheWatchKindExistsInEvents(t *testing.T) {
|
||||
require.Empty(t, cmp.Diff(resource.(types.Resource153UnwrapperT[*summaryv1.InferenceSecret]).UnwrapT(), uw.UnwrapT(), protocmp.Transform()))
|
||||
case types.Resource153UnwrapperT[*summaryv1.InferencePolicy]:
|
||||
require.Empty(t, cmp.Diff(resource.(types.Resource153UnwrapperT[*summaryv1.InferencePolicy]).UnwrapT(), uw.UnwrapT(), protocmp.Transform()))
|
||||
case types.Resource153UnwrapperT[*summaryv1.Classifier]:
|
||||
require.Empty(t, cmp.Diff(resource.(types.Resource153UnwrapperT[*summaryv1.Classifier]).UnwrapT(), uw.UnwrapT(), protocmp.Transform()))
|
||||
case types.Resource153UnwrapperT[*summaryv1.RetrievalModel]:
|
||||
require.Empty(t, cmp.Diff(resource.(types.Resource153UnwrapperT[*summaryv1.RetrievalModel]).UnwrapT(), uw.UnwrapT(), protocmp.Transform()))
|
||||
case types.Resource153UnwrapperT[*subcav1.CertAuthorityOverride]:
|
||||
|
||||
Vendored
+9
@@ -157,6 +157,7 @@ type collections struct {
|
||||
inferenceModels *collection[*summarizerv1.InferenceModel, inferenceModelIndex]
|
||||
inferenceSecrets *collection[*summarizerv1.InferenceSecret, inferenceSecretIndex]
|
||||
inferencePolicies *collection[*summarizerv1.InferencePolicy, inferencePolicyIndex]
|
||||
classifiers *collection[*summarizerv1.Classifier, classifierIndex]
|
||||
retrievalModels *collection[*summarizerv1.RetrievalModel, retrievalModelIndex]
|
||||
certAuthorityOverrides *collection[*subcav1.CertAuthorityOverride, certAuthorityOverrideIndex]
|
||||
}
|
||||
@@ -838,6 +839,14 @@ func setupCollections(c Config) (*collections, error) {
|
||||
|
||||
out.inferencePolicies = collect
|
||||
out.byKind[resourceKind] = out.inferencePolicies
|
||||
case types.KindClassifier:
|
||||
collect, err := newClassifierCollection(c.Summarizer, watch)
|
||||
if err != nil {
|
||||
return nil, trace.Wrap(err)
|
||||
}
|
||||
|
||||
out.classifiers = collect
|
||||
out.byKind[resourceKind] = out.classifiers
|
||||
case types.KindRetrievalModel:
|
||||
collect, err := newRetrievalModelCollection(c.Summarizer, watch)
|
||||
if err != nil {
|
||||
|
||||
Vendored
+93
@@ -255,6 +255,99 @@ func (c *Cache) AllInferencePolicies(ctx context.Context) iter.Seq2[*summarizerv
|
||||
}
|
||||
}
|
||||
|
||||
type classifierIndex struct{}
|
||||
|
||||
var classifierNameIndex = classifierIndex{}
|
||||
|
||||
func newClassifierCollection(upstream services.Summarizer, w types.WatchKind) (*collection[*summarizerv1.Classifier, classifierIndex], error) {
|
||||
if upstream == nil {
|
||||
return nil, trace.BadParameter("missing parameter Summarizer")
|
||||
}
|
||||
|
||||
return &collection[*summarizerv1.Classifier, classifierIndex]{
|
||||
store: newStore(
|
||||
types.KindClassifier,
|
||||
proto.CloneOf[*summarizerv1.Classifier],
|
||||
map[classifierIndex]func(*summarizerv1.Classifier) string{
|
||||
classifierNameIndex: func(r *summarizerv1.Classifier) string {
|
||||
return r.GetMetadata().GetName()
|
||||
},
|
||||
}),
|
||||
fetcher: func(ctx context.Context, loadSecrets bool) ([]*summarizerv1.Classifier, error) {
|
||||
out, err := stream.Collect(clientutils.Resources(ctx, upstream.ListClassifiers))
|
||||
return out, trace.Wrap(err)
|
||||
},
|
||||
headerTransform: func(hdr *types.ResourceHeader) *summarizerv1.Classifier {
|
||||
return summarizerv1.Classifier_builder{
|
||||
Kind: hdr.Kind,
|
||||
Version: hdr.Version,
|
||||
Metadata: headerv1.Metadata_builder{
|
||||
Name: hdr.Metadata.Name,
|
||||
}.Build(),
|
||||
}.Build()
|
||||
},
|
||||
watch: w,
|
||||
}, nil
|
||||
}
|
||||
|
||||
func (c *Cache) GetClassifier(ctx context.Context, name string) (*summarizerv1.Classifier, error) {
|
||||
ctx, span := c.Tracer.Start(ctx, "cache/GetClassifier")
|
||||
defer span.End()
|
||||
|
||||
getter := genericGetter[*summarizerv1.Classifier, classifierIndex]{
|
||||
cache: c,
|
||||
collection: c.collections.classifiers,
|
||||
index: classifierNameIndex,
|
||||
upstreamGet: c.Config.Summarizer.GetClassifier,
|
||||
}
|
||||
out, err := getter.get(ctx, name)
|
||||
return out, trace.Wrap(err)
|
||||
}
|
||||
|
||||
func (c *Cache) ListClassifiers(ctx context.Context, pageSize int, pageToken string) ([]*summarizerv1.Classifier, string, error) {
|
||||
ctx, span := c.Tracer.Start(ctx, "cache/ListClassifiers")
|
||||
defer span.End()
|
||||
|
||||
lister := genericLister[*summarizerv1.Classifier, classifierIndex]{
|
||||
cache: c,
|
||||
collection: c.collections.classifiers,
|
||||
index: classifierNameIndex,
|
||||
upstreamList: c.Config.Summarizer.ListClassifiers,
|
||||
nextToken: func(t *summarizerv1.Classifier) string {
|
||||
return t.GetMetadata().GetName()
|
||||
},
|
||||
}
|
||||
out, next, err := lister.list(ctx, pageSize, pageToken)
|
||||
return out, next, trace.Wrap(err)
|
||||
}
|
||||
|
||||
func (c *Cache) RangeClassifiers(ctx context.Context, start, end string) iter.Seq2[*summarizerv1.Classifier, error] {
|
||||
lister := genericLister[*summarizerv1.Classifier, classifierIndex]{
|
||||
cache: c,
|
||||
collection: c.collections.classifiers,
|
||||
index: classifierNameIndex,
|
||||
upstreamList: c.Config.Summarizer.ListClassifiers,
|
||||
nextToken: func(t *summarizerv1.Classifier) string {
|
||||
return t.GetMetadata().GetName()
|
||||
},
|
||||
}
|
||||
|
||||
return func(yield func(*summarizerv1.Classifier, error) bool) {
|
||||
ctx, span := c.Tracer.Start(ctx, "cache/RangeClassifiers")
|
||||
defer span.End()
|
||||
|
||||
for classifier, err := range lister.Range(ctx, start, end) {
|
||||
if !yield(classifier, err) {
|
||||
return
|
||||
}
|
||||
|
||||
if err != nil {
|
||||
return
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
type retrievalModelIndex struct{}
|
||||
|
||||
var retrievalModelNameIndex = retrievalModelIndex{}
|
||||
|
||||
Vendored
+138
@@ -18,12 +18,16 @@ package cache
|
||||
|
||||
import (
|
||||
"context"
|
||||
"iter"
|
||||
"testing"
|
||||
"testing/synctest"
|
||||
|
||||
"github.com/stretchr/testify/require"
|
||||
|
||||
summarizerv1 "github.com/gravitational/teleport/api/gen/proto/go/teleport/summarizer/v1"
|
||||
"github.com/gravitational/teleport/api/types"
|
||||
"github.com/gravitational/teleport/api/types/summarizer"
|
||||
"github.com/gravitational/teleport/lib/itertools/stream"
|
||||
)
|
||||
|
||||
func TestInferenceModelCache(t *testing.T) {
|
||||
@@ -95,6 +99,140 @@ func TestInferencePolicyCache(t *testing.T) {
|
||||
})
|
||||
}
|
||||
|
||||
func TestClassifierCache(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
p := newTestPack(t, ForAuth)
|
||||
t.Cleanup(p.Close)
|
||||
|
||||
testResources153(t, p, testFuncs[*summarizerv1.Classifier]{
|
||||
newResource: func(name string) (*summarizerv1.Classifier, error) {
|
||||
return summarizer.NewClassifier(name, summarizerv1.ClassifierSpec_builder{
|
||||
Kinds: []string{string(types.SSHSessionKind)},
|
||||
Criteria: "sessions that touch production data",
|
||||
}.Build()), nil
|
||||
},
|
||||
create: func(ctx context.Context, item *summarizerv1.Classifier) error {
|
||||
_, err := p.summarizer.CreateClassifier(ctx, item)
|
||||
return err
|
||||
},
|
||||
list: p.summarizer.ListClassifiers,
|
||||
cacheList: p.cache.ListClassifiers,
|
||||
deleteAll: p.summarizer.DeleteAllClassifiers,
|
||||
})
|
||||
}
|
||||
|
||||
// collectClassifiers drains a Classifier range into a slice, failing the test
|
||||
// on error.
|
||||
func collectClassifiers(t require.TestingT, it iter.Seq2[*summarizerv1.Classifier, error]) []*summarizerv1.Classifier {
|
||||
out, err := stream.Collect(it)
|
||||
require.NoError(t, err)
|
||||
return out
|
||||
}
|
||||
|
||||
// createClassifiers creates classifiers with the given names in the backend
|
||||
// and blocks until the cache has observed all of them. Names must be unique.
|
||||
func createClassifiers(t *testing.T, ctx context.Context, p *testPack, names []string) {
|
||||
t.Helper()
|
||||
for _, name := range names {
|
||||
classifier := summarizer.NewClassifier(name, summarizerv1.ClassifierSpec_builder{
|
||||
Kinds: []string{string(types.SSHSessionKind)},
|
||||
Criteria: "sessions that touch production data",
|
||||
}.Build())
|
||||
_, err := p.summarizer.CreateClassifier(ctx, classifier)
|
||||
require.NoError(t, err, "failed to create Classifier %q", name)
|
||||
}
|
||||
|
||||
// Wait for the cache to observe all created classifiers.
|
||||
synctest.Wait()
|
||||
results := collectClassifiers(t, p.cache.RangeClassifiers(ctx, "", ""))
|
||||
require.Len(t, results, len(names))
|
||||
}
|
||||
|
||||
// TestClassifierCacheRange tests that RangeClassifiers iterates in name order
|
||||
// and honors range bounds.
|
||||
func TestClassifierCacheRange(t *testing.T) {
|
||||
synctest.Test(t, func(t *testing.T) {
|
||||
ctx := t.Context()
|
||||
|
||||
p := newTestPack(t, ForAuth)
|
||||
t.Cleanup(p.Close)
|
||||
|
||||
createClassifiers(t, ctx, p, []string{
|
||||
"test-classifier-2",
|
||||
"test-classifier-1",
|
||||
"test-classifier-3",
|
||||
})
|
||||
|
||||
names := func(in []*summarizerv1.Classifier) []string {
|
||||
out := make([]string, len(in))
|
||||
for i, c := range in {
|
||||
out[i] = c.GetMetadata().GetName()
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
// Full range is ascending by name. (t.Run subtests are not permitted
|
||||
// inside a synctest bubble, so the cases are inlined.)
|
||||
got := collectClassifiers(t, p.cache.RangeClassifiers(ctx, "", ""))
|
||||
require.Equal(t, []string{
|
||||
"test-classifier-1",
|
||||
"test-classifier-2",
|
||||
"test-classifier-3",
|
||||
}, names(got))
|
||||
|
||||
// Bounded range is exclusive of end.
|
||||
got = collectClassifiers(t, p.cache.RangeClassifiers(ctx, "test-classifier-2", "test-classifier-3"))
|
||||
require.Equal(t, []string{
|
||||
"test-classifier-2",
|
||||
}, names(got))
|
||||
|
||||
// Open-ended range starts at the given name.
|
||||
got = collectClassifiers(t, p.cache.RangeClassifiers(ctx, "test-classifier-2", ""))
|
||||
require.Equal(t, []string{
|
||||
"test-classifier-2",
|
||||
"test-classifier-3",
|
||||
}, names(got))
|
||||
})
|
||||
}
|
||||
|
||||
// TestClassifierCacheRangeUpstreamFallback tests that RangeClassifiers falls
|
||||
// back to the upstream service when the cache is unhealthy.
|
||||
func TestClassifierCacheRangeUpstreamFallback(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
ctx := t.Context()
|
||||
|
||||
// neverOK keeps the cache permanently unhealthy, forcing reads to fall back
|
||||
// to the upstream service.
|
||||
p := newTestPack(t, func(cfg Config) Config {
|
||||
cfg = ForAuth(cfg)
|
||||
cfg.neverOK = true
|
||||
return cfg
|
||||
})
|
||||
t.Cleanup(p.Close)
|
||||
|
||||
want := []string{
|
||||
"test-classifier-2",
|
||||
"test-classifier-1",
|
||||
"test-classifier-3",
|
||||
}
|
||||
for _, name := range want {
|
||||
_, err := p.summarizer.CreateClassifier(ctx, summarizer.NewClassifier(name, summarizerv1.ClassifierSpec_builder{
|
||||
Kinds: []string{string(types.SSHSessionKind)},
|
||||
Criteria: "sessions that touch production data",
|
||||
}.Build()))
|
||||
require.NoError(t, err, "failed to create Classifier %q", name)
|
||||
}
|
||||
|
||||
got := collectClassifiers(t, p.cache.RangeClassifiers(ctx, "", ""))
|
||||
gotNames := make([]string, len(got))
|
||||
for i, c := range got {
|
||||
gotNames[i] = c.GetMetadata().GetName()
|
||||
}
|
||||
require.ElementsMatch(t, want, gotNames)
|
||||
}
|
||||
|
||||
func TestRetrievalModelCache(t *testing.T) {
|
||||
synctest.Test(t, func(t *testing.T) {
|
||||
p := newTestPack(t, ForAuth)
|
||||
|
||||
@@ -291,6 +291,8 @@ func (e *EventsService) NewWatcher(ctx context.Context, watch types.Watch) (type
|
||||
parser = newInferencePolicyParser()
|
||||
case types.KindInferenceSecret:
|
||||
parser = newInferenceSecretParser()
|
||||
case types.KindClassifier:
|
||||
parser = newClassifierParser()
|
||||
case types.KindRetrievalModel:
|
||||
parser = newRetrievalModelParser()
|
||||
case types.KindCertAuthorityOverride:
|
||||
|
||||
@@ -36,6 +36,7 @@ type SummarizerService struct {
|
||||
modelService *generic.ServiceWrapper[*summarizerv1.InferenceModel]
|
||||
secretService *generic.ServiceWrapper[*summarizerv1.InferenceSecret]
|
||||
policyService *generic.ServiceWrapper[*summarizerv1.InferencePolicy]
|
||||
classifierService *generic.ServiceWrapper[*summarizerv1.Classifier]
|
||||
retrievalModelService *generic.ServiceWrapper[*summarizerv1.RetrievalModel]
|
||||
}
|
||||
|
||||
@@ -233,6 +234,75 @@ func (s *SummarizerService) AllInferencePolicies(
|
||||
return s.policyService.Resources(ctx, "", "")
|
||||
}
|
||||
|
||||
// CreateClassifier creates a new session summarization classifier in the
|
||||
// backend.
|
||||
func (s *SummarizerService) CreateClassifier(
|
||||
ctx context.Context, classifier *summarizerv1.Classifier,
|
||||
) (*summarizerv1.Classifier, error) {
|
||||
res, err := s.classifierService.CreateResource(ctx, classifier)
|
||||
return res, trace.Wrap(err)
|
||||
}
|
||||
|
||||
// DeleteClassifier deletes a session summarization classifier from the
|
||||
// backend by name.
|
||||
func (s *SummarizerService) DeleteClassifier(
|
||||
ctx context.Context, name string,
|
||||
) error {
|
||||
return trace.Wrap(s.classifierService.DeleteResource(ctx, name))
|
||||
}
|
||||
|
||||
// GetClassifier retrieves a session summarization classifier from the backend
|
||||
// by name.
|
||||
func (s *SummarizerService) GetClassifier(
|
||||
ctx context.Context, name string,
|
||||
) (*summarizerv1.Classifier, error) {
|
||||
res, err := s.classifierService.GetResource(ctx, name)
|
||||
return res, trace.Wrap(err)
|
||||
}
|
||||
|
||||
// ListClassifiers lists session summarization classifiers in the backend with
|
||||
// pagination support. Returns a slice of classifiers and a next page token.
|
||||
func (s *SummarizerService) ListClassifiers(
|
||||
ctx context.Context, size int, pageToken string,
|
||||
) ([]*summarizerv1.Classifier, string, error) {
|
||||
res, nextToken, err := s.classifierService.ListResources(ctx, size, pageToken)
|
||||
return res, nextToken, trace.Wrap(err)
|
||||
}
|
||||
|
||||
// DeleteAllClassifiers deletes all session summarization classifiers from the
|
||||
// backend. This should only be used by the cache.
|
||||
func (s *SummarizerService) DeleteAllClassifiers(ctx context.Context) error {
|
||||
return trace.Wrap(s.classifierService.DeleteAllResources(ctx))
|
||||
}
|
||||
|
||||
// UpdateClassifier updates an existing session summarization classifier in
|
||||
// the backend.
|
||||
func (s *SummarizerService) UpdateClassifier(
|
||||
ctx context.Context, classifier *summarizerv1.Classifier,
|
||||
) (*summarizerv1.Classifier, error) {
|
||||
res, err := s.classifierService.ConditionalUpdateResource(ctx, classifier)
|
||||
return res, trace.Wrap(err)
|
||||
}
|
||||
|
||||
// UpsertClassifier creates or updates a session summarization classifier in
|
||||
// the backend. If the classifier already exists, it will be updated.
|
||||
func (s *SummarizerService) UpsertClassifier(
|
||||
ctx context.Context, classifier *summarizerv1.Classifier,
|
||||
) (*summarizerv1.Classifier, error) {
|
||||
res, err := s.classifierService.UpsertResource(ctx, classifier)
|
||||
return res, trace.Wrap(err)
|
||||
}
|
||||
|
||||
// RangeClassifiers returns an iterator that retrieves session summarization
|
||||
// classifiers from the backend, without pagination, starting with the
|
||||
// resource named start and ending before the resource named end. Empty bounds
|
||||
// iterate from the beginning and/or to the end of the collection.
|
||||
func (s *SummarizerService) RangeClassifiers(
|
||||
ctx context.Context, start, end string,
|
||||
) iter.Seq2[*summarizerv1.Classifier, error] {
|
||||
return s.classifierService.Resources(ctx, start, end)
|
||||
}
|
||||
|
||||
// CreateRetrievalModel creates the search model in the backend.
|
||||
// Only one RetrievalModel can exist per cluster.
|
||||
func (s *SummarizerService) CreateRetrievalModel(
|
||||
@@ -278,6 +348,7 @@ const (
|
||||
inferenceModelPrefix = "inference_models"
|
||||
inferenceSecretPrefix = "inference_secrets"
|
||||
inferencePolicyPrefix = "inference_policies"
|
||||
classifierPrefix = "classifiers"
|
||||
retrievalModelPrefix = "retrieval_model"
|
||||
)
|
||||
|
||||
@@ -353,6 +424,19 @@ func NewSummarizerService(cfg SummarizerServiceConfig) (*SummarizerService, erro
|
||||
return nil, trace.Wrap(err)
|
||||
}
|
||||
|
||||
classifierService, err := generic.NewServiceWrapper(
|
||||
generic.ServiceConfig[*summarizerv1.Classifier]{
|
||||
Backend: cfg.Backend,
|
||||
ResourceKind: types.KindClassifier,
|
||||
BackendPrefix: backend.NewKey(classifierPrefix),
|
||||
MarshalFunc: services.MarshalProtoResource[*summarizerv1.Classifier],
|
||||
UnmarshalFunc: services.UnmarshalProtoResource[*summarizerv1.Classifier],
|
||||
ValidateFunc: services.ValidateClassifier,
|
||||
})
|
||||
if err != nil {
|
||||
return nil, trace.Wrap(err)
|
||||
}
|
||||
|
||||
validateRetrievalModel := func(m *summarizerv1.RetrievalModel) error {
|
||||
err := summarizer.ValidateRetrievalModel(m)
|
||||
if err != nil {
|
||||
@@ -379,6 +463,7 @@ func NewSummarizerService(cfg SummarizerServiceConfig) (*SummarizerService, erro
|
||||
modelService: modelService,
|
||||
secretService: secretService,
|
||||
policyService: policyService,
|
||||
classifierService: classifierService,
|
||||
retrievalModelService: retrievalModelService,
|
||||
}, nil
|
||||
}
|
||||
@@ -504,6 +589,46 @@ func (p *inferenceSecretParser) parse(event backend.Event) (types.Resource, erro
|
||||
}
|
||||
}
|
||||
|
||||
func newClassifierParser() *classifierParser {
|
||||
return &classifierParser{
|
||||
baseParser: newBaseParser(backend.NewKey(classifierPrefix)),
|
||||
}
|
||||
}
|
||||
|
||||
type classifierParser struct {
|
||||
baseParser
|
||||
}
|
||||
|
||||
func (p *classifierParser) parse(event backend.Event) (types.Resource, error) {
|
||||
switch event.Type {
|
||||
case types.OpDelete:
|
||||
components := event.Item.Key.Components()
|
||||
if len(components) != 2 {
|
||||
return nil, trace.NotFound("failed parsing %v: expected 2 components, got %d", event.Item.Key.String(), len(components))
|
||||
}
|
||||
name := components[1]
|
||||
return &types.ResourceHeader{
|
||||
Kind: types.KindClassifier,
|
||||
Version: types.V1,
|
||||
Metadata: types.Metadata{
|
||||
Name: name,
|
||||
},
|
||||
}, nil
|
||||
case types.OpPut:
|
||||
classifier, err := services.UnmarshalProtoResource[*summarizerv1.Classifier](
|
||||
event.Item.Value,
|
||||
services.WithExpires(event.Item.Expires),
|
||||
services.WithRevision(event.Item.Revision),
|
||||
)
|
||||
if err != nil {
|
||||
return nil, trace.Wrap(err)
|
||||
}
|
||||
return types.Resource153ToLegacy(classifier), nil
|
||||
default:
|
||||
return nil, trace.BadParameter("event %v is not supported", event.Type)
|
||||
}
|
||||
}
|
||||
|
||||
func newRetrievalModelParser() *retrievalModelParser {
|
||||
return &retrievalModelParser{
|
||||
baseParser: newBaseParser(backend.NewKey(retrievalModelPrefix)),
|
||||
|
||||
@@ -837,6 +837,288 @@ func TestSummarizerService_AllInferencePolicies(t *testing.T) {
|
||||
}
|
||||
}
|
||||
|
||||
func newClassifier(name string) *summarizerv1.Classifier {
|
||||
return summarizer.NewClassifier(name, summarizerv1.ClassifierSpec_builder{
|
||||
Kinds: []string{string(types.SSHSessionKind)},
|
||||
Criteria: "sessions that touch production data",
|
||||
}.Build())
|
||||
}
|
||||
|
||||
func TestSummarizerService_CreateClassifier(t *testing.T) {
|
||||
ctx, service := setupSummarizerTest(t)
|
||||
|
||||
t.Run("ok", func(t *testing.T) {
|
||||
want := newClassifier("dummy-classifier")
|
||||
got, err := service.CreateClassifier(
|
||||
ctx,
|
||||
// Clone to avoid Marshaling modifying want
|
||||
proto.Clone(want).(*summarizerv1.Classifier),
|
||||
)
|
||||
require.NoError(t, err)
|
||||
assert.NotEmpty(t, got.GetMetadata().GetRevision())
|
||||
assert.Empty(t, cmp.Diff(
|
||||
want,
|
||||
got,
|
||||
protocmp.Transform(),
|
||||
protocmp.IgnoreFields(&headerv1.Metadata{}, "revision"),
|
||||
))
|
||||
})
|
||||
t.Run("invalid criteria", func(t *testing.T) {
|
||||
c := newClassifier("invalid-classifier")
|
||||
c.GetSpec().SetCriteria("")
|
||||
_, err := service.CreateClassifier(
|
||||
ctx,
|
||||
proto.Clone(c).(*summarizerv1.Classifier),
|
||||
)
|
||||
require.Error(t, err)
|
||||
assert.ErrorContains(t, err, "spec.criteria is required")
|
||||
})
|
||||
t.Run("invalid filter", func(t *testing.T) {
|
||||
c := newClassifier("invalid-classifier")
|
||||
c.GetSpec().SetFilter("$%^@$")
|
||||
_, err := service.CreateClassifier(
|
||||
ctx,
|
||||
proto.Clone(c).(*summarizerv1.Classifier),
|
||||
)
|
||||
require.Error(t, err)
|
||||
assert.ErrorContains(t, err, "spec.filter has to be a valid predicate")
|
||||
})
|
||||
t.Run("no upsert", func(t *testing.T) {
|
||||
res := newClassifier("no-upsert")
|
||||
_, err := service.CreateClassifier(
|
||||
ctx,
|
||||
// Clone to avoid Marshaling modifying want
|
||||
proto.Clone(res).(*summarizerv1.Classifier),
|
||||
)
|
||||
require.NoError(t, err)
|
||||
_, err = service.CreateClassifier(
|
||||
ctx,
|
||||
// Clone to avoid Marshaling modifying want
|
||||
proto.Clone(res).(*summarizerv1.Classifier),
|
||||
)
|
||||
require.Error(t, err)
|
||||
assert.True(t, trace.IsAlreadyExists(err))
|
||||
})
|
||||
}
|
||||
|
||||
func TestSummarizerService_UpsertClassifier(t *testing.T) {
|
||||
ctx, service := setupSummarizerTest(t)
|
||||
|
||||
want := newClassifier("dummy-classifier")
|
||||
got, err := service.UpsertClassifier(
|
||||
ctx,
|
||||
// Clone to avoid Marshaling modifying want
|
||||
proto.Clone(want).(*summarizerv1.Classifier),
|
||||
)
|
||||
require.NoError(t, err)
|
||||
assert.NotEmpty(t, got.GetMetadata().GetRevision())
|
||||
assert.Empty(t, cmp.Diff(
|
||||
want,
|
||||
got,
|
||||
protocmp.Transform(),
|
||||
protocmp.IgnoreFields(&headerv1.Metadata{}, "revision"),
|
||||
))
|
||||
|
||||
// Ensure we can upsert over an existing resource
|
||||
_, err = service.UpsertClassifier(
|
||||
ctx,
|
||||
// Clone to avoid Marshaling modifying want
|
||||
proto.Clone(want).(*summarizerv1.Classifier),
|
||||
)
|
||||
require.NoError(t, err)
|
||||
}
|
||||
|
||||
func TestSummarizerService_ListClassifiers(t *testing.T) {
|
||||
ctx, service := setupSummarizerTest(t)
|
||||
// Create entities to list
|
||||
createdObjects := []*summarizerv1.Classifier{}
|
||||
// Create 49 entities to test an incomplete page at the end.
|
||||
for i := range 49 {
|
||||
created, err := service.CreateClassifier(
|
||||
ctx,
|
||||
newClassifier(fmt.Sprintf("classifier-%d", i)),
|
||||
)
|
||||
require.NoError(t, err)
|
||||
createdObjects = append(createdObjects, created)
|
||||
}
|
||||
t.Run("default page size", func(t *testing.T) {
|
||||
page, nextToken, err := service.ListClassifiers(ctx, 0, "")
|
||||
require.NoError(t, err)
|
||||
assert.Len(t, page, 49)
|
||||
assert.Empty(t, nextToken)
|
||||
|
||||
// Expect that we get all the things we have created
|
||||
for _, created := range createdObjects {
|
||||
assert.True(t, slices.ContainsFunc(page, func(classifier *summarizerv1.Classifier) bool {
|
||||
return proto.Equal(created, classifier)
|
||||
}))
|
||||
}
|
||||
})
|
||||
t.Run("pagination", func(t *testing.T) {
|
||||
fetched := []*summarizerv1.Classifier{}
|
||||
token := ""
|
||||
iterations := 0
|
||||
for {
|
||||
iterations++
|
||||
page, nextToken, err := service.ListClassifiers(ctx, 10, token)
|
||||
require.NoError(t, err)
|
||||
fetched = append(fetched, page...)
|
||||
if nextToken == "" {
|
||||
break
|
||||
}
|
||||
token = nextToken
|
||||
}
|
||||
assert.Equal(t, 5, iterations)
|
||||
|
||||
assert.Len(t, fetched, 49)
|
||||
// Expect that we get all the things we have created
|
||||
for _, created := range createdObjects {
|
||||
assert.True(t, slices.ContainsFunc(fetched, func(classifier *summarizerv1.Classifier) bool {
|
||||
return proto.Equal(created, classifier)
|
||||
}))
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
func TestSummarizerService_GetClassifier(t *testing.T) {
|
||||
ctx, service := setupSummarizerTest(t)
|
||||
|
||||
t.Run("ok", func(t *testing.T) {
|
||||
want := newClassifier("dummy-classifier")
|
||||
_, err := service.CreateClassifier(
|
||||
ctx,
|
||||
// Clone to avoid Marshaling modifying want
|
||||
proto.Clone(want).(*summarizerv1.Classifier),
|
||||
)
|
||||
require.NoError(t, err)
|
||||
got, err := service.GetClassifier(ctx, "dummy-classifier")
|
||||
require.NoError(t, err)
|
||||
assert.NotEmpty(t, got.GetMetadata().GetRevision())
|
||||
assert.Empty(t, cmp.Diff(
|
||||
want,
|
||||
got,
|
||||
protocmp.Transform(),
|
||||
protocmp.IgnoreFields(&headerv1.Metadata{}, "revision"),
|
||||
))
|
||||
})
|
||||
t.Run("not found", func(t *testing.T) {
|
||||
_, err := service.GetClassifier(ctx, "foobar")
|
||||
require.Error(t, err)
|
||||
assert.True(t, trace.IsNotFound(err))
|
||||
})
|
||||
}
|
||||
|
||||
func TestSummarizerService_DeleteClassifier(t *testing.T) {
|
||||
ctx, service := setupSummarizerTest(t)
|
||||
|
||||
t.Run("ok", func(t *testing.T) {
|
||||
_, err := service.CreateClassifier(
|
||||
ctx,
|
||||
newClassifier("dummy-classifier"),
|
||||
)
|
||||
require.NoError(t, err)
|
||||
|
||||
_, err = service.GetClassifier(ctx, "dummy-classifier")
|
||||
require.NoError(t, err)
|
||||
|
||||
err = service.DeleteClassifier(ctx, "dummy-classifier")
|
||||
require.NoError(t, err)
|
||||
|
||||
_, err = service.GetClassifier(ctx, "dummy-classifier")
|
||||
require.Error(t, err)
|
||||
assert.True(t, trace.IsNotFound(err))
|
||||
})
|
||||
t.Run("not found", func(t *testing.T) {
|
||||
err := service.DeleteClassifier(ctx, "foobar")
|
||||
require.Error(t, err)
|
||||
assert.True(t, trace.IsNotFound(err))
|
||||
})
|
||||
}
|
||||
|
||||
func TestSummarizerService_UpdateClassifier(t *testing.T) {
|
||||
ctx, service := setupSummarizerTest(t)
|
||||
|
||||
t.Run("ok", func(t *testing.T) {
|
||||
// Create resource for us to Update since we can't update a non-existent resource.
|
||||
created, err := service.CreateClassifier(
|
||||
ctx,
|
||||
newClassifier("dummy-classifier"),
|
||||
)
|
||||
require.NoError(t, err)
|
||||
want := proto.Clone(created).(*summarizerv1.Classifier)
|
||||
want.GetSpec().SetKinds([]string{string(types.SSHSessionKind), string(types.DatabaseSessionKind)})
|
||||
|
||||
updated, err := service.UpdateClassifier(
|
||||
ctx,
|
||||
// Clone to avoid Marshaling modifying want
|
||||
proto.Clone(want).(*summarizerv1.Classifier),
|
||||
)
|
||||
require.NoError(t, err)
|
||||
assert.NotEqual(t, created.GetMetadata().GetRevision(), updated.GetMetadata().GetRevision())
|
||||
assert.Empty(t, cmp.Diff(
|
||||
want,
|
||||
updated,
|
||||
protocmp.Transform(),
|
||||
protocmp.IgnoreFields(&headerv1.Metadata{}, "revision"),
|
||||
))
|
||||
|
||||
got, err := service.GetClassifier(ctx, "dummy-classifier")
|
||||
require.NoError(t, err)
|
||||
assert.Empty(t, cmp.Diff(
|
||||
want,
|
||||
got,
|
||||
protocmp.Transform(),
|
||||
protocmp.IgnoreFields(&headerv1.Metadata{}, "revision"),
|
||||
))
|
||||
assert.Equal(t, updated.GetMetadata().GetRevision(), got.GetMetadata().GetRevision())
|
||||
})
|
||||
t.Run("no create", func(t *testing.T) {
|
||||
_, err := service.UpdateClassifier(
|
||||
ctx,
|
||||
newClassifier("non-existing-classifier"),
|
||||
)
|
||||
require.Error(t, err)
|
||||
})
|
||||
}
|
||||
|
||||
func TestSummarizerService_RangeClassifiers(t *testing.T) {
|
||||
ctx, service := setupSummarizerTest(t)
|
||||
// Create entities to retrieve
|
||||
createdObjects := []*summarizerv1.Classifier{}
|
||||
for i := range 5 {
|
||||
created, err := service.CreateClassifier(
|
||||
ctx,
|
||||
newClassifier(fmt.Sprintf("classifier-%d", i)),
|
||||
)
|
||||
require.NoError(t, err)
|
||||
createdObjects = append(createdObjects, created)
|
||||
}
|
||||
|
||||
fetched := []*summarizerv1.Classifier{}
|
||||
for classifier, err := range service.RangeClassifiers(ctx, "", "") {
|
||||
require.NoError(t, err)
|
||||
fetched = append(fetched, classifier)
|
||||
}
|
||||
assert.Len(t, fetched, 5)
|
||||
|
||||
// Expect that we get all the things we have created
|
||||
for _, created := range createdObjects {
|
||||
assert.True(t, slices.ContainsFunc(fetched, func(classifier *summarizerv1.Classifier) bool {
|
||||
return proto.Equal(created, classifier)
|
||||
}))
|
||||
}
|
||||
|
||||
// A bounded range returns only the resources within it, excluding the end.
|
||||
bounded := []*summarizerv1.Classifier{}
|
||||
for classifier, err := range service.RangeClassifiers(ctx, "classifier-1", "classifier-3") {
|
||||
require.NoError(t, err)
|
||||
bounded = append(bounded, classifier)
|
||||
}
|
||||
require.Len(t, bounded, 2)
|
||||
assert.Equal(t, "classifier-1", bounded[0].GetMetadata().GetName())
|
||||
assert.Equal(t, "classifier-2", bounded[1].GetMetadata().GetName())
|
||||
}
|
||||
|
||||
func newRetrievalModel() *summarizerv1.RetrievalModel {
|
||||
return summarizer.NewRetrievalModel(summarizerv1.RetrievalModelSpec_builder{
|
||||
Openai: summarizerv1.OpenAIProvider_builder{
|
||||
|
||||
@@ -225,6 +225,7 @@ func NewPresetEditorRole() types.Role {
|
||||
types.NewRule(types.KindInferenceModel, RW()),
|
||||
types.NewRule(types.KindInferenceSecret, RW()),
|
||||
types.NewRule(types.KindInferencePolicy, RW()),
|
||||
types.NewRule(types.KindClassifier, RW()),
|
||||
types.NewRule(types.KindRetrievalModel, RW()),
|
||||
types.NewRule(types.KindClientIPRestriction, RW()),
|
||||
types.NewRule(access.KindScopedRole, RW()),
|
||||
@@ -880,6 +881,7 @@ func NewPresetTerraformProviderRole() types.Role {
|
||||
types.NewRule(types.KindInferenceModel, RW()),
|
||||
types.NewRule(types.KindInferenceSecret, RW()),
|
||||
types.NewRule(types.KindInferencePolicy, RW()),
|
||||
types.NewRule(types.KindClassifier, RW()),
|
||||
types.NewRule(types.KindRetrievalModel, RW()),
|
||||
types.NewRule(types.KindSAMLIdPServiceProvider, RW()),
|
||||
types.NewRule(types.KindScopedToken, RW()),
|
||||
|
||||
@@ -777,6 +777,7 @@ func TestAddRoleDefaults(t *testing.T) {
|
||||
types.NewRule(types.KindInferenceModel, RW()),
|
||||
types.NewRule(types.KindInferenceSecret, RW()),
|
||||
types.NewRule(types.KindInferencePolicy, RW()),
|
||||
types.NewRule(types.KindClassifier, RW()),
|
||||
types.NewRule(types.KindRetrievalModel, RW()),
|
||||
types.NewRule(types.KindScopedToken, RW()),
|
||||
types.NewRule(access.KindScopedRole, RW()),
|
||||
|
||||
@@ -303,6 +303,8 @@ func ParseShortcut(in string) (string, error) {
|
||||
return types.KindInferenceSecret, nil
|
||||
case types.KindInferencePolicy, "inference_policies":
|
||||
return types.KindInferencePolicy, nil
|
||||
case types.KindClassifier, types.KindClassifier + "s":
|
||||
return types.KindClassifier, nil
|
||||
case types.KindRetrievalModel:
|
||||
return types.KindRetrievalModel, nil
|
||||
case types.KindRelayServer, types.KindRelayServer + "s":
|
||||
|
||||
@@ -19,6 +19,7 @@ package services
|
||||
import (
|
||||
"context"
|
||||
"iter"
|
||||
"slices"
|
||||
"strings"
|
||||
|
||||
"github.com/gravitational/trace"
|
||||
@@ -78,6 +79,19 @@ type Summarizer interface {
|
||||
// backend by name.
|
||||
DeleteInferencePolicy(ctx context.Context, name string) error
|
||||
|
||||
// CreateClassifier creates a new session summarization classifier in the
|
||||
// backend.
|
||||
CreateClassifier(ctx context.Context, classifier *summarizerv1.Classifier) (*summarizerv1.Classifier, error)
|
||||
// UpdateClassifier updates an existing session summarization classifier in
|
||||
// the backend.
|
||||
UpdateClassifier(ctx context.Context, classifier *summarizerv1.Classifier) (*summarizerv1.Classifier, error)
|
||||
// UpsertClassifier creates or updates a session summarization classifier
|
||||
// in the backend. If the classifier already exists, it will be updated.
|
||||
UpsertClassifier(ctx context.Context, classifier *summarizerv1.Classifier) (*summarizerv1.Classifier, error)
|
||||
// DeleteClassifier deletes a session summarization classifier from the
|
||||
// backend by name.
|
||||
DeleteClassifier(ctx context.Context, name string) error
|
||||
|
||||
// CreateRetrievalModel creates the search model in the backend.
|
||||
// Only one RetrievalModel can exist per cluster.
|
||||
CreateRetrievalModel(ctx context.Context, model *summarizerv1.RetrievalModel) (*summarizerv1.RetrievalModel, error)
|
||||
@@ -122,6 +136,19 @@ type SummarizerServiceGetter interface {
|
||||
// summary inference policies from the backend, without pagination.
|
||||
AllInferencePolicies(ctx context.Context) iter.Seq2[*summarizerv1.InferencePolicy, error]
|
||||
|
||||
// GetClassifier retrieves a session summarization classifier from the
|
||||
// backend by name.
|
||||
GetClassifier(ctx context.Context, name string) (*summarizerv1.Classifier, error)
|
||||
// ListClassifiers lists session summarization classifiers in the backend
|
||||
// with pagination support. Returns a slice of classifiers and a next page
|
||||
// token.
|
||||
ListClassifiers(ctx context.Context, size int, pageToken string) ([]*summarizerv1.Classifier, string, error)
|
||||
// RangeClassifiers returns an iterator that retrieves session summarization
|
||||
// classifiers from the backend, without pagination, starting with the
|
||||
// resource named start and ending before the resource named end. Empty
|
||||
// bounds iterate from the beginning and/or to the end of the collection.
|
||||
RangeClassifiers(ctx context.Context, start, end string) iter.Seq2[*summarizerv1.Classifier, error]
|
||||
|
||||
// GetRetrievalModel retrieves the search model from the backend.
|
||||
// Since only one RetrievalModel can exist per cluster, no name is required.
|
||||
GetRetrievalModel(ctx context.Context) (*summarizerv1.RetrievalModel, error)
|
||||
@@ -235,6 +262,80 @@ func (ctx *InferencePolicyMatchingContext) ExtendWithSessionEnd(sessionEnd event
|
||||
ctx.Resource = rebuildResourceFromSessionEndEvent(sessionEnd)
|
||||
}
|
||||
|
||||
// MatchingClassifiers returns the classifiers from the given sequence that
|
||||
// apply to a session of the given kind and matching context. Classifiers are
|
||||
// matched by session kind and filter expression the same way inference
|
||||
// policies are, except that all matching classifiers are returned rather than
|
||||
// the first one. The sequence is typically
|
||||
// [SummarizerServiceGetter.RangeClassifiers].
|
||||
func MatchingClassifiers(
|
||||
classifiers iter.Seq2[*summarizerv1.Classifier, error],
|
||||
sessionKind types.SessionKind,
|
||||
matchingCtx *InferencePolicyMatchingContext,
|
||||
) ([]*summarizerv1.Classifier, error) {
|
||||
parser, err := NewWhereParser(matchingCtx)
|
||||
if err != nil {
|
||||
return nil, trace.Wrap(err)
|
||||
}
|
||||
|
||||
var matched []*summarizerv1.Classifier
|
||||
for c, err := range classifiers {
|
||||
if err != nil {
|
||||
return nil, trace.Wrap(err)
|
||||
}
|
||||
if !slices.Contains(c.GetSpec().GetKinds(), string(sessionKind)) {
|
||||
continue
|
||||
}
|
||||
|
||||
if filter := c.GetSpec().GetFilter(); filter != "" {
|
||||
parseResult, err := parser.Parse(filter)
|
||||
if err != nil {
|
||||
return nil, trace.Wrap(err)
|
||||
}
|
||||
|
||||
pred, ok := parseResult.(predicate.BoolPredicate)
|
||||
if !ok {
|
||||
return nil, trace.BadParameter("unsupported type: %T", parseResult)
|
||||
}
|
||||
|
||||
if !pred() {
|
||||
continue
|
||||
}
|
||||
}
|
||||
|
||||
matched = append(matched, c)
|
||||
}
|
||||
return matched, nil
|
||||
}
|
||||
|
||||
// ValidateClassifier validates a classifier, including checking filter
|
||||
// syntax. This function wraps [apisummarizer.ValidateClassifier], as no
|
||||
// function in the api/types tree can depend on the lib/services package.
|
||||
func ValidateClassifier(c *summarizerv1.Classifier) error {
|
||||
err := apisummarizer.ValidateClassifier(c)
|
||||
if err != nil {
|
||||
return trace.Wrap(err)
|
||||
}
|
||||
|
||||
s := c.GetSpec()
|
||||
if s.GetFilter() != "" {
|
||||
parser, err := NewWhereParser(&InferencePolicyMatchingContext{})
|
||||
if err != nil {
|
||||
return trace.Wrap(err)
|
||||
}
|
||||
|
||||
parseResult, err := parser.Parse(s.GetFilter())
|
||||
if err != nil {
|
||||
return trace.Wrap(err, "spec.filter has to be a valid predicate")
|
||||
}
|
||||
if _, ok := parseResult.(predicate.BoolPredicate); !ok {
|
||||
return trace.BadParameter("spec.filter has to be a boolean expression")
|
||||
}
|
||||
}
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
// ValidateInferencePolicy validates an inference policy, including checking
|
||||
// filter syntax. This function wraps [apisummarizer.ValidateInferencePolicy],
|
||||
// as no function in the api/types tree can depend on the lib/services package.
|
||||
|
||||
@@ -318,3 +318,110 @@ func TestValidateInferencePolicy(t *testing.T) {
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func newTestClassifier(name string, kinds []string, filter string) *summarizerv1.Classifier {
|
||||
return apisummarizer.NewClassifier(name, summarizerv1.ClassifierSpec_builder{
|
||||
Kinds: kinds,
|
||||
Filter: filter,
|
||||
Criteria: "sessions that touch production data",
|
||||
}.Build())
|
||||
}
|
||||
|
||||
func TestMatchingClassifiers(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
user, err := types.NewUser("alice")
|
||||
require.NoError(t, err)
|
||||
|
||||
server, err := types.NewServer("server-1", types.KindNode, types.ServerSpecV2{Hostname: "host-1"})
|
||||
require.NoError(t, err)
|
||||
server.SetStaticLabels(map[string]string{"env": "prod"})
|
||||
|
||||
matchingCtx := &InferencePolicyMatchingContext{
|
||||
User: user,
|
||||
Resource: server,
|
||||
}
|
||||
|
||||
all := []*summarizerv1.Classifier{
|
||||
newTestClassifier("ssh-any", []string{"ssh"}, ""),
|
||||
newTestClassifier("ssh-prod", []string{"ssh"}, `equals(resource.metadata.labels["env"], "prod")`),
|
||||
newTestClassifier("ssh-dev", []string{"ssh"}, `equals(resource.metadata.labels["env"], "dev")`),
|
||||
newTestClassifier("ssh-alice", []string{"ssh"}, `equals(user.metadata.name, "alice")`),
|
||||
newTestClassifier("db-any", []string{"db"}, ""),
|
||||
}
|
||||
classifiers := func(yield func(*summarizerv1.Classifier, error) bool) {
|
||||
for _, c := range all {
|
||||
if !yield(c, nil) {
|
||||
return
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
matched, err := MatchingClassifiers(classifiers, types.SSHSessionKind, matchingCtx)
|
||||
require.NoError(t, err)
|
||||
names := make([]string, 0, len(matched))
|
||||
for _, c := range matched {
|
||||
names = append(names, c.GetMetadata().GetName())
|
||||
}
|
||||
assert.Equal(t, []string{"ssh-any", "ssh-prod", "ssh-alice"}, names)
|
||||
|
||||
matched, err = MatchingClassifiers(classifiers, types.DatabaseSessionKind, matchingCtx)
|
||||
require.NoError(t, err)
|
||||
require.Len(t, matched, 1)
|
||||
assert.Equal(t, "db-any", matched[0].GetMetadata().GetName())
|
||||
|
||||
matched, err = MatchingClassifiers(classifiers, types.KubernetesSessionKind, matchingCtx)
|
||||
require.NoError(t, err)
|
||||
assert.Empty(t, matched)
|
||||
|
||||
iterErr := trace.BadParameter("backend failure")
|
||||
failing := func(yield func(*summarizerv1.Classifier, error) bool) {
|
||||
yield(nil, iterErr)
|
||||
}
|
||||
_, err = MatchingClassifiers(failing, types.SSHSessionKind, matchingCtx)
|
||||
require.ErrorIs(t, err, iterErr)
|
||||
|
||||
invalid := []*summarizerv1.Classifier{
|
||||
newTestClassifier("bad-filter", []string{"ssh"}, "$%^@$"),
|
||||
}
|
||||
invalidSeq := func(yield func(*summarizerv1.Classifier, error) bool) {
|
||||
for _, c := range invalid {
|
||||
if !yield(c, nil) {
|
||||
return
|
||||
}
|
||||
}
|
||||
}
|
||||
_, err = MatchingClassifiers(invalidSeq, types.SSHSessionKind, matchingCtx)
|
||||
require.Error(t, err)
|
||||
}
|
||||
|
||||
func TestValidateClassifier(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
c := newTestClassifier(
|
||||
"my-classifier",
|
||||
[]string{"ssh", "k8s", "db"},
|
||||
`equals(resource.metadata.labels["env"], "prod") || equals(user.metadata.name, "admin")`,
|
||||
)
|
||||
require.NoError(t, ValidateClassifier(c))
|
||||
|
||||
// Empty filter should also be valid.
|
||||
c.GetSpec().SetFilter("")
|
||||
require.NoError(t, ValidateClassifier(c))
|
||||
|
||||
// Filter syntax errors are rejected.
|
||||
c.GetSpec().SetFilter("$%^@$")
|
||||
err := ValidateClassifier(c)
|
||||
require.Error(t, err)
|
||||
assert.ErrorContains(t, err, "spec.filter has to be a valid predicate")
|
||||
|
||||
c.GetSpec().SetFilter("user.metadata.name")
|
||||
err = ValidateClassifier(c)
|
||||
require.Error(t, err)
|
||||
assert.ErrorContains(t, err, "spec.filter has to be a boolean expression")
|
||||
|
||||
// Errors from the api-level validation propagate.
|
||||
c.GetSpec().SetFilter("")
|
||||
c.GetSpec().SetCriteria("")
|
||||
assert.ErrorContains(t, ValidateClassifier(c), "spec.criteria is required")
|
||||
}
|
||||
|
||||
@@ -138,6 +138,8 @@ type UserACL struct {
|
||||
InferencePolicy ResourceAccess `json:"inferencePolicy"`
|
||||
// InferenceSecret defines access to session summaries inference secret.
|
||||
InferenceSecret ResourceAccess `json:"inferenceSecret"`
|
||||
// Classifier defines access to session summarization classifiers.
|
||||
Classifier ResourceAccess `json:"classifier"`
|
||||
// AutoUpdateConfig defines access to autoupdate config.
|
||||
AutoUpdateConfig ResourceAccess `json:"autoUpdateConfig"`
|
||||
// AutoUpdateVersion defines access to autoupdate version.
|
||||
@@ -322,6 +324,7 @@ func NewUserACL(user types.User, userRoles RoleSet, features proto.Features, des
|
||||
InferenceModel: newAccess(userRoles, ctx, types.KindInferenceModel),
|
||||
InferencePolicy: newAccess(userRoles, ctx, types.KindInferencePolicy),
|
||||
InferenceSecret: newAccess(userRoles, ctx, types.KindInferenceSecret),
|
||||
Classifier: newAccess(userRoles, ctx, types.KindClassifier),
|
||||
AutoUpdateConfig: autoUpdateConfig,
|
||||
AutoUpdateVersion: autoUpdateVersion,
|
||||
AutoUpdateAgentRollout: autoUpdateAgentRollout,
|
||||
|
||||
@@ -0,0 +1,176 @@
|
||||
// Teleport
|
||||
// Copyright (C) 2026 Gravitational, Inc.
|
||||
//
|
||||
// This program is free software: you can redistribute it and/or modify
|
||||
// it under the terms of the GNU Affero General Public License as published by
|
||||
// the Free Software Foundation, either version 3 of the License, or
|
||||
// (at your option) any later version.
|
||||
//
|
||||
// This program is distributed in the hope that it will be useful,
|
||||
// but WITHOUT ANY WARRANTY; without even the implied warranty of
|
||||
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
|
||||
// GNU Affero General Public License for more details.
|
||||
//
|
||||
// You should have received a copy of the GNU Affero General Public License
|
||||
// along with this program. If not, see <http://www.gnu.org/licenses/>.
|
||||
|
||||
package resources
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"io"
|
||||
"strings"
|
||||
|
||||
"github.com/gravitational/trace"
|
||||
|
||||
summarizerv1 "github.com/gravitational/teleport/api/gen/proto/go/teleport/summarizer/v1"
|
||||
"github.com/gravitational/teleport/api/types"
|
||||
"github.com/gravitational/teleport/api/utils/clientutils"
|
||||
"github.com/gravitational/teleport/lib/asciitable"
|
||||
"github.com/gravitational/teleport/lib/auth/authclient"
|
||||
"github.com/gravitational/teleport/lib/itertools/stream"
|
||||
"github.com/gravitational/teleport/lib/services"
|
||||
)
|
||||
|
||||
// classifierCollection is a collection of Classifier resources that can be
|
||||
// written to an io.Writer in a human-readable format.
|
||||
type classifierCollection []*summarizerv1.Classifier
|
||||
|
||||
func (c classifierCollection) Resources() []types.Resource {
|
||||
out := make([]types.Resource, 0, len(c))
|
||||
for _, item := range c {
|
||||
out = append(out, types.ProtoResource153ToLegacy(item))
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
func (c classifierCollection) WriteText(w io.Writer, verbose bool) error {
|
||||
headers := []string{"Name", "Description", "Kinds"}
|
||||
var rows [][]string
|
||||
for _, item := range c {
|
||||
meta := item.GetMetadata()
|
||||
spec := item.GetSpec()
|
||||
rows = append(rows, []string{
|
||||
meta.GetName(),
|
||||
meta.GetDescription(),
|
||||
strings.Join(spec.GetKinds(), ", "),
|
||||
})
|
||||
}
|
||||
var t asciitable.Table
|
||||
if verbose {
|
||||
t = asciitable.MakeTable(headers, rows...)
|
||||
} else {
|
||||
t = asciitable.MakeTableWithTruncatedColumn(headers, rows, "Description")
|
||||
}
|
||||
|
||||
// stable sort by name.
|
||||
t.SortRowsBy([]int{0}, true)
|
||||
_, err := t.AsBuffer().WriteTo(w)
|
||||
return trace.Wrap(err)
|
||||
}
|
||||
|
||||
func classifierHandler() Handler {
|
||||
return Handler{
|
||||
getHandler: getClassifiers,
|
||||
createHandler: createClassifier,
|
||||
updateHandler: updateClassifier,
|
||||
deleteHandler: deleteClassifier,
|
||||
singleton: false,
|
||||
mfaRequired: false,
|
||||
description: "Specifies a classifier applied to sessions during summarization",
|
||||
}
|
||||
}
|
||||
|
||||
func createClassifier(
|
||||
ctx context.Context, clt *authclient.Client, raw services.UnknownResource, opts CreateOpts,
|
||||
) error {
|
||||
classifier, err := services.UnmarshalProtoResource[*summarizerv1.Classifier](
|
||||
raw.Raw, services.DisallowUnknown())
|
||||
if err != nil {
|
||||
return trace.Wrap(err)
|
||||
}
|
||||
|
||||
sclt := clt.SummarizerServiceClient()
|
||||
if opts.Force {
|
||||
req := summarizerv1.UpsertClassifierRequest_builder{
|
||||
Classifier: classifier,
|
||||
}.Build()
|
||||
_, err = sclt.UpsertClassifier(ctx, req)
|
||||
} else {
|
||||
req := summarizerv1.CreateClassifierRequest_builder{
|
||||
Classifier: classifier,
|
||||
}.Build()
|
||||
_, err = sclt.CreateClassifier(ctx, req)
|
||||
}
|
||||
if err != nil {
|
||||
return trace.Wrap(err)
|
||||
}
|
||||
|
||||
fmt.Printf("classifier %q has been created\n", classifier.GetMetadata().GetName())
|
||||
return nil
|
||||
}
|
||||
|
||||
func updateClassifier(
|
||||
ctx context.Context, clt *authclient.Client, raw services.UnknownResource, opts CreateOpts,
|
||||
) error {
|
||||
classifier, err := services.UnmarshalProtoResource[*summarizerv1.Classifier](
|
||||
raw.Raw, services.DisallowUnknown())
|
||||
if err != nil {
|
||||
return trace.Wrap(err)
|
||||
}
|
||||
req := summarizerv1.UpdateClassifierRequest_builder{
|
||||
Classifier: classifier,
|
||||
}.Build()
|
||||
if _, err := clt.SummarizerServiceClient().UpdateClassifier(ctx, req); err != nil {
|
||||
return trace.Wrap(err)
|
||||
}
|
||||
fmt.Printf("classifier %q has been updated\n", classifier.GetMetadata().GetName())
|
||||
return nil
|
||||
}
|
||||
|
||||
func getClassifiers(
|
||||
ctx context.Context, clt *authclient.Client, ref services.Ref, opts GetOpts,
|
||||
) (Collection, error) {
|
||||
ssclt := clt.SummarizerServiceClient()
|
||||
if ref.Name != "" {
|
||||
req := summarizerv1.GetClassifierRequest_builder{
|
||||
Name: ref.Name,
|
||||
}.Build()
|
||||
res, err := ssclt.GetClassifier(ctx, req)
|
||||
if err != nil {
|
||||
return nil, trace.Wrap(err)
|
||||
}
|
||||
return classifierCollection{res.GetClassifier()}, nil
|
||||
}
|
||||
|
||||
items, err := stream.Collect(clientutils.Resources(
|
||||
ctx,
|
||||
func(
|
||||
ctx context.Context, _ int, nextPageToken string,
|
||||
) ([]*summarizerv1.Classifier, string, error) {
|
||||
resp, err := ssclt.ListClassifiers(
|
||||
ctx,
|
||||
summarizerv1.ListClassifiersRequest_builder{
|
||||
PageToken: nextPageToken,
|
||||
}.Build())
|
||||
return resp.GetClassifiers(), resp.GetNextPageToken(), trace.Wrap(err)
|
||||
}))
|
||||
if err != nil {
|
||||
return nil, trace.Wrap(err)
|
||||
}
|
||||
return classifierCollection(items), nil
|
||||
}
|
||||
|
||||
func deleteClassifier(
|
||||
ctx context.Context, clt *authclient.Client, ref services.Ref,
|
||||
) error {
|
||||
req := summarizerv1.DeleteClassifierRequest_builder{
|
||||
Name: ref.Name,
|
||||
}.Build()
|
||||
if _, err := clt.SummarizerServiceClient().DeleteClassifier(ctx, req); err != nil {
|
||||
return trace.Wrap(err)
|
||||
}
|
||||
fmt.Printf("classifier %q has been deleted\n", ref.Name)
|
||||
return nil
|
||||
}
|
||||
@@ -54,6 +54,7 @@ func Handlers() map[string]Handler {
|
||||
types.KindBotInstance: botInstanceHandler(),
|
||||
types.KindCertAuthority: certAuthorityHandler(),
|
||||
types.KindCertAuthorityOverride: certAuthorityOverrideHandler(),
|
||||
types.KindClassifier: classifierHandler(),
|
||||
types.KindClusterAuthPreference: authPreferenceHandler(),
|
||||
types.KindClusterMaintenanceConfig: clusterMaintenanceConfigHandler(),
|
||||
types.KindClusterNetworkingConfig: networkingConfigHandler(),
|
||||
|
||||
Reference in New Issue
Block a user