|
| 1 | +# frozen_string_literal: true |
| 2 | + |
| 3 | +require "digest" |
| 4 | +require "base64" |
| 5 | +require_relative "../../errors" |
| 6 | +require_relative "../../server/provider" |
| 7 | +require_relative "../../models" |
| 8 | + |
| 9 | +module MCP |
| 10 | + module Auth |
| 11 | + module Server |
| 12 | + module Handlers |
| 13 | + class BaseRequest |
| 14 | + attr_reader :grant_type, :client_id, :client_secret |
| 15 | + |
| 16 | + def initialize( |
| 17 | + grant_type:, |
| 18 | + client_id:, |
| 19 | + client_secret: nil |
| 20 | + ) |
| 21 | + @grant_type = grant_type |
| 22 | + @client_id = client_id |
| 23 | + @client_secret = client_secret |
| 24 | + end |
| 25 | + end |
| 26 | + |
| 27 | + class AuthorizationCodeRequest < BaseRequest |
| 28 | + attr_reader :code, :code_verifier, :redirect_uri |
| 29 | + |
| 30 | + def initialize( |
| 31 | + code:, |
| 32 | + code_verifier:, |
| 33 | + redirect_uri: nil, |
| 34 | + **base_kwargs |
| 35 | + ) |
| 36 | + super(**base_kwargs) |
| 37 | + |
| 38 | + @code = code |
| 39 | + @code_verifier = code_verifier |
| 40 | + @redirect_uri = redirect_uri |
| 41 | + |
| 42 | + if @grant_type != "authorization_code" |
| 43 | + raise Errors::AuthorizationError.invalid_request("grant_type must be authorization_code") |
| 44 | + end |
| 45 | + end |
| 46 | + end |
| 47 | + |
| 48 | + class TokenHandler |
| 49 | + def initialize( |
| 50 | + auth_server_provider:, |
| 51 | + request_parser: |
| 52 | + ) |
| 53 | + @auth_server_provider = auth_server_provider |
| 54 | + @request_parser = request_parser |
| 55 | + end |
| 56 | + |
| 57 | + def handle(request) |
| 58 | + params_h = @request_parser.parse_query_params(request) |
| 59 | + request = AuthorizationCodeRequest.new(**params_h) |
| 60 | + |
| 61 | + client_info = @auth_server_provider.get_client(request.client_id) |
| 62 | + validate_request_client!(request:, client_info:) |
| 63 | + |
| 64 | + auth_code = @auth_server_provider.load_authorization_code(request.code) |
| 65 | + validate_auth_code!(request:, auth_code:) |
| 66 | + validate_pkce!(request:, auth_code:) |
| 67 | + tokens = @auth_server_provider.exchange_authorization_code(auth_code) |
| 68 | + |
| 69 | + [200, {}, tokens] |
| 70 | + rescue Errors::ClientAuthenticationError => e |
| 71 | + bad_request_error( |
| 72 | + error_code: Errors::AuthorizationError::UNAUTHORIZED_CLIENT, |
| 73 | + error_description: e.message, |
| 74 | + ) |
| 75 | + rescue Errors::AuthorizationError => e |
| 76 | + bad_request_error( |
| 77 | + error_codea: e.error_code, |
| 78 | + error_description: e.message, |
| 79 | + ) |
| 80 | + rescue |
| 81 | + bad_request_error( |
| 82 | + error_code: Errors::AuthorizationError::SERVER_ERROR, |
| 83 | + error_description: "unexpected error", |
| 84 | + ) |
| 85 | + end |
| 86 | + |
| 87 | + private |
| 88 | + |
| 89 | + def validate_request_client!(request:, client_info: nil) |
| 90 | + if client_info.nil? |
| 91 | + raise Errors::AuthorizationError.invalid_request("invalid client_id") |
| 92 | + end |
| 93 | + |
| 94 | + client_info.authenticate!( |
| 95 | + request_client_id: request.client_id, |
| 96 | + request_client_secret: request.client_secret, |
| 97 | + ) |
| 98 | + unless client_info.valid_grant_type?(request.grant_type) |
| 99 | + raise Errors::AuthorizationError.new( |
| 100 | + error_code: Errors::AuthorizationError::UNSUPPORTED_GRANT_TYPE, |
| 101 | + message: "supported grant type are #{client_info.grant_types}", |
| 102 | + ) |
| 103 | + end |
| 104 | + end |
| 105 | + |
| 106 | + def validate_auth_code!(request:, auth_code:) |
| 107 | + if auth_code.nil? || !auth_code.belongs_to_client?(request.client_id) |
| 108 | + # If auth code is for another client, pretend it's not there |
| 109 | + raise Errors::AuthorizationError.invalid_grant("authorization code does not exist") |
| 110 | + end |
| 111 | + |
| 112 | + if auth_code.expired? |
| 113 | + raise Errors::AuthorizationError.invalid_grant("authorization code expired") |
| 114 | + end |
| 115 | + |
| 116 | + # verify redirect_uri doesn't change between /authorize and /tokens |
| 117 | + # see https://datatracker.ietf.org/doc/html/rfc6749#section-10.6 |
| 118 | + authorize_request_redirect_uri = auth_code.redirect_uri_provided_explicitly ? auth_code.redirect_uri : nil |
| 119 | + if request.redirect_uri != authorize_request_redirect_uri |
| 120 | + raise Errors::AuthorizationError.invalid_request("redirect_uri did not match the one when creating auth code") |
| 121 | + end |
| 122 | + end |
| 123 | + |
| 124 | + def validate_pkce!(request:, auth_code:) |
| 125 | + sha256 = Digest::SHA256.digest(request.code_verifier.encode) |
| 126 | + request_code_challenge = Base64.urlsafe_encode64(sha256).tr("=", "") |
| 127 | + |
| 128 | + unless auth_code.code_challenge_match?(request_code_challenge) |
| 129 | + # see https://datatracker.ietf.org/doc/html/rfc7636#section-4.6 |
| 130 | + raise Errors::AuthorizationError.invalid_grant("incorrect code_verifier") |
| 131 | + end |
| 132 | + end |
| 133 | + |
| 134 | + def bad_request_error( |
| 135 | + error_code:, |
| 136 | + error_description: |
| 137 | + ) |
| 138 | + body = { error: error_code, error_description: } |
| 139 | + |
| 140 | + [400, { "Cache-Control": "no-store", "Pragma": "no-cache" }, body] |
| 141 | + end |
| 142 | + end |
| 143 | + end |
| 144 | + end |
| 145 | + end |
| 146 | +end |
0 commit comments