herringbone 0.5.0 → 0.6.0

data/bin/herringbone

@@ -10,6 +10,13 @@
 #                                               and open it in the browser (printed when stdout is redirected)
 #   add --verify-checksums to any inspect form to read page bodies and check their CRCs
 # `inspect` reads only the footer, page headers and page indexes: no values are decompressed.
+#
+# Encrypted files: both commands take --key=KEY (a key matched by its fingerprint id, or used for
+# the footer when it is the only one), --key=KEY_ID=KEY, --footer-key=KEY and
+# --column-key=PATH=KEY (all repeatable but --footer-key), and --aad-prefix=PREFIX. KEY is hex, or
+# base64:..., or raw:... for the bytes as typed (a 16, 24 or 32-character String that is not
+# valid hex is taken as is too). Keys that are needed but not given are asked for on stdin
+# (without echo in a terminal; an empty answer skips the key), unless --no-prompt is given.
 $LOAD_PATH.unshift File.expand_path("../lib", __dir__)
 require "herringbone"
 require "json"
@@ -18,6 +25,7 @@ require "json"
 rescue LoadError
 end
 require "tmpdir"
+require "io/console"
 
 def jsonable(v)
   case v
@@ -43,12 +51,81 @@ def open_with_default_app(path)
   !!system(*command, out: File::NULL, err: File::NULL, exception: false)
 end
 
-USAGE = "usage: herringbone cat FILE [N]\n" \
-  "       herringbone inspect FILE [--format=text|json|html] [--pages] [--verify-checksums]\n" \
-  "       (--pages only with --format=text)"
+# A key as given on the command line or at the prompt: hex, base64:..., raw:..., or the bytes as
+# typed when they are not valid hex but have the size of a key
+# @return [String] the key
+# @raise [ArgumentError] when it is none of those, or not 16, 24 or 32 bytes long
+def parse_key(text, what)
+  text = text.strip
+  bytes = if text.start_with?("base64:")
+    text.delete_prefix("base64:").unpack1("m0")
+  elsif text.start_with?("raw:")
+    text.delete_prefix("raw:")
+  elsif text.match?(/\A(\h\h)+\z/) && Herringbone::Encryption::KEY_SIZES.include?(text.size / 2)
+    [text].pack("H*")
+  else
+    text
+  end
+  Herringbone::Encryption.check_key!(bytes, what)
+end
+
+# Asks for a key on stderr and reads it from stdin, without echo when stdin is a terminal
+# @return [String, nil] the key, nil when the answer is empty or stdin is exhausted
+def ask_key(metadata, owner)
+  what = (owner == :footer) ? "the footer" : "column #{owner}"
+  $stderr.print "Key for #{what}#{" (key metadata #{metadata.inspect})" if metadata}, hex: "
+  line = $stdin.tty? ? $stdin.noecho(&:gets) : $stdin.gets
+  $stderr.puts
+  return nil if line.nil? || line.strip.empty?
+  parse_key(line, "the key for #{what}")
+rescue ArgumentError => e
+  raise unless $stdin.tty?
+  warn "herringbone: #{e.message}"
+  retry
+end
+
+# The decryption: option from the key flags
+# @return [Hash{Symbol => Object}]
+def decryption_from(key_flags, prompt)
+  by_metadata = {}
+  plain = []
+  decryption = {columns: {}}
+  key_flags.each do |flag|
+    name, value = flag.split("=", 2)
+    case name
+    when "--footer-key" then decryption[:footer_key] = parse_key(value.to_s, "--footer-key")
+    when "--aad-prefix" then decryption[:aad_prefix] = value.to_s
+    when "--key"
+      target, key = value.to_s.split("=", 2)
+      if key.nil? || value.start_with?("base64:", "raw:")
+        plain << Herringbone::Key.new(parse_key(value.to_s, "--key"))
+        by_metadata[plain.last.id] = plain.last.bytes
+      else
+        by_metadata[target] = parse_key(key, "--key=#{target}")
+      end
+    when "--column-key"
+      target, key = value.to_s.split("=", 2)
+      raise ArgumentError, "--column-key expects --column-key=PATH=KEY" unless key
+      decryption[:columns][target] = parse_key(key, "--column-key=#{target}")
+    end
+  end
+  decryption[:footer_key] ||= plain.first.bytes if plain.size == 1
+  decryption[:keys] = ->(metadata, owner) { by_metadata.fetch(metadata) { prompt ? ask_key(metadata, owner) : nil } }
+  decryption
+end
+
+USAGE = "usage: herringbone cat FILE [N] [KEY OPTIONS]\n" \
+  "       herringbone inspect FILE [--format=text|json|html] [--pages] [--verify-checksums] [KEY OPTIONS]\n" \
+  "       (--pages only with --format=text)\n" \
+  "key options, for encrypted files: --key=KEY --key=KEY_ID=KEY --footer-key=KEY --column-key=PATH=KEY\n" \
+  "       --aad-prefix=PREFIX --no-prompt (KEY is hex, base64:... or raw:...)"
+
+KEY_FLAGS = %w[--footer-key --column-key --key --aad-prefix].freeze
 
 command, *rest = ARGV
 flags, args = rest.partition { |a| a.start_with?("--") }
+key_flags, flags = flags.partition { |f| KEY_FLAGS.include?(f.split("=", 2).first) }
+prompt = !flags.delete("--no-prompt")
 path, limit = args
 valid = case command
 when "cat" then flags.empty? && path && (limit.nil? || limit.match?(/\A\d+\z/))
@@ -64,14 +141,15 @@ unless valid
   exit 1
 end
 
-# Parquet and file errors are reported in one line, without a backtrace
+# Parquet, file and key errors are reported in one line, without a backtrace
 begin
+  decryption = decryption_from(key_flags, prompt)
   File.open(path, "rb") do |io|
     if command == "cat"
       n = limit&.to_i
-      Herringbone::Reader.new(io).each_row(limit: n) { |row| puts JSON.generate(jsonable(row)) }
+      Herringbone::Reader.new(io, decryption: decryption).each_row(limit: n) { |row| puts JSON.generate(jsonable(row)) }
     else
-      inspector = Herringbone::Inspector.new(io)
+      inspector = Herringbone::Inspector.new(io, decryption: decryption)
       inspector.verify_checksums if flags.include?("--verify-checksums")
       case format
       when "json"
@@ -90,7 +168,8 @@ begin
       end
     end
   end
-rescue Herringbone::Error, SystemCallError => e
+rescue Herringbone::Error, SystemCallError, ArgumentError => e
   warn "herringbone: #{e.message}"
+  warn "herringbone: give keys with --footer-key=KEY, --column-key=PATH=KEY or --key=KEY_METADATA=KEY" if e.is_a?(Herringbone::DecryptionError)
   exit 1
 end

data/lib/herringbone/bloom_filter.rb

@@ -300,8 +300,6 @@ module Herringbone
       "#<#{self.class.name} #{num_bytes} bytes#{" for #{@column.dotted_path}" if @column}>"
     end
 
-    private
-
     # Header for this filter: split block algorithm, XXH64 hash, uncompressed
     #
     # @return [Format::BloomFilterHeader] the header
@@ -314,57 +312,4 @@ module Herringbone
       )
     end
   end
-
-  class Reader
-    # Internal (used by reads with where:): the bloom filter of a column chunk. +column+ is a
-    # dotted path ("a.b"), an Array path or a Schema::Column. Returns a BloomFilter, or nil when
-    # the chunk has none (or one of an unknown kind).
-    #
-    # @param row_group_index [Integer] index of the row group
-    # @param column [String, Array<String, Symbol>, Symbol, Schema::Column] the leaf column
-    # @return [BloomFilter, nil] the filter, or nil when there is none or it is unsupported
-    # @raise [IndexError] when there is no such row group
-    # @raise [ArgumentError] when there is no such column
-    # @raise [FormatError] when the filter is truncated or its header cannot be decoded
-    def bloom_filter(row_group_index, column)
-      col = bloom_filter_column(column)
-      rg = row_groups.fetch(row_group_index) { raise IndexError, "No row group #{row_group_index}" }
-      meta = rg.columns.fetch(col.index).meta_data
-      offset = meta&.bloom_filter_offset
-      return nil unless offset
-      length = meta.bloom_filter_length
-      @io.seek(offset)
-      if length
-        buf = @io.read(length)
-        raise FormatError, "Truncated bloom filter" if buf.nil? || buf.bytesize < length
-        return BloomFilter.decode(buf, column: col)
-      end
-      # Without a length (older writers), read the header first, then the bitset it announces
-      head = @io.read(256) || "".b
-      reader = Thrift::Reader.new(head)
-      header = reader.read_struct(Format::BloomFilterHeader)
-      return nil unless BloomFilter.supported_header?(header)
-      @io.seek(offset + reader.pos)
-      bitset = @io.read(header.num_bytes)
-      raise FormatError, "Truncated bloom filter" if bitset.nil? || bitset.bytesize != header.num_bytes
-      BloomFilter.new(bitset: bitset, column: col)
-    rescue Thrift::Error => e
-      raise FormatError, "Corrupt bloom filter header for #{col.dotted_path}: #{e.message}"
-    end
-
-    private
-
-    # Resolves the +column+ argument of #bloom_filter to a leaf column
-    #
-    # @param column [String, Array<String, Symbol>, Symbol, Schema::Column] dotted path, path
-    #   Array or column
-    # @return [Schema::Column] the column
-    # @raise [ArgumentError] when the schema has no such column
-    def bloom_filter_column(column)
-      return column if column.is_a?(Schema::Column)
-      col = schema.column(column.is_a?(Array) ? column.map(&:to_s) : column.to_s)
-      raise ArgumentError, "No column #{column.inspect}" unless col
-      col
-    end
-  end
 end

data/lib/herringbone/compression.rb

@@ -83,6 +83,28 @@ module Herringbone
     # Codec name => codec id, the inverse of NAMES
     CODECS_BY_NAME = NAMES.invert.freeze
 
+    # Codec id => compression levels it accepts, for the writer's compression_level: option
+    LEVELS = Ractor.make_shareable({
+      Format::Codec::ZSTD => (-(1 << 17))..22,
+      Format::Codec::GZIP => 0..9,
+      Format::Codec::BROTLI => 0..11
+    })
+
+    # Checks that +codec+ takes a compression level and that +level+ is in its range
+    #
+    # @param codec [Integer] Format::Codec id
+    # @param level [Integer, nil] compression level, nil for the codec's default
+    # @return [Integer, nil] the level
+    # @raise [ArgumentError] when the codec has no levels or +level+ is out of its range
+    def check_level!(codec, level)
+      return if level.nil?
+      name = NAMES.fetch(codec, codec).inspect
+      range = LEVELS.fetch(codec) { raise ArgumentError, "#{name} compression has no levels" }
+      level = Integer(level)
+      raise ArgumentError, "#{name} compression level must be in #{range}, got #{level}" unless range.cover?(level)
+      level
+    end
+
     # Codec id for a codec name; Integers are taken to be ids already and returned unchecked
     #
     # @param name [Integer, Symbol, String] codec id, or a name from NAMES (case-insensitive)
@@ -129,18 +151,19 @@ module Herringbone
     #
     # @param codec [Integer] Format::Codec id
     # @param data [String] bytes to compress
+    # @param level [Integer, nil] compression level checked with check_level!, nil for the default
     # @return [String] compressed bytes (binary)
     # @raise [UnsupportedError] for a codec herringbone does not implement
     # @raise [MissingCodecError] when the codec's gem is not loaded
-    def compress(codec, data)
+    def compress(codec, data, level = nil)
       case codec
       when Format::Codec::UNCOMPRESSED then data
       when Format::Codec::SNAPPY then Codecs::Snappy.compress(data)
-      when Format::Codec::GZIP then Zlib.gzip(data)
+      when Format::Codec::GZIP then level ? Zlib.gzip(data, level: level) : Zlib.gzip(data)
       when Format::Codec::LZ4_RAW then Codecs::LZ4.compress_block(data)
       when Format::Codec::LZ4 then Codecs::LZ4.compress_hadoop(data)
-      when Format::Codec::ZSTD then library(codec).compress(data)
-      when Format::Codec::BROTLI then library(codec).deflate(data)
+      when Format::Codec::ZSTD then level ? library(codec).compress(data, level: level) : library(codec).compress(data)
+      when Format::Codec::BROTLI then level ? library(codec).deflate(data, quality: level) : library(codec).deflate(data)
       else
         raise UnsupportedError, "Unsupported compression codec #{Format::Codec::NAMES[codec] || codec}"
       end.b