Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
110 changes: 110 additions & 0 deletions authbridge/authlib/listener/extproc/mcpreject_test.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,110 @@
package extproc

import (
"encoding/json"
"testing"

corev3 "github.com/envoyproxy/go-control-plane/envoy/config/core/v3"
extprocv3 "github.com/envoyproxy/go-control-plane/envoy/service/ext_proc/v3"
typev3 "github.com/envoyproxy/go-control-plane/envoy/type/v3"

"github.com/kagenti/kagenti-extensions/authbridge/authlib/pipeline"
)

// pctxWithMCP builds the minimum pctx needed to trigger the JSON-RPC
// rendering path on the extproc listener.
func pctxWithMCP(id any) *pipeline.Context {
return &pipeline.Context{
Extensions: pipeline.Extensions{
MCP: &pipeline.MCPExtension{
Method: "tools/call",
RPCID: id,
},
},
}
}

func ibacBlocked() pipeline.Action {
return pipeline.DenyStatus(403, "ibac.blocked", "intent does not align")
}

// immediateBody fishes the ImmediateResponse body out of a
// ProcessingResponse — saves boilerplate in every assertion below.
func immediateBody(t *testing.T, resp *extprocv3.ProcessingResponse) (typev3.StatusCode, []byte, []*corev3.HeaderValueOption) {
t.Helper()
imm, ok := resp.GetResponse().(*extprocv3.ProcessingResponse_ImmediateResponse)
if !ok {
t.Fatalf("response is not ImmediateResponse: %T", resp.GetResponse())
}
var hs []*corev3.HeaderValueOption
if imm.ImmediateResponse.Headers != nil {
hs = imm.ImmediateResponse.Headers.SetHeaders
}
return imm.ImmediateResponse.Status.Code, imm.ImmediateResponse.Body, hs
}

// TestRejectFromActionForRequest_MCPRequest verifies the extproc
// listener emits an HTTP 200 ImmediateResponse with a JSON-RPC error
// body when the rejected request was MCP — the envoy-sidecar twin of
// the forwardproxy fix.
func TestRejectFromActionForRequest_MCPRequest(t *testing.T) {
resp := rejectFromActionForRequest(ibacBlocked(), pctxWithMCP(float64(7)))
status, body, headers := immediateBody(t, resp)

if status != typev3.StatusCode_OK {
t.Fatalf("status = %v, want OK (200) — JSON-RPC errors travel over a 200 transport", status)
}
gotCT := false
for _, h := range headers {
if h.Header.Key == "content-type" && string(h.Header.RawValue) == "application/json" {
gotCT = true
}
}
if !gotCT {
t.Errorf("content-type header missing or wrong: %v", headers)
}
var parsed map[string]any
if err := json.Unmarshal(body, &parsed); err != nil {
t.Fatalf("body is not JSON: %v\n%s", err, body)
}
if parsed["jsonrpc"] != "2.0" {
t.Errorf("jsonrpc = %v, want 2.0", parsed["jsonrpc"])
}
if id, _ := parsed["id"].(float64); id != 7 {
t.Errorf("id = %v, want 7", parsed["id"])
}
errObj, _ := parsed["error"].(map[string]any)
if errObj["code"] != float64(-32000) {
t.Errorf("error.code = %v, want -32000", errObj["code"])
}
if errObj["message"] != "intent does not align" {
t.Errorf("error.message = %v, want IBAC reason", errObj["message"])
}
}

// TestRejectFromActionForRequest_NonMCPFallsBack: non-MCP outbound
// denials keep today's transport-level error shape. Asserts the
// status mirrors the violation, not 200.
func TestRejectFromActionForRequest_NonMCPFallsBack(t *testing.T) {
resp := rejectFromActionForRequest(ibacBlocked(), &pipeline.Context{})
status, _, _ := immediateBody(t, resp)
if status != typev3.StatusCode_Forbidden {
t.Fatalf("status = %v, want Forbidden (403) — non-MCP request should keep HTTP-level shape", status)
}
}

// TestRejectFromActionForRequest_NotificationFallsBack: JSON-RPC
// notifications (id == nil) get no JSON-RPC response by spec, so we
// fall through to the HTTP-level shape rather than echoing a null id.
func TestRejectFromActionForRequest_NotificationFallsBack(t *testing.T) {
pctx := &pipeline.Context{
Extensions: pipeline.Extensions{
MCP: &pipeline.MCPExtension{Method: "tools/call", RPCID: nil},
},
}
resp := rejectFromActionForRequest(ibacBlocked(), pctx)
status, _, _ := immediateBody(t, resp)
if status != typev3.StatusCode_Forbidden {
t.Fatalf("status = %v, want Forbidden (403) — notification should keep HTTP-level shape", status)
}
}
29 changes: 27 additions & 2 deletions authbridge/authlib/listener/extproc/server.go
Original file line number Diff line number Diff line change
Expand Up @@ -22,6 +22,7 @@ import (
"google.golang.org/grpc/status"

"github.com/kagenti/kagenti-extensions/authbridge/authlib/auth"
"github.com/kagenti/kagenti-extensions/authbridge/authlib/listener/httpx"
"github.com/kagenti/kagenti-extensions/authbridge/authlib/listener/internal/sseframe"
"github.com/kagenti/kagenti-extensions/authbridge/authlib/listener/skiphost"
"github.com/kagenti/kagenti-extensions/authbridge/authlib/pipeline"
Expand Down Expand Up @@ -499,7 +500,7 @@ func (s *Server) handleOutbound(stream extprocv3.ExternalProcessor_ProcessServer
if action.Type == pipeline.Reject {
s.recordOutboundReject(pctx, action)
s.OutboundPipeline.RunFinish(ctx, pctx, pipeline.OutcomeFromContext(pctx))
return rejectFromAction(action), nil
return rejectFromActionForRequest(action, pctx), nil
}

s.recordOutboundSession(pctx)
Expand Down Expand Up @@ -551,7 +552,7 @@ func (s *Server) handleOutboundBody(stream extprocv3.ExternalProcessor_ProcessSe
if action.Type == pipeline.Reject {
s.recordOutboundReject(pctx, action)
s.OutboundPipeline.RunFinish(ctx, pctx, pipeline.OutcomeFromContext(pctx))
return rejectFromAction(action), nil
return rejectFromActionForRequest(action, pctx), nil
}

s.recordOutboundSession(pctx)
Expand Down Expand Up @@ -851,6 +852,30 @@ func replaceTokenResponse(token string) *extprocv3.ProcessingResponse {
}
}

// rejectFromActionForRequest is the MCP-aware sibling of rejectFromAction.
// When pctx carries an MCP JSON-RPC request shape (Method + non-nil RPCID),
// the response is an HTTP 200 carrying a JSON-RPC 2.0 error frame so the
// caller's MCP client surfaces this as one failed tool call rather than a
// transport break. All other shapes fall through to rejectFromAction.
func rejectFromActionForRequest(action pipeline.Action, pctx *pipeline.Context) *extprocv3.ProcessingResponse {
if pctx != nil && pctx.Extensions.MCP != nil &&

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

suggestion — The MCP-detection predicate here (MCP != nil && Method != "" && RPCID != nil) is reimplemented inline, while httpx.shouldRenderMCPError encodes the identical condition. They're logically equal today but can drift independently. Consider exporting the httpx predicate (e.g. httpx.ShouldRenderMCPError(pctx)) and calling it from both listeners so the "is this an MCP request?" rule lives in one place.

pctx.Extensions.MCP.Method != "" && pctx.Extensions.MCP.RPCID != nil {
body := httpx.MarshalMCPRejectionBody(action, pctx.Extensions.MCP.RPCID)
return &extprocv3.ProcessingResponse{
Response: &extprocv3.ProcessingResponse_ImmediateResponse{
ImmediateResponse: &extprocv3.ImmediateResponse{
Status: &typev3.HttpStatus{Code: typev3.StatusCode(http.StatusOK)},
Body: body,
Headers: &extprocv3.HeaderMutation{SetHeaders: []*corev3.HeaderValueOption{{
Header: &corev3.HeaderValue{Key: "content-type", RawValue: []byte("application/json")},
}}},
},
},
}
}
return rejectFromAction(action)
}

// rejectFromAction turns a pipeline Reject into an Envoy ImmediateResponse,
// preserving the plugin's status/headers/body. Replaces the old
// denyResponse helper which hardcoded {"error":...,"message":...} at each
Expand Down
14 changes: 12 additions & 2 deletions authbridge/authlib/listener/forwardproxy/server.go
Original file line number Diff line number Diff line change
Expand Up @@ -260,7 +260,12 @@ func (s *Server) handleRequest(w http.ResponseWriter, r *http.Request) {

if action.Type == pipeline.Reject {
s.recordOutboundReject(pctx, action)
httpx.WriteRejection(w, action)
// Render as a JSON-RPC error frame when the rejected
// request was MCP JSON-RPC, so the agent's MCP client
// surfaces this as one failed tool call rather than a
// transport break. Falls through to plain HTTP-level
// rejection for non-MCP traffic.
httpx.WriteRejectionForRequest(w, action, pctx)
return
}
}
Expand Down Expand Up @@ -748,7 +753,12 @@ func (s *Server) handleConnect(w http.ResponseWriter, r *http.Request) {
action := s.OutboundPipeline.Run(r.Context(), pctx)
if action.Type == pipeline.Reject {
s.recordOutboundReject(pctx, action)
httpx.WriteRejection(w, action)
// Render as a JSON-RPC error frame when the rejected
// request was MCP JSON-RPC, so the agent's MCP client
// surfaces this as one failed tool call rather than a
// transport break. Falls through to plain HTTP-level
// rejection for non-MCP traffic.
httpx.WriteRejectionForRequest(w, action, pctx)
return
}
}
Expand Down
126 changes: 126 additions & 0 deletions authbridge/authlib/listener/httpx/render.go
Original file line number Diff line number Diff line change
Expand Up @@ -4,6 +4,7 @@
package httpx

import (
"encoding/json"
"net/http"

"github.com/kagenti/kagenti-extensions/authbridge/authlib/pipeline"
Expand All @@ -28,3 +29,128 @@ func WriteRejection(w http.ResponseWriter, action pipeline.Action) {
w.WriteHeader(status)
_, _ = w.Write(body)
}

// WriteRejectionForRequest renders a Reject the same way as
// WriteRejection EXCEPT when the rejected request was a JSON-RPC
// request that an MCP-aware parser already classified — in that case
// it writes a JSON-RPC 2.0 error frame at HTTP 200 with the original
// id echoed back, so the caller's MCP client surfaces this as a
// failed tool call rather than a transport break.
//
// The MCP-shape detection is conservative: we only rewrite when
// pctx.Extensions.MCP is populated with a non-empty Method AND a
// non-nil RPCID. JSON-RPC notifications (no id) deliberately fall
// through to plain WriteRejection — by spec the client expects no
// response, so emitting a JSON-RPC error frame would violate the
// notification contract.
//
// All non-MCP requests fall through to WriteRejection, so this is a
// safe drop-in replacement at any call site.
func WriteRejectionForRequest(w http.ResponseWriter, action pipeline.Action, pctx *pipeline.Context) {
if !shouldRenderMCPError(pctx) {
WriteRejection(w, action)
return
}
writeMCPRejection(w, action, pctx.Extensions.MCP.RPCID)
}

func shouldRenderMCPError(pctx *pipeline.Context) bool {
if pctx == nil || pctx.Extensions.MCP == nil {
return false
}
mcp := pctx.Extensions.MCP
if mcp.Method == "" || mcp.RPCID == nil {
return false
}
return true
}

// JSON-RPC 2.0 error code for application errors. -32000..-32099 is
// the "implementation-defined server-error" range reserved by the
// spec; -32000 is the conventional generic value used when there's
// no protocol-level reason for a more specific code. Authbridge's
// denials are policy decisions outside the JSON-RPC layer, so the
// generic server-error code fits — operators read the human reason
// and structured details, not the numeric code, to understand what
// happened.
const jsonRPCServerError = -32000

func writeMCPRejection(w http.ResponseWriter, action pipeline.Action, id any) {
body := MarshalMCPRejectionBody(action, id)
w.Header().Set("Content-Type", "application/json")
w.WriteHeader(http.StatusOK)
_, _ = w.Write(body)
}

// mcpRejectionFallback is a precomputed JSON-RPC 2.0 error frame used as
// the last-resort body when both the full and minimal marshal attempts
// fail. It is guaranteed to round-trip through json.Unmarshal so MCP
// clients always see a parseable frame on a 200 application/json
// response.
var mcpRejectionFallback = []byte(`{"jsonrpc":"2.0","id":null,"error":{"code":-32000,"message":"request rejected"}}`)

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

nitmcpRejectionFallback hardcodes -32000 in the raw string literal instead of referencing the jsonRPCServerError const just below. If the const ever changes, this last-resort frame silently diverges. A // keep in sync with jsonRPCServerError comment (or building it via fmt.Sprintf at init) would close the gap.

// MarshalMCPRejectionBody renders a Reject Action as a JSON-RPC 2.0
// error frame body. It is guaranteed to return a non-empty, parseable
// JSON byte slice so callers can safely send it on a 200
// application/json response: marshal the full frame first; on error
// (e.g. a plugin populated Violation.Details with an unmarshalable
// value, or id is unmarshalable), retry with a minimal frame that omits
// optional data and falls back to id=null; on further error, return a
// constant precomputed frame.
func MarshalMCPRejectionBody(action pipeline.Action, id any) []byte {
v := action.Violation
message := "request rejected"
var data map[string]any
if v != nil {
if v.Reason != "" {
message = v.Reason
}
data = map[string]any{}
if v.Code != "" {
data["error"] = v.Code
}
if v.PluginName != "" {
data["plugin"] = v.PluginName
}
if v.Description != "" {
data["description"] = v.Description
}
if len(v.Details) > 0 {
data["details"] = v.Details
}
if len(data) == 0 {
data = nil
}
}
if body, err := json.Marshal(map[string]any{
"jsonrpc": "2.0",
"id": id,
"error": map[string]any{
"code": jsonRPCServerError,
"message": message,
"data": data,

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Nit: "data": data is included unconditionally, so when action.Violation is nil (or all Violation fields are empty and data stays nil), the wire response contains "data":null. JSON-RPC 2.0 §5.1 says the data member SHOULD be omitted when there are no additional details. Consider a conditional include:

errorObj := map[string]any{
    "code":    jsonRPCServerError,
    "message": message,
}
if data != nil {
    errorObj["data"] = data
}

Not a blocker — the current shape is spec-valid — but it avoids a superfluous null field that some strict MCP clients might warn on.

},
}); err == nil {
return body
}
// Full frame failed to marshal — retry without optional data, and
// keep the original id only if it survives marshaling on its own
// (otherwise drop to null per JSON-RPC 2.0 §5.1).
safeID := any(nil)
if id != nil {
if _, err := json.Marshal(id); err == nil {
safeID = id
}
}
if body, err := json.Marshal(map[string]any{
"jsonrpc": "2.0",
"id": safeID,
"error": map[string]any{
"code": jsonRPCServerError,
"message": message,
},
}); err == nil {
return body
}
return mcpRejectionFallback
}
Loading
Loading