#!/usr/bin/env perl

# Copyright (C) Yichun Zhang (agentzh)
# Copyright (C) Guanlan Dai

# TODO: port this script into the nginx core for greater flexibility
# and better performance.

use strict;
use warnings;

our $VERSION = '0.19';

use File::Spec ();
use FindBin ();
use List::Util qw( max );
use Getopt::Long qw( GetOptions :config no_ignore_case require_order);
use File::Temp qw( tempdir );
use POSIX qw( WNOHANG );
use Time::HiRes qw( sleep );

sub resolve_includes {
    my $type = shift;
    my @paths = map {
        my $abs_path = File::Spec->rel2abs($_);
        unless (-f $abs_path) {
          die "Could not find $type include '$abs_path'";
        }
        "include $abs_path;";
    } @_;
    return join("\n", @paths);
}

sub format_directives {
    my $name = shift;
    my @directives = map {
        "$name $_;";
    } @_;
    return join("\n", @directives);
}

sub check_option_format {
    my ($opt_name, $regex, $opt_format, $args_ref) = @_;
    for (@$args_ref) {
        unless (m/$regex/) {
            die "Invalid value for $opt_name option. Expected: $opt_format\n";
        }
    }
}

my @all_args = @ARGV;

my (@http_confs, @http_includes, @main_includes, @shdicts, @ns, @src_a);

my $errlog_level = "warn";

GetOptions("c=i",             \(my $conns_num),
           "e=s",             \@src_a,
           "gdb",             \(my $use_gdb),
           "h|help",          \(my $help),
           "errlog-level=s",  \$errlog_level,
           "http-conf=s",     \@http_confs,
           "http-include=s",  \@http_includes,
           "shdict=s",        \@shdicts,
           "I=s@",            \(my $Inc),
           "main-include=s",  \@main_includes,
           "nginx=s",         \(my $nginx_path),
           "valgrind",        \(my $use_valgrind),
           "valgrind-opts=s", \(my $valgrind_opts),
           "ns=s",            \@ns,
           "resolve-ipv6",    \(my $resolve_ipv6),
           "rr",              \(my $use_rr),
           "V|v",             \(my $version))
   or usage(1);

check_option_format('--shdict', qr/^ [_a-z]+ \s+ \d+ (?i) [km]? $/x, 'NAME SIZE', \@shdicts);
check_option_format('--ns',     qr/^ (?: \d{1,3} (?: \. \d{1,3} ){3} | \[ [0-9a-fA-F:]+ \] ) $/x, 'IP', \@ns);

my $src;
if (@src_a) {
    $src = join('; ', @src_a);
}

if ($help) {
    usage(0);
}

if (!$nginx_path) {
    use Config;
    my $ext = $Config{_exe};
    if (!$ext) {
        if ($^O eq 'msys') {
            $ext = '.exe';
        } else {
            $ext = '';
        }
    }
    $nginx_path = File::Spec->catfile($FindBin::RealBin, "..", "nginx", "sbin", "nginx$ext");
    if (!-f $nginx_path) {
        $nginx_path = File::Spec->catfile($FindBin::RealBin, "nginx$ext");
        if (!-f $nginx_path) {
            $nginx_path = "nginx";  # find in PATH
        }
    }
}

#warn $nginx_path;

if ($version) {
    warn "resty $VERSION\n";
    my $cmd = "$nginx_path -V";
    exec $cmd or die "Failed to run command \"$cmd\": $!\n";
}

my $lua_package_path_config = '';
if (defined $Inc) {
    my $package_path = "";
    my $package_cpath = "";
    for my $dir (@$Inc) {
        if (!-d $dir) {
            die "Search directory $dir is not found.\n";
        }
        $package_path .= File::Spec->catfile($dir, "?.lua;");
        $package_cpath .= File::Spec->catfile($dir, "?.so;");
    }
    $lua_package_path_config = <<_EOC_;
    lua_package_path "$package_path;";
    lua_package_cpath "$package_cpath;";
_EOC_
}

my $luafile = shift;
if (!defined $src and !defined $luafile) {
    die qq{Neither Lua input file nor -e "" option specified.\n};
}

my $conns = $conns_num || 64;

my @nameservers;

if (@ns > 0) {
    unshift @nameservers, @ns;

} else {
    # try to read the nameservers used by the system resolver:
    if (open my $in, "/etc/resolv.conf") {
        while (<$in>) {
            if (/^\s*nameserver\s+(\d+(?:\.\d+){3})(?:\s+|$)/) {
                unshift @nameservers, $1;
                if (@nameservers > 10) {
                    last;
                }
            }
        }
        close $in;
    }
}

if (!@nameservers) {
    # default to Google's open DNS servers
    unshift @nameservers, "8.8.8.8", "8.8.4.4";
}

if (!$resolve_ipv6) {
    push @nameservers, "ipv6=off";
}

#warn "@nameservers\n";

my $prefix_dir;
if ($^O eq 'msys') {
    # to work around a bug in msys perl (at least 5.8.8 msys 64int)
    $prefix_dir = "resty_cli_temp";
    if (-d $prefix_dir) {
        system("rm -rf $prefix_dir") == 0 or die $!;
    }
    mkdir $prefix_dir or die "failed to mkdir $prefix_dir: $!";

} else {
    $prefix_dir = tempdir(CLEANUP => 1);
    if ($^O eq 'MSWin32') {
        require Win32;
        $prefix_dir = Win32::GetLongPathName($prefix_dir);
    }
}
#warn "prefix dir: $prefix_dir\n";

my $logs_dir = File::Spec->catfile($prefix_dir, "logs");
mkdir $logs_dir or die "failed to mkdir $logs_dir: $!";

my $conf_dir = File::Spec->catfile($prefix_dir, "conf");
mkdir $conf_dir or die "failed to mkdir $conf_dir: $!";

my $inline_lua = '';
my $quoted_luafile;
if (defined $src) {
    my $file = File::Spec->catfile($conf_dir, "a.lua");
    open my $out, ">$file"
        or die "Cannot open $file for writing: $!\n";
    print $out $src;
    close $out;
    my $chunk_name = "=(command line -e)";
    $quoted_luafile = quote_as_lua_str($file);

    $inline_lua = <<"_EOC_";
                local fname = $quoted_luafile
                local f = assert(io.open(fname, "r"))
                local chunk = f:read("*a")
                local inline_gen = assert(loadstring(chunk, "$chunk_name"))
_EOC_
}

my $file_lua = '';
if (defined $luafile) {
    if (!-f $luafile) {
        die "Lua input file $luafile not found.\n";
    }

    my $chunk_name = quote_as_lua_str("\@$luafile");
    $quoted_luafile = quote_as_lua_str($luafile);
    $file_lua = <<"_EOC_";
                local fname = $quoted_luafile
                local f = assert(io.open(fname, "r"))
                local chunk = f:read("*a")
                local file_gen = assert(loadstring(chunk, $chunk_name))
_EOC_
}

my @user_args =  @ARGV;
my $args = gen_lua_code_for_args(\@user_args, \@all_args);

my $loader = <<_EOC_;
            local gen
            do
                $args
$inline_lua
$file_lua

                gen = function()
                  if inline_gen then inline_gen() end
                  if file_gen then file_gen() end
                end
            end
_EOC_

my $env_list = '';
for my $var (sort keys %ENV) {
    #warn $var;
    $env_list .= "env $var;\n";
}

my $main_include_directives = resolve_includes('main', @main_includes);
my $http_include_directives = resolve_includes('http', @http_includes);
my $lua_shared_dicts = format_directives('lua_shared_dict', @shdicts);

my $http_conf_lines = join "", map { "    $_\n" } @http_confs;

my $conf_file = File::Spec->catfile($conf_dir, "nginx.conf");
open my $out, ">$conf_file"
    or die "Cannot open $conf_file for writing: $!\n";

print $out <<_EOC_;
daemon off;
master_process off;
worker_processes 1;
pid logs/nginx.pid;

$env_list

error_log stderr $errlog_level;
#error_log stderr debug;

events {
    worker_connections $conns;
}

$main_include_directives

http {
    access_log off;
    lua_socket_log_errors off;
    resolver @nameservers;
    lua_regex_cache_max_entries 40960;
    $lua_shared_dicts
$lua_package_path_config
$http_conf_lines
    $http_include_directives
    init_by_lua_block {
        local stdout = io.stdout
        local ngx_null = ngx.null
        local maxn = table.maxn
        local unpack = unpack
        local concat = table.concat

        local expand_table
        function expand_table(src, inplace)
            local n = maxn(src)
            local dst = inplace and src or {}
            for i = 1, n do
                local arg = src[i]
                local typ = type(arg)
                if arg == nil then
                    dst[i] = "nil"

                elseif typ == "boolean" then
                    if arg then
                        dst[i] = "true"
                    else
                        dst[i] = "false"
                    end

                elseif arg == ngx_null then
                    dst[i] = "null"

                elseif typ == "table" then
                    dst[i] = expand_table(arg, false)

                elseif typ ~= "string" then
                    dst[i] = tostring(arg)

                else
                    dst[i] = arg
                end
            end
            return concat(dst)
        end

        local function output(...)
            local args = {...}

            return stdout:write(expand_table(args, true))
        end

        ngx.print = output
        ngx.say = function (...)
                local ok, err = output(...)
                if ok then
                    return output("\\n")
                end
                return ok, err
            end
        print = ngx.say

        ngx.flush = function (...) return stdout:flush() end
        -- we cannot close stdout here due to a bug in Lua:
        ngx.eof = function (...) return true end
        ngx.exit = os.exit
    }

    init_worker_by_lua_block {
        local exit = os.exit
        local stderr = io.stderr
        local ffi = require "ffi"

        local function handle_err(err)
            if err then
                err = string.gsub(err, "^init_worker_by_lua:%d+: ", "")
                stderr:write("ERROR: ", err, "\\n")
            end
            return exit(1)
        end

        local ok, err = pcall(function ()
            if not ngx.config
               or not ngx.config.ngx_lua_version
               or ngx.config.ngx_lua_version < 10009
            then
                error("at least ngx_lua 0.10.9 is required")
            end

            local signal_graceful_exit = require("ngx.process").signal_graceful_exit
            if not signal_graceful_exit then
                error("lua-resty-core library is too old; "
                      .. "missing the signal_graceful_exit() function in ngx.process")
            end

$loader
            -- print("calling timer.at...")
            local ok, err = ngx.timer.at(0, function ()
                -- io.stderr:write("timer firing")
                local ok, err = xpcall(gen, function (err)
                    -- level 3: we skip this function and the
                    -- error() call itself in our stacktrace
                    local trace = debug.traceback(err, 3)
                    return handle_err(trace)
                end)
                if not ok then
                    return handle_err(err)
                end
                if ffi.abi("win") then
                    return exit(0)
                end
                signal_graceful_exit()
            end)
            if not ok then
                return handle_err(err)
            end
            -- print("timer created")
        end)

        if not ok then
            return handle_err(err)
        end
    }
}
_EOC_

close $out;

my @cmd = ($nginx_path, '-p', "$prefix_dir/", '-c', "conf/nginx.conf");

if ($use_gdb) {
    if ($use_valgrind) {
        die "ERROR: options --gdb and --valgrind cannot be specified at the ",
            "same time.\n";
    }

    if ($use_rr) {
        die "ERROR: options --gdb and --rr cannot be specified at the ",
            "same time.\n";
    }

    unshift @cmd, "gdb", "--args",

} elsif ($use_rr) {
    if ($use_valgrind) {
        die "ERROR: options --rr and --valgrind cannot be specified at the ",
            "same time.\n";
    }

    unshift @cmd, "rr", "record",

} elsif ($use_valgrind) {
    my @new = ('valgrind');
    if ($valgrind_opts) {
        $valgrind_opts =~ s/^\s+|\s+$//g;
        push @new, split /\s+/, $valgrind_opts;
    }

    unshift @cmd, @new;
}

my $child_pid;

sub forward_signal {
    my $signame = shift;

    if ($child_pid) {
        #warn "killing $child_pid\n";

        kill $signame => $child_pid
            or die "failed to send $signame to process $child_pid: $!";

        my $nap = 0.001;
        my $total_nap = 0;
        while ($total_nap < 0.1) {
            #warn "sleeping $nap";
            $nap *= 1.5;
            $total_nap += $nap;
            sleep $nap;

            my $ret = waitpid $child_pid, WNOHANG;
            next if $ret == 0;
            last if $ret < 0;

            my $rc = 0;
            if ($signame eq 'INT') {
                $rc = 130;

            } elsif ($signame eq 'TERM') {
                $rc = 143;
            }

            if ($?) {
                $rc = ($? >> 8);
                if ($rc == 0) {
                    $rc = $?;
                }
            }
            exit $rc;
        }
    }
}

for my $sig (qw/ INT TERM QUIT HUP USR1 USR2 WINCH /) {
    $SIG{$sig} = \&forward_signal;
}

my $pid = fork();

if (!defined $pid) {
    die "fork() failed: $!\n";
}

if ($pid == 0) {  # child process
    #warn "exec @cmd...";
    exec(@cmd)
        or die "Failed to run command \"@cmd\": $!\n";

} else {
    $child_pid = $pid;
    waitpid($child_pid, 0);
    my $rc = 0;
    if ($?) {
        $rc = ($? >> 8);
        if ($rc == 0) {
            $rc = $?;
        }
    }
    exit $rc;
}

sub usage {
    my $rc = shift;
    my $msg = <<_EOC_;
resty [options] [lua-file [args]]

Options:
    -c NUM              Set maximal connection count (default: 64).
    -e PROG             Run the inlined Lua code in "prog".
    --gdb               Use GDB to run the underlying C process.
    --help              Print this help.

    --errlog-level LEVEL
                        Set nginx error_log level.
                        Can be debug, info, notice, warn, error, crit, alert, or emerg.

    --http-conf CONF    Specifies nginx.conf snippet inserted into the http {}
                        configuration block (multiple instances are supported).

    --http-include PATH Include the specified file in the nginx http configuration block
                        (multiple instances are supported).

    -I DIR              Add dir to the search paths for Lua libraries.

    --main-include PATH Include the specified file in the nginx main configuration block
                        (multiple instances are supported).

    --ns IP             Specify a custom name server (multiple instances are supported).

    --resolve-ipv6      Make the nginx resolver lookup both IPv4 and IPv6 addresses.

    --rr                Use Mozilla rr to record the execution of the underlying C process.

    --shdict 'NAME SIZE'
                        Create the specified lua shared dicts in the http
                        configuration block (multiple instances are supported).

    --nginx             Specify the nginx path (this option might be removed in the future).
    -V                  Print version numbers and nginx configurations.
    --valgrind          Use valgrind to run nginx.
    --valgrind-opts     Pass extra options to valgrind.

For bug reporting instructions, please see:

    <https://openresty.org/en/community.html>

Copyright (C) Yichun Zhang (agentzh). All rights reserved.
_EOC_
    if ($rc == 0) {
        print $msg;
        exit(0);
    }

    warn $msg;
    exit($rc);
}

sub get_bracket_level {
    my %bracket_levels;
    my $bracket_level = 0;
    my $max_level = 0;

    # scan all args and store level of closing brackets
    for my $arg (@_) {
        while ($arg =~ /\](=*)\]/g) {
            my $level = length($1);
            if ($level > $max_level) {
                $max_level = $level;
            }
            $bracket_levels{$level} = 1;
        }
    }

    # if args contain closing bracket
    if (%bracket_levels) {
        # find the shortest form of the long brackets accordingly
        for (my $i = 1; $i < $max_level; $i++) {
            if (!exists $bracket_levels{$i}) {
                $bracket_level = $i;
                last;
            }
        }

        if ($bracket_level == 0) {
            $bracket_level = $max_level + 1;
        }
        return $bracket_level;
    }
    return 1;
}

sub quote_as_lua_str {
    my ($str) = @_;
    my $bracket_level = get_bracket_level($str);
    my $left_bracket = "[" . "=" x $bracket_level . "[";
    my $right_bracket = "]" . "=" x $bracket_level . "]";

    return $left_bracket . $str . $right_bracket;
}

sub gen_lua_code_for_args {
    my ($user_args, $all_args) = @_;

    my $luasrc = "arg = {}\n";

    # args[n] (n = 0)
    $luasrc .= "arg[0] = $quoted_luafile\n";

    # args[n] (n > 0)
    for my $i (0 .. $#user_args) {
        my $index = $i + 1;
        my $quoted_arg = quote_as_lua_str($user_args[$i]);
        $luasrc .= "arg[$index] = $quoted_arg\n";
    }

    # args[n] (n < 0)
    my $left_num = $#all_args - $#user_args;
    for my $i (0 .. $left_num - 2) {
        my $index = 0 - $left_num + $i + 1;
        my $quoted_arg = quote_as_lua_str($all_args[$i]);
        $luasrc .= "arg[$index] = $quoted_arg\n";
    }

    # args[n] (n = the index of resty-cli itself)
    my $index = 0 - $left_num;
    my $quoted_arg = quote_as_lua_str($0);
    $luasrc .= "arg[$index] = $quoted_arg\n";

    #warn $luasrc;
    return $luasrc;
}
