openresty/util/resty

#!/usr/bin/env perl

# Copyright (C) Yichun Zhang (agentzh)

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

use strict;
use warnings;

use File::Spec ();
use FindBin ();
use Getopt::Std qw( getopts );
use File::Temp qw( tempdir );
use POSIX qw( WNOHANG );

my %opts;
getopts("he:c:VI:", \%opts)
    or usage(1);

if ($opts{h}) {
    usage(0);
}

my $nginx_path = File::Spec->catfile($FindBin::Bin, "..", "nginx", "sbin", "nginx");
#warn $nginx_path;
if (!-f $nginx_path) {
    $nginx_path = "nginx";  # find in PATH
}

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

my $lua_package_path_config = '';
if ($opts{I}) {
    my $dir = $opts{I};
    if (!-d $dir) {
        die "Search directory $dir is not found.\n";
    }
    $lua_package_path_config = <<_EOC_;
    lua_package_path "$dir/?.lua;;";
    lua_package_cpath "$dir/?.so;;";
_EOC_
}

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

if (@ARGV) {
    die "unknown arguments: @ARGV\n";
}

my $conns = $opts{c} || 64;

my @nameservers;

# 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+|$)/) {
            push @nameservers, $1;
            if (@nameservers > 10) {
                last;
            }
        }
    }
    close $in;
}

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

#warn "@nameservers\n";

my $prefix_dir = tempdir(CLEANUP => 1);
#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 $chunk_name;
if (defined $opts{e}) {
    $luafile = File::Spec->catfile($conf_dir, "a.lua");
    open my $out, ">$luafile"
        or die "Cannot open $luafile for writing: $!\n";
    print $out $opts{e};
    close $out;
    $chunk_name = "=(command line -e)";

} else {
    $chunk_name = "\@$luafile";
}

my $loader = <<_EOC_;
            local gen
            do
                local fname = "$luafile"
                local f = assert(io.open(fname, "r"))
                local chunk = f:read("*a")
                gen = assert(loadstring(chunk, "$chunk_name"))
            end
_EOC_

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

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;

$env_list

error_log /dev/stderr warn;
#error_log /dev/stderr debug;

events {
    worker_connections $conns;
}

http {
    access_log off;
    lua_socket_log_errors off;
    resolver @nameservers;
$lua_package_path_config
    init_by_lua '
        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 '
        local exit = os.exit
        local stderr = io.stderr

        local function handle_err(err)
            if err then
                err = string.gsub(err, "^init_worker_by_lua:%d+: ", "")
                stderr:write(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 < 9011
            then
                error("at least ngx_lua 0.9.12 is required")
            end

$loader
            -- print("calling timer.at...")
            local ok, err = ngx.timer.at(0, function ()
                -- io.stderr:write("timer firing")
                local ok, err = pcall(gen)
                if not ok then
                    return handle_err(err)
                end
                local rc = err
                if rc and type(rc) ~= "number" then
                    return handle_err("bad return value of type " .. type(rc))
                end
                return exit(rc)
            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/";

my $child_pid;

sub sigint {
    $SIG{INT} = \&sigint;
    if ($child_pid) {
        kill INT => $child_pid;
    }
}
$SIG{INT} = \&sigint;

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 (defined $?) {
        $rc = ($? >> 8);
    }
    exit($rc);
}

sub usage {
    my $rc = shift;
    my $msg = <<_EOC_;
resty [-I dir] [-h] [-c num] [-e prog] [-V] [lua-file]

Options:
    -c num      set maximal connection count (default: 64).
    -e prog     run the inlined Lua code in "prog".
    -h          print this help.
    -I dir      Add dir to the search paths for Lua libraries.
    -V          print the underlying nginx version and configurations.

For bug reporting instructions, please see:
<http://openresty.org/#Community>
_EOC_
    if ($rc == 0) {
        print $msg;
        exit(0);
    }

    warn $msg;
    exit($rc);
}