Commit f76b9dd

mo khan <mo@mokhan.ca>
2026-09-07 03:19:37
feat: restore OAuth support dropped during the reorg
lib/elelem/mcp/oauth.rb and its spec were lost when MCP was reorganized into this standalone gem (the port only carried over client.rb, token_storage.rb, and their specs). Ports the PKCE authorization-code flow verbatim, fixing `class MCP` to `module MCP` to match this gem's actual nesting, and replaces the removed net-hippie block-based http client with a small Net::HTTP-backed HttpAdapter matching the same get/post(&block) shape the spec mocks. Wires HttpServer#post to retry once through OAuth on 401 instead of raising "not yet supported", and restores the Authorization header and its request_headers spec example that were dropped on the (incorrect) assumption OAuth had been deliberately cut rather than just not yet ported. Claude-Session: https://claude.ai/code/session_01DUuj4amvRrPxDnHPHBkvz5
1 parent 2b8eedf
lib/elelem/mcp/client.rb
@@ -157,6 +157,7 @@ module Elelem
         @open_timeout = open_timeout
         @id = 0
         @session_id = nil
+        @access_token = nil
         handshake!
       end
 
@@ -185,7 +186,7 @@ module Elelem
         post(msg)
       end
 
-      def post(msg)
+      def post(msg, retry_auth: true)
         request = Net::HTTP::Post.new(@uri)
         request["content-type"] = "application/json"
         request_headers.each { |k, v| request[k] = v }
@@ -204,7 +205,10 @@ module Elelem
               return parse_response(response)
             when Net::HTTPUnauthorized
               response.body
-              raise "MCP: server requires OAuth (not yet supported)"
+              raise "Authorization failed" unless retry_auth
+
+              @access_token = OAuth.new(@uri.to_s).token
+              return post(msg, retry_auth: false)
             else
               raise "HTTP #{response.code}: #{response.body}"
             end
@@ -215,6 +219,7 @@ module Elelem
       def request_headers
         base = { "Accept" => "application/json, text/event-stream" }
         base["Mcp-Session-Id"] = @session_id if @session_id
+        base["Authorization"] = "Bearer #{@access_token}" if @access_token
         @headers.merge(base)
       end
 
lib/elelem/mcp/oauth.rb
@@ -0,0 +1,237 @@
+# frozen_string_literal: true
+
+module Elelem
+  module MCP
+    class OAuth
+      CALLBACK_PORT = 18273
+      REDIRECT_URI = "http://127.0.0.1:#{CALLBACK_PORT}/callback"
+
+      def initialize(resource_url, http: HttpAdapter.new)
+        @resource_url = resource_url
+        @http = http
+        @storage = TokenStorage.new
+      end
+
+      def token
+        stored = @storage.load(@resource_url)
+        return stored[:access_token] if stored && !expired?(stored)
+        return refresh(stored[:refresh_token]) if stored&.dig(:refresh_token)
+
+        authorize
+      end
+
+      private
+
+      def expired?(stored)
+        return false unless stored[:expires_at]
+
+        Time.now.to_i >= stored[:expires_at] - 60
+      end
+
+      def authorize
+        metadata = discover_auth_server
+        client = load_or_register_client(metadata)
+        verifier, challenge = generate_pkce
+        state = SecureRandom.hex(16)
+
+        auth_url = build_auth_url(metadata, client, challenge, state)
+        open_browser(auth_url)
+
+        code = wait_for_callback(state)
+        tokens = exchange_code(metadata, client, code, verifier)
+
+        @storage.save(
+          @resource_url,
+          access_token: tokens["access_token"],
+          refresh_token: tokens["refresh_token"],
+          expires_in: tokens["expires_in"]
+        )
+
+        tokens["access_token"]
+      end
+
+      def refresh(refresh_token)
+        metadata = discover_auth_server
+        client = load_or_register_client(metadata)
+        uri = URI.parse(metadata["token_endpoint"])
+
+        body = {
+          grant_type: "refresh_token",
+          refresh_token: refresh_token,
+          client_id: client[:client_id]
+        }
+
+        response = post_form(uri, body)
+        tokens = JSON.parse(response.body)
+
+        @storage.save(
+          @resource_url,
+          access_token: tokens["access_token"],
+          refresh_token: tokens["refresh_token"] || refresh_token,
+          expires_in: tokens["expires_in"]
+        )
+
+        tokens["access_token"]
+      rescue StandardError => e
+        warn "Token refresh failed: #{e.message}"
+        authorize
+      end
+
+      def discover_auth_server
+        resource_uri = URI.parse(@resource_url)
+        metadata_url = "#{resource_uri.scheme}://#{resource_uri.host}/.well-known/oauth-protected-resource"
+
+        resource_metadata = fetch_json(metadata_url)
+        auth_server_url = resource_metadata["authorization_servers"]&.first
+        raise "No authorization server found" unless auth_server_url
+
+        auth_metadata_url = "#{auth_server_url}/.well-known/oauth-authorization-server"
+        fetch_json(auth_metadata_url)
+      end
+
+      def load_or_register_client(metadata)
+        stored = @storage.load_client(@resource_url)
+        return stored if stored
+
+        client = register_client(metadata)
+        @storage.save_client(@resource_url, client)
+        @storage.load_client(@resource_url)
+      end
+
+      def register_client(metadata)
+        endpoint = metadata["registration_endpoint"]
+        raise "Dynamic registration not supported" unless endpoint
+
+        body = {
+          client_name: "elelem",
+          redirect_uris: [REDIRECT_URI],
+          grant_types: %w[authorization_code refresh_token],
+          response_types: ["code"],
+          token_endpoint_auth_method: "none"
+        }
+
+        response = post_json(endpoint, body)
+        JSON.parse(response.body)
+      end
+
+      def generate_pkce
+        verifier = SecureRandom.urlsafe_base64(32)
+        challenge = Base64.urlsafe_encode64(
+          Digest::SHA256.digest(verifier),
+          padding: false
+        )
+        [verifier, challenge]
+      end
+
+      def build_auth_url(metadata, client, challenge, state)
+        params = {
+          response_type: "code",
+          client_id: client[:client_id],
+          redirect_uri: REDIRECT_URI,
+          scope: metadata["scopes_supported"]&.join(" ") || "openid",
+          state: state,
+          code_challenge: challenge,
+          code_challenge_method: "S256"
+        }
+
+        "#{metadata["authorization_endpoint"]}?#{URI.encode_www_form(params)}"
+      end
+
+      def open_browser(url)
+        commands = ["xdg-open", "open", "start"]
+        commands.each do |cmd|
+          return if system(cmd, url, out: File::NULL, err: File::NULL)
+        end
+        warn "Open this URL in your browser: #{url}"
+      end
+
+      def wait_for_callback(expected_state)
+        code = nil
+        @server = WEBrick::HTTPServer.new(
+          Port: CALLBACK_PORT,
+          Logger: WEBrick::Log.new(File::NULL),
+          AccessLog: []
+        )
+
+        at_exit { @server&.shutdown }
+
+        @server.mount_proc("/callback") do |req, res|
+          state = req.query["state"]
+          raise "State mismatch" unless state == expected_state
+
+          code = req.query["code"]
+          res.content_type = "text/html"
+          res.body = "<html><body><h1>Authorization complete</h1><p>You can close this window.</p></body></html>"
+          @server.shutdown
+        end
+
+        Timeout.timeout(120) { @server.start }
+        code
+      rescue Timeout::Error
+        @server.shutdown
+        raise "OAuth callback timed out"
+      end
+
+      def exchange_code(metadata, client, code, verifier)
+        uri = URI.parse(metadata["token_endpoint"])
+
+        body = {
+          grant_type: "authorization_code",
+          code: code,
+          redirect_uri: REDIRECT_URI,
+          client_id: client[:client_id],
+          code_verifier: verifier
+        }
+
+        response = post_form(uri, body)
+        JSON.parse(response.body)
+      end
+
+      def fetch_json(url)
+        response = nil
+        @http.get(url) { |r| response = r }
+        JSON.parse(response.body)
+      end
+
+      def post_json(url, body)
+        response = nil
+        @http.post(
+          url,
+          headers: { "Content-Type" => "application/json" },
+          body: body.to_json
+        ) { |r| response = r }
+        response
+      end
+
+      def post_form(uri, body)
+        response = nil
+        @http.post(
+          uri.to_s,
+          headers: { "Content-Type" => "application/x-www-form-urlencoded" },
+          body: URI.encode_www_form(body)
+        ) { |r| response = r }
+        response
+      end
+
+      class HttpAdapter
+        def get(url)
+          uri = URI.parse(url)
+          Net::HTTP.start(uri.host, uri.port, use_ssl: uri.scheme == "https") do |conn|
+            yield conn.get(uri.request_uri)
+          end
+        end
+
+        def post(url, headers: {}, body: nil)
+          uri = URI.parse(url)
+          request = Net::HTTP::Post.new(uri)
+          headers.each { |k, v| request[k] = v }
+          request.body = body
+
+          Net::HTTP.start(uri.host, uri.port, use_ssl: uri.scheme == "https") do |conn|
+            yield conn.request(request)
+          end
+        end
+      end
+    end
+  end
+end
lib/elelem/mcp.rb
@@ -1,15 +1,20 @@
 # frozen_string_literal: true
 
+require "base64"
 require "digest"
 require "elelem"
 require "fileutils"
 require "json"
 require "net/http"
 require "open3"
+require "securerandom"
+require "timeout"
 require "uri"
+require "webrick"
 
 require_relative "mcp/version"
 require_relative "mcp/token_storage"
+require_relative "mcp/oauth"
 require_relative "mcp/client"
 
 Elelem::Plugins.register(:mcp) do |agent|
spec/elelem/mcp/http_server_spec.rb
@@ -72,6 +72,7 @@ RSpec.describe Elelem::MCP::HttpServer do
       server = described_class.allocate
       server.instance_variable_set(:@headers, {})
       server.instance_variable_set(:@session_id, nil)
+      server.instance_variable_set(:@access_token, nil)
       server
     end
 
@@ -85,5 +86,11 @@ RSpec.describe Elelem::MCP::HttpServer do
       headers = subject.send(:request_headers)
       expect(headers["Mcp-Session-Id"]).to eq("abc123")
     end
+
+    it "includes authorization when access_token set" do
+      subject.instance_variable_set(:@access_token, "token123")
+      headers = subject.send(:request_headers)
+      expect(headers["Authorization"]).to eq("Bearer token123")
+    end
   end
 end
spec/elelem/mcp/oauth_spec.rb
@@ -0,0 +1,92 @@
+# frozen_string_literal: true
+
+RSpec.describe Elelem::MCP::OAuth do
+  let(:resource_url) { "https://mcp.example.com/sse" }
+  let(:http) { double("http") }
+  let(:storage) { instance_double(Elelem::MCP::TokenStorage) }
+
+  subject { described_class.new(resource_url, http: http) }
+
+  before do
+    allow(Elelem::MCP::TokenStorage).to receive(:new).and_return(storage)
+  end
+
+  describe "#token" do
+    context "with valid cached token" do
+      it "returns cached access_token" do
+        allow(storage).to receive(:load).with(resource_url).and_return({
+          access_token: "cached_token",
+          expires_at: Time.now.to_i + 3600
+        })
+
+        expect(subject.token).to eq("cached_token")
+      end
+    end
+
+    context "with expired token and refresh_token" do
+      let(:auth_metadata) do
+        {
+          "authorization_endpoint" => "https://auth.example.com/authorize",
+          "token_endpoint" => "https://auth.example.com/token",
+          "registration_endpoint" => "https://auth.example.com/register"
+        }
+      end
+
+      it "refreshes using refresh_token" do
+        allow(storage).to receive(:load).with(resource_url).and_return({
+          access_token: "old_token",
+          refresh_token: "refresh_abc",
+          expires_at: Time.now.to_i - 100
+        })
+
+        allow(storage).to receive(:load_client).with(resource_url).and_return({
+          client_id: "elelem-123"
+        })
+
+        resource_response = double(body: { "authorization_servers" => ["https://auth.example.com"] }.to_json)
+        auth_response = double(body: auth_metadata.to_json)
+        token_response = double(body: { "access_token" => "new_token", "expires_in" => 3600 }.to_json)
+
+        allow(http).to receive(:get).with("https://mcp.example.com/.well-known/oauth-protected-resource").and_yield(resource_response)
+        allow(http).to receive(:get).with("https://auth.example.com/.well-known/oauth-authorization-server").and_yield(auth_response)
+        allow(http).to receive(:post).and_yield(token_response)
+        allow(storage).to receive(:save)
+
+        expect(subject.token).to eq("new_token")
+      end
+    end
+  end
+
+  describe "PKCE generation" do
+    it "generates valid verifier and challenge" do
+      verifier, challenge = subject.send(:generate_pkce)
+
+      expect(verifier.length).to be >= 43
+      expect(challenge.length).to be >= 43
+      expect(challenge).not_to include("+", "/", "=")
+
+      expected_challenge = Base64.urlsafe_encode64(
+        Digest::SHA256.digest(verifier),
+        padding: false
+      )
+      expect(challenge).to eq(expected_challenge)
+    end
+  end
+
+  describe "expiration check" do
+    it "considers token expired when within 60 seconds of expiry" do
+      stored = { expires_at: Time.now.to_i + 30 }
+      expect(subject.send(:expired?, stored)).to be true
+    end
+
+    it "considers token valid when more than 60 seconds remain" do
+      stored = { expires_at: Time.now.to_i + 120 }
+      expect(subject.send(:expired?, stored)).to be false
+    end
+
+    it "considers token valid when no expires_at" do
+      stored = { expires_at: nil }
+      expect(subject.send(:expired?, stored)).to be false
+    end
+  end
+end
elelem-mcp.gemspec
@@ -25,11 +25,15 @@ Gem::Specification.new do |spec|
   end
   spec.require_paths = ["lib"]
 
+  spec.add_dependency "base64", "~> 0.3"
   spec.add_dependency "digest", "~> 3.2"
   spec.add_dependency "elelem", "~> 0.1"
   spec.add_dependency "fileutils", "~> 1.8"
   spec.add_dependency "json", "~> 2.21"
   spec.add_dependency "net-http", "~> 0.9"
   spec.add_dependency "open3", "~> 0.2"
+  spec.add_dependency "securerandom", "~> 0.4"
+  spec.add_dependency "timeout", "~> 0.6"
   spec.add_dependency "uri", "~> 1.1"
+  spec.add_dependency "webrick", "~> 1.9"
 end
Gemfile.lock
@@ -2,13 +2,17 @@ PATH
   remote: .
   specs:
     elelem-mcp (0.1.0)
+      base64 (~> 0.3)
       digest (~> 3.2)
       elelem (~> 0.1)
       fileutils (~> 1.8)
       json (~> 2.21)
       net-http (~> 0.9)
       open3 (~> 0.2)
+      securerandom (~> 0.4)
+      timeout (~> 0.6)
       uri (~> 1.1)
+      webrick (~> 1.9)
 
 GEM
   remote: https://rubygems.org/