#!/usr/bin/env bash
# =============================================================================
#  start-squid.sh - install (offline) and start a Squid proxy cache daemon
#                   on RHEL 9.6, ready for performance testing
# =============================================================================
#
#  USAGE (as root)
#      ./start-squid.sh            install if needed, configure, start
#      ./start-squid.sh start      same as above
#      ./start-squid.sh stop       stop the test Squid and the test web server
#      ./start-squid.sh restart    stop, then start again (same settings)
#      ./start-squid.sh status     show whether everything is running
#      ./start-squid.sh cleanup    undo every change made by this kit
#
#  WHAT IS STARTED
#      * squid-perf-test.service    the Squid proxy under test, 127.0.0.1:3401
#      * squid-perf-origin.service  a small test web server, 127.0.0.1:8009.
#                                   It plays the part of the internet web
#                                   sites the proxy fetches pages from.
#
#      The normal squid.service (port 3128) is NOT touched: the kit uses its
#      own configuration, cache, log folder and systemd service.
#
#  WHAT "start" DOES, STEP BY STEP
#      1. Checks the operating system (RHEL 9 expected).
#      2. Installs squid + httpd-tools (for the "ab" load tool) if missing,
#         WITHOUT internet:
#           a) from RPM files in ./rpms/  (see download-rpms.sh), or
#           b) from a local dnf repository already set up on the host
#              (for example the mounted RHEL 9.6 DVD / ISO).
#      3. Checks that the ports are free and that SELinux allows Squid to
#         use its port.
#      4. Creates the folders of the test Squid (configuration, cache, logs).
#      5. Writes the test configuration /etc/squid/perf-test/squid.conf.
#      6. Installs and starts the test web server (squid-perf-origin).
#      7. Writes the systemd service squid-perf-test.service.
#      8. Checks the configuration ("squid -k parse") and prepares the disk
#         cache folders ("squid -z").
#      9. Starts Squid, waits until the first page is served through it, and
#         checks that Squid really caches.
#
#  "cleanup" undoes steps 3 to 9 (the RPM packages stay installed).
# =============================================================================

set -euo pipefail

# "start" applies the values of settings.conf (and the command line);
# the other actions use the values the running test Squid was started with.
if [[ ${1:-start} == start ]]; then
    USE_RUNNING_SETTINGS=no
fi
source "$(dirname -- "${BASH_SOURCE[0]}")/lib/common.sh"

SQUID_UNIT_FILE="/etc/systemd/system/$SERVICE_NAME.service"
ORIGIN_UNIT_FILE="/etc/systemd/system/$ORIGIN_SERVICE_NAME.service"

# Remembers changes that "cleanup" must undo (for example an SELinux port).
KIT_STATE_FILE="$KIT_DIR/.kit-state"


# ----------------------------------------------------------------------------
#  Step 1 - operating system check
# ----------------------------------------------------------------------------
check_operating_system() {
    # /etc/os-release defines NAME, VERSION_ID, ...
    source /etc/os-release

    if [[ ${VERSION_ID%%.*} != 9 ]]; then
        die "This kit is made for RHEL 9. Found: $PRETTY_NAME"
    fi
    if [[ $ID != rhel || $VERSION_ID != 9.6 ]]; then
        warn "Target is RHEL 9.6; this host is '$PRETTY_NAME'. Continuing."
    fi
    log "Operating system: $PRETTY_NAME"

    command -v python3 >/dev/null || die "python3 is missing (needed for the test web server)."
    command -v curl    >/dev/null || die "curl is missing."
}


# ----------------------------------------------------------------------------
#  Step 2 - install Squid without internet access
# ----------------------------------------------------------------------------
install_squid_offline() {
    if rpm -q squid httpd-tools >/dev/null 2>&1; then
        log "Squid is already installed: $(rpm -q squid)"
        return
    fi

    local rpm_dir="$KIT_DIR/rpms"

    if compgen -G "$rpm_dir/*.rpm" >/dev/null; then
        # a) RPM files shipped next to this script. All other repositories
        #    are disabled so dnf never tries to reach the internet.
        log "Installing Squid from RPM files in $rpm_dir"
        dnf install -y --disablerepo='*' "$rpm_dir"/*.rpm
    else
        # b) A local repository already configured on this host (RHEL DVD).
        log "Installing Squid from the local dnf repositories"
        if ! dnf install -y squid httpd-tools; then
            die "Could not install Squid offline.
       Either copy RPMs into $rpm_dir (see download-rpms.sh),
       or mount the RHEL 9.6 DVD and configure it as a local repository."
        fi
    fi
    ok "Installed $(rpm -q squid) and $(rpm -q httpd-tools)"
}


# ----------------------------------------------------------------------------
#  Step 3 - port checks
# ----------------------------------------------------------------------------

# Name of the program listening on TCP port $1, or nothing if the port is free.
port_owner() {
    ss -Hltnp "sport = :$1" 2>/dev/null |
        grep -o 'users:(("[^"]*"' | head -1 | cut -d'"' -f2 || true
}

# SELinux type of TCP port $1 (for example squid_port_t), or nothing.
selinux_port_type() {
    local port="$1"
    semanage port -l 2>/dev/null |
        awk -v port="$port" '
            $2 == "tcp" {
                # The rest of the line is a list such as "3128, 3401, 4827".
                for (i = 3; i <= NF; i++) {
                    gsub(/,/, "", $i)
                    split($i, range, "-")
                    low = range[1]; high = (range[2] == "" ? range[1] : range[2])
                    if (port >= low + 0 && port <= high + 0) { print $1; exit }
                }
            }'
}

check_ports() {
    local owner

    # The proxy port: free, or used by the kit's own (running) test Squid.
    owner="$(port_owner "$SQUID_PORT")"
    if [[ -n $owner ]] && ! squid_is_running; then
        die "Port $SQUID_PORT is already used by '$owner'.
       Choose another SQUID_PORT in settings.conf (for example 4827)."
    fi

    # The origin port: free, or used by the kit's own test web server.
    owner="$(port_owner "$ORIGIN_PORT")"
    if [[ -n $owner ]] && ! origin_is_serving; then
        die "Port $ORIGIN_PORT is already used by '$owner'.
       Choose another ORIGIN_PORT in settings.conf (for example 8008)."
    fi

    log "Ports $SQUID_PORT (proxy) and $ORIGIN_PORT (test web server) are available"
}

# SELinux lets Squid listen only on ports of certain types. 3128, 3401 and
# 4827 are squid_port_t already. Another port is added to squid_port_t, and
# "cleanup" removes it again.
allow_squid_port_in_selinux() {
    if ! selinux_is_enabled || ! command -v semanage >/dev/null; then
        return
    fi

    local port_type
    port_type="$(selinux_port_type "$SQUID_PORT")"
    case "$port_type" in
        squid_port_t|http_cache_port_t)
            log "SELinux: port $SQUID_PORT is $port_type (Squid may use it)"
            ;;
        "")
            log "SELinux: adding port $SQUID_PORT to squid_port_t"
            semanage port -a -t squid_port_t -p tcp "$SQUID_PORT"
            echo "selinux_port_added=$SQUID_PORT" >> "$KIT_STATE_FILE"
            ;;
        *)
            die "SELinux: port $SQUID_PORT is already labelled $port_type.
       Choose another SQUID_PORT in settings.conf (for example 3401)."
            ;;
    esac

    # Squid may connect to web ports (http_port_t) always, and to any port
    # while the boolean squid_connect_any is on (the RHEL default).
    if [[ $(getsebool squid_connect_any 2>/dev/null | awk '{print $3}') == off &&
          $(selinux_port_type "$ORIGIN_PORT") != http_port_t ]]; then
        warn "SELinux: squid_connect_any is off and port $ORIGIN_PORT is not a web port;"
        warn "Squid will not reach the test web server. Use ORIGIN_PORT=8009."
    fi
}


# ----------------------------------------------------------------------------
#  Step 4 - folders of the test Squid
# ----------------------------------------------------------------------------
create_folders() {
    log "Creating $CONF_DIR, $CACHE_DIR and $LOG_DIR"
    mkdir -p "$CONF_DIR" "$CACHE_DIR" "$LOG_DIR"

    # Squid works as the user "squid"; it must own its cache and log folders.
    chown squid:squid "$CACHE_DIR" "$LOG_DIR"
    chmod 750 "$CACHE_DIR" "$LOG_DIR"

    # Give the folders the standard Squid SELinux labels.
    if command -v restorecon >/dev/null; then
        restorecon -R "$CONF_DIR" "$CACHE_DIR" "$LOG_DIR"
    fi
}


# ----------------------------------------------------------------------------
#  Step 5 - Squid configuration for the test
# ----------------------------------------------------------------------------
write_squid_config() {
    local access_log_line="access_log none"
    if [[ $TEST_ACCESS_LOG == on ]]; then
        access_log_line="access_log daemon:$LOG_DIR/access.log squid"
    fi

    log "Writing $SQUID_CONF"
    cat > "$SQUID_CONF" <<EOF
# Created by the Squid performance kit ($KIT_DIR/start-squid.sh).
# Values come from settings.conf. "./start-squid.sh cleanup" removes this file.
# This is a separate Squid instance; the normal /etc/squid/squid.conf is not used.

# --- Where Squid listens ------------------------------------------------------
# Only on this machine, so the test proxy cannot be used from the network.
http_port $SQUID_HOST:$SQUID_PORT

# Name shown in error pages and in the "X-Cache: HIT from ..." header.
visible_hostname $SERVICE_NAME

# --- Files of this instance -----------------------------------------------------
pid_filename $PID_FILE
cache_log    $LOG_DIR/cache.log
coredump_dir $CACHE_DIR
cache_store_log none

# One line per request costs speed; "TEST_ACCESS_LOG=on" in settings.conf
# switches it on, as in RHEL's default squid.conf.
$access_log_line

# --- Who may use the proxy ------------------------------------------------------
# "localhost" = 127.0.0.1 and ::1. The cache manager pages (live statistics
# at http://$SQUID_HOST:$SQUID_PORT/squid-internal-mgr/...) are for this machine only.
http_access allow localhost manager
http_access deny  manager
http_access allow localhost
http_access deny  all

# --- Memory cache ---------------------------------------------------------------
# RAM for cached objects. Squid needs more RAM than this in total (indexes,
# connections, buffers).
cache_mem $CACHE_MEM_MB MB
# Objects up to this size are kept in RAM; larger ones only on disk.
maximum_object_size_in_memory $MAX_OBJECT_IN_MEMORY_KB KB

# --- Disk cache -----------------------------------------------------------------
# Type, folder, size in MB, first-level and second-level sub-folders.
cache_dir $CACHE_DIR_TYPE $CACHE_DIR $CACHE_DIR_MB 16 256
# Largest object Squid caches at all.
maximum_object_size $MAX_OBJECT_SIZE_MB MB

# --- How long objects stay fresh --------------------------------------------------
# Answers with their own caching headers (Cache-Control, Expires) follow those.
# The rest: as in RHEL's default squid.conf.
refresh_pattern ^ftp:             1440  20%  10080
refresh_pattern -i (/cgi-bin/|\?)    0   0%      0
refresh_pattern .                    0  20%   4320

# --- Limits -----------------------------------------------------------------------
# Every client connection and every connection to a web server is an open file.
max_filedescriptors $OPEN_FILES_LIMIT

# On stop, wait at most this long for open requests (Squid default 30 s).
shutdown_lifetime $SHUTDOWN_LIFETIME_SECONDS seconds
EOF

    # Remember the values the test Squid runs with (see lib/common.sh).
    cat > "$RUNNING_SETTINGS_FILE" <<EOF
# Written by start-squid.sh: the settings the test Squid was started with.
# test-squid.sh and "start-squid.sh stop|status|cleanup" read them back.
SQUID_PORT="$SQUID_PORT"
ORIGIN_PORT="$ORIGIN_PORT"
ORIGIN_PROCESSES="$ORIGIN_PROCESSES"
CACHE_MEM_MB="$CACHE_MEM_MB"
MAX_OBJECT_IN_MEMORY_KB="$MAX_OBJECT_IN_MEMORY_KB"
MAX_OBJECT_SIZE_MB="$MAX_OBJECT_SIZE_MB"
CACHE_DIR_TYPE="$CACHE_DIR_TYPE"
CACHE_DIR_MB="$CACHE_DIR_MB"
OPEN_FILES_LIMIT="$OPEN_FILES_LIMIT"
TEST_ACCESS_LOG="$TEST_ACCESS_LOG"
EOF

    if command -v restorecon >/dev/null; then
        restorecon "$SQUID_CONF" "$RUNNING_SETTINGS_FILE"
    fi
}


# ----------------------------------------------------------------------------
#  Step 6 - the test web server ("origin server")
# ----------------------------------------------------------------------------
#  The script is copied out of the kit folder, because the kit may sit in a
#  home folder that the unprivileged user "nobody" cannot read.
install_origin_server() {
    log "Installing the test web server in $ORIGIN_INSTALL_DIR"
    mkdir -p "$ORIGIN_INSTALL_DIR"
    install -m 755 "$KIT_DIR/lib/origin-server.py" "$ORIGIN_SCRIPT"
    command -v restorecon >/dev/null && restorecon -R "$ORIGIN_INSTALL_DIR"

    if has_systemd; then
        log "Writing $ORIGIN_UNIT_FILE"
        cat > "$ORIGIN_UNIT_FILE" <<EOF
# Created by $KIT_DIR/start-squid.sh - "./start-squid.sh cleanup" removes it.
[Unit]
Description=Test web server of the Squid performance kit ($ORIGIN)
Documentation=file://$KIT_DIR/README.md

[Service]
Type=simple
User=nobody
Group=nobody
ExecStart=/usr/bin/python3 $ORIGIN_SCRIPT --address $ORIGIN_HOST --port $ORIGIN_PORT --processes $(origin_process_count)
LimitNOFILE=$OPEN_FILES_LIMIT
# It needs no privileges and no files.
NoNewPrivileges=yes
ProtectSystem=strict
ProtectHome=yes
PrivateTmp=yes
EOF
        systemctl daemon-reload
    fi

    # (Re)start it, so a changed port or process count is used.
    origin_stop
    origin_start

    local waited=0
    while ! origin_is_serving && (( waited < 100 )); do
        sleep 0.1
        waited=$(( waited + 1 ))
    done
    origin_is_serving || die "The test web server did not answer on $HEALTH_URL.
       Look at:  journalctl -u $ORIGIN_SERVICE_NAME   (or $ORIGIN_LOG without systemd)"
    log "Test web server is answering on $ORIGIN_URL ($(origin_process_count) processes)"
}


# ----------------------------------------------------------------------------
#  Step 7 - systemd service of the test Squid
# ----------------------------------------------------------------------------
#  A copy of RHEL's squid.service, pointed at the test configuration.
write_squid_unit() {
    if ! has_systemd; then
        return      # without systemd, squid_start runs Squid directly
    fi
    log "Writing $SQUID_UNIT_FILE"
    cat > "$SQUID_UNIT_FILE" <<EOF
# Created by $KIT_DIR/start-squid.sh - "./start-squid.sh cleanup" removes it.
[Unit]
Description=Squid proxy of the performance test kit ($PROXY)
Documentation=file://$KIT_DIR/README.md man:squid(8)
After=network.target $ORIGIN_SERVICE_NAME.service

[Service]
Type=notify
NotifyAccess=all
PIDFile=$PID_FILE
LimitNOFILE=$OPEN_FILES_LIMIT
ExecStart=$SQUID_BIN --foreground -f $SQUID_CONF
ExecReload=/usr/bin/kill -HUP \$MAINPID
KillMode=mixed
EOF
    systemctl daemon-reload
}


# ----------------------------------------------------------------------------
#  Step 8 - configuration check and disk cache folders
# ----------------------------------------------------------------------------
check_config() {
    log "Checking the Squid configuration (squid -k parse)"
    local output
    if ! output="$("$SQUID_BIN" -k parse -f "$SQUID_CONF" 2>&1)"; then
        echo "$output" | grep -iE 'error|fatal' >&2 || echo "$output" >&2
        die "The Squid configuration has errors (see messages above)."
    fi
    # Show warnings, but do not stop for them.
    echo "$output" | grep -iE 'WARNING|ERROR' | sed 's/^/    /' || true
}

# "squid -z" creates the 16 x 256 sub-folders of the disk cache. It is only
# needed once; RHEL's own squid.service does the same (cache_swap.sh).
prepare_disk_cache() {
    if [[ -d $CACHE_DIR/00 ]]; then
        return
    fi
    log "Creating the disk cache folders in $CACHE_DIR (squid -z)"
    "$SQUID_BIN" --foreground -z -f "$SQUID_CONF" >> "$LOG_DIR/squid-z.out" 2>&1 ||
        die "squid -z failed. See $LOG_DIR/squid-z.out and $CACHE_LOG"
    command -v restorecon >/dev/null && restorecon -R "$CACHE_DIR"
    return 0
}


# ----------------------------------------------------------------------------
#  Step 9 - start the daemon and check that it caches
# ----------------------------------------------------------------------------
start_daemon() {
    if squid_is_running; then
        log "The test Squid is running - restarting it to apply the configuration"
        squid_stop
    fi
    prepare_disk_cache

    log "Starting the test Squid on $PROXY"
    squid_start
    if ! wait_until_serving 60; then
        die "No page came through Squid within 60 s.
       Look at:  $CACHE_LOG   and   journalctl -u $SERVICE_NAME"
    fi
}

# Fetch one cacheable page twice: the second answer must come from the cache.
check_caching() {
    local url="$ORIGIN_URL/cacheable/1k/start-check-$(date +%s)"
    local first second
    first="$(cache_status_of "$url")"
    second="$(cache_status_of "$url")"
    if [[ $first == MISS && $second == HIT ]]; then
        log "Caching works: first request $first, second request $second"
    else
        warn "Caching check: first request '$first', second '$second' (expected MISS, then HIT)."
        warn "The cache tests will fail. Look at $CACHE_LOG"
    fi
}

print_summary() {
    local selinux
    selinux="$(getenforce 2>/dev/null || echo 'not available')"

    echo
    ok "The test Squid is running and caching."
    echo "    Version        : Squid $(squid_version)"
    echo "    Proxy address  : $PROXY   (for example: curl -x http://$PROXY $ORIGIN_URL/health)"
    echo "    Master PID     : $(squid_master_pid)"
    echo "    Worker PID     : $(squid_worker_pids | tr '\n' ' ')"
    echo "    Memory cache   : $CACHE_MEM_MB MB (objects up to $MAX_OBJECT_IN_MEMORY_KB KB)"
    echo "    Disk cache     : $CACHE_DIR_TYPE, $CACHE_DIR_MB MB in $CACHE_DIR"
    echo "    Statistics     : $MANAGER_URL/info"
    echo "    Log            : $CACHE_LOG"
    echo "    Test web server: $ORIGIN_URL"
    echo "    SELinux        : $selinux"
    if has_systemd; then
        echo "    Managed by     : systemd  (systemctl status $SERVICE_NAME $ORIGIN_SERVICE_NAME)"
    else
        echo "    Managed by     : squid -k (no systemd found, e.g. inside a container)"
    fi
    echo
    echo "    Next step      : ./test-squid.sh"
}


# ----------------------------------------------------------------------------
#  Other actions
# ----------------------------------------------------------------------------
show_status() {
    local healthy=yes

    if origin_is_serving; then
        ok "Test web server is answering on $ORIGIN_URL"
    else
        warn "Test web server is NOT answering on $ORIGIN_URL"
        healthy=no
    fi

    if squid_is_serving; then
        ok "Test Squid is running (master PID $(squid_master_pid)) and serving through $PROXY"
        # Selected lines of Squid's own statistics; never fail on them.
        squid_manager_page info |
            grep -E 'Squid Object Cache|UP Time|HTTP requests received|Hits as % of all|Storage Mem size|Storage Swap size|file desc currently in use|Maximum Resident' |
            sed 's/^[[:space:]]*/    /' || true
    elif squid_is_running; then
        warn "Test Squid processes exist, but no page comes through $PROXY. See $CACHE_LOG"
        healthy=no
    else
        warn "Test Squid is not running."
        healthy=no
    fi

    [[ $healthy == yes ]] || exit 1
}

stop_everything() {
    squid_stop
    origin_stop
    ok "Test Squid and test web server stopped."
}

cleanup_test_setup() {
    log "Stopping the test Squid and the test web server"
    squid_stop || true
    origin_stop

    local unit_file
    for unit_file in "$SQUID_UNIT_FILE" "$ORIGIN_UNIT_FILE"; do
        if [[ -f $unit_file ]]; then
            log "Removing $unit_file"
            rm -f "$unit_file"
        fi
    done
    has_systemd && systemctl daemon-reload

    log "Removing $CONF_DIR, $CACHE_DIR, $LOG_DIR and $ORIGIN_INSTALL_DIR"
    rm -rf "$CONF_DIR" "$CACHE_DIR" "$LOG_DIR" "$ORIGIN_INSTALL_DIR"
    rm -f "$PID_FILE" "$ORIGIN_PID_FILE"

    if [[ -f $KIT_STATE_FILE ]]; then
        local port
        port="$(awk -F= '$1 == "selinux_port_added" {print $2}' "$KIT_STATE_FILE" | tail -1)"
        if [[ -n $port ]]; then
            log "SELinux: removing port $port from squid_port_t"
            semanage port -d -t squid_port_t -p tcp "$port" || true
        fi
        rm -f "$KIT_STATE_FILE"
    fi

    ok "Cleanup finished. (The squid and httpd-tools packages stay installed;"
    echo "    remove them with 'dnf remove squid httpd-tools' if no longer needed.)"
}


# ----------------------------------------------------------------------------
#  Main
# ----------------------------------------------------------------------------
main() {
    local action="${1:-start}"
    require_root "$@"

    case "$action" in
        start)
            check_operating_system
            install_squid_offline
            check_ports
            allow_squid_port_in_selinux
            create_folders
            write_squid_config
            install_origin_server
            write_squid_unit
            check_config
            start_daemon
            check_caching
            print_summary
            ;;
        stop)
            stop_everything
            ;;
        restart)
            origin_is_serving || origin_start
            squid_restart
            wait_until_serving 60 || die "The test Squid did not come back. See $CACHE_LOG"
            ok "Test Squid restarted."
            ;;
        status)
            show_status
            ;;
        cleanup)
            cleanup_test_setup
            ;;
        *)
            echo "Usage: $0 [start|stop|restart|status|cleanup]"
            exit 2
            ;;
    esac
}

main "$@"