#!/bin/sh

set -eu

PROXY_HOST="${PROXY_HOST:-proxy}"
PROXY_PORT="${PROXY_PORT:-3128}"
SOCKS_PORT="${SOCKS_PORT:-1080}"

TEST_HTTP_URL="${TEST_HTTP_URL:-http://api.ipify.org}"
TEST_HTTPS_URL="${TEST_HTTPS_URL:-https://api.ipify.org}"
TEST_ATTEMPTS="${TEST_ATTEMPTS:-20}"
TEST_RETRY_DELAY="${TEST_RETRY_DELAY:-1}"
TEST_TIMEOUT="${TEST_TIMEOUT:-15}"

EXPECTED_EXIT_IP="${EXPECTED_EXIT_IP:-}"
ASSERT_SAME_PROXY_EXIT="${ASSERT_SAME_PROXY_EXIT:-true}"
ASSERT_DIFFERENT_FROM_DIRECT="${ASSERT_DIFFERENT_FROM_DIRECT:-true}"

normalize_response() {
    printf '%s' "$1" | tr -d '[:space:]'
}

request() {
    label="$1"
    proxy_url="$2"
    target_url="$3"
    attempt=1

    while [ "$attempt" -le "$TEST_ATTEMPTS" ]; do
        printf 'Testing %s, attempt %s/%s...\n' \
            "$label" \
            "$attempt" \
            "$TEST_ATTEMPTS" >&2

        if [ -n "$proxy_url" ]; then
            if response="$(
                curl -q \
                    --fail \
                    --silent \
                    --show-error \
                    --location \
                    --connect-timeout 5 \
                    --max-time "$TEST_TIMEOUT" \
                    --noproxy "" \
                    --proxy "$proxy_url" \
                    "$target_url" \
                    2>/tmp/curl-error
            )"; then
                response="$(normalize_response "$response")"

                if [ -n "$response" ]; then
                    printf '%s' "$response"
                    return 0
                fi
            fi
        else
            if response="$(
                curl -q \
                    --fail \
                    --silent \
                    --show-error \
                    --location \
                    --connect-timeout 5 \
                    --max-time "$TEST_TIMEOUT" \
                    --noproxy "*" \
                    "$target_url" \
                    2>/tmp/curl-error
            )"; then
                response="$(normalize_response "$response")"

                if [ -n "$response" ]; then
                    printf '%s' "$response"
                    return 0
                fi
            fi
        fi

        attempt=$((attempt + 1))

        if [ "$attempt" -le "$TEST_ATTEMPTS" ]; then
            sleep "$TEST_RETRY_DELAY"
        fi
    done

    printf '%s failed after %s attempts.\n' \
        "$label" \
        "$TEST_ATTEMPTS" >&2

    if [ -s /tmp/curl-error ]; then
        cat /tmp/curl-error >&2
    fi

    return 1
}

assert_expected_ip() {
    label="$1"
    actual="$2"

    if [ "$actual" != "$EXPECTED_EXIT_IP" ]; then
        printf 'FAIL: %s used an unexpected exit IP.\n' "$label" >&2
        printf 'Expected: %s\n' "$EXPECTED_EXIT_IP" >&2
        printf 'Actual:   %s\n' "$actual" >&2
        exit 1
    fi
}

printf '%s\n' 'Starting 3proxy smoke tests...'

http_ip="$(
    request \
        'HTTP listener with an HTTP target' \
        "http://${PROXY_HOST}:${PROXY_PORT}" \
        "$TEST_HTTP_URL"
)"

https_ip="$(
    request \
        'HTTP listener with an HTTPS target' \
        "http://${PROXY_HOST}:${PROXY_PORT}" \
        "$TEST_HTTPS_URL"
)"

socks_ip="$(
    request \
        'SOCKS listener with remote DNS' \
        "socks5h://${PROXY_HOST}:${SOCKS_PORT}" \
        "$TEST_HTTPS_URL"
)"

printf 'HTTP target exit IP:  %s\n' "$http_ip"
printf 'HTTPS target exit IP: %s\n' "$https_ip"
printf 'SOCKS exit IP:        %s\n' "$socks_ip"

if [ -n "$EXPECTED_EXIT_IP" ]; then
    assert_expected_ip 'HTTP target' "$http_ip"
    assert_expected_ip 'HTTPS target' "$https_ip"
    assert_expected_ip 'SOCKS target' "$socks_ip"
elif [ "$ASSERT_SAME_PROXY_EXIT" = 'true' ]; then
    if [ "$http_ip" != "$https_ip" ] || [ "$http_ip" != "$socks_ip" ]; then
        printf '%s\n' \
            'FAIL: proxy listeners returned different exit IPs.' >&2
        exit 1
    fi
fi

if [ "$ASSERT_DIFFERENT_FROM_DIRECT" = 'true' ]; then
    direct_ip="$(request 'direct connection' '' "$TEST_HTTPS_URL")"

    printf 'Direct exit IP:       %s\n' "$direct_ip"

    if [ "$http_ip" = "$direct_ip" ] || \
       [ "$https_ip" = "$direct_ip" ] || \
       [ "$socks_ip" = "$direct_ip" ]; then
        printf '%s\n' \
            'FAIL: at least one proxied request used the direct exit IP.' >&2
        exit 1
    fi
fi

printf '%s\n' 'PASS: all proxy smoke tests succeeded.'