Adds authorization callback handler

Author: mazinesy

lib/mcp.rb

@@ -24,13 +24,17 @@
 require_relative "mcp/auth/server/provider"
 require_relative "mcp/auth/server/settings"
 require_relative "mcp/auth/server/uri_helper"
-require_relative "mcp/auth/server/client_registry"
-require_relative "mcp/auth/server/state_registry"
+require_relative "mcp/auth/server/registries/client_registry"
+require_relative "mcp/auth/server/registries/state_registry"
+require_relative "mcp/auth/server/registries/auth_code_registry"
+require_relative "mcp/auth/server/registries/token_registry"
+require_relative "mcp/auth/server/registries/in_memory_registry"
 require_relative "mcp/auth/server/request_parser"
 require_relative "mcp/auth/server/providers/mcp_authorization_server_provider"
 require_relative "mcp/auth/server/handlers/metadata_handler"
 require_relative "mcp/auth/server/handlers/registration_handler"
 require_relative "mcp/auth/server/handlers/authorization_handler"
+require_relative "mcp/auth/server/handlers/callback_handler"
 
 module MCP
   class << self

lib/mcp/auth/server/client_registry.rb

@@ -94,7 +94,6 @@ def handle(request)
                 redirect_uri_provided_explicitly: auth_request.redirect_uri_provided?,
                 response_type: auth_request.response_type,
               )
-
               location = @auth_server_provider.authorize(client_info:, auth_params:)
               headers = {
                 "Cache-Control": "no-store",

lib/mcp/auth/server/handlers/authorization_handler.rb

@@ -0,0 +1,51 @@
+# frozen_string_literal: true
+
+require_relative "../../errors"
+require_relative "../../server/provider"
+
+module MCP
+  module Auth
+    module Server
+      module Handlers
+        class CallbackHandler
+          def initialize(
+            auth_server_provider:,
+            request_parser:
+          )
+            @auth_server_provider = auth_server_provider
+            @request_parser = request_parser
+          end
+
+          def handle(request)
+            params_h = @request_parser.parse_query_params(request)
+            code = params_h[:code]
+            state = params_h[:state]
+            if code.nil? || state.nil?
+              return bad_request_error(error: "invalid_request", error_description: "missing code or state parameter")
+            end
+
+            begin
+              redirect_uri = @auth_server_provider.authorize_callback(code:, state:)
+              headers = { Location: redirect_uri }
+
+              [302, headers, nil]
+            rescue Errors::AuthorizationError => e
+              bad_request_error(
+                error: e.error_code,
+                error_description: e.message,
+              )
+            rescue
+              [500, {}, { error: "server_error" }]
+            end
+          end
+
+          private
+
+          def bad_request_error(error:, error_description:)
+            [400, {}, { error:, error_description: }]
+          end
+        end
+      end
+    end
+  end
+end

lib/mcp/auth/server/handlers/callback_handler.rb

@@ -2,7 +2,6 @@
 
 require "securerandom"
 require "time"
-require "json"
 require_relative "../../../serialization_utils"
 require_relative "../../errors"
 require_relative "../../models"

lib/mcp/auth/server/handlers/registration_handler.rb

@@ -153,6 +153,22 @@ def authorize(client_info:, auth_params:)
           raise NotImplementedError, "#{self.class.name}#authorize is not implemented"
         end
 
+        # Handles the callback from the OAuth provider after the user has authorized
+        # the application. This is called when the OAuth provider redirects back to
+        # the MCP server's callback endpoint.
+        #
+        # Implementations should validate the code and state parameters, exchange the
+        # authorization code with the OAuth provider if needed, and return a URL to
+        # redirect the user back to the original client application.
+        #
+        # @param code [String] The authorization code from the OAuth provider
+        # @param state [String] The state parameter that was passed to the authorize endpoint
+        # @return [String] The URL to redirect the user to (typically the client's redirect_uri)
+        # @raise [MCP::Auth::Errors::AuthorizeError] If the callback parameters are invalid
+        def authorize_callback(code:, state:)
+          raise NotImplementedError, "#{self.class.name}#handle_callback is not implemented"
+        end
+
         # Loads an AuthorizationCode by its code string.
         #
         # @param client_info [MCP::Auth::Server::OAuthClientInformationFull] The client that requested the authorization code.

lib/mcp/auth/server/provider.rb

@@ -1,9 +1,8 @@
 # frozen_string_literal: true
 
 require "securerandom"
+require "net/http"
 require_relative "../../../serialization_utils"
-require_relative "../client_registry"
-require_relative "../state_registry"
 require_relative "../provider"
 
 module MCP
@@ -39,28 +38,33 @@ class McpAuthorizationServerProvider
           include OAuthAuthorizationServerProvider
           include SerializationUtils
 
+          FIVE_MINUTES_IN_SECONDS = 300
+
           def initialize(
             auth_server_settings:,
-            client_registry: nil,
-            state_registry: nil
+            client_registry:,
+            state_registry:,
+            auth_code_registry:,
+            token_registry:
           )
             @settings = auth_server_settings
-            @client_registry = client_registry || InMemoryClientRegistry.new
-            @state_registry = state_registry || InMemoryStateRegistry.new
+            @client_registry = client_registry
+            @state_registry = state_registry
+            @auth_code_registry = auth_code_registry
+            @token_registry = token_registry
           end
 
           def get_client(client_id)
-            @client_registry.find_by_id(client_id)
+            @client_registry.find_client(client_id)
           end
 
           def register_client(client_info)
-            @client_registry.create(client_info)
+            @client_registry.create_client(client_info)
           end
 
           def authorize(client_info:, auth_params:)
             state = auth_params.state || SecureRandom.hex(16)
-
-            @state_registry.create(state, to_h(auth_params))
+            @state_registry.create_state(state, to_h(auth_params))
 
             auth_url = URI(@settings.auth_server_authorization_endpoint)
             auth_url.query = URI.encode_www_form([
@@ -73,6 +77,67 @@ def authorize(client_info:, auth_params:)
 
             auth_url.to_s
           end
+
+          def authorize_callback(code:, state:)
+            state_data = @state_registry.find_state(state)
+            raise Errors::AuthorizationError.invalid_request("invalid state parameter") if state_data.nil?
+
+            access_token = query_access_token!({
+              client_id: @settings.client_id,
+              client_secret: @settings.client_secret,
+              code:,
+              redirect_uri: @settings.mcp_callback_endpoint,
+              grant_type: "authorization_code",
+            })
+
+            mcp_auth_code = "mcp_#{SecureRandom.hex(16)}"
+            auth_code = MCP::Auth::Server::AuthorizationCode.new(
+              code: mcp_auth_code,
+              client_id: state_data[:client_id],
+              redirect_uri: state_data[:redirect_uri],
+              redirect_uri_provided_explicitly: state_data[:redirect_uri_provided_explicitly],
+              expires_at: Time.now.to_i + FIVE_MINUTES_IN_SECONDS,
+              scopes: ["mcp:user"],
+              code_challenge: state_data[:code_challenge],
+            )
+            @auth_code_registry.create_auth_code(mcp_auth_code, auth_code)
+            @token_registry.create_token(mcp_auth_code, AccessToken.new(
+              token: access_token,
+              client_id: state_data[:client_id],
+              scopes: @settings.auth_server_scopes,
+              expires_at: nil,
+            ))
+
+            redirect_uri = URI(state_data[:redirect_uri])
+            redirect_uri.query = URI.encode_www_form([
+              ["code", mcp_auth_code],
+              ["state", state],
+            ])
+            @state_registry.delete_state(state)
+
+            redirect_uri
+          end
+
+          private
+
+          def query_access_token!(data)
+            uri = URI(@settings.auth_server_token_endpoint)
+            response = Net::HTTP.post_form(uri, stringify_keys(data))
+            raise Errors::AuthorizationError.new(
+              error_code: Errors::AuthorizationError::INVALID_REQUEST,
+              message: "failed to exchange code for token",
+            ) unless response.is_a?(Net::HTTPOK)
+
+            data = JSON.parse(response.body)
+            if data.key?("error")
+              raise Errors::AuthorizationError.new(
+                error_code: data["error"],
+                message: data["error_description"] || data["error"],
+              )
+            end
+
+            data["access_token"]
+          end
         end
       end
     end

lib/mcp/auth/server/providers/mcp_authorization_server_provider.rb

@@ -0,0 +1,23 @@
+# frozen_string_literal: true
+
+module MCP
+  module Auth
+    module Server
+      module Registries
+        module AuthCodeRegistry
+          def create_auth_code(code_id, data)
+            raise NotImplementedError, "#{self.class.name}#create_auth_code is not implemented"
+          end
+
+          def find_auth_code(code_id)
+            raise NotImplementedError, "#{self.class.name}#find_auth_code is not implemented"
+          end
+
+          def delete_auth_code(code_id)
+            raise NotImplementedError, "#{self.class.name}#delete_auth_code is not implemented"
+          end
+        end
+      end
+    end
+  end
+end

lib/mcp/auth/server/registries/auth_code_registry.rb

@@ -0,0 +1,19 @@
+# frozen_string_literal: true
+
+module MCP
+  module Auth
+    module Server
+      module Registries
+        module ClientRegistry
+          def create_client(client_info)
+            raise NotImplementedError, "#{self.class.name}#create_client is not implemented"
+          end
+
+          def find_client(client_id)
+            raise NotImplementedError, "#{self.class.name}#find_client is not implemented"
+          end
+        end
+      end
+    end
+  end
+end

lib/mcp/auth/server/registries/client_registry.rb

@@ -0,0 +1,80 @@
+# frozen_string_literal: true
+
+require_relative "auth_code_registry"
+require_relative "client_registry"
+require_relative "state_registry"
+require_relative "token_registry"
+
+module MCP
+  module Auth
+    module Server
+      module Registries
+        class InMemoryRegistry
+          include AuthCodeRegistry
+          include ClientRegistry
+          include StateRegistry
+          include TokenRegistry
+
+          def initialize
+            @codes = {}
+            @clients = {}
+            @states = {}
+            @tokens = {}
+          end
+
+          def create_client(client_info)
+            raise ArgumentError, "Client '#{client_info.client_id}' already exists" if @clients.key?(client_info.client_id)
+
+            @clients[client_info.client_id] = client_info
+          end
+
+          def find_client(client_id)
+            @clients[client_id]
+          end
+
+          def create_auth_code(code_id, data)
+            raise ArgumentError, "Code with id '#{code}' already exists" if @codes.key?(code_id)
+
+            @codes[code_id] = data
+          end
+
+          def find_code(code_id)
+            @codes[code_id]
+          end
+
+          def delete_code(code_id)
+            @codes.delete(code_id)
+          end
+
+          def create_state(state_id, state)
+            raise ArgumentError, "State with id '#{state_id}' already exists" if @states.key?(state_id)
+
+            @states[state_id] = state
+          end
+
+          def find_state(state_id)
+            @states[state_id]
+          end
+
+          def delete_state(state_id)
+            @states.delete(state_id)
+          end
+
+          def create_token(token_id, token)
+            raise ArgumentError, "Token with id '#{token_id}' already exists" if @tokens.key?(token_id)
+
+            @tokens[token_id] = token
+          end
+
+          def find_token(token_id)
+            @tokens[token_id]
+          end
+
+          def delete_token(token_id)
+            @tokens.delete(token_id)
+          end
+        end
+      end
+    end
+  end
+end