Improve tests and refactor

lib/rubrik.rb

@@ -3,13 +3,13 @@
 
 require "sorbet-runtime"
 require "pdf-reader"
-require "irb"
 
 module Rubrik
   class Error < StandardError; end
 end
 
 require_relative "rubrik/document"
 require_relative "rubrik/document/increment"
+require_relative "rubrik/document/serialize_object"
 require_relative "rubrik/fill_signature"
 require_relative "rubrik/sign"

lib/rubrik/document.rb

@@ -12,19 +12,18 @@ class Document
     SIGNATURE_SIZE = 8_192
 
     sig {returns(T.any(File, Tempfile, StringIO))}
-    attr_reader :io
+    attr_accessor :io
 
     sig {returns(PDF::Reader::ObjectHash)}
-    attr_reader :objects
+    attr_accessor :objects
 
     sig {returns(T::Array[{id: PDF::Reader::Reference, value: T.untyped}])}
-    attr_reader :modified_objects
-
-    sig {returns(PDF::Reader::Reference)}
-    attr_reader :interactive_form_id
+    attr_accessor :modified_objects
 
     sig {returns(Integer)}
-    attr_reader :last_object_id
+    attr_accessor :last_object_id
+
+    private :io=, :objects=, :modified_objects=, :last_object_id=
 
     sig {params(input: T.any(File, Tempfile, StringIO)).void}
     def initialize(input)
@@ -36,9 +35,11 @@ def initialize(input)
       fetch_or_create_interactive_form!
     end
 
-    sig {void}
+    sig {returns(PDF::Reader::Reference)}
+    # Returns the reference of the Signature Value dictionary.
     def add_signature_field
-      # create signature value dictionary
+      # To add an signature to the PDF, we need the following structure
+      # Interactive Form -> Signature Field -> Signature Value
       signature_value_id = assign_new_object_id!
       modified_objects << {
         id: signature_value_id,
@@ -63,7 +64,7 @@ def add_signature_field
           V: signature_value_id,
           Type: :Annot,
           Subtype: :Widget,
-          Rect: [20, 20, 120, 120],
+          Rect: [0, 0, 0, 0],
           F: 4,
           P: first_page_reference
         }
@@ -75,6 +76,8 @@ def add_signature_field
       modified_objects << {id: first_page_reference, value: modified_page}
 
       (interactive_form[:Fields] ||= []) << signature_field_id
+
+      signature_value_id
     end
 
     private
@@ -87,7 +90,7 @@ def interactive_form
     sig {void}
     def fetch_or_create_interactive_form!
       root_ref = objects.trailer[:Root]
-      root = T.let(objects.fetch(root_ref), Hash)
+      root = T.let(objects.fetch(root_ref), T::Hash[Symbol, T.untyped])
 
       if root.key?(:AcroForm)
         form_id = root[:AcroForm]
@@ -111,17 +114,5 @@ def fetch_or_create_interactive_form!
     def assign_new_object_id!
       PDF::Reader::Reference.new(self.last_object_id += 1, 0)
     end
-
-    sig {params(io: T.any(File, Tempfile, StringIO)).returns(T.any(File, Tempfile, StringIO))}
-    attr_writer :io
-
-    sig {params(objects: PDF::Reader::ObjectHash).returns(PDF::Reader::ObjectHash)}
-    attr_writer :objects
-
-    sig {params(modified_objects: T::Array[{id: PDF::Reader::Reference, value: T.untyped}]).returns(T::Array[{id: PDF::Reader::Reference, value: T.untyped}])}
-    attr_writer :modified_objects
-
-    sig {params(last_object_id: Integer).returns(Integer)}
-    attr_writer :last_object_id
   end
 end

lib/rubrik/document/increment.rb

@@ -9,19 +9,22 @@ module Increment
       extend T::Sig
       extend self
 
-      sig {params(document: Rubrik::Document, io: T.any(File, Tempfile, StringIO)).returns(T.any(File, Tempfile, StringIO))}
+      sig {params(document: Rubrik::Document, io: T.any(File, Tempfile, StringIO)).void}
       def call(document, io:)
         document.io.rewind
-        IO.copy_stream(T.unsafe(document.io), T.unsafe(io))
+        IO.copy_stream(document.io, io)
 
         io << "\n"
-        new_xref = Array.new
+
+        new_xref = T.let([], T::Array[T::Hash[Symbol, Integer]])
+        new_xref << {id: 0}
 
         document.modified_objects.each do |object|
           integer_id = T.let(object[:id].to_i, Integer)
           new_xref << {id: integer_id, offset: io.pos}
 
-          io << "#{integer_id} 0 obj\n" "#{serialize(object[:value])}\n" "endobj\n\n"
+          value = object[:value]
+          io << "#{integer_id} 0 obj\n" "#{SerializeObject[value]}\n" "endobj\n\n"
         end
 
         updated_trailer = document.objects.trailer.dup
@@ -31,72 +34,34 @@ def call(document, io:)
         new_xref_pos = io.pos
 
         new_xref_subsections = new_xref
-          .sort_by { _1[:id] }
-          .chunk_while { _1[:id] + 1 == _2[:id] }
+          .sort_by { |entry| entry.fetch(:id) }
+          .chunk_while {  _1.fetch(:id) + 1 == _2[:id] }
 
         io << "xref\n"
-        io << "0 1\n"
-        io << "0000000000 65535 f\n"
 
         new_xref_subsections.each do |subsection|
-          starting_id = subsection.first[:id]
+          starting_id = T.must(subsection.first).fetch(:id)
           length = subsection.length
 
           io << "#{starting_id} #{length}\n"
+
+          if starting_id.zero?
+            io << "0000000000 65535 f\n"
+            subsection.shift
+          end
+
           subsection.each { |entry| io << "#{format("%010d", entry[:offset])} 00000 n\n" }
         end
 
         io << "trailer\n"
-        io << "#{serialize(updated_trailer)}\n"
+        io << "#{SerializeObject[updated_trailer]}\n"
         io << "startxref\n"
         io << "#{new_xref_pos.to_s}\n"
         io << "%%EOF\n"
-
-        io.rewind
-        io
       end
 
       private
 
-      sig {params(obj: T.untyped).returns(String)}
-      def serialize(obj)
-        case obj
-        when Hash
-          serialized_objs = obj.flatten.map { |e| serialize(e) }
-          "<<#{serialized_objs.join(" ")}>>"
-        when Symbol
-          "/#{obj}"
-        when Array
-          serialized_objs = obj.map { |e| serialize(e) }
-          "[#{serialized_objs.join(" ")}]"
-        when PDF::Reader::Reference
-          "#{obj.id} #{obj.gen} R"
-        when String
-          "(#{obj})"
-        when TrueClass
-          "true"
-        when FalseClass
-          "false"
-        when Document::CONTENTS_PLACEHOLDER
-          "<#{"0" * Document::SIGNATURE_SIZE}>"
-        when Document::BYTE_RANGE_PLACEHOLDER
-          "[0 0000000000 0000000000 0000000000]"
-        when Numeric
-          obj.to_s
-        when NilClass
-          "null"
-        when PDF::Reader::Stream
-          <<~OBJECT.chomp
-            #{serialize(obj.hash)}
-            stream
-            #{obj.data}
-            endstream
-          OBJECT
-        else
-          raise NotImplementedError.new("Don't know how to serialize #{obj}")
-        end
-      end
-
       sig {params(document: Rubrik::Document).returns(Integer)}
       def last_xref_pos(document)
         PDF::Reader::Buffer.new(document.io, seek: 0).find_first_xref_offset

lib/rubrik/document/serialize_object.rb

@@ -0,0 +1,53 @@
+# typed: true
+# frozen_string_literal: true
+
+module Rubrik
+  class Document
+    module SerializeObject
+      include Kernel
+      extend T::Sig
+      extend self
+
+      sig {params(obj: T.untyped).returns(String)}
+      def [](obj)
+        case obj
+        when Hash
+          serialized_objs = obj.flatten.map { |e| SerializeObject[e] }
+          "<<#{serialized_objs.join(" ")}>>"
+        when Symbol
+          "/#{obj}"
+        when Array
+          serialized_objs = obj.map { |e| SerializeObject[e] }
+          "[#{serialized_objs.join(" ")}]"
+        when PDF::Reader::Reference
+          "#{obj.id} #{obj.gen} R"
+        when String
+          "(#{obj})"
+        when TrueClass
+          "true"
+        when FalseClass
+          "false"
+        when Document::CONTENTS_PLACEHOLDER
+          "<#{"0" * Document::SIGNATURE_SIZE}>"
+        when Document::BYTE_RANGE_PLACEHOLDER
+          "[0 0000000000 0000000000 0000000000]"
+        when Float, Integer
+          obj.to_s
+        when NilClass
+          "null"
+        when PDF::Reader::Stream
+          <<~OBJECT.chomp
+            #{SerializeObject[obj.hash]}
+            stream
+            #{obj.data}
+            endstream
+          OBJECT
+        else
+          raise "Don't know how to serialize #{obj}"
+        end
+      end
+
+      alias call []
+    end
+  end
+end

lib/rubrik/fill_signature.rb

@@ -15,7 +15,7 @@ module FillSignature
           private_key: OpenSSL::PKey::RSA,
           public_key: OpenSSL::X509::Certificate,
           certificate_chain: T::Array[OpenSSL::X509::Certificate])
-        .returns(T.any(File, StringIO, Tempfile))}
+        .void}
 
     FIRST_OFFSET = 0
 
@@ -29,8 +29,7 @@ def call(io, signature_value_ref:, private_key:, public_key:, certificate_chain:
       io.gets("<")
 
       first_length = io.pos - 1
-      # we need to double the SIGNATURE_SIZE because the hex encoding double the data size
-      # we also need to sum +2 to account for "<" and ">" of the hex string
+      # We need to sum +2 to account for "<" and ">" of the hex string
       second_offset = first_length + Document::SIGNATURE_SIZE + 2
       second_length = io.size - second_offset
 

lib/rubrik/sign.rb

@@ -13,15 +13,15 @@ module Sign
            certificate_chain: T::Array[OpenSSL::X509::Certificate])
          .void}
     def self.call(input, output, private_key:, public_key:, certificate_chain: [])
+      input.binmode
+      output.reopen(T.unsafe(output), "wb+") if !output.is_a?(StringIO)
+
       document = Rubrik::Document.new(input)
 
-      document.add_signature_field
+      signature_value_ref = document.add_signature_field
 
       Document::Increment.call(document, io: output)
 
-      signature_value = T.must(document.modified_objects.find { _1.dig(:value, :Type) == :Sig })
-
-      signature_value_ref = T.let(signature_value[:id], PDF::Reader::Reference)
       FillSignature.call(output, signature_value_ref:, private_key:, public_key:, certificate_chain:)
     end
   end