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
14 changes: 12 additions & 2 deletions internal/stackql/util/plan_utility.go
Original file line number Diff line number Diff line change
Expand Up @@ -238,12 +238,18 @@ func TransformSQLRawParameters(input map[string]interface{}, ignoreTuples bool)
for k, v := range input {
switch v := v.(type) {
case *sqlparser.FuncExpr:
logging.GetLogger().Infof("%v\n", v)
// JSON() values pass through intact for query param transposition
// (issue #700); other FuncExprs are evaluated elsewhere (issue #689).
if isJSONFuncExpr(v) {
rv[k] = v
}
continue
case parserutil.ParameterMetadata:
switch t := v.GetVal().(type) { //nolint:gocritic // understandable
case *sqlparser.FuncExpr:
logging.GetLogger().Infof("%v\n", t)
if isJSONFuncExpr(t) {
rv[k] = t
}
continue
}
}
Expand All @@ -259,6 +265,10 @@ func TransformSQLRawParameters(input map[string]interface{}, ignoreTuples bool)
return rv, nil
}

func isJSONFuncExpr(f *sqlparser.FuncExpr) bool {
return strings.EqualFold(f.Name.GetRawVal(), "JSON")
}

func extractRaw(raw interface{}, ignoreTuples bool) (interface{}, error) {
switch r := raw.(type) {
case sqlparser.BoolVal:
Expand Down
31 changes: 31 additions & 0 deletions internal/stackql/util/plan_utility_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -4,6 +4,7 @@ import (
"testing"

"github.com/stackql/stackql-parser/go/vt/sqlparser"
"github.com/stackql/stackql/internal/stackql/parserutil"
)

func mustParse(t *testing.T, sql string) sqlparser.Statement {
Expand Down Expand Up @@ -92,3 +93,33 @@ func TestExtractSQLNodeParams_NotInIsNotFannedOut(t *testing.T) {
t.Errorf("NOT IN must not supply the parameter, got %v", got[0])
}
}

func TestTransformSQLRawParameters_JSONFuncExprPassesThrough(t *testing.T) {
// Issue #700: JSON() values must survive to the request preparator for
// query param transposition; other FuncExprs are still discarded here.
jsonFunc := &sqlparser.FuncExpr{Name: sqlparser.NewColIdent("JSON")}
wrappedJSONFunc := &sqlparser.FuncExpr{Name: sqlparser.NewColIdent("json")}
otherFunc := &sqlparser.FuncExpr{Name: sqlparser.NewColIdent("strftime")}
input := map[string]interface{}{
"InstanceRequirements": jsonFunc,
"WrappedRequirements": parserutil.NewComparisonParameterMetadata(nil, wrappedJSONFunc, 0),
"Evaluated": otherFunc,
"Scalar": sqlparser.NewStrVal([]byte("x86_64")),
}
got, err := TransformSQLRawParameters(input, true)
if err != nil {
t.Fatalf("TransformSQLRawParameters: %v", err)
}
if got["InstanceRequirements"] != jsonFunc {
t.Errorf("bare JSON() not passed through: %v", got["InstanceRequirements"])
}
if got["WrappedRequirements"] != wrappedJSONFunc {
t.Errorf("metadata-wrapped JSON() not passed through: %v", got["WrappedRequirements"])
}
if _, present := got["Evaluated"]; present {
t.Errorf("non-JSON FuncExpr must be discarded, got %v", got["Evaluated"])
}
if got["Scalar"] != "x86_64" {
t.Errorf("scalar mishandled: %v", got["Scalar"])
}
}
56 changes: 3 additions & 53 deletions pkg/mcp_server/server.go
Original file line number Diff line number Diff line change
Expand Up @@ -6,7 +6,6 @@ import (
"fmt"
"io"
"net/http"
"os"
"sync"
"time"

Expand Down Expand Up @@ -641,63 +640,14 @@ func (s *simpleMCPServer) run(ctx context.Context) error {
case serverTransportSSE:
return fmt.Errorf("SSE transport obsoleted; use streamable HTTP transport instead")
case serverTransportStdIO:
// Windows text-mode pipes CRLF-terminate JSON-RPC lines; the SDK's
// ndjson decoder treats the bare carriage return before the newline
// as trailing garbage and kills the session (issue #668). A raw CR
// is illegal inside a JSON string (control characters must be
// escaped), so stripping every CR from the inbound stream is
// lossless for spec-compliant traffic.
if filtered, filterErr := newCRStrippedStdin(); filterErr == nil {
os.Stdin = filtered
} else {
s.logger.Warnf("could not install CRLF-tolerant stdin filter: %v", filterErr)
}
return s.server.Run(ctx, &mcp.StdioTransport{})
// Handles CRLF framing (issue #668) and malformed frames (issue
// #701); see resilientStdioTransport.
return s.server.Run(ctx, newResilientStdioTransport(s.logger))
default:
return fmt.Errorf("unsupported transport: %s", s.config.Server.Transport)
}
}

// crFilterReader removes carriage-return bytes from the wrapped stream.
type crFilterReader struct {
r io.Reader
}

func (c *crFilterReader) Read(p []byte) (int, error) {
for {
n, err := c.r.Read(p)
kept := 0
for i := 0; i < n; i++ {
if p[i] == '\r' {
continue
}
p[kept] = p[i]
kept++
}
if kept > 0 || err != nil || len(p) == 0 {
return kept, err
}
// Every byte read was a CR; retry rather than returning (0, nil).
}
}

// newCRStrippedStdin returns the read end of a pipe fed by a goroutine that
// copies the process's real stdin with all carriage returns removed. The
// pump goroutine lives for the remainder of the process, which matches the
// lifetime of a stdio MCP session.
func newCRStrippedStdin() (*os.File, error) {
pr, pw, err := os.Pipe()
if err != nil {
return nil, err
}
stdin := os.Stdin
go func() {
defer pw.Close()
io.Copy(pw, &crFilterReader{r: stdin}) //nolint:errcheck // EOF/close terminates the session anyway
}()
return pr, nil
}

// Stop gracefully stops the MCP server and all transports.
func (s *simpleMCPServer) Stop() error {
s.mu.Lock()
Expand Down
232 changes: 232 additions & 0 deletions pkg/mcp_server/stdio_transport.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,232 @@
package mcp_server //nolint:revive // package name is established

import (
"bufio"
"bytes"
"context"
"encoding/json"
"errors"
"fmt"
"io"
"os"
"sync"

"github.com/modelcontextprotocol/go-sdk/jsonrpc"
"github.com/modelcontextprotocol/go-sdk/mcp"
"github.com/sirupsen/logrus"
)

const (
// maxStdioFrameBytes bounds a single inbound newline-delimited frame so a
// peer cannot exhaust memory with one unterminated line (issue #701).
maxStdioFrameBytes = 8 << 20 // 8 MiB
// loggedFrameSampleBytes bounds how much untrusted client input is logged.
loggedFrameSampleBytes = 256
stdioReadBufferBytes = 64 << 10

// The id is null because the request id cannot be trusted from a frame
// that failed to decode.
jsonRPCParseErrorLine = `{"jsonrpc":"2.0","id":null,"error":{"code":-32700,"message":"Parse error"}}`
jsonRPCInvalidRequestLine = `{"jsonrpc":"2.0","id":null,"error":{"code":-32600,"message":"Invalid Request"}}`
)

// resilientStdioTransport replaces mcp.StdioTransport, whose read loop
// treats any JSON decode failure on stdin as fatal to the session (issue
// #701; unchanged as of SDK v1.7.0). Frames are read here instead and
// undecodable ones answered per JSON-RPC 2.0 (-32700 / -32600) without
// terminating the session. Inbound batch arrays get -32600 rather than
// fan-out: batching was removed in MCP 2025-06-18 and the SDK kills
// modern-protocol sessions on batch frames anyway.
type resilientStdioTransport struct {
logger *logrus.Logger
}

func newResilientStdioTransport(logger *logrus.Logger) *resilientStdioTransport {
return &resilientStdioTransport{logger: logger}
}

func (t *resilientStdioTransport) Connect(_ context.Context) (mcp.Connection, error) {
return newResilientStdioConn(os.Stdin, os.Stdout, t.logger, maxStdioFrameBytes), nil
}

type stdioMsgOrErr struct {
msg jsonrpc.Message
err error
}

// resilientStdioConn implements mcp.Connection over a newline-delimited
// JSON-RPC stream. All writes (server responses and this type's own error
// responses) share writeMu so frames cannot interleave on stdout.
type resilientStdioConn struct {
out io.Writer
logger *logrus.Logger
incoming <-chan stdioMsgOrErr

writeMu sync.Mutex

closers []io.Closer
closeOnce sync.Once
closed chan struct{}
closeErr error
}

func newResilientStdioConn(
in io.ReadCloser,
out io.WriteCloser,
logger *logrus.Logger,
maxFrameBytes int,
) *resilientStdioConn {
incoming := make(chan stdioMsgOrErr)
conn := &resilientStdioConn{
out: out,
logger: logger,
incoming: incoming,
closers: []io.Closer{in, out},
closed: make(chan struct{}),
}
// A dedicated read goroutine lets Read select against the closed channel,
// mirroring the SDK transport (including its unavoidable leak if a
// blocked os.Stdin read never unblocks).
go conn.readLoop(bufio.NewReaderSize(in, stdioReadBufferBytes), maxFrameBytes, incoming)
return conn
}

func (c *resilientStdioConn) readLoop(
reader *bufio.Reader,
maxFrameBytes int,
incoming chan<- stdioMsgOrErr,
) {
for {
frame, overflowed, err := readBoundedFrame(reader, maxFrameBytes)
if msg, ok := c.admitFrame(frame, overflowed); ok {
select {
case incoming <- stdioMsgOrErr{msg: msg}:
case <-c.closed:
return
}
}
if err != nil {
select {
case incoming <- stdioMsgOrErr{err: err}:
case <-c.closed:
}
return
}
}
}

// admitFrame vets one inbound frame, answering undeliverable ones with a
// JSON-RPC error response; it returns true only for frames the SDK's
// message layer is guaranteed to accept.
func (c *resilientStdioConn) admitFrame(frame []byte, overflowed bool) (jsonrpc.Message, bool) {
if overflowed {
c.rejectFrame(jsonRPCParseErrorLine, "oversized", frame)
return nil, false
}
// Windows text-mode pipes CRLF-terminate JSON-RPC lines (issue #668); a
// raw CR is illegal inside a JSON string, so stripping every CR is
// lossless for spec-compliant traffic.
frame = bytes.TrimSpace(bytes.ReplaceAll(frame, []byte{'\r'}, nil))
if len(frame) == 0 {
return nil, false
}
if !json.Valid(frame) {
c.rejectFrame(jsonRPCParseErrorLine, "malformed JSON", frame)
return nil, false
}
msg, decErr := jsonrpc.DecodeMessage(frame)
if decErr != nil {
c.rejectFrame(jsonRPCInvalidRequestLine, "non-JSON-RPC", frame)
return nil, false
}
return msg, true
}

func (c *resilientStdioConn) rejectFrame(response string, reason string, frame []byte) {
if c.logger != nil {
sample := frame
if len(sample) > loggedFrameSampleBytes {
sample = sample[:loggedFrameSampleBytes]
}
c.logger.Warnf(
"stdio transport: rejecting %s frame (%d bytes, sample %q)",
reason, len(frame), sample,
)
}
c.writeMu.Lock()
defer c.writeMu.Unlock()
if _, err := io.WriteString(c.out, response+"\n"); err != nil && c.logger != nil {
c.logger.Warnf("stdio transport: failed to write error response: %v", err)
}
}

// readBoundedFrame reads one newline-terminated frame, discarding (but fully
// consuming) any frame longer than maxFrameBytes. The returned frame
// excludes the terminator; overflowed reports whether the cap was exceeded.
func readBoundedFrame(reader *bufio.Reader, maxFrameBytes int) ([]byte, bool, error) {
var frame []byte
overflowed := false
for {
chunk, err := reader.ReadSlice('\n')
if !overflowed {
if len(frame)+len(chunk) > maxFrameBytes {
overflowed = true
frame = nil
} else {
frame = append(frame, chunk...)
}
}
if errors.Is(err, bufio.ErrBufferFull) {
continue
}
frame = bytes.TrimSuffix(frame, []byte{'\n'})
return frame, overflowed, err
}
}

func (c *resilientStdioConn) SessionID() string { return "" }

func (c *resilientStdioConn) Read(ctx context.Context) (jsonrpc.Message, error) {
select {
case <-ctx.Done():
return nil, ctx.Err()
case v := <-c.incoming:
if v.err != nil {
return nil, v.err
}
return v.msg, nil
case <-c.closed:
return nil, io.EOF
}
}

func (c *resilientStdioConn) Write(ctx context.Context, msg jsonrpc.Message) error {
select {
case <-ctx.Done():
return ctx.Err()
default:
}
data, err := jsonrpc.EncodeMessage(msg)
if err != nil {
return fmt.Errorf("marshaling message: %w", err)
}
data = append(data, '\n') // newline delimited
c.writeMu.Lock()
defer c.writeMu.Unlock()
_, err = c.out.Write(data)
return err
}

func (c *resilientStdioConn) Close() error {
c.closeOnce.Do(func() {
var errs []error
for _, closer := range c.closers {
if err := closer.Close(); err != nil {
errs = append(errs, err)
}
}
c.closeErr = errors.Join(errs...)
close(c.closed)
})
return c.closeErr
}
Loading
Loading