feat(discovery+openshell): DNS-AID agent discovery with OpenShell policy enforcement #469
There are no files selected for viewing
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,59 @@ | ||
| // Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. | ||
| // | ||
| // Licensed under the Apache License, Version 2.0 (the "License"); | ||
| // you may not use this file except in compliance with the License. | ||
| // You may obtain a copy of the License at | ||
| // | ||
| // http://www.apache.org/licenses/LICENSE-2.0 | ||
| // | ||
| // Unless required by applicable law or agreed to in writing, software | ||
| // distributed under the License is distributed on an "AS IS" BASIS, | ||
| // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. | ||
| // See the License for the specific language governing permissions and | ||
| // limitations under the License. | ||
|
|
||
| // Package testdns provides shared test utilities for spinning up an in-process | ||
| // DNS server. It lives in internal/ so it is only importable from this module | ||
| // and can be reused across pkg/discovery, pkg/api, and any future consumers | ||
| // without each package keeping its own copy. | ||
| package testdns | ||
|
|
||
| import ( | ||
| "context" | ||
| "net" | ||
| "testing" | ||
|
|
||
| "github.com/miekg/dns" | ||
| ) | ||
|
|
||
| // Start launches a UDP DNS server on 127.0.0.1 with an ephemeral port that | ||
| // dispatches to the supplied handler. It returns the host:port the caller | ||
| // should configure their DNS client with. The server is shut down via | ||
| // t.Cleanup when the test finishes. | ||
| func Start(t *testing.T, handler dns.HandlerFunc) string { | ||
| t.Helper() | ||
|
|
||
| // noctx lint requires the context-aware ListenConfig variant rather | ||
| // than the bare net.ListenPacket. Background() is fine here — test | ||
| // servers don't need a real deadline. | ||
| lc := &net.ListenConfig{} | ||
| pc, err := lc.ListenPacket(context.Background(), "udp", "127.0.0.1:0") | ||
| if err != nil { | ||
| t.Fatalf("testdns: listen: %v", err) | ||
| } | ||
|
|
||
| server := &dns.Server{ | ||
| PacketConn: pc, | ||
| Handler: handler, | ||
| } | ||
|
|
||
| go func() { | ||
| _ = server.ActivateAndServe() | ||
| }() | ||
|
|
||
| t.Cleanup(func() { | ||
| _ = server.Shutdown() | ||
| }) | ||
|
|
||
| return pc.LocalAddr().String() | ||
| } |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,159 @@ | ||
| // Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. | ||
| // | ||
| // Licensed under the Apache License, Version 2.0 (the "License"); | ||
| // you may not use this file except in compliance with the License. | ||
| // You may obtain a copy of the License at | ||
| // | ||
| // http://www.apache.org/licenses/LICENSE-2.0 | ||
| // | ||
| // Unless required by applicable law or agreed to in writing, software | ||
| // distributed under the License is distributed on an "AS IS" BASIS, | ||
| // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. | ||
| // See the License for the specific language governing permissions and | ||
| // limitations under the License. | ||
|
|
||
| package api | ||
|
|
||
| import ( | ||
| "context" | ||
| "net/http" | ||
| "regexp" | ||
|
|
||
| "log/slog" | ||
|
|
||
| "golang.org/x/sync/errgroup" | ||
|
|
||
| "github.com/NVIDIA/aicr/pkg/defaults" | ||
| "github.com/NVIDIA/aicr/pkg/discovery" | ||
| "github.com/NVIDIA/aicr/pkg/errors" | ||
| "github.com/NVIDIA/aicr/pkg/openshell" | ||
| "github.com/NVIDIA/aicr/pkg/serializer" | ||
| "github.com/NVIDIA/aicr/pkg/server" | ||
| ) | ||
|
|
||
| // maxPolicyCheckConcurrency bounds the number of in-flight policy fetches | ||
| // per /v1/agents request. Sequential fetches with a 3s PolicyFetchTimeout | ||
| // would otherwise risk N×3s latency under fanout. | ||
| const maxPolicyCheckConcurrency = 8 | ||
|
|
||
| // dnsNameRegex validates DNS domain names: labels separated by dots, optional trailing dot. | ||
| // Each label: 1-63 chars of [a-z0-9_-], total <= 253 chars. | ||
| var dnsNameRegex = regexp.MustCompile(`^[a-z0-9_][a-z0-9_\-]{0,62}(\.[a-z0-9_][a-z0-9_\-]{0,62})*\.?$`) | ||
|
Comment on lines
+40
to
+41
There was a problem hiding this comment. 🧹 Nitpick | 🔵 Trivial | 💤 Low value Clarify underscore allowance in DNS label validation. The regex permits underscores ( 🤖 Prompt for AI Agents |
||
|
|
||
| // agentResponse represents a single agent in the /v1/agents JSON response. | ||
| type agentResponse struct { | ||
| Name string `json:"name"` | ||
| Protocol string `json:"protocol"` | ||
| Endpoint string `json:"endpoint"` | ||
| Port uint16 `json:"port"` | ||
| Priority uint16 `json:"priority"` | ||
| Params agentParamsSummary `json:"params,omitempty"` | ||
| PolicyResult *policyResultSummary `json:"policyResult,omitempty"` | ||
| } | ||
|
|
||
| // agentParamsSummary is a subset of SvcParams relevant for the API response. | ||
| type agentParamsSummary struct { | ||
| Realm string `json:"realm,omitempty"` | ||
| BAP []string `json:"bap,omitempty"` | ||
| Policy string `json:"policy,omitempty"` | ||
| ConnectClass string `json:"connectClass,omitempty"` | ||
| } | ||
|
|
||
| // policyResultSummary is the OpenShell policy evaluation result for an agent. | ||
| type policyResultSummary struct { | ||
| Allowed bool `json:"allowed"` | ||
| Violations []string `json:"violations,omitempty"` | ||
| } | ||
|
|
||
| // handleAgents returns an http.HandlerFunc that discovers all agents in a domain, | ||
| // evaluates OpenShell policy for each, and returns them as JSON. The domain is | ||
| // taken from the "domain" query parameter and defaults to defaultDomain. | ||
| // If guard is nil, policy evaluation is skipped. | ||
| func handleAgents(disc *discovery.Discoverer, guard *openshell.Guard, defaultDomain string) http.HandlerFunc { | ||
| return func(w http.ResponseWriter, r *http.Request) { | ||
| if r.Method != http.MethodGet { | ||
| w.Header().Set("Allow", http.MethodGet) | ||
| server.WriteError(w, r, http.StatusMethodNotAllowed, errors.ErrCodeMethodNotAllowed, | ||
| "Method not allowed", false, nil) | ||
| return | ||
| } | ||
|
|
||
| ctx, cancel := context.WithTimeout(r.Context(), defaults.DiscoveryHandlerTimeout) | ||
| defer cancel() | ||
|
|
||
| domain := r.URL.Query().Get("domain") | ||
| if domain == "" { | ||
| domain = defaultDomain | ||
| } | ||
| if len(domain) > 253 || !dnsNameRegex.MatchString(domain) { | ||
| server.WriteError(w, r, http.StatusBadRequest, errors.ErrCodeInvalidRequest, | ||
|
IngmarVG-IB marked this conversation as resolved.
|
||
| "invalid domain parameter", false, map[string]any{"domain": domain}) | ||
| return | ||
| } | ||
|
|
||
| records, err := disc.DiscoverAll(ctx, domain) | ||
| if err != nil { | ||
| server.WriteErrorFromErr(w, r, err, "agent discovery failed", nil) | ||
| return | ||
| } | ||
|
|
||
| // Build the response shells first so per-agent policy checks can fan | ||
| // out concurrently into stable indices. | ||
| agents := make([]agentResponse, len(records)) | ||
| for i := range records { | ||
| rec := &records[i] | ||
| agents[i] = agentResponse{ | ||
| Name: rec.Name, | ||
| Protocol: string(rec.Protocol), | ||
| Endpoint: rec.Endpoint, | ||
| Port: rec.Port, | ||
| Priority: rec.Priority, | ||
| Params: agentParamsSummary{ | ||
| Realm: rec.Params.Realm, | ||
| BAP: rec.Params.BAP, | ||
| Policy: rec.Params.Policy, | ||
| ConnectClass: rec.Params.ConnectClass, | ||
| }, | ||
| } | ||
| } | ||
|
|
||
| // Evaluate OpenShell policy concurrently. Sequential checks would | ||
| // scale as N×PolicyFetchTimeout (3s) on a cold cache and risk | ||
| // blowing past DiscoveryHandlerTimeout. | ||
| if guard != nil { | ||
| g, gctx := errgroup.WithContext(ctx) | ||
| g.SetLimit(maxPolicyCheckConcurrency) | ||
| for i := range records { | ||
| if records[i].Params.Policy == "" { | ||
| continue | ||
| } | ||
| g.Go(func() error { | ||
| rec := &records[i] | ||
| result, err := guard.Check(gctx, rec) | ||
| if err != nil { | ||
| slog.Warn("openshell guard check failed", | ||
| "agent", rec.Name, "error", err) | ||
| return nil // skip; do not fail the whole request | ||
| } | ||
| summary := policyResultSummary{Allowed: result.Allowed} | ||
| for _, v := range result.Violations { | ||
| summary.Violations = append(summary.Violations, v.Rule+": "+v.Detail) | ||
| } | ||
| agents[i].PolicyResult = &summary | ||
| return nil | ||
| }) | ||
| } | ||
| // Goroutines never return errors today, but Wait() also surfaces | ||
| // gctx cancellation if the request deadline trips first. | ||
| _ = g.Wait() | ||
| } | ||
|
|
||
| resp := map[string]any{ | ||
| "domain": domain, | ||
| "agents": agents, | ||
| "count": len(agents), | ||
| } | ||
|
|
||
| serializer.RespondJSON(w, http.StatusOK, resp) | ||
| } | ||
| } | ||
Oops, something went wrong.
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
Uh oh!
There was an error while loading. Please reload this page.