♻️ uniformize with other sdks

Steepfile

@@ -18,10 +18,12 @@ target :lib do
   # Use the stubs located at sig/custom/<library_name>.rbs instead.
   library 'date'
   library 'fileutils'
+  library 'ipaddr'
   library 'logger'
   library 'json'
   library 'openssl'
   library 'pathname'
+  library 'resolv'
   library 'tempfile'
   library 'time'
   library 'uri'

lib/mindee/http.rb

@@ -1,3 +1,4 @@
 # frozen_string_literal: true
 
+require_relative 'http/cancellation_token'
 require_relative 'http/http_error_handler'

lib/mindee/http/cancellation_token.rb

@@ -0,0 +1,24 @@
+# frozen_string_literal: true
+
+module Mindee
+  module HTTP
+    # Custom cancellation token class for polling.
+    class CancellationToken
+      private attr_reader :is_canceled
+
+      def initialize
+        @is_canceled = false
+      end
+
+      # Cancel the token.
+      def cancel
+        @is_canceled = true
+      end
+
+      # Check if the token is canceled.
+      def canceled?
+        @is_canceled
+      end
+    end
+  end
+end

lib/mindee/input/sources/url_input_source.rb

@@ -3,21 +3,36 @@
 require 'net/http'
 require 'uri'
 require 'fileutils'
+require 'ipaddr'
+require 'resolv'
 require_relative '../../logging'
 
 module Mindee
   module Input
     module Source
       # Load a remote document from a file url.
       class URLInputSource
+        # IP ranges that must not be the target of a source URL (SSRF protection).
+        DISALLOWED_RANGES = [
+          IPAddr.new('127.0.0.0/8'), # IPv4 loopback
+          IPAddr.new('::1/128'), # IPv6 loopback
+          IPAddr.new('169.254.0.0/16'), # IPv4 link-local
+          IPAddr.new('fe80::/10'), # IPv6 link-local
+          IPAddr.new('10.0.0.0/8'),     # RFC 1918
+          IPAddr.new('172.16.0.0/12'),  # RFC 1918
+          IPAddr.new('192.168.0.0/16'), # RFC 1918
+          IPAddr.new('0.0.0.0/8'),      # "this" network
+          IPAddr.new('224.0.0.0/4'),    # IPv4 multicast
+          IPAddr.new('ff00::/8'),        # IPv6 multicast
+          IPAddr.new('fc00::/7'),        # IPv6 unique-local
+          IPAddr.new('100.64.0.0/10'), # Carrier-grade NAT (RFC 6598)
+        ].freeze
         # @return [String]
         attr_reader :url
 
         def initialize(url)
-          raise Error::MindeeInputError, 'URL must be HTTPS' unless url.start_with? 'https://'
-
+          validate_url!(url)
           logger.debug("URL input: #{url}")
-
           @url = url
         end
 
@@ -85,6 +100,62 @@ def fetch_file_content(username: nil, password: nil, token: nil, max_redirects:
 
         private
 
+        # Validates the URL against SSRF risks before it is stored or fetched.
+        # Checks scheme, embedded credentials, loopback hostnames, and resolved addresses.
+        # @param url [String]
+        # @raise [Error::MindeeInputError]
+        # rubocop:disable Metrics/CyclomaticComplexity, Metrics/PerceivedComplexity
+        def validate_url!(url)
+          uri = URI.parse(url)
+
+          raise Error::MindeeInputError, 'Only HTTPS source URLs are allowed' unless uri.scheme&.downcase == 'https'
+
+          userinfo = uri.userinfo
+          raise Error::MindeeInputError, 'Source URLs must not embed user credentials' if userinfo && !userinfo.empty?
+
+          host = uri.hostname
+          raise Error::MindeeInputError, 'Source URL is missing a host' if host.nil? || host.empty?
+
+          lower_host = host.downcase
+          if lower_host == 'localhost' || lower_host.end_with?('.localhost') ||
+             lower_host == 'ip6-localhost' || lower_host == 'ip6-loopback'
+            raise Error::MindeeInputError, "Loopback hostnames are not allowed: #{host}"
+          end
+
+          resolved_addresses(host).each do |addr|
+            if DISALLOWED_RANGES.any? { |range| range.include?(addr) }
+              raise Error::MindeeInputError,
+                    "Source URL host resolves to a disallowed address: #{addr}"
+            end
+          end
+        rescue URI::InvalidURIError
+          raise Error::MindeeInputError, "Invalid URL: #{url}"
+        end
+        # rubocop:enable Metrics/CyclomaticComplexity, Metrics/PerceivedComplexity
+
+        # Resolves a host to an array of IPAddr objects.
+        # Handles literal IP addresses directly without a DNS round-trip.
+        # @param host [String]
+        # @return [Array<IPAddr>]
+        # @raise [Error::MindeeInputError] if the host cannot be resolved.
+        def resolved_addresses(host)
+          return [IPAddr.new(host)] if literal_ip?(host)
+
+          addrs = Resolv.getaddresses(host)
+          raise Error::MindeeInputError, "Unable to resolve source URL host: #{host}" if addrs.empty?
+
+          addrs.map { |a| IPAddr.new(a) }
+        end
+
+        # Returns true if the string is a valid literal IPv4 or IPv6 address.
+        # @param host [String]
+        def literal_ip?(host)
+          IPAddr.new(host)
+          true
+        rescue IPAddr::InvalidAddressError, IPAddr::AddressFamilyError
+          false
+        end
+
         def extract_filename_from_url(uri)
           filename = File.basename(uri.path.to_s)
           filename.empty? ? '' : filename
@@ -105,6 +176,7 @@ def make_request(uri, request, max_redirects)
               location = response['location']
               raise Error::MindeeInputError, 'No location in redirection header.' if location.nil?
 
+              validate_url!(location)
               new_uri = URI.parse(location)
               request = Net::HTTP::Get.new(new_uri)
               make_request(new_uri, request, max_redirects - 1)

lib/mindee/v1/client.rb

@@ -219,6 +219,7 @@ def parse_queued(job_id, product_class, endpoint: nil)
         Mindee::V1::Parsing::Common::ApiResponse.new(product_class, prediction, raw_http)
       end
 
+      # rubocop:disable Metrics/CyclomaticComplexity
       # Enqueue a document for async parsing and automatically try to retrieve it
       #
       # @param input_source [Mindee::Input::Source::LocalInputSource, Mindee::Input::Source::URLInputSource]
@@ -247,13 +248,23 @@ def parse_queued(job_id, product_class, endpoint: nil)
       #   * `:delay_sec` [Numeric] Delay between polling attempts. Defaults to 1.5.
       #   * `:max_retries` [Integer] Maximum number of retries. Defaults to 80.
       # @param endpoint [Mindee::V1::HTTP::Endpoint] Endpoint of the API.
+      # @param cancellation_token [Mindee::HTTP::CancellationToken, nil] Token for cancellation.
       # @return [Mindee::V1::Parsing::Common::ApiResponse]
-      def enqueue_and_parse(input_source, product_class, endpoint, options)
+      # rubocop:disable Metrics/PerceivedComplexity
+      def enqueue_and_parse(
+        input_source,
+        product_class,
+        endpoint,
+        options,
+        cancellation_token = nil
+      )
         validate_async_params(options.initial_delay_sec, options.delay_sec, options.max_retries)
         enqueue_res = enqueue(input_source, product_class, endpoint: endpoint, options: options)
         job = enqueue_res.job or raise Error::MindeeAPIError, 'Expected job to be present'
         job_id = job.id
 
+        raise Mindee::Error::MindeeError, 'Enqueueing of the document was canceled.' if cancellation_token&.canceled?
+
         sleep(options.initial_delay_sec)
         polling_attempts = 1
         logger.debug("Successfully enqueued document with job id: '#{job_id}'")
@@ -266,6 +277,8 @@ def enqueue_and_parse(input_source, product_class, endpoint, options)
         # @type var valid_statuses: Array[(:waiting | :processing | :completed | :failed)]
         while valid_statuses.include?(queue_res_job.status) && polling_attempts < options.max_retries
           logger.debug("Polling server for parsing result with job id: '#{job_id}'. Attempt #{polling_attempts}")
+          raise Mindee::Error::MindeeError, 'Enqueueing of the document was canceled.' if cancellation_token&.canceled?
+
           sleep(options.delay_sec)
           queue_res = parse_queued(job_id, product_class, endpoint: endpoint)
           queue_res_job = queue_res.job or raise Error::MindeeAPIError, 'Expected job to be present'
@@ -280,7 +293,9 @@ def enqueue_and_parse(input_source, product_class, endpoint, options)
 
         queue_res
       end
+      # rubocop:enable Metrics/PerceivedComplexity
 
+      # rubocop:enable Metrics/CyclomaticComplexity
       # Sends a document to a workflow.
       #
       # Accepts options either as a Hash or as a WorkflowOptions struct.

lib/mindee/v2/client.rb

@@ -58,12 +58,15 @@ def enqueue(
       #   The source of the input document (local file or URL).
       # @param params [Hash, Input::BaseParameters] Parameters for the inference.
       # @param polling_options [Hash, PollingOptions, nil] Parameters for polling.
+      # @param cancellation_token [Mindee::HTTP::CancellationToken, nil] Token for cancellation.
       # @return [Parsing::BaseResponse]
+      # rubocop:disable Metrics/CyclomaticComplexity, Metrics/PerceivedComplexity
       def enqueue_and_get_result(
         product,
         input_source,
         params,
-        polling_options = nil
+        polling_options = nil,
+        cancellation_token = nil
       )
         enqueue_response = enqueue(product, input_source, params)
         normalized_params = normalize_parameters(product.params_type, params, polling_options: polling_options)
@@ -77,6 +80,8 @@ def enqueue_and_get_result(
         job_id = enqueue_response.job.id
         logger.debug("Successfully enqueued document with job id: #{job_id}.")
 
+        raise Mindee::Error::MindeeError, 'Enqueueing of the document was canceled.' if cancellation_token&.canceled?
+
         sleep(normalized_params.polling_options.initial_delay_sec)
         retry_counter = 1
         poll_results = get_job(job_id)
@@ -94,6 +99,8 @@ def enqueue_and_get_result(
             "Job status: #{poll_results.job.status}."
           )
 
+          raise Mindee::Error::MindeeError, 'Enqueueing of the document was canceled.' if cancellation_token&.canceled?
+
           sleep(normalized_params.polling_options.delay_sec)
           poll_results = get_job(job_id)
           retry_counter += 1
@@ -110,6 +117,7 @@ def enqueue_and_get_result(
         raise Mindee::Error::MindeeError,
               "Asynchronous parsing request timed out after #{sec_count} seconds"
       end
+      # rubocop:enable Metrics/CyclomaticComplexity, Metrics/PerceivedComplexity
 
       # Searches for a list of available models for the given API key.
       # @param model_name [String]

lib/mindee/v2/parsing/field/field_confidence.rb

@@ -87,8 +87,8 @@ def >=(other)
             end
           end
 
-          # less than or equality of two FieldConfidence instances.
-          # # @param other [String, Integer, FieldConfidence] The other confidence to compare.
+          # Less than or equality of two FieldConfidence instances.
+          # @param other [String, Integer, FieldConfidence] The other confidence to compare.
           def <=(other)
             if other.is_a?(FieldConfidence)
               to_i <= val_to_i(other.value)
@@ -101,12 +101,47 @@ def <=(other)
             end
           end
 
+          # Greater than comparison of two FieldConfidence instances.
+          # @param other [String, Integer, FieldConfidence] The other confidence to compare.
+          def >(other)
+            if other.is_a?(FieldConfidence)
+              to_i > val_to_i(other.value)
+            elsif other.is_a?(String)
+              to_i > val_to_i(other)
+            elsif other.is_a?(Integer)
+              to_i > other
+            else
+              raise ArgumentError, "Invalid type: #{other.class}"
+            end
+          end
+
+          # Less than comparison of two FieldConfidence instances.
+          # @param other [String, Integer, FieldConfidence] The other confidence to compare.
+          def <(other)
+            if other.is_a?(FieldConfidence)
+              to_i < val_to_i(other.value)
+            elsif other.is_a?(String)
+              to_i < val_to_i(other)
+            elsif other.is_a?(Integer)
+              to_i < other
+            else
+              raise ArgumentError, "Invalid type: #{other.class}"
+            end
+          end
+
           # rubocop:enable Style/CaseLikeIf
 
           # Aliases for the comparison operators.
           alias eql? ==
+          alias equal_value? ==
           alias gteql? >=
+          alias greater_than_or_equal? >=
           alias lteql? <=
+          alias less_than_or_equal? <=
+          alias gt? >
+          alias greater_than? >
+          alias lt? <
+          alias less_than? <
 
           protected
 

sig/mindee/http/cancellation_token.rbs

@@ -0,0 +1,15 @@
+# lib/http/cancellation_token.rb
+
+module Mindee
+  module HTTP
+    class CancellationToken
+      @is_canceled: bool
+
+      attr_reader is_canceled: bool
+
+      def canceled?: -> bool
+      def initialize: -> void
+      def cancel: -> void
+    end
+  end
+end

sig/mindee/input/sources/url_input_source.rbs

@@ -3,6 +3,8 @@ module Mindee
   module Input
     module Source
       class URLInputSource
+        DISALLOWED_RANGES: Array[IPAddr]
+
         def url: -> String
         def logger: () -> Logger
         def initialize: (String) -> void
@@ -14,6 +16,9 @@ module Mindee
         def make_request: (URI::Generic, Net::HTTP::Get, Integer) -> Net::HTTPRedirection
         def get_file_extension: (String) -> String?
         def generate_file_name: (?extension: String?) -> String
+        def validate_url!: (String) -> void
+        def resolved_addresses: (String) -> Array[IPAddr]
+        def literal_ip?: (String) -> bool
       end
     end
   end

sig/mindee/v2/parsing/field/field_confidence.rbs

@@ -20,6 +20,15 @@ module Mindee
           def lteql? : (String | Integer | FieldConfidence) -> bool
           def >=: (String | Integer | FieldConfidence) -> bool
           def gteql?: (String | Integer | FieldConfidence) -> bool
+          def greater_than_or_equal?: (String | Integer | FieldConfidence) -> bool
+          def >: (String | Integer | FieldConfidence) -> bool
+          def gt?: (String | Integer | FieldConfidence) -> bool
+          def greater_than?: (String | Integer | FieldConfidence) -> bool
+          def <: (String | Integer | FieldConfidence) -> bool
+          def lt?: (String | Integer | FieldConfidence) -> bool
+          def less_than?: (String | Integer | FieldConfidence) -> bool
+          def less_than_or_equal?: (String | Integer | FieldConfidence) -> bool
+          def equal_value?: (String | Integer | FieldConfidence) -> bool
           def inspect: -> String
 
           def val_to_i: (String) -> Integer

spec/http/cancellation_token_spec.rb

@@ -0,0 +1,23 @@
+# frozen_string_literal: true
+
+require 'mindee'
+
+describe Mindee::HTTP::CancellationToken do
+  it 'starts in a not-cancelled state' do
+    token = described_class.new
+    expect(token.canceled?).to be false
+  end
+
+  it 'is cancelled after calling #cancel' do
+    token = described_class.new
+    token.cancel
+    expect(token.canceled?).to be true
+  end
+
+  it 'remains cancelled once cancel has been called' do
+    token = described_class.new
+    token.cancel
+    token.cancel
+    expect(token.canceled?).to be true
+  end
+end