diff --git a/docs/server/tools.md b/docs/server/tools.md index 9ac1ee1..656d7dc 100644 --- a/docs/server/tools.md +++ b/docs/server/tools.md @@ -8,8 +8,9 @@ - `handle` — synchronous logic from input to `ToolResult`. - `handleWithHeaders` — synchronous logic that also receives the request headers. - `serverLogic` — effectful logic, with the request headers. + - `handleSecured` (synchronous) or `securedServerLogic` (effectful) — logic that also receives the principal made by the server's security logic; usable only on a secured server — see [transport security](transport.md). - A tool that pushes to the client while running (progress, logging) instead uses `streamingServerLogic` — see [server capabilities](capabilities.md). + A tool that pushes to the client while running (progress, logging) instead uses `streamingServerLogic` — see [server capabilities](capabilities.md). A secured, streaming tool uses `securedStreamingServerLogic`, which gives both the principal and the streaming context — see [transport security](transport.md). - Assemble tools into an `McpServer` and call `.endpoint(path)` to create a Tapir endpoint. ```scala mdoc:compile-only diff --git a/docs/server/transport.md b/docs/server/transport.md index 5a33b68..c63f21a 100644 --- a/docs/server/transport.md +++ b/docs/server/transport.md @@ -106,6 +106,82 @@ object SecuredMcpServer: NettySyncServer().port(8080).addEndpoint(securedEndpoint).startAndWait() ``` -The result of the security logic does not reach the tool logic. If a tool needs data from the caller, read the request headers with `handleWithHeaders` (or `serverLogic` with headers). +The result of the security logic of `prependSecurity` does not reach the tool logic. Use `prependSecurity` if the tools do not need data from the caller. If a tool needs such data, read the request headers with `handleWithHeaders` (or `serverLogic` with headers), or give the tools a principal, as below. For all the security inputs - API keys, basic and bearer authorization, OAuth2 flows - see the [Tapir endpoint security documentation](https://tapir.softwaremill.com/en/latest/endpoint/security.html). + +### Giving the security result to the tools + +To validate the caller one time and give the result to the tool logic, use `serverSecurityLogic` (or `serverSecurityLogicPure`, if the logic needs no effect). It takes the same security input and error output as `prependSecurity`, and makes a principal - a value of your own type, such as the identity of the caller. + +The server gives the principal to the logic of each tool which you add to it. Define such a tool with `handleSecured`, or with `securedServerLogic` if the logic needs an effect. Tools which do not need the principal keep their usual logic. The other builders of `McpServer` stay available, so you can configure the server before or after you add the security logic: + +```scala mdoc:compile-only +import chimp.server.* +import sttp.model.StatusCode +import sttp.shared.Identity +import sttp.tapir.* +import sttp.tapir.server.netty.sync.NettySyncServer + +case class User(email: String) + +object McpServerWithPrincipal: + def main(args: Array[String]): Unit = + val echo = tool("echo").input[String].handle(message => ToolResult.text(message)) + val whoAmI = tool("whoAmI").input[String].handleSecured[User]((_, user) => ToolResult.text(user.email)) + + val securedEndpoint = McpServer[Identity]() + .serverSecurityLogicPure( + auth.bearer[String](), + statusCode(StatusCode.Unauthorized).and(stringBody) + )(token => if token == "s3cret" then Right(User("employee@example.com")) else Left("Invalid token")) + .name("my-mcp-server") + .addTools(echo, whoAmI) + .endpoint(List("mcp")) + + NettySyncServer().port(8080).addEndpoint(securedEndpoint).startAndWait() +``` + +The security logic runs one time for each request, before the server handles the MCP message. If it gives a rejection, the server sends the error output and no tool logic runs. This is necessary if the client must get an HTTP status code, because a tool which rejects a call can only give a JSON-RPC error with status code 200. + +### Combining security with streaming + +`.streaming` on a secured server gives a `SecuredStreamingMcpServer`, which accepts streaming tools alongside the plain secured ones. Define such a tool with `securedStreamingServerLogic`, which gives both the principal and the `StreamingServerContext`: + +```scala mdoc:compile-only +import chimp.protocol.LoggingLevel +import chimp.server.* +import chimp.server.ox.OxServerHttpTransport +import chimp.server.transport.SecuredServerStreamingHttpTransport +import io.circe.Json +import sttp.model.StatusCode +import sttp.shared.Identity +import sttp.tapir.* +import sttp.tapir.server.netty.sync.NettySyncServer + +case class User(email: String) + +object McpServerWithPrincipalAndStreaming: + def main(args: Array[String]): Unit = + val whoAmI = tool("whoAmI") + .input[String] + .securedStreamingServerLogic[Identity, User]: (_, user, ctx, _) => + ctx.log(LoggingLevel.Info, Json.fromString(s"called by ${user.email}")) + ToolResult.text(user.email) + + val securedStreamingServer = McpServer[Identity]() + .serverSecurityLogicPure( + auth.bearer[String](), + statusCode(StatusCode.Unauthorized).and(stringBody) + )(token => if token == "s3cret" then Right(User("employee@example.com")) else Left("Invalid token")) + .streaming + .addStreamingTool(whoAmI) + + // `OxServerHttpTransport` is one effect backend's streaming machinery; substitute your own (Pekko, ZIO, ...). + val backend = OxServerHttpTransport(List("mcp")) + val securedEndpoint = SecuredServerStreamingHttpTransport(List("mcp"), backend).serve(securedStreamingServer) + + NettySyncServer().port(8080).addEndpoint(securedEndpoint).startAndWait() +``` + +`SecuredServerStreamingHttpTransport` takes an existing `ServerStreamingHttpTransport` as its source of streaming machinery, rather than extending it, so the same effect-specific backend instance - `OxServerHttpTransport`, `PekkoServerHttpTransport`, `ZioServerHttpTransport` - serves both a plain `StreamingMcpServer` and a `SecuredStreamingMcpServer`. diff --git a/server-streaming/server-ox/src/test/scala/chimp/server/ox/OxMcpServerHttpSpec.scala b/server-streaming/server-ox/src/test/scala/chimp/server/ox/OxMcpServerHttpSpec.scala index de51b3b..724be68 100644 --- a/server-streaming/server-ox/src/test/scala/chimp/server/ox/OxMcpServerHttpSpec.scala +++ b/server-streaming/server-ox/src/test/scala/chimp/server/ox/OxMcpServerHttpSpec.scala @@ -1,20 +1,39 @@ package chimp.server.ox -import chimp.client.transport.ClientTransport +import chimp.client.transport.{ClientHttpTransport, ClientTransport} import chimp.client.transport.ox.OxClientHttpTransport import chimp.client.{BidirectionalMcpClient, McpClient} import chimp.protocol.{Implementation, ProtocolVersion} -import chimp.server.{McpServer, McpServerStreamingTests, McpServerTests, StreamingMcpServer, SyncToFuture} +import chimp.server.transport.SecuredServerStreamingHttpTransport +import chimp.server.{ + McpServer, + McpServerStreamingTests, + McpServerTests, + SecuredMcpServer, + SecuredMcpServerStreamingTests, + SecuredMcpServerTests, + SecuredStreamingMcpServer, + StreamingMcpServer, + SyncToFuture +} import org.scalatest.Assertion -import ox.supervised -import sttp.client4.DefaultSyncBackend +import ox.{supervised, Ox} +import sttp.capabilities.WebSockets +import sttp.client4.{DefaultSyncBackend, SyncBackend} +import sttp.model.Header import sttp.model.Uri.UriContext import sttp.shared.Identity -import sttp.tapir.server.netty.sync.NettySyncServer +import sttp.tapir.server.netty.sync.{NettySyncServer, OxStreams} +import sttp.tapir.server.ServerEndpoint import scala.concurrent.Future -class OxMcpServerHttpSpec extends McpServerTests[Identity] with McpServerStreamingTests[Identity] with SyncToFuture: +class OxMcpServerHttpSpec + extends McpServerTests[Identity] + with McpServerStreamingTests[Identity] + with SecuredMcpServerTests[Identity] + with SecuredMcpServerStreamingTests[Identity] + with SyncToFuture: private val clientInfo = Implementation("chimp-server-test", "0.0.1") override protected def withServer(server: McpServer[Identity])(test: McpClient[Identity] => Identity[Assertion]): Future[Assertion] = @@ -23,21 +42,51 @@ class OxMcpServerHttpSpec extends McpServerTests[Identity] with McpServerStreami override protected def withStreamingServer( server: StreamingMcpServer[Identity] )(test: BidirectionalMcpClient[Identity] => Identity[Assertion]): Future[Assertion] = + withHttpServer(OxServerHttpTransport(List("mcp")).serve(server)): (port, backend, ox) => + given Ox = ox + val transport = + OxClientHttpTransport(backend, uri"http://localhost:$port/mcp", ProtocolVersion.Latest, ClientTransport.defaultTimeout) + try test(McpClient.bidirectional(transport, clientInfo)) + finally transport.close() + + override protected def withSecuredServer( + server: SecuredMcpServer[Identity, String, String, User], + token: String + )(test: McpClient[Identity] => Identity[Assertion]): Future[Assertion] = + withHttpServer(server.endpoint(List("mcp"))): (port, backend, _) => + val transport = ClientHttpTransport[Identity]( + backend, + uri"http://localhost:$port/mcp", + headers = List(Header.authorization("Bearer", token)) + ) + try test(McpClient(transport, clientInfo)) + finally transport.close() + + override protected def withSecuredStreamingServer( + server: SecuredStreamingMcpServer[Identity, String, String, User], + token: String + )(test: BidirectionalMcpClient[Identity] => Identity[Assertion]): Future[Assertion] = + val endpoint = SecuredServerStreamingHttpTransport(List("mcp"), OxServerHttpTransport(List("mcp"))).serve(server) + withHttpServer(endpoint): (port, backend, ox) => + given Ox = ox + val transport = OxClientHttpTransport( + backend, + uri"http://localhost:$port/mcp", + ProtocolVersion.Latest, + ClientTransport.defaultTimeout, + headers = List(Header.authorization("Bearer", token)) + ) + try test(McpClient.bidirectional(transport, clientInfo)) + finally transport.close() + + private def withHttpServer( + endpoint: ServerEndpoint[OxStreams & WebSockets, Identity] + )(test: (Int, SyncBackend, Ox) => Assertion): Future[Assertion] = toFuture: supervised: - val endpoint = OxServerHttpTransport(List("mcp")).serve(server) val binding = NettySyncServer().port(0).addEndpoint(endpoint).start() try val backend = DefaultSyncBackend() - try - val transport = - OxClientHttpTransport( - backend, - uri"http://localhost:${binding.port}/mcp", - ProtocolVersion.Latest, - ClientTransport.defaultTimeout - ) - try test(McpClient.bidirectional(transport, clientInfo)) - finally transport.close() + try test(binding.port, backend, summon[Ox]) finally backend.close() finally binding.stop() diff --git a/server-streaming/server-pekko/src/test/scala/chimp/server/pekko/PekkoMcpServerHttpSpec.scala b/server-streaming/server-pekko/src/test/scala/chimp/server/pekko/PekkoMcpServerHttpSpec.scala index 016c7c8..f3c0944 100644 --- a/server-streaming/server-pekko/src/test/scala/chimp/server/pekko/PekkoMcpServerHttpSpec.scala +++ b/server-streaming/server-pekko/src/test/scala/chimp/server/pekko/PekkoMcpServerHttpSpec.scala @@ -3,17 +3,36 @@ package chimp.server.pekko import chimp.client.transport.pekko.PekkoClientHttpTransport import chimp.client.{BidirectionalMcpClient, McpClient} import chimp.protocol.Implementation -import chimp.server.{McpServer, McpServerStreamingTests, McpServerTests, StreamingMcpServer} +import chimp.server.transport.SecuredServerStreamingHttpTransport +import chimp.server.{ + McpServer, + McpServerStreamingTests, + McpServerTests, + SecuredMcpServer, + SecuredMcpServerStreamingTests, + SecuredMcpServerTests, + SecuredStreamingMcpServer, + StreamingMcpServer +} import org.apache.pekko.http.scaladsl.Http import org.scalatest.Assertion +import sttp.capabilities.WebSockets +import sttp.capabilities.pekko.PekkoStreams import sttp.client4.pekkohttp.PekkoHttpBackend +import sttp.model.Header import sttp.model.Uri.UriContext import sttp.tapir.server.pekkohttp.PekkoHttpServerInterpreter +import sttp.tapir.server.ServerEndpoint import scala.concurrent.duration.DurationInt import scala.concurrent.{ExecutionContext, Future} -class PekkoMcpServerHttpSpec extends McpServerTests[Future] with McpServerStreamingTests[Future] with PekkoToFuture: +class PekkoMcpServerHttpSpec + extends McpServerTests[Future] + with McpServerStreamingTests[Future] + with SecuredMcpServerTests[Future] + with SecuredMcpServerStreamingTests[Future] + with PekkoToFuture: private val clientInfo = Implementation("chimp-server-test", "0.0.1") override protected def withServer(server: McpServer[Future])(test: McpClient[Future] => Future[Assertion]): Future[Assertion] = @@ -22,17 +41,37 @@ class PekkoMcpServerHttpSpec extends McpServerTests[Future] with McpServerStream override protected def withStreamingServer( server: StreamingMcpServer[Future] )(test: BidirectionalMcpClient[Future] => Future[Assertion]): Future[Assertion] = + withHttpServer(PekkoServerHttpTransport(List("mcp")).serve(server), Nil): transport => + McpClient.bidirectional(transport, clientInfo).flatMap(test) + + override protected def withSecuredServer( + server: SecuredMcpServer[Future, String, String, User], + token: String + )(test: McpClient[Future] => Future[Assertion]): Future[Assertion] = + withHttpServer(server.endpoint(List("mcp")), List(Header.authorization("Bearer", token))): transport => + McpClient(transport, clientInfo).flatMap(test) + + override protected def withSecuredStreamingServer( + server: SecuredStreamingMcpServer[Future, String, String, User], + token: String + )(test: BidirectionalMcpClient[Future] => Future[Assertion]): Future[Assertion] = + val endpoint = SecuredServerStreamingHttpTransport(List("mcp"), PekkoServerHttpTransport(List("mcp"))).serve(server) + withHttpServer(endpoint, List(Header.authorization("Bearer", token))): transport => + McpClient.bidirectional(transport, clientInfo).flatMap(test) + + private def withHttpServer( + endpoint: ServerEndpoint[PekkoStreams & WebSockets, Future], + headers: List[Header] + )(test: PekkoClientHttpTransport => Future[Assertion]): Future[Assertion] = given ExecutionContext = actorSystem.dispatcher - val endpoint = PekkoServerHttpTransport(List("mcp")).serve(server) Http() .newServerAt("localhost", 0) .bind(PekkoHttpServerInterpreter().toRoute(endpoint)) .flatMap: binding => val backend = PekkoHttpBackend.usingActorSystem(actorSystem) - val transport = PekkoClientHttpTransport(backend, uri"http://localhost:${binding.localAddress.getPort}/mcp") - McpClient - .bidirectional(transport, clientInfo) - .flatMap(test) + val transport = + PekkoClientHttpTransport(backend, uri"http://localhost:${binding.localAddress.getPort}/mcp", headers = headers) + test(transport) .transformWith: result => transport .close() diff --git a/server-streaming/server-zio/src/test/scala/chimp/server/zio/ZioMcpServerHttpSpec.scala b/server-streaming/server-zio/src/test/scala/chimp/server/zio/ZioMcpServerHttpSpec.scala index 32525b8..8bd995a 100644 --- a/server-streaming/server-zio/src/test/scala/chimp/server/zio/ZioMcpServerHttpSpec.scala +++ b/server-streaming/server-zio/src/test/scala/chimp/server/zio/ZioMcpServerHttpSpec.scala @@ -4,17 +4,36 @@ import chimp.client.transport.ClientTransport import chimp.client.transport.zio.ZioClientHttpTransport import chimp.client.{BidirectionalMcpClient, McpClient} import chimp.protocol.{Implementation, ProtocolVersion} -import chimp.server.{McpServer, McpServerStreamingTests, McpServerTests, StreamingMcpServer} +import chimp.server.transport.SecuredServerStreamingHttpTransport +import chimp.server.{ + McpServer, + McpServerStreamingTests, + McpServerTests, + SecuredMcpServer, + SecuredMcpServerStreamingTests, + SecuredMcpServerTests, + SecuredStreamingMcpServer, + StreamingMcpServer +} import org.scalatest.Assertion +import sttp.capabilities.WebSockets +import sttp.capabilities.zio.ZioStreams import sttp.client4.* import sttp.client4.httpclient.zio.HttpClientZioBackend +import sttp.model.Header import sttp.tapir.server.ziohttp.ZioHttpInterpreter +import sttp.tapir.server.ServerEndpoint import zio.http.Server import zio.{Scope, Task, ZIO} import scala.concurrent.Future -class ZioMcpServerHttpSpec extends McpServerTests[Task] with McpServerStreamingTests[Task] with ZioToFuture: +class ZioMcpServerHttpSpec + extends McpServerTests[Task] + with McpServerStreamingTests[Task] + with SecuredMcpServerTests[Task] + with SecuredMcpServerStreamingTests[Task] + with ZioToFuture: private val clientInfo = Implementation("chimp-server-test", "0.0.1") override protected def withServer(server: McpServer[Task])(test: McpClient[Task] => Task[Assertion]): Future[Assertion] = @@ -23,15 +42,42 @@ class ZioMcpServerHttpSpec extends McpServerTests[Task] with McpServerStreamingT override protected def withStreamingServer( server: StreamingMcpServer[Task] )(test: BidirectionalMcpClient[Task] => Task[Assertion]): Future[Assertion] = + withHttpServer(ZioServerHttpTransport(List("mcp")).serve(server), Nil): transport => + McpClient.bidirectional(transport, clientInfo).flatMap(test) + + override protected def withSecuredServer( + server: SecuredMcpServer[Task, String, String, User], + token: String + )(test: McpClient[Task] => Task[Assertion]): Future[Assertion] = + withHttpServer(server.endpoint(List("mcp")), List(Header.authorization("Bearer", token))): transport => + McpClient(transport, clientInfo).flatMap(test) + + override protected def withSecuredStreamingServer( + server: SecuredStreamingMcpServer[Task, String, String, User], + token: String + )(test: BidirectionalMcpClient[Task] => Task[Assertion]): Future[Assertion] = + val endpoint = SecuredServerStreamingHttpTransport(List("mcp"), ZioServerHttpTransport(List("mcp"))).serve(server) + withHttpServer(endpoint, List(Header.authorization("Bearer", token))): transport => + McpClient.bidirectional(transport, clientInfo).flatMap(test) + + private def withHttpServer( + endpoint: ServerEndpoint[ZioStreams & WebSockets, Task], + headers: List[Header] + )(test: ZioClientHttpTransport => Task[Assertion]): Future[Assertion] = toFuture: - val routes = ZioHttpInterpreter().toHttp(ZioServerHttpTransport(List("mcp")).serve(server)) + val routes = ZioHttpInterpreter().toHttp(endpoint) ZIO.scoped: (for port <- Server.install(routes) result <- HttpClientZioBackend().flatMap: backend => ZioClientHttpTransport - .scoped(backend, uri"http://localhost:$port/mcp", ProtocolVersion.Latest, ClientTransport.defaultTimeout) - .flatMap(transport => McpClient.bidirectional(transport, clientInfo)) - .flatMap(client => test(client)) + .scoped( + backend, + uri"http://localhost:$port/mcp", + ProtocolVersion.Latest, + ClientTransport.defaultTimeout, + headers = headers + ) + .flatMap(test) .ensuring(backend.close().ignore) yield result).provideSome[Scope](Server.defaultWithPort(0)) diff --git a/server/src/main/scala/chimp/server/McpServer.scala b/server/src/main/scala/chimp/server/McpServer.scala index 75fca77..49f380c 100644 --- a/server/src/main/scala/chimp/server/McpServer.scala +++ b/server/src/main/scala/chimp/server/McpServer.scala @@ -1,8 +1,10 @@ package chimp.server import chimp.protocol.* -import chimp.server.transport.ServerHttpTransport +import chimp.server.transport.{SecuredServerHttpTransport, ServerHttpTransport} +import sttp.monad.MonadError import sttp.tapir.server.ServerEndpoint +import sttp.tapir.{EndpointInput, EndpointOutput} type CompletionHandler[F[_]] = (CompleteRef, CompleteArgument, Option[CompleteContext]) => F[Completion] @@ -91,6 +93,20 @@ case class McpServer[F[_]]( def endpoint(path: List[String]): ServerEndpoint[Any, F] = ServerHttpTransport(path).serve(this) + /** Adds the security input, the error output which describes a rejection, and the logic which validates the security input and makes the + * principal. The principal is given to the logic of the tools which are added to the returned server. + */ + def serverSecurityLogic[S, E, P](securityInput: EndpointInput[S], errorOutput: EndpointOutput[E])( + logic: S => F[Either[E, P]] + ): SecuredMcpServer[F, S, E, P] = + SecuredMcpServer(this, securityInput, errorOutput, _ => logic) + + /** The same as [[serverSecurityLogic]], but for security logic which needs no effect. */ + def serverSecurityLogicPure[S, E, P](securityInput: EndpointInput[S], errorOutput: EndpointOutput[E])( + logic: S => Either[E, P] + ): SecuredMcpServer[F, S, E, P] = + SecuredMcpServer(this, securityInput, errorOutput, monad => input => monad.unit(logic(input))) + def streaming: StreamingMcpServer[F] = StreamingMcpServer( name, @@ -174,3 +190,166 @@ case class StreamingMcpServer[F[_]]( def withSubscriptions(handler: ResourceSubscriptions[F]): StreamingMcpServer[F] = copy(subscriptions = Some(handler)) + +/** An [[McpServer]] with security logic, which runs before the server handles an MCP message. The result of the security logic, the + * principal, is given to the logic of the tools which are added to this server. Tools of the initial server, which do not need the + * principal, are kept. + * + * @tparam S + * The type of the security input, for example a bearer token. + * @tparam E + * The type of the error output, which the server sends if the security logic gives a rejection. + * @tparam P + * The type of the principal, which the security logic makes from the security input. + */ +case class SecuredMcpServer[F[_], S, E, P]( + server: McpServer[F], + securityInput: EndpointInput[S], + errorOutput: EndpointOutput[E], + securityLogic: MonadError[F] => S => F[Either[E, P]], + securedTools: List[ServerTool[?, ?, F, SecuredServerContext[F, P]]] = Nil +) extends McpServerDef[F, SecuredServerContext[F, P]]: + def name: String = server.name + def version: String = server.version + def instructions: Option[String] = server.instructions + def showJsonSchemaMetadata: Boolean = server.showJsonSchemaMetadata + def originCheck: OriginCheck = server.originCheck + def prompts: List[ServerPrompt[F]] = server.prompts + def resources: List[ServerResource[F]] = server.resources + def resourceTemplates: List[ServerResourceTemplate[F]] = server.resourceTemplates + def completion: Option[CompletionHandler[F]] = server.completion + def loggingLevel: Option[SetLoggingLevelHandler[F]] = server.loggingLevel + def subscriptions: Option[ResourceSubscriptions[F]] = server.subscriptions + + def tools: List[ServerTool[?, ?, F, SecuredServerContext[F, P]]] = server.tools ++ securedTools + + def name(value: String): SecuredMcpServer[F, S, E, P] = + copy(server = server.name(value)) + + def version(value: String): SecuredMcpServer[F, S, E, P] = + copy(server = server.version(value)) + + def instructions(value: String): SecuredMcpServer[F, S, E, P] = + copy(server = server.instructions(value)) + + def withJsonSchemaMetadata(value: Boolean): SecuredMcpServer[F, S, E, P] = + copy(server = server.withJsonSchemaMetadata(value)) + + def withOriginCheck(value: OriginCheck): SecuredMcpServer[F, S, E, P] = + copy(server = server.withOriginCheck(value)) + + def addTool(tool: ServerTool[?, ?, F, SecuredServerContext[F, P]]): SecuredMcpServer[F, S, E, P] = + copy(securedTools = securedTools :+ tool) + + def addTools(tools: ServerTool[?, ?, F, SecuredServerContext[F, P]]*): SecuredMcpServer[F, S, E, P] = + copy(securedTools = this.securedTools ++ tools) + + def addPrompt(prompt: ServerPrompt[F]): SecuredMcpServer[F, S, E, P] = + copy(server = server.addPrompt(prompt)) + + def addPrompts(prompts: ServerPrompt[F]*): SecuredMcpServer[F, S, E, P] = + copy(server = server.addPrompts(prompts*)) + + def addResource(resource: ServerResource[F]): SecuredMcpServer[F, S, E, P] = + copy(server = server.addResource(resource)) + + def addResources(resources: ServerResource[F]*): SecuredMcpServer[F, S, E, P] = + copy(server = server.addResources(resources*)) + + def addResourceTemplate(resourceTemplate: ServerResourceTemplate[F]): SecuredMcpServer[F, S, E, P] = + copy(server = server.addResourceTemplate(resourceTemplate)) + + def addResourceTemplates(resourceTemplates: ServerResourceTemplate[F]*): SecuredMcpServer[F, S, E, P] = + copy(server = server.addResourceTemplates(resourceTemplates*)) + + def withCompletion(handler: CompletionHandler[F]): SecuredMcpServer[F, S, E, P] = + copy(server = server.withCompletion(handler)) + + def withLoggingLevel(handler: SetLoggingLevelHandler[F]): SecuredMcpServer[F, S, E, P] = + copy(server = server.withLoggingLevel(handler)) + + def withSubscriptions(handler: ResourceSubscriptions[F]): SecuredMcpServer[F, S, E, P] = + copy(server = server.withSubscriptions(handler)) + + def endpoint(path: List[String]): ServerEndpoint[Any, F] = SecuredServerHttpTransport[F, S, E, P](path).serve(this) + + def streaming: SecuredStreamingMcpServer[F, S, E, P] = SecuredStreamingMcpServer(this) + +/** A [[SecuredMcpServer]] which also accepts streaming tools, which are given a [[SecuredStreamingServerContext]] combining the principal + * with the [[StreamingServerContext]]. Tools of the initial secured server, which need only the principal, are kept. + */ +case class SecuredStreamingMcpServer[F[_], S, E, P]( + server: SecuredMcpServer[F, S, E, P], + streamingTools: List[ServerTool[?, ?, F, SecuredStreamingServerContext[F, P]]] = Nil +) extends McpServerDef[F, SecuredStreamingServerContext[F, P]]: + def name: String = server.name + def version: String = server.version + def instructions: Option[String] = server.instructions + def showJsonSchemaMetadata: Boolean = server.showJsonSchemaMetadata + def originCheck: OriginCheck = server.originCheck + def prompts: List[ServerPrompt[F]] = server.prompts + def resources: List[ServerResource[F]] = server.resources + def resourceTemplates: List[ServerResourceTemplate[F]] = server.resourceTemplates + def completion: Option[CompletionHandler[F]] = server.completion + def loggingLevel: Option[SetLoggingLevelHandler[F]] = server.loggingLevel + def subscriptions: Option[ResourceSubscriptions[F]] = server.subscriptions + + def securityInput: EndpointInput[S] = server.securityInput + def errorOutput: EndpointOutput[E] = server.errorOutput + def securityLogic: MonadError[F] => S => F[Either[E, P]] = server.securityLogic + + def tools: List[ServerTool[?, ?, F, SecuredStreamingServerContext[F, P]]] = server.tools ++ streamingTools + + def name(value: String): SecuredStreamingMcpServer[F, S, E, P] = + copy(server = server.name(value)) + + def version(value: String): SecuredStreamingMcpServer[F, S, E, P] = + copy(server = server.version(value)) + + def instructions(value: String): SecuredStreamingMcpServer[F, S, E, P] = + copy(server = server.instructions(value)) + + def withJsonSchemaMetadata(value: Boolean): SecuredStreamingMcpServer[F, S, E, P] = + copy(server = server.withJsonSchemaMetadata(value)) + + def withOriginCheck(value: OriginCheck): SecuredStreamingMcpServer[F, S, E, P] = + copy(server = server.withOriginCheck(value)) + + def addTool(tool: ServerTool[?, ?, F, SecuredServerContext[F, P]]): SecuredStreamingMcpServer[F, S, E, P] = + copy(server = server.addTool(tool)) + + def addTools(tools: ServerTool[?, ?, F, SecuredServerContext[F, P]]*): SecuredStreamingMcpServer[F, S, E, P] = + copy(server = server.addTools(tools*)) + + def addStreamingTool(tool: ServerTool[?, ?, F, SecuredStreamingServerContext[F, P]]): SecuredStreamingMcpServer[F, S, E, P] = + copy(streamingTools = streamingTools :+ tool) + + def addStreamingTools(tools: ServerTool[?, ?, F, SecuredStreamingServerContext[F, P]]*): SecuredStreamingMcpServer[F, S, E, P] = + copy(streamingTools = this.streamingTools ++ tools) + + def addPrompt(prompt: ServerPrompt[F]): SecuredStreamingMcpServer[F, S, E, P] = + copy(server = server.addPrompt(prompt)) + + def addPrompts(prompts: ServerPrompt[F]*): SecuredStreamingMcpServer[F, S, E, P] = + copy(server = server.addPrompts(prompts*)) + + def addResource(resource: ServerResource[F]): SecuredStreamingMcpServer[F, S, E, P] = + copy(server = server.addResource(resource)) + + def addResources(resources: ServerResource[F]*): SecuredStreamingMcpServer[F, S, E, P] = + copy(server = server.addResources(resources*)) + + def addResourceTemplate(resourceTemplate: ServerResourceTemplate[F]): SecuredStreamingMcpServer[F, S, E, P] = + copy(server = server.addResourceTemplate(resourceTemplate)) + + def addResourceTemplates(resourceTemplates: ServerResourceTemplate[F]*): SecuredStreamingMcpServer[F, S, E, P] = + copy(server = server.addResourceTemplates(resourceTemplates*)) + + def withCompletion(handler: CompletionHandler[F]): SecuredStreamingMcpServer[F, S, E, P] = + copy(server = server.withCompletion(handler)) + + def withLoggingLevel(handler: SetLoggingLevelHandler[F]): SecuredStreamingMcpServer[F, S, E, P] = + copy(server = server.withLoggingLevel(handler)) + + def withSubscriptions(handler: ResourceSubscriptions[F]): SecuredStreamingMcpServer[F, S, E, P] = + copy(server = server.withSubscriptions(handler)) diff --git a/server/src/main/scala/chimp/server/ServerContext.scala b/server/src/main/scala/chimp/server/ServerContext.scala index c764818..ebcfa60 100644 --- a/server/src/main/scala/chimp/server/ServerContext.scala +++ b/server/src/main/scala/chimp/server/ServerContext.scala @@ -10,6 +10,14 @@ trait ServerContext[F[_]] object ServerContext: def noop[F[_]]: ServerContext[F] = new ServerContext[F] {} +/** A context which also gives the principal that the security logic of a [[SecuredMcpServer]] made from the request. */ +trait SecuredServerContext[F[_], +P] extends ServerContext[F]: + def principal: P + +object SecuredServerContext: + def apply[F[_], P](value: P): SecuredServerContext[F, P] = new SecuredServerContext[F, P]: + def principal: P = value + trait StreamingServerContext[F[_]] extends ServerContext[F]: def reportProgress(progress: Double, total: Option[Double] = None, message: Option[String] = None): F[Unit] def log(level: LoggingLevel, data: Json, logger: Option[String] = None): F[Unit] @@ -32,3 +40,19 @@ private[server] final class SinkStreamingServerContext[F[_]](sink: OutboundSink[ sink.send( JSONRPCMessage.Notification(method = "notifications/message", params = Some(LoggingMessageParams(level, data, logger).asJson)) ) + +/** A context which gives both the principal that the security logic of a [[SecuredMcpServer]] made from the request, and the streaming + * capabilities of a [[StreamingServerContext]]. + */ +trait SecuredStreamingServerContext[F[_], +P] extends SecuredServerContext[F, P] with StreamingServerContext[F] + +private[server] final class SinkSecuredStreamingServerContext[F[_], P]( + sink: OutboundSink[F], + progressToken: Option[ProgressToken], + val principal: P +)(using MonadError[F]) + extends SecuredStreamingServerContext[F, P]: + private val delegate = SinkStreamingServerContext[F](sink, progressToken) + def reportProgress(progress: Double, total: Option[Double] = None, message: Option[String] = None): F[Unit] = + delegate.reportProgress(progress, total, message) + def log(level: LoggingLevel, data: Json, logger: Option[String] = None): F[Unit] = delegate.log(level, data, logger) diff --git a/server/src/main/scala/chimp/server/Tool.scala b/server/src/main/scala/chimp/server/Tool.scala index 1e4f489..e6aecd3 100644 --- a/server/src/main/scala/chimp/server/Tool.scala +++ b/server/src/main/scala/chimp/server/Tool.scala @@ -103,12 +103,39 @@ case class Tool[I, O]( def serverLogic[F[_]](logic: (I, Seq[Header]) => F[ToolResult[O]]): ServerTool[I, O, F, ServerContext[F]] = ServerTool(name, description, inputSchema, inputDecoder, outputSchema, annotations, (input, _, headers) => logic(input, headers)) + /** Attaches effectful logic, with access to the principal; usable only on a [[SecuredMcpServer]]. */ + def securedServerLogic[F[_], P](logic: (I, P, Seq[Header]) => F[ToolResult[O]]): ServerTool[I, O, F, SecuredServerContext[F, P]] = + ServerTool( + name, + description, + inputSchema, + inputDecoder, + outputSchema, + annotations, + (input, context, headers) => logic(input, context.principal, headers) + ) + /** Attaches effectful logic with access to the [[StreamingServerContext]]; usable only on a streaming server. */ def streamingServerLogic[F[_]]( logic: (I, StreamingServerContext[F], Seq[Header]) => F[ToolResult[O]] ): ServerTool[I, O, F, StreamingServerContext[F]] = ServerTool(name, description, inputSchema, inputDecoder, outputSchema, annotations, logic) + /** Attaches effectful logic with access to the principal and the [[StreamingServerContext]]; usable only on a secured streaming server. + */ + def securedStreamingServerLogic[F[_], P]( + logic: (I, P, StreamingServerContext[F], Seq[Header]) => F[ToolResult[O]] + ): ServerTool[I, O, F, SecuredStreamingServerContext[F, P]] = + ServerTool( + name, + description, + inputSchema, + inputDecoder, + outputSchema, + annotations, + (input, context, headers) => logic(input, context.principal, context, headers) + ) + /** Attaches synchronous logic that also receives the request headers. */ def handleWithHeaders(logic: (I, Seq[Header]) => ToolResult[O]): ServerTool[I, O, Identity, ServerContext[Identity]] = ServerTool(name, description, inputSchema, inputDecoder, outputSchema, annotations, (i, _, headers) => logic(i, headers)) @@ -117,6 +144,10 @@ case class Tool[I, O]( def handle(logic: I => ToolResult[O]): ServerTool[I, O, Identity, ServerContext[Identity]] = handleWithHeaders((i, _) => logic(i)) + /** Attaches synchronous logic over the decoded input and the principal; usable only on a [[SecuredMcpServer]]. */ + def handleSecured[P](logic: (I, P) => ToolResult[O]): ServerTool[I, O, Identity, SecuredServerContext[Identity, P]] = + securedServerLogic[Identity, P]((i, principal, _) => logic(i, principal)) + /** A fully-defined tool: its metadata plus the logic handling a call, in effect `F` with context `C`. */ case class ServerTool[I, O, F[_], -C <: ServerContext[F]]( name: String, diff --git a/server/src/main/scala/chimp/server/transport/JsonRpcHttpResponses.scala b/server/src/main/scala/chimp/server/transport/JsonRpcHttpResponses.scala new file mode 100644 index 0000000..b89b9ec --- /dev/null +++ b/server/src/main/scala/chimp/server/transport/JsonRpcHttpResponses.scala @@ -0,0 +1,26 @@ +package chimp.server.transport + +import chimp.server.{McpResponse, OriginCheck, OutboundSink} +import io.circe.Json +import sttp.model.{Header, HeaderNames, StatusCode} +import sttp.monad.MonadError +import sttp.monad.syntax.* + +private[transport] def originAllowed(originCheck: OriginCheck, headers: Seq[Header]): Boolean = + val host = headers.find(_.name.equalsIgnoreCase(HeaderNames.Host)).map(_.value) + val origin = headers.find(_.name.equalsIgnoreCase(HeaderNames.Origin)).map(_.value) + originCheck.validate(host, origin) + +private[transport] def respondToJsonRpc[F[_]](originCheck: OriginCheck, headers: Seq[Header])(handle: => F[McpResponse])(using + m: MonadError[F] +): F[(StatusCode, Option[Json])] = + if !originAllowed(originCheck, headers) then m.unit((StatusCode.Forbidden, None)) + else handle.map(response => (response.statusCode, response.body)) + +private[transport] def respondWithEventStream[F[_], Caps]( + originCheck: OriginCheck, + headers: Seq[Header], + backend: ServerStreamingHttpTransport[F, Caps] +)(handle: OutboundSink[F] => F[Option[Json]])(using m: MonadError[F]): F[(StatusCode, backend.EventStream)] = + if !originAllowed(originCheck, headers) then m.unit((StatusCode.Forbidden, backend.emptyStream)) + else backend.eventStream(handle).map(events => (StatusCode.Ok, events)) diff --git a/server/src/main/scala/chimp/server/transport/SecuredServerHttpTransport.scala b/server/src/main/scala/chimp/server/transport/SecuredServerHttpTransport.scala new file mode 100644 index 0000000..ebc452f --- /dev/null +++ b/server/src/main/scala/chimp/server/transport/SecuredServerHttpTransport.scala @@ -0,0 +1,41 @@ +package chimp.server.transport + +import chimp.server.* +import io.circe.Json +import sttp.model.Header +import sttp.monad.MonadError +import sttp.monad.syntax.* +import sttp.tapir.* +import sttp.tapir.json.circe.* +import sttp.tapir.server.ServerEndpoint + +/** Implementation of unidirectional MCP server using Streamable HTTP, protected by the security logic of a [[SecuredMcpServer]]. The + * security logic runs before any MCP message is handled. If it gives a rejection, the server sends the error output and no tool logic + * runs. If it gives a principal, the principal goes to the tool logic. + * + * @param path + * The MCP endpoint path. + */ +final case class SecuredServerHttpTransport[F[_], S, E, P](path: List[String]): + def serve(server: SecuredMcpServer[F, S, E, P]): ServerEndpoint[Any, F] = + val handler = new McpHandler(server) + val mcpEndpoint = endpoint.post + .securityIn(server.securityInput) + .in(path.foldLeft(emptyInput)((inputSoFar, pathComponent) => inputSoFar / pathComponent)) + .in(extractFromRequest(_.headers)) + .in(jsonBody[Json]) + .errorOut(server.errorOutput) + .out(statusCode) + .out(jsonBody[Option[Json]]) + + ServerEndpoint( + mcpEndpoint, + server.securityLogic, + me => { (principal: P) => (input: (Seq[Header], Json)) => + val (headers, json) = input + given MonadError[F] = me + respondToJsonRpc(server.originCheck, headers) { + handler.handleJsonRpc(json, headers, _ => SecuredServerContext[F, P](principal)) + }.map(Right(_)) + } + ) diff --git a/server/src/main/scala/chimp/server/transport/SecuredServerStreamingHttpTransport.scala b/server/src/main/scala/chimp/server/transport/SecuredServerStreamingHttpTransport.scala new file mode 100644 index 0000000..4a6faa7 --- /dev/null +++ b/server/src/main/scala/chimp/server/transport/SecuredServerStreamingHttpTransport.scala @@ -0,0 +1,53 @@ +package chimp.server.transport + +import chimp.protocol.ProgressToken +import chimp.server.* +import io.circe.Json +import sttp.model.Header +import sttp.monad.MonadError +import sttp.monad.syntax.* +import sttp.tapir.* +import sttp.tapir.json.circe.* +import sttp.tapir.server.ServerEndpoint + +/** Implementation of bidirectional MCP server using Streamable HTTP, protected by the security logic of a [[SecuredMcpServer]]. The + * security logic runs before any MCP message is handled. If it gives a rejection, the server sends the error output and no tool logic + * runs. If it gives a principal, the principal goes to the tool logic, together with the [[StreamingServerContext]]. + * + * Takes an existing [[ServerStreamingHttpTransport]] as the source of the streaming machinery for the effect backend, rather than + * extending it, so that the same backend instance - for example an `chimp.server.ox.OxServerHttpTransport` - serves both a plain + * [[StreamingMcpServer]] and a [[SecuredStreamingMcpServer]]. Its own `path` plays no part here; only its streaming machinery is used. + * + * @param path + * The MCP endpoint path. + * @param backend + * Supplies the streaming machinery for the effect type `F` and the streaming capability `Caps`. + */ +final case class SecuredServerStreamingHttpTransport[F[_], Caps, S, E, P]( + path: List[String], + backend: ServerStreamingHttpTransport[F, Caps] +): + def serve(server: SecuredStreamingMcpServer[F, S, E, P]): ServerEndpoint[Caps, F] = + val handler = new McpHandler[F, SecuredStreamingServerContext[F, P]](server) + val mcpEndpoint = endpoint.post + .securityIn(server.securityInput) + .in(path.foldLeft(emptyInput)((inputSoFar, pathComponent) => inputSoFar / pathComponent)) + .in(extractFromRequest(_.headers)) + .in(jsonBody[Json]) + .errorOut(server.errorOutput) + .out(statusCode) + .out(backend.sseBody) + + ServerEndpoint( + mcpEndpoint, + server.securityLogic, + me => { (principal: P) => (input: (Seq[Header], Json)) => + val (headers, json) = input + given MonadError[F] = me + respondWithEventStream(server.originCheck, headers, backend) { sink => + val makeContext: Option[ProgressToken] => SecuredStreamingServerContext[F, P] = + token => SinkSecuredStreamingServerContext(sink, token, principal) + handler.handleJsonRpc(json, headers, makeContext).map(_.body) + }.map(Right(_)) + } + ) diff --git a/server/src/main/scala/chimp/server/transport/ServerHttpTransport.scala b/server/src/main/scala/chimp/server/transport/ServerHttpTransport.scala index 1b91bda..44be5a5 100644 --- a/server/src/main/scala/chimp/server/transport/ServerHttpTransport.scala +++ b/server/src/main/scala/chimp/server/transport/ServerHttpTransport.scala @@ -2,7 +2,7 @@ package chimp.server.transport import chimp.server.* import io.circe.Json -import sttp.model.{Header, HeaderNames, StatusCode} +import sttp.model.Header import sttp.monad.MonadError import sttp.monad.syntax.* import sttp.tapir.* @@ -30,12 +30,6 @@ final case class ServerHttpTransport[F[_]](path: List[String]) extends ServerTra me => { (input: (Seq[Header], Json)) => val (headers, json) = input given MonadError[F] = me - val host = headers.find(_.name.equalsIgnoreCase(HeaderNames.Host)).map(_.value) - val origin = headers.find(_.name.equalsIgnoreCase(HeaderNames.Origin)).map(_.value) - if !server.originCheck.validate(host, origin) then me.unit(Right((StatusCode.Forbidden, None))) - else - handler - .handleJsonRpc(json, headers) - .map(response => Right((response.statusCode, response.body))) + respondToJsonRpc(server.originCheck, headers)(handler.handleJsonRpc(json, headers)).map(Right(_)) } ) diff --git a/server/src/main/scala/chimp/server/transport/ServerStreamingHttpTransport.scala b/server/src/main/scala/chimp/server/transport/ServerStreamingHttpTransport.scala index 4dae664..8ae7503 100644 --- a/server/src/main/scala/chimp/server/transport/ServerStreamingHttpTransport.scala +++ b/server/src/main/scala/chimp/server/transport/ServerStreamingHttpTransport.scala @@ -4,7 +4,7 @@ import chimp.protocol.ProgressToken import chimp.server.* import io.circe.Json import sttp.capabilities.Streams -import sttp.model.{Header, HeaderNames, StatusCode} +import sttp.model.Header import sttp.monad.MonadError import sttp.monad.syntax.* import sttp.tapir.* @@ -14,20 +14,22 @@ import sttp.tapir.server.ServerEndpoint /** Abstract base for bidirectional MCP server using Streamable HTTP. Responds to JSON-RPC messages from an MCP client with a * Server-Sent-Event stream. Messages in the stream are interleaved with the final response on that stream. * - * The extra type parameter `S` carries the streaming capability evidence required by the Tapir [[sttp.tapir.server.ServerEndpoint]] to - * produce asynchronous stream of Server-Sent Events as response. + * The extra type parameter `Caps` carries the streaming capability evidence required by the Tapir [[sttp.tapir.server.ServerEndpoint]] to + * produce asynchronous stream of Server-Sent Events as response. An instance also serves as the streaming machinery which + * [[SecuredServerStreamingHttpTransport]] needs from the same effect backend - its `path` plays no part in that, so the same instance, or + * another one for the same `F`/`Caps`, can back both a plain [[StreamingMcpServer]] and a [[SecuredStreamingMcpServer]]. * * @param path * The MCP endpoint path. */ -abstract class ServerStreamingHttpTransport[F[_], S](path: List[String]) extends StreamingServerTransport[F, ServerEndpoint[S, F]]: - val streams: Streams[S] +abstract class ServerStreamingHttpTransport[F[_], Caps](path: List[String]) extends StreamingServerTransport[F, ServerEndpoint[Caps, F]]: + val streams: Streams[Caps] type EventStream - def sseBody: StreamBodyIO[streams.BinaryStream, EventStream, S] + def sseBody: StreamBodyIO[streams.BinaryStream, EventStream, Caps] def emptyStream: EventStream def eventStream(handle: OutboundSink[F] => F[Option[Json]]): F[EventStream] - final def serve(server: StreamingMcpServer[F]): ServerEndpoint[S, F] = + final def serve(server: StreamingMcpServer[F]): ServerEndpoint[Caps, F] = val handler = new McpHandler[F, StreamingServerContext[F]](server) val endpoint = infallibleEndpoint.post .in(path.foldLeft(emptyInput)((inputSoFar, pathComponent) => inputSoFar / pathComponent)) @@ -41,14 +43,10 @@ abstract class ServerStreamingHttpTransport[F[_], S](path: List[String]) extends me => { (input: (Seq[Header], Json)) => val (headers, json) = input given MonadError[F] = me - val host = headers.find(_.name.equalsIgnoreCase(HeaderNames.Host)).map(_.value) - val origin = headers.find(_.name.equalsIgnoreCase(HeaderNames.Origin)).map(_.value) - if !server.originCheck.validate(host, origin) then me.unit(Right((StatusCode.Forbidden, emptyStream))) - else - eventStream { sink => - val makeContext: Option[ProgressToken] => StreamingServerContext[F] = - token => SinkStreamingServerContext(sink, token) - handler.handleJsonRpc(json, headers, makeContext).map(_.body) - }.map(events => Right((StatusCode.Ok, events))) + respondWithEventStream(server.originCheck, headers, this) { sink => + val makeContext: Option[ProgressToken] => StreamingServerContext[F] = + token => SinkStreamingServerContext(sink, token) + handler.handleJsonRpc(json, headers, makeContext).map(_.body) + }.map(Right(_)) } ) diff --git a/server/src/test/scala/chimp/server/SecuredMcpServerStreamingTests.scala b/server/src/test/scala/chimp/server/SecuredMcpServerStreamingTests.scala new file mode 100644 index 0000000..6dc12c1 --- /dev/null +++ b/server/src/test/scala/chimp/server/SecuredMcpServerStreamingTests.scala @@ -0,0 +1,85 @@ +package chimp.server + +import chimp.client.{BidirectionalMcpClient, McpAuthorizationException} +import chimp.client.notifications.ServerNotification +import chimp.protocol.* +import io.circe.{Codec, Json} +import org.scalatest.{Assertion, RecoverMethods} +import org.scalatest.flatspec.AsyncFlatSpec +import org.scalatest.matchers.should.Matchers +import sttp.model.StatusCode +import sttp.monad.syntax.* +import sttp.tapir.* + +import java.util.concurrent.ConcurrentLinkedQueue +import scala.concurrent.Future +import scala.jdk.CollectionConverters.* + +trait SecuredMcpServerTestFixtures[F[_]]: + protected case class User(email: String) + + protected val validToken = "s3cret" + +trait SecuredMcpServerStreamingTests[F[_]] extends AsyncFlatSpec with Matchers with RecoverMethods with SecuredMcpServerTestFixtures[F]: + this: ToFuture[F] => + + /** @param token + * The bearer token the client authenticates with; defaults to [[validToken]] so most tests need not pass it. + */ + protected def withSecuredStreamingServer(server: SecuredStreamingMcpServer[F, String, String, User], token: String = validToken)( + test: BidirectionalMcpClient[F] => F[Assertion] + ): Future[Assertion] + + private case class NoInput() derives Codec, Schema + + private def securityLogic(token: String): Either[String, User] = + if token == validToken then Right(User("employee@example.com")) else Left("Invalid token") + + private def securityLogicEffectful(token: String): F[Either[String, User]] = + monad.unit(securityLogic(token)) + + private def whoAmITool: ServerTool[NoInput, NoStructuredOutput, F, SecuredStreamingServerContext[F, User]] = + tool("whoAmI") + .description("Logs the caller's email, then returns it") + .input[NoInput] + .securedStreamingServerLogic[F, User]: (_, user, ctx, _) => + ctx.log(LoggingLevel.Info, Json.fromString(user.email)).map(_ => ToolResult.text(user.email)) + + protected def securedStreamingServer: SecuredStreamingMcpServer[F, String, String, User] = + McpServer[F]() + .serverSecurityLogicPure(auth.bearer[String](), statusCode(StatusCode.Unauthorized).and(stringBody))(securityLogic) + .withLoggingLevel(_ => monad.unit(())) + .streaming + .addStreamingTool(whoAmITool) + + protected def securedStreamingServerWithEffectfulSecurityLogic: SecuredStreamingMcpServer[F, String, String, User] = + McpServer[F]() + .serverSecurityLogic(auth.bearer[String](), statusCode(StatusCode.Unauthorized).and(stringBody))(securityLogicEffectful) + .withLoggingLevel(_ => monad.unit(())) + .streaming + .addStreamingTool(whoAmITool) + + private def assertWhoAmIDeliversPrincipal(client: BidirectionalMcpClient[F]): F[Assertion] = + val messages = ConcurrentLinkedQueue[Json]() + val listener: ServerNotification => F[Unit] = { + case ServerNotification.LoggingMessage(params) => messages.add(params.data); monad.unit(()) + case _ => monad.unit(()) + } + client + .onServerNotification(notification => listener(notification)) + .flatMap(_ => client.callTool("whoAmI", Json.obj())) + .flatMap: result => + waitUntil(messages.size >= 1).map: _ => + result.content shouldBe List(ToolContent.Text("text", "employee@example.com")) + messages.asScala.toList shouldBe List(Json.fromString("employee@example.com")) + + "a secured streaming MCP server" should "give the principal to a streaming tool, and deliver its log notifications" in + withSecuredStreamingServer(securedStreamingServer)(assertWhoAmIDeliversPrincipal) + + it should "give the principal to a streaming tool when the security logic is effectful" in + withSecuredStreamingServer(securedStreamingServerWithEffectfulSecurityLogic)(assertWhoAmIDeliversPrincipal) + + it should "reject an invalid security input with the error output, before any tool logic runs" in + recoverToExceptionIf[McpAuthorizationException] { + Future(withSecuredStreamingServer(securedStreamingServer, token = "wrong")(_ => monad.unit(succeed))).flatten + }.map(_.statusCode shouldBe StatusCode.Unauthorized.code) diff --git a/server/src/test/scala/chimp/server/SecuredMcpServerTests.scala b/server/src/test/scala/chimp/server/SecuredMcpServerTests.scala new file mode 100644 index 0000000..3a4208e --- /dev/null +++ b/server/src/test/scala/chimp/server/SecuredMcpServerTests.scala @@ -0,0 +1,120 @@ +package chimp.server + +import chimp.client.{McpAuthorizationException, McpClient} +import chimp.protocol.{Implementation, ResourceContents, ToolContent} +import io.circe.{Codec, Json} +import org.scalatest.{Assertion, RecoverMethods} +import org.scalatest.flatspec.AsyncFlatSpec +import org.scalatest.matchers.should.Matchers +import sttp.monad.syntax.* +import sttp.model.StatusCode +import sttp.tapir.* + +import scala.concurrent.Future + +trait SecuredMcpServerTests[F[_]] extends AsyncFlatSpec with Matchers with RecoverMethods with SecuredMcpServerTestFixtures[F]: + this: ToFuture[F] => + + protected def withSecuredServer(server: SecuredMcpServer[F, String, String, User], token: String = validToken)( + test: McpClient[F] => F[Assertion] + ): Future[Assertion] + + private case class EchoInput(message: String) derives Codec, Schema + + private def securityLogic(token: String): Either[String, User] = + if token == validToken then Right(User("employee@example.com")) else Left("Invalid token") + + private def securityLogicEffectful(token: String): F[Either[String, User]] = + monad.unit(securityLogic(token)) + + private def echoTool: ServerTool[EchoInput, NoStructuredOutput, F, ServerContext[F]] = + tool("echo") + .description("Echoes a message.") + .input[EchoInput] + .serverLogic[F]((in, _) => monad.unit(ToolResult.text(in.message))) + + private def whoAmITool: ServerTool[EchoInput, NoStructuredOutput, F, SecuredServerContext[F, User]] = + tool("whoAmI") + .description("Echoes a message and the caller's email.") + .input[EchoInput] + .securedServerLogic[F, User]((in, user, _) => monad.unit(ToolResult.text(s"${in.message} ${user.email}"))) + + private def greetingResource: ServerResource[F] = + resource("test://greeting") + .name("greeting") + .mimeType("text/plain") + .serverLogic[F](_ => + monad.unit(Right(List(ResourceContents.Text(uri = "test://greeting", text = "hello", mimeType = Some("text/plain"))))) + ) + + private def securedServer: SecuredMcpServer[F, String, String, User] = + McpServer[F]() + .addTool(echoTool) + .serverSecurityLogicPure(auth.bearer[String](), statusCode(StatusCode.Unauthorized).and(stringBody))(securityLogic) + .addTool(whoAmITool) + + private def securedServerWithEffectfulSecurityLogic: SecuredMcpServer[F, String, String, User] = + McpServer[F]() + .addTool(echoTool) + .serverSecurityLogic(auth.bearer[String](), statusCode(StatusCode.Unauthorized).and(stringBody))(securityLogicEffectful) + .addTool(whoAmITool) + + private def serverConfiguredAfterSecurity: SecuredMcpServer[F, String, String, User] = + McpServer[F]() + .serverSecurityLogicPure(auth.bearer[String](), statusCode(StatusCode.Unauthorized).and(stringBody))(securityLogic) + .name("secured-server") + .version("2.0.0") + .addResource(greetingResource) + .addTool(whoAmITool) + + private def assertWhoAmIDeliversPrincipal(client: McpClient[F]): F[Assertion] = + client + .callTool("whoAmI", Json.obj("message" -> Json.fromString("hi"))) + .map: result => + result.isError shouldBe false + result.content shouldBe List(ToolContent.Text("text", "hi employee@example.com")) + + "a secured MCP server" should "give the principal to the tool logic" in + withSecuredServer(securedServer)(assertWhoAmIDeliversPrincipal) + + it should "give the principal to the tool logic when the security logic is effectful" in + withSecuredServer(securedServerWithEffectfulSecurityLogic)(assertWhoAmIDeliversPrincipal) + + it should "also serve tools which do not need the principal" in + withSecuredServer(securedServer): client => + client + .listTools() + .flatMap: tools => + tools.tools.map(_.name) should contain allOf ("echo", "whoAmI") + client + .callTool("echo", Json.obj("message" -> Json.fromString("hi"))) + .map: result => + result.content shouldBe List(ToolContent.Text("text", "hi")) + + it should "use the identity which the builders set after the security logic" in + withSecuredServer(serverConfiguredAfterSecurity): client => + monad.unit(client.serverInfo shouldBe Implementation("secured-server", "2.0.0")) + + it should "serve a resource which was added after the security logic" in + withSecuredServer(serverConfiguredAfterSecurity): client => + client.serverCapabilities.resources shouldBe defined + client + .listResources() + .flatMap: listed => + listed.resources.map(_.uri) shouldBe List("test://greeting") + client + .readResource("test://greeting") + .map: result => + result.contents.head match + case ResourceContents.Text(_, text, _, _) => text shouldBe "hello" + case other => fail(s"expected text contents, got $other") + + it should "reject an invalid security input with HTTP 401 before any tool logic runs" in + recoverToExceptionIf[McpAuthorizationException] { + Future(withSecuredServer(securedServer, token = "wrong")(_ => monad.unit(succeed))).flatten + }.map(_.statusCode shouldBe StatusCode.Unauthorized.code) + + it should "reject invalid security input with HTTP 401 when the security logic is effectful" in + recoverToExceptionIf[McpAuthorizationException] { + Future(withSecuredServer(securedServerWithEffectfulSecurityLogic, token = "wrong")(_ => monad.unit(succeed))).flatten + }.map(_.statusCode shouldBe StatusCode.Unauthorized.code)