Skip to content

Commit a6bc755

Browse files
araujofaraujof
committed
Merge remote-tracking branch 'origin/main' into feat/cpex_plugin
2 parents 0295322 + 1cd2a12 commit a6bc755

5 files changed

Lines changed: 522 additions & 4 deletions

File tree

‎authbridge/authlib/listener/extproc/mcpreject_test.go‎

Lines changed: 110 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,110 @@
1+
package extproc
2+
3+
import (
4+
"encoding/json"
5+
"testing"
6+
7+
corev3 "github.com/envoyproxy/go-control-plane/envoy/config/core/v3"
8+
extprocv3 "github.com/envoyproxy/go-control-plane/envoy/service/ext_proc/v3"
9+
typev3 "github.com/envoyproxy/go-control-plane/envoy/type/v3"
10+
11+
"github.com/kagenti/kagenti-extensions/authbridge/authlib/pipeline"
12+
)
13+
14+
// pctxWithMCP builds the minimum pctx needed to trigger the JSON-RPC
15+
// rendering path on the extproc listener.
16+
func pctxWithMCP(id any) *pipeline.Context {
17+
return &pipeline.Context{
18+
Extensions: pipeline.Extensions{
19+
MCP: &pipeline.MCPExtension{
20+
Method: "tools/call",
21+
RPCID: id,
22+
},
23+
},
24+
}
25+
}
26+
27+
func ibacBlocked() pipeline.Action {
28+
return pipeline.DenyStatus(403, "ibac.blocked", "intent does not align")
29+
}
30+
31+
// immediateBody fishes the ImmediateResponse body out of a
32+
// ProcessingResponse — saves boilerplate in every assertion below.
33+
func immediateBody(t *testing.T, resp *extprocv3.ProcessingResponse) (typev3.StatusCode, []byte, []*corev3.HeaderValueOption) {
34+
t.Helper()
35+
imm, ok := resp.GetResponse().(*extprocv3.ProcessingResponse_ImmediateResponse)
36+
if !ok {
37+
t.Fatalf("response is not ImmediateResponse: %T", resp.GetResponse())
38+
}
39+
var hs []*corev3.HeaderValueOption
40+
if imm.ImmediateResponse.Headers != nil {
41+
hs = imm.ImmediateResponse.Headers.SetHeaders
42+
}
43+
return imm.ImmediateResponse.Status.Code, imm.ImmediateResponse.Body, hs
44+
}
45+
46+
// TestRejectFromActionForRequest_MCPRequest verifies the extproc
47+
// listener emits an HTTP 200 ImmediateResponse with a JSON-RPC error
48+
// body when the rejected request was MCP — the envoy-sidecar twin of
49+
// the forwardproxy fix.
50+
func TestRejectFromActionForRequest_MCPRequest(t *testing.T) {
51+
resp := rejectFromActionForRequest(ibacBlocked(), pctxWithMCP(float64(7)))
52+
status, body, headers := immediateBody(t, resp)
53+
54+
if status != typev3.StatusCode_OK {
55+
t.Fatalf("status = %v, want OK (200) — JSON-RPC errors travel over a 200 transport", status)
56+
}
57+
gotCT := false
58+
for _, h := range headers {
59+
if h.Header.Key == "content-type" && string(h.Header.RawValue) == "application/json" {
60+
gotCT = true
61+
}
62+
}
63+
if !gotCT {
64+
t.Errorf("content-type header missing or wrong: %v", headers)
65+
}
66+
var parsed map[string]any
67+
if err := json.Unmarshal(body, &parsed); err != nil {
68+
t.Fatalf("body is not JSON: %v\n%s", err, body)
69+
}
70+
if parsed["jsonrpc"] != "2.0" {
71+
t.Errorf("jsonrpc = %v, want 2.0", parsed["jsonrpc"])
72+
}
73+
if id, _ := parsed["id"].(float64); id != 7 {
74+
t.Errorf("id = %v, want 7", parsed["id"])
75+
}
76+
errObj, _ := parsed["error"].(map[string]any)
77+
if errObj["code"] != float64(-32000) {
78+
t.Errorf("error.code = %v, want -32000", errObj["code"])
79+
}
80+
if errObj["message"] != "intent does not align" {
81+
t.Errorf("error.message = %v, want IBAC reason", errObj["message"])
82+
}
83+
}
84+
85+
// TestRejectFromActionForRequest_NonMCPFallsBack: non-MCP outbound
86+
// denials keep today's transport-level error shape. Asserts the
87+
// status mirrors the violation, not 200.
88+
func TestRejectFromActionForRequest_NonMCPFallsBack(t *testing.T) {
89+
resp := rejectFromActionForRequest(ibacBlocked(), &pipeline.Context{})
90+
status, _, _ := immediateBody(t, resp)
91+
if status != typev3.StatusCode_Forbidden {
92+
t.Fatalf("status = %v, want Forbidden (403) — non-MCP request should keep HTTP-level shape", status)
93+
}
94+
}
95+
96+
// TestRejectFromActionForRequest_NotificationFallsBack: JSON-RPC
97+
// notifications (id == nil) get no JSON-RPC response by spec, so we
98+
// fall through to the HTTP-level shape rather than echoing a null id.
99+
func TestRejectFromActionForRequest_NotificationFallsBack(t *testing.T) {
100+
pctx := &pipeline.Context{
101+
Extensions: pipeline.Extensions{
102+
MCP: &pipeline.MCPExtension{Method: "tools/call", RPCID: nil},
103+
},
104+
}
105+
resp := rejectFromActionForRequest(ibacBlocked(), pctx)
106+
status, _, _ := immediateBody(t, resp)
107+
if status != typev3.StatusCode_Forbidden {
108+
t.Fatalf("status = %v, want Forbidden (403) — notification should keep HTTP-level shape", status)
109+
}
110+
}

‎authbridge/authlib/listener/extproc/server.go‎

Lines changed: 27 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -22,6 +22,7 @@ import (
2222
"google.golang.org/grpc/status"
2323

2424
"github.com/kagenti/kagenti-extensions/authbridge/authlib/auth"
25+
"github.com/kagenti/kagenti-extensions/authbridge/authlib/listener/httpx"
2526
"github.com/kagenti/kagenti-extensions/authbridge/authlib/listener/internal/sseframe"
2627
"github.com/kagenti/kagenti-extensions/authbridge/authlib/listener/skiphost"
2728
"github.com/kagenti/kagenti-extensions/authbridge/authlib/pipeline"
@@ -499,7 +500,7 @@ func (s *Server) handleOutbound(stream extprocv3.ExternalProcessor_ProcessServer
499500
if action.Type == pipeline.Reject {
500501
s.recordOutboundReject(pctx, action)
501502
s.OutboundPipeline.RunFinish(ctx, pctx, pipeline.OutcomeFromContext(pctx))
502-
return rejectFromAction(action), nil
503+
return rejectFromActionForRequest(action, pctx), nil
503504
}
504505

505506
s.recordOutboundSession(pctx)
@@ -551,7 +552,7 @@ func (s *Server) handleOutboundBody(stream extprocv3.ExternalProcessor_ProcessSe
551552
if action.Type == pipeline.Reject {
552553
s.recordOutboundReject(pctx, action)
553554
s.OutboundPipeline.RunFinish(ctx, pctx, pipeline.OutcomeFromContext(pctx))
554-
return rejectFromAction(action), nil
555+
return rejectFromActionForRequest(action, pctx), nil
555556
}
556557

557558
s.recordOutboundSession(pctx)
@@ -851,6 +852,30 @@ func replaceTokenResponse(token string) *extprocv3.ProcessingResponse {
851852
}
852853
}
853854

855+
// rejectFromActionForRequest is the MCP-aware sibling of rejectFromAction.
856+
// When pctx carries an MCP JSON-RPC request shape (Method + non-nil RPCID),
857+
// the response is an HTTP 200 carrying a JSON-RPC 2.0 error frame so the
858+
// caller's MCP client surfaces this as one failed tool call rather than a
859+
// transport break. All other shapes fall through to rejectFromAction.
860+
func rejectFromActionForRequest(action pipeline.Action, pctx *pipeline.Context) *extprocv3.ProcessingResponse {
861+
if pctx != nil && pctx.Extensions.MCP != nil &&
862+
pctx.Extensions.MCP.Method != "" && pctx.Extensions.MCP.RPCID != nil {
863+
body := httpx.MarshalMCPRejectionBody(action, pctx.Extensions.MCP.RPCID)
864+
return &extprocv3.ProcessingResponse{
865+
Response: &extprocv3.ProcessingResponse_ImmediateResponse{
866+
ImmediateResponse: &extprocv3.ImmediateResponse{
867+
Status: &typev3.HttpStatus{Code: typev3.StatusCode(http.StatusOK)},
868+
Body: body,
869+
Headers: &extprocv3.HeaderMutation{SetHeaders: []*corev3.HeaderValueOption{{
870+
Header: &corev3.HeaderValue{Key: "content-type", RawValue: []byte("application/json")},
871+
}}},
872+
},
873+
},
874+
}
875+
}
876+
return rejectFromAction(action)
877+
}
878+
854879
// rejectFromAction turns a pipeline Reject into an Envoy ImmediateResponse,
855880
// preserving the plugin's status/headers/body. Replaces the old
856881
// denyResponse helper which hardcoded {"error":...,"message":...} at each

‎authbridge/authlib/listener/forwardproxy/server.go‎

Lines changed: 12 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -260,7 +260,12 @@ func (s *Server) handleRequest(w http.ResponseWriter, r *http.Request) {
260260

261261
if action.Type == pipeline.Reject {
262262
s.recordOutboundReject(pctx, action)
263-
httpx.WriteRejection(w, action)
263+
// Render as a JSON-RPC error frame when the rejected
264+
// request was MCP JSON-RPC, so the agent's MCP client
265+
// surfaces this as one failed tool call rather than a
266+
// transport break. Falls through to plain HTTP-level
267+
// rejection for non-MCP traffic.
268+
httpx.WriteRejectionForRequest(w, action, pctx)
264269
return
265270
}
266271
}
@@ -748,7 +753,12 @@ func (s *Server) handleConnect(w http.ResponseWriter, r *http.Request) {
748753
action := s.OutboundPipeline.Run(r.Context(), pctx)
749754
if action.Type == pipeline.Reject {
750755
s.recordOutboundReject(pctx, action)
751-
httpx.WriteRejection(w, action)
756+
// Render as a JSON-RPC error frame when the rejected
757+
// request was MCP JSON-RPC, so the agent's MCP client
758+
// surfaces this as one failed tool call rather than a
759+
// transport break. Falls through to plain HTTP-level
760+
// rejection for non-MCP traffic.
761+
httpx.WriteRejectionForRequest(w, action, pctx)
752762
return
753763
}
754764
}

‎authbridge/authlib/listener/httpx/render.go‎

Lines changed: 126 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -4,6 +4,7 @@
44
package httpx
55

66
import (
7+
"encoding/json"
78
"net/http"
89

910
"github.com/kagenti/kagenti-extensions/authbridge/authlib/pipeline"
@@ -28,3 +29,128 @@ func WriteRejection(w http.ResponseWriter, action pipeline.Action) {
2829
w.WriteHeader(status)
2930
_, _ = w.Write(body)
3031
}
32+
33+
// WriteRejectionForRequest renders a Reject the same way as
34+
// WriteRejection EXCEPT when the rejected request was a JSON-RPC
35+
// request that an MCP-aware parser already classified — in that case
36+
// it writes a JSON-RPC 2.0 error frame at HTTP 200 with the original
37+
// id echoed back, so the caller's MCP client surfaces this as a
38+
// failed tool call rather than a transport break.
39+
//
40+
// The MCP-shape detection is conservative: we only rewrite when
41+
// pctx.Extensions.MCP is populated with a non-empty Method AND a
42+
// non-nil RPCID. JSON-RPC notifications (no id) deliberately fall
43+
// through to plain WriteRejection — by spec the client expects no
44+
// response, so emitting a JSON-RPC error frame would violate the
45+
// notification contract.
46+
//
47+
// All non-MCP requests fall through to WriteRejection, so this is a
48+
// safe drop-in replacement at any call site.
49+
func WriteRejectionForRequest(w http.ResponseWriter, action pipeline.Action, pctx *pipeline.Context) {
50+
if !shouldRenderMCPError(pctx) {
51+
WriteRejection(w, action)
52+
return
53+
}
54+
writeMCPRejection(w, action, pctx.Extensions.MCP.RPCID)
55+
}
56+
57+
func shouldRenderMCPError(pctx *pipeline.Context) bool {
58+
if pctx == nil || pctx.Extensions.MCP == nil {
59+
return false
60+
}
61+
mcp := pctx.Extensions.MCP
62+
if mcp.Method == "" || mcp.RPCID == nil {
63+
return false
64+
}
65+
return true
66+
}
67+
68+
// JSON-RPC 2.0 error code for application errors. -32000..-32099 is
69+
// the "implementation-defined server-error" range reserved by the
70+
// spec; -32000 is the conventional generic value used when there's
71+
// no protocol-level reason for a more specific code. Authbridge's
72+
// denials are policy decisions outside the JSON-RPC layer, so the
73+
// generic server-error code fits — operators read the human reason
74+
// and structured details, not the numeric code, to understand what
75+
// happened.
76+
const jsonRPCServerError = -32000
77+
78+
func writeMCPRejection(w http.ResponseWriter, action pipeline.Action, id any) {
79+
body := MarshalMCPRejectionBody(action, id)
80+
w.Header().Set("Content-Type", "application/json")
81+
w.WriteHeader(http.StatusOK)
82+
_, _ = w.Write(body)
83+
}
84+
85+
// mcpRejectionFallback is a precomputed JSON-RPC 2.0 error frame used as
86+
// the last-resort body when both the full and minimal marshal attempts
87+
// fail. It is guaranteed to round-trip through json.Unmarshal so MCP
88+
// clients always see a parseable frame on a 200 application/json
89+
// response.
90+
var mcpRejectionFallback = []byte(`{"jsonrpc":"2.0","id":null,"error":{"code":-32000,"message":"request rejected"}}`)
91+
92+
// MarshalMCPRejectionBody renders a Reject Action as a JSON-RPC 2.0
93+
// error frame body. It is guaranteed to return a non-empty, parseable
94+
// JSON byte slice so callers can safely send it on a 200
95+
// application/json response: marshal the full frame first; on error
96+
// (e.g. a plugin populated Violation.Details with an unmarshalable
97+
// value, or id is unmarshalable), retry with a minimal frame that omits
98+
// optional data and falls back to id=null; on further error, return a
99+
// constant precomputed frame.
100+
func MarshalMCPRejectionBody(action pipeline.Action, id any) []byte {
101+
v := action.Violation
102+
message := "request rejected"
103+
var data map[string]any
104+
if v != nil {
105+
if v.Reason != "" {
106+
message = v.Reason
107+
}
108+
data = map[string]any{}
109+
if v.Code != "" {
110+
data["error"] = v.Code
111+
}
112+
if v.PluginName != "" {
113+
data["plugin"] = v.PluginName
114+
}
115+
if v.Description != "" {
116+
data["description"] = v.Description
117+
}
118+
if len(v.Details) > 0 {
119+
data["details"] = v.Details
120+
}
121+
if len(data) == 0 {
122+
data = nil
123+
}
124+
}
125+
if body, err := json.Marshal(map[string]any{
126+
"jsonrpc": "2.0",
127+
"id": id,
128+
"error": map[string]any{
129+
"code": jsonRPCServerError,
130+
"message": message,
131+
"data": data,
132+
},
133+
}); err == nil {
134+
return body
135+
}
136+
// Full frame failed to marshal — retry without optional data, and
137+
// keep the original id only if it survives marshaling on its own
138+
// (otherwise drop to null per JSON-RPC 2.0 §5.1).
139+
safeID := any(nil)
140+
if id != nil {
141+
if _, err := json.Marshal(id); err == nil {
142+
safeID = id
143+
}
144+
}
145+
if body, err := json.Marshal(map[string]any{
146+
"jsonrpc": "2.0",
147+
"id": safeID,
148+
"error": map[string]any{
149+
"code": jsonRPCServerError,
150+
"message": message,
151+
},
152+
}); err == nil {
153+
return body
154+
}
155+
return mcpRejectionFallback
156+
}

0 commit comments

Comments
 (0)