#!/usr/bin/env ruby
# frozen_string_literal: true

# Example usage:
#   MODULE_ROOT=/path/to/puppetlabs-stdlib-folder puppetserver ruby pw_hash_compare.rb
#
#   Optionally adjust THREADS (default 10) and ITERATIONS (default 20) to increase the load
#
require 'puppet'
require 'tmpdir'
require 'fileutils'

# XXX: Small hack to be able to pass shell env vars, `puppetserver ruby` wrapper
# is not reliably passing, so we accept KEY=VALUE too
ARGV.each { |a| k, v = a.split('=', 2); ENV[k] = v if v }

THREADS    = Integer(ENV.fetch('THREADS', '10'))
ITERATIONS = Integer(ENV.fetch('ITERATIONS', '20'))
PASSWORD   = ENV.fetch('PASSWORD', 'password')
HASH_TYPE  = ENV.fetch('HASH_TYPE', 'bcrypt')
# bcrypt salt validates: <cost>$<22 chars>.
SALT       = ENV.fetch('SALT', '05$salt.salt.salt.salt.sa')

module_root = ENV['MODULE_ROOT'] or abort 'set MODULE_ROOT to the stdlib version dir (containing lib/)'
fn = File.join(module_root, 'lib/puppet/parser/functions/pw_hash.rb')
abort "no pw_hash at #{fn}" unless File.exist?(fn)

Puppet.initialize_settings rescue nil # rubocop:disable Style/RescueModifier

# Isolate a single stdlib version: so the loader cannot pick a different one.
tmp = Dir.mktmpdir('pw_hash_cmp')
FileUtils.ln_s(File.expand_path(module_root), File.join(tmp, 'stdlib'))
at_exit { FileUtils.remove_entry(tmp) }

environment = Puppet::Node::Environment.create(:production, [tmp])
loaders     = Puppet::Pops::Loaders.new(environment)
Puppet.push_context(current_environment: environment, loaders: loaders)

node     = Puppet::Node.new('pw_hash-compare', environment: environment)
compiler = Puppet::Parser::Compiler.new(node)
scope    = compiler.topscope

def hash_once(scope, env, loaders, password, type, salt)
  # Re-establish the Puppet context inside each thread; the context stack is per-thread
  Puppet.override(current_environment: env, loaders: loaders) do
    scope.call_function('pw_hash', [password, type, salt])
  end
end

# Sequential reference. Single-threaded, should not be corrupted.
reference = hash_once(scope, environment, loaders, PASSWORD, HASH_TYPE, SALT)

puts "module:    #{File.basename(module_root)}"
puts "file md5:  #{Digest::MD5.file(fn).hexdigest}" if defined?(Digest)
puts "platform:  #{RUBY_PLATFORM}"
puts "call:      pw_hash(#{PASSWORD.inspect}, #{HASH_TYPE.inspect}, #{SALT.inspect})"
puts "reference: #{reference.inspect}"
puts "load:      #{THREADS} threads x #{ITERATIONS} iterations"
puts

jruby     = RUBY_PLATFORM == 'java'
mismatch  = jruby ? java.util.concurrent.atomic.AtomicInteger.new(0) : [0]
errored   = jruby ? java.util.concurrent.atomic.AtomicInteger.new(0) : [0]
sample_lock = Mutex.new
sample = nil
bump = lambda do |counter|
  jruby ? counter.incrementAndGet : sample_lock.synchronize { counter[0] += 1 }
end
read = ->(counter) { jruby ? counter.get : counter[0] }

# Don't spam per-thread stack traces
Thread.report_on_exception = false if Thread.respond_to?(:report_on_exception=)

workers = Array.new(THREADS) do
  Thread.new do
    ITERATIONS.times do
      begin
        result = hash_once(scope, environment, loaders, PASSWORD, HASH_TYPE, SALT)
        next if result == reference

        bump.call(mismatch)
        sample_lock.synchronize { sample ||= "wrong hash: #{result.inspect}" }
      rescue StandardError => e
        # The race also surfaces as an exception: a corrupted crypt(3) buffer
        # makes the enhanced-salt probe fail, raising 'system does not support
        # enhanced salts' even though the platform does support them.
        bump.call(errored)
        sample_lock.synchronize { sample ||= "#{e.class}: #{e.message}" }
      end
    end
  end
end
workers.each(&:join)

bad_mismatch = read.call(mismatch)
bad_errored  = read.call(errored)
total        = THREADS * ITERATIONS
corrupt      = bad_mismatch + bad_errored

puts "wrong-hash: #{bad_mismatch}/#{total}"
puts "errored:    #{bad_errored}/#{total}"
puts "  sample: #{sample}" if sample
puts

if corrupt.zero?
  puts "OK: #{total} concurrent calls, all correct. No corruption found."
  exit 0
else
  puts "FAIL: #{corrupt}/#{total} concurrent calls corrupted!"
  exit 1
end