PranitInitial commit439f83c3 months agoHistory
1628 lines · 59.3 KB · Zig
RawDownload
const std = @import("std");const mcp_auth = @import("mcp_auth.zig");const mcp_contract = @import("mcp_contract.zig");const operation_control = @import("operation_control.zig");const protocol_negotiation = @import("protocol_negotiation.zig");const io_mod = @import("../shared/io.zig"); const Allocator = std.mem.Allocator; pub const protocol_version = protocol_negotiation.modern_protocol_version;const max_sse_events: usize = 1024;const max_tool_schema_depth: usize = 64;const max_json_safe_integer: i64 = 9_007_199_254_740_991; pub const Control = struct {    deadline: std.Io.Clock.Timestamp,    cancel_flag: ?*std.atomic.Value(bool) = null,    lifecycle_cancel_flag: ?*const std.atomic.Value(bool) = null,     pub fn cancellation(self: Control) operation_control.CancellationSources {        return .{            .caller = self.cancel_flag,            .runtime = self.lifecycle_cancel_flag,        };    }}; pub const PostOptions = struct {    url: []const u8,    static_headers: []const mcp_contract.McpHttpHeader = &.{},    request_body: []const u8,    input_schema_json: ?[]const u8 = null,    max_response_bytes: usize,    max_event_bytes: usize,    allow_version_error: bool = false,    allow_discovery_mismatch_status: bool = false,    progress: ?mcp_contract.ProgressSink = null,    notifications: ?mcp_contract.NotificationSink = null,    long_lived: bool = false,    precommit: ?*mcp_contract.TransportPrecommit = null,    control: Control,}; pub const PostResponse = struct {    body: []u8,    version_error: bool = false,    discovery_mismatch_status: bool = false,    auth_rejection: AuthRejection = .none,    www_authenticate: ?[]u8 = null,     pub fn deinit(self: *PostResponse, alloc: Allocator) void {        alloc.free(self.body);        if (self.www_authenticate) |value| alloc.free(value);        self.* = undefined;    }}; pub const AuthRejection = enum {    none,    unauthorized,    insufficient_scope,}; pub const EndpointError = error{    InvalidEndpoint,    InsecureEndpoint,}; pub const HeaderError = error{    InvalidHeaderName,    InvalidHeaderValue,    DuplicateHeader,    ReservedHeader,}; pub fn validateEndpoint(url: []const u8) EndpointError!void {    const uri = std.Uri.parse(url) catch return error.InvalidEndpoint;    if (uri.user != null or uri.password != null or uri.fragment != null) {        return error.InvalidEndpoint;    }    const host_component = uri.host orelse return error.InvalidEndpoint;     if (std.ascii.eqlIgnoreCase(uri.scheme, "https")) return;    if (!std.ascii.eqlIgnoreCase(uri.scheme, "http") or uri.port == null) {        return error.InsecureEndpoint;    }     var host_buf: [std.Io.net.HostName.max_len]u8 = undefined;    const host = host_component.toRaw(&host_buf) catch return error.InvalidEndpoint;    if (std.mem.eql(u8, host, "127.0.0.1") or        std.ascii.eqlIgnoreCase(host, "localhost") or        std.mem.eql(u8, host, "[::1]"))    {        return;    }    return error.InsecureEndpoint;} pub fn validateStaticHeaders(headers: []const mcp_contract.McpHttpHeader) HeaderError!void {    for (headers, 0..) |header, index| {        if (!isHttpToken(header.name)) return error.InvalidHeaderName;        try validateHeaderValue(header.value);        if (isReservedHeader(header.name)) return error.ReservedHeader;         for (headers[0..index]) |previous| {            if (std.ascii.eqlIgnoreCase(previous.name, header.name)) {                return error.DuplicateHeader;            }        }    }} pub fn validateHeaderValue(value: []const u8) HeaderError!void {    if (!isValidHeaderValue(value)) return error.InvalidHeaderValue;} fn isHttpToken(value: []const u8) bool {    if (value.len == 0) return false;    for (value) |byte| {        if (std.ascii.isAlphanumeric(byte)) continue;        switch (byte) {            '!', '#', '$', '%', '&', '\'', '*', '+', '-', '.', '^', '_', '`', '|', '~' => {},            else => return false,        }    }    return true;} pub fn validateToolInputSchemaHeaders(alloc: Allocator, schema_json: []const u8) !void {    var parsed = std.json.parseFromSlice(std.json.Value, alloc, schema_json, .{}) catch {        return error.InvalidToolInputSchema;    };    defer parsed.deinit();    try validateToolInputSchemaValue(alloc, parsed.value);} pub fn post(alloc: Allocator, options: PostOptions) !PostResponse {    try validateEndpoint(options.url);    try validateStaticHeaders(options.static_headers);    if (options.max_response_bytes == 0 or options.max_event_bytes == 0) {        return error.InvalidResponseLimit;    }     const cancellation = options.control.cancellation();    if (cancellation.cancelled()) return error.Cancelled;     const now = std.Io.Clock.Timestamp.now(io_mod.getIo(), .awake);    if (!std.Io.Clock.Timestamp.compare(now, .lt, options.control.deadline)) {        return error.McpRequestTimedOut;    }     const Operation = struct {        alloc: Allocator,        options: PostOptions,         fn run(self: @This()) anyerror!PostResponse {            return postCore(self.alloc, self.options);        }    };    const Event = union(enum) {        request: anyerror!PostResponse,        cancelled: anyerror!void,        deadline: anyerror!void,    };    const Cleanup = struct {        fn drain(result_alloc: Allocator, select: *std.Io.Select(Event)) void {            while (select.cancel()) |item| switch (item) {                .request => |result| {                    var response = result catch continue;                    response.deinit(result_alloc);                },                .cancelled, .deadline => {},            };        }    };     var select_buffer: [3]Event = undefined;    var select: std.Io.Select(Event) = .init(io_mod.getIo(), &select_buffer);    select.concurrent(.cancelled, waitForCancellation, .{cancellation}) catch |err| return err;    select.concurrent(.deadline, waitForDeadline, .{options.control.deadline}) catch |err| {        select.cancelDiscard();        return err;    };    select.concurrent(.request, Operation.run, .{Operation{        .alloc = alloc,        .options = options,    }}) catch |err| {        select.cancelDiscard();        return err;    };     const event = select.await() catch |err| {        Cleanup.drain(alloc, &select);        return err;    };    switch (event) {        .request => |result| {            Cleanup.drain(alloc, &select);            if (cancellation.cancelled()) {                var response = result catch return error.Cancelled;                response.deinit(alloc);                return error.Cancelled;            }            return result;        },        .cancelled => |result| {            result catch |err| {                Cleanup.drain(alloc, &select);                return err;            };            Cleanup.drain(alloc, &select);            return error.Cancelled;        },        .deadline => |result| {            result catch |err| {                Cleanup.drain(alloc, &select);                return err;            };            Cleanup.drain(alloc, &select);            if (cancellation.cancelled()) return error.Cancelled;            return error.McpRequestTimedOut;        },    }} const RequestInfo = struct {    id: i64,    method: []const u8,    name: ?[]const u8,    params: std.json.ObjectMap,}; const OwnedHeaders = struct {    items: std.ArrayList(std.http.Header) = .empty,    owned: std.ArrayList([]u8) = .empty,     fn deinit(self: *OwnedHeaders, alloc: Allocator) void {        for (self.owned.items) |value| alloc.free(value);        self.owned.deinit(alloc);        self.items.deinit(alloc);        self.* = undefined;    }     fn own(self: *OwnedHeaders, alloc: Allocator, value: []const u8) ![]const u8 {        const copy = try alloc.dupe(u8, value);        errdefer alloc.free(copy);        try self.owned.append(alloc, copy);        return copy;    }     fn append(self: *OwnedHeaders, alloc: Allocator, name: []const u8, value: []const u8) !void {        try self.items.append(alloc, .{ .name = name, .value = value });    }     fn appendOwned(self: *OwnedHeaders, alloc: Allocator, name: []const u8, value: []const u8) !void {        const owned_name = try self.own(alloc, name);        const owned_value = try self.own(alloc, value);        try self.append(alloc, owned_name, owned_value);    }}; const PreparedRequest = struct {    headers: OwnedHeaders,    id: i64,    is_discovery: bool,     fn deinit(self: *PreparedRequest, alloc: Allocator) void {        self.headers.deinit(alloc);        self.* = undefined;    }}; fn postCore(alloc: Allocator, options: PostOptions) !PostResponse {    var prepared = try prepareRequest(        alloc,        options.static_headers,        options.request_body,        options.input_schema_json,    );    defer prepared.deinit(alloc);     const uri = std.Uri.parse(options.url) catch return error.InvalidEndpoint;    var client: std.http.Client = .{        .allocator = alloc,        .io = io_mod.getIo(),    };    defer client.deinit();     var request = try client.request(.POST, uri, .{        .redirect_behavior = .unhandled,        .headers = .{            .content_type = .{ .override = "application/json" },            .accept_encoding = .omit,        },        .extra_headers = prepared.headers.items.items,    });    defer request.deinit();    defer if (options.precommit) |precommit| precommit.release();    if (options.precommit) |precommit| try precommit.acquire();    try request.sendBodyComplete(@constCast(options.request_body));    if (options.precommit) |precommit| precommit.release();     var response = request.receiveHead(&.{}) catch |err| switch (err) {        error.HttpContentEncodingUnsupported => return error.UnsupportedContentEncoding,        else => |e| return e,    };    const status = response.head.status;    if (status.class() == .redirect) return error.RedirectNotAllowed;     if (status == .unauthorized or status == .forbidden) {        const authenticate = try mcp_auth.collectAuthenticateHeader(            alloc,            response.head,        );        defer if (authenticate) |value| alloc.free(value);        if (status == .unauthorized or authenticate != null) {            return .{                .body = try alloc.dupe(u8, ""),                .auth_rejection = if (status == .unauthorized)                    .unauthorized                else                    .insufficient_scope,                .www_authenticate = if (authenticate) |value|                    try alloc.dupe(u8, value)                else                    null,            };        }    }     const version_error = options.allow_version_error and status == .bad_request;    const discovery_mismatch_status = options.allow_discovery_mismatch_status and        (status == .not_found or status == .method_not_allowed);    if (status != .ok and !version_error and !discovery_mismatch_status) {        return error.UnexpectedHttpStatus;    }    if (response.head.content_encoding != .identity) {        return error.UnsupportedContentEncoding;    }    if (discovery_mismatch_status) {        return .{            .body = try alloc.dupe(u8, ""),            .discovery_mismatch_status = true,        };    }     const content_type = response.head.content_type orelse return error.MissingContentType;    const media_type = parseMediaType(content_type) orelse return error.UnsupportedContentType;    if (version_error and media_type != .json) return error.UnsupportedContentType;     var transfer_buffer: [16 * 1024]u8 = undefined;    const reader = response.reader(&transfer_buffer);    const body = switch (media_type) {        .json => try readJsonResponse(            alloc,            reader,            prepared.id,            options.max_response_bytes,            prepared.is_discovery,        ),        .sse => try readSseResponse(            alloc,            reader,            prepared.id,            if (options.long_lived) std.math.maxInt(usize) else options.max_response_bytes,            options.max_event_bytes,            if (options.long_lived) std.math.maxInt(usize) else max_sse_events,            options.progress,            options.notifications,            options.control.cancellation(),        ),    };    return .{        .body = body,        .version_error = version_error,        .discovery_mismatch_status = discovery_mismatch_status,    };} fn prepareRequest(    alloc: Allocator,    static_headers: []const mcp_contract.McpHttpHeader,    request_body: []const u8,    input_schema_json: ?[]const u8,) !PreparedRequest {    var parsed = std.json.parseFromSlice(std.json.Value, alloc, request_body, .{}) catch {        return error.InvalidModernRequest;    };    defer parsed.deinit();    const request_info = try modernRequestInfo(parsed.value);    const is_discovery = std.mem.eql(u8, request_info.method, "server/discover");     var headers: OwnedHeaders = .{};    errdefer headers.deinit(alloc);    try headers.append(alloc, "Accept", "application/json, text/event-stream");    try headers.append(alloc, "MCP-Protocol-Version", protocol_version);    try headers.appendOwned(alloc, "Mcp-Method", request_info.method);    if (request_info.name) |name| {        const encoded_name = try encodeHeaderValue(alloc, name);        defer alloc.free(encoded_name);        try headers.appendOwned(alloc, "Mcp-Name", encoded_name);    }    for (static_headers) |header| {        try headers.append(alloc, header.name, header.value);    }     if (std.mem.eql(u8, request_info.method, "tools/call")) {        const schema_json = input_schema_json orelse return error.MissingToolInputSchema;        try appendProjectedHeaders(            alloc,            &headers,            schema_json,            request_info.params,        );    } else if (input_schema_json != null) {        return error.UnexpectedToolInputSchema;    }    return .{ .headers = headers, .id = request_info.id, .is_discovery = is_discovery };} fn modernRequestInfo(value: std.json.Value) !RequestInfo {    if (value != .object) return error.InvalidModernRequest;    const jsonrpc = value.object.get("jsonrpc") orelse return error.InvalidModernRequest;    if (jsonrpc != .string or !std.mem.eql(u8, jsonrpc.string, "2.0")) {        return error.InvalidModernRequest;    }    const id_value = value.object.get("id") orelse return error.InvalidModernRequest;    if (id_value != .integer) return error.InvalidModernRequest;    const method_value = value.object.get("method") orelse return error.InvalidModernRequest;    if (method_value != .string or method_value.string.len == 0) {        return error.InvalidModernRequest;    }    const params_value = value.object.get("params") orelse return error.InvalidModernRequest;    if (params_value != .object) return error.InvalidModernRequest;    const meta_value = params_value.object.get("_meta") orelse return error.InvalidModernRequest;    if (meta_value != .object) return error.InvalidModernRequest;    const version_value = meta_value.object.get("io.modelcontextprotocol/protocolVersion") orelse {        return error.InvalidModernRequest;    };    if (version_value != .string or !std.mem.eql(u8, version_value.string, protocol_version)) {        return error.InvalidModernRequest;    }    const client_info = meta_value.object.get("io.modelcontextprotocol/clientInfo") orelse {        return error.InvalidModernRequest;    };    if (client_info != .object) return error.InvalidModernRequest;    const client_name = client_info.object.get("name") orelse return error.InvalidModernRequest;    const client_version = client_info.object.get("version") orelse return error.InvalidModernRequest;    if (client_name != .string or client_name.string.len == 0 or        client_version != .string or client_version.string.len == 0)    {        return error.InvalidModernRequest;    }    const capabilities = meta_value.object.get("io.modelcontextprotocol/clientCapabilities") orelse {        return error.InvalidModernRequest;    };    if (capabilities != .object) return error.InvalidModernRequest;     const method = method_value.string;    const name = if (std.mem.eql(u8, method, "tools/call")) blk: {        const name_value = params_value.object.get("name") orelse return error.InvalidModernRequest;        if (name_value != .string) return error.InvalidModernRequest;        break :blk name_value.string;    } else null;     return .{        .id = id_value.integer,        .method = method,        .name = name,        .params = params_value.object,    };} fn appendProjectedHeaders(    alloc: Allocator,    headers: *OwnedHeaders,    schema_json: []const u8,    request_params: std.json.ObjectMap,) !void {    var schema = std.json.parseFromSlice(std.json.Value, alloc, schema_json, .{}) catch {        return error.InvalidToolInputSchema;    };    defer schema.deinit();    try validateToolInputSchemaValue(alloc, schema.value);     const arguments_value = request_params.get("arguments") orelse return error.InvalidToolArguments;    if (arguments_value != .object) return error.InvalidToolArguments;    if (schema.value != .object) return error.InvalidToolInputSchema;    const properties_value = schema.value.object.get("properties") orelse return;    if (properties_value != .object) return error.InvalidToolInputSchema;     try appendProjectedProperties(        alloc,        headers,        properties_value.object,        arguments_value.object,        0,    );} 

Lines 1–500 of 1628.