#!/usr/bin/env bash
set -euo pipefail

usage() {
  cat <<'EOF'
Usage: Scripts/burrow-shadowsocks-singmux-interop

Run controlled Burrow Shadowsocks top-level sing-mux interop tests against the
same github.com/metacubex/sing-mux Go module used by Mihomo. By default this
checks smux, yamux, and h2mux with and without Mihomo's padded sing-mux
preface/framing. The test does not change system proxy settings or use the
user's Burrow database.
EOF
}

if [[ "${1:-}" == "-h" || "${1:-}" == "--help" ]]; then
  usage
  exit 0
fi

repo_root="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)"
burrow_bin="${BURROW_BIN:-$repo_root/target/debug/burrow}"
sing_mux_dir="${SING_MUX_DIR:-/Users/jettchen/go/pkg/mod/github.com/metacubex/sing-mux@v0.3.9}"
sing_dir="${SING_DIR:-/Users/jettchen/go/pkg/mod/github.com/metacubex/sing@v0.5.7}"

if [[ -z "${BURROW_SINGMUX_INTEROP_SINGLE:-}" ]]; then
  for protocol in ${BURROW_SINGMUX_PROTOCOLS:-smux yamux h2mux}; do
    for padding in ${BURROW_SINGMUX_PADDING_MODES:-false true}; do
      echo "=== Mihomo sing-mux interop: $protocol padding=$padding ==="
      BURROW_SINGMUX_INTEROP_SINGLE=1 \
        BURROW_SINGMUX_PROTOCOL="$protocol" \
        BURROW_SINGMUX_PADDING="$padding" \
        "$0"
    done
  done
  exit 0
fi

protocol="${BURROW_SINGMUX_PROTOCOL:-smux}"
case "$protocol" in
  smux | yamux | h2mux) ;;
  *)
    echo "unsupported sing-mux protocol: $protocol" >&2
    exit 2
    ;;
esac

padding="${BURROW_SINGMUX_PADDING:-false}"
case "$padding" in
  false | true) ;;
  *)
    echo "unsupported sing-mux padding mode: $padding" >&2
    exit 2
    ;;
esac

node_udp="${BURROW_SINGMUX_NODE_UDP:-false}"
case "$node_udp" in
  false | true) ;;
  *)
    echo "unsupported Shadowsocks node udp flag: $node_udp" >&2
    exit 2
    ;;
esac

if [[ -z "${BURROW_BIN:-}" ]]; then
  (cd "$repo_root" && cargo build -p burrow)
elif [[ ! -x "$burrow_bin" ]]; then
  echo "BURROW_BIN is not executable: $burrow_bin" >&2
  exit 2
fi

if ! command -v go >/dev/null 2>&1; then
  echo "go is required for the Mihomo sing-mux interop peer" >&2
  exit 2
fi

if ! command -v uv >/dev/null 2>&1; then
  echo "uv is required for result assertions" >&2
  exit 2
fi

if [[ ! -d "$sing_mux_dir" ]]; then
  echo "sing-mux module directory not found: $sing_mux_dir" >&2
  exit 2
fi

if [[ ! -d "$sing_dir" ]]; then
  echo "sing module directory not found: $sing_dir" >&2
  exit 2
fi

tmpdir="$(mktemp -d "${TMPDIR:-/tmp}/burrow-singmux-interop.XXXXXX")"
daemon_pid=""
server_pid=""
status=0

cleanup() {
  if [[ -n "$server_pid" ]]; then
    kill "$server_pid" >/dev/null 2>&1 || true
  fi
  if [[ -n "$daemon_pid" ]]; then
    kill "$daemon_pid" >/dev/null 2>&1 || true
  fi
  if [[ "${BURROW_SELFTEST_KEEP:-}" == "1" || "$status" -ne 0 ]]; then
    echo "preserving sing-mux interop artifacts: $tmpdir" >&2
  else
    rm -rf "$tmpdir"
  fi
}
trap 'status=$?; cleanup' EXIT

socket="$tmpdir/burrow.sock"
port_file="$tmpdir/singmux.port"
result_file="$tmpdir/singmux-result.json"
payload="$tmpdir/proxy-payload.json"
server_log="$tmpdir/singmux-server.log"
daemon_log="$tmpdir/daemon.log"
go_dir="$tmpdir/go-peer"
mkdir -p "$go_dir"

cat >"$go_dir/go.mod" <<EOF
module burrow-singmux-interop

go 1.23

require (
	github.com/metacubex/sing v0.5.7
	github.com/metacubex/sing-mux v0.3.9
)

replace github.com/metacubex/sing => $sing_dir
replace github.com/metacubex/sing-mux => $sing_mux_dir
EOF

cat >"$go_dir/main.go" <<'EOF'
package main

import (
	"context"
	"encoding/json"
	"fmt"
	"io"
	"net"
	"os"
	"strings"
	"sync"
	"time"

	mux "github.com/metacubex/sing-mux"
	"github.com/metacubex/sing/common/buf"
	"github.com/metacubex/sing/common/bufio"
	"github.com/metacubex/sing/common/logger"
	M "github.com/metacubex/sing/common/metadata"
	N "github.com/metacubex/sing/common/network"
)

type noopLogger struct{}

func (noopLogger) Trace(args ...any)                       {}
func (noopLogger) Debug(args ...any)                       {}
func (noopLogger) Info(args ...any)                        {}
func (noopLogger) Warn(args ...any)                        {}
func (noopLogger) Error(args ...any)                       {}
func (noopLogger) Fatal(args ...any)                       {}
func (noopLogger) Panic(args ...any)                       {}
func (noopLogger) TraceContext(context.Context, ...any)    {}
func (noopLogger) DebugContext(context.Context, ...any)    {}
func (noopLogger) InfoContext(context.Context, ...any)     {}
func (noopLogger) WarnContext(context.Context, ...any)     {}
func (noopLogger) ErrorContext(context.Context, ...any)    {}
func (noopLogger) FatalContext(context.Context, ...any)    {}
func (noopLogger) PanicContext(context.Context, ...any)    {}

var _ logger.ContextLogger = noopLogger{}

type observed struct {
	MuxTarget        map[string]any `json:"mux_target,omitempty"`
	TCPDestination   map[string]any `json:"tcp_destination,omitempty"`
	TCPFirstLine     string         `json:"tcp_first_line,omitempty"`
	UDPDestination   map[string]any `json:"udp_destination,omitempty"`
	UDPPayload       string         `json:"udp_payload,omitempty"`
	Protocol         string         `json:"protocol,omitempty"`
	Padding          bool           `json:"padding,omitempty"`
	PacketReplyCount int            `json:"packet_reply_count,omitempty"`
}

type handler struct {
	mu       sync.Mutex
	result   observed
	done     chan struct{}
	doneOnce sync.Once
}

func (h *handler) markDoneIfReady() {
	if h.result.TCPFirstLine != "" && h.result.UDPPayload != "" {
		h.doneOnce.Do(func() { close(h.done) })
	}
}

func (h *handler) NewConnection(ctx context.Context, conn net.Conn, metadata M.Metadata) error {
	defer conn.Close()
	request, err := readHTTPHeaders(conn)
	if err != nil {
		return err
	}
	firstLine := strings.SplitN(string(request), "\r\n", 2)[0]
	body := []byte("burrow sing-mux tcp ok\n")
	response := fmt.Sprintf(
		"HTTP/1.1 200 OK\r\nContent-Length: %d\r\nConnection: close\r\nContent-Type: text/plain\r\n\r\n%s",
		len(body),
		body,
	)
	if _, err := conn.Write([]byte(response)); err != nil {
		return err
	}
	h.mu.Lock()
	h.result.TCPDestination = socksaddrJSON(metadata.Destination)
	h.result.TCPFirstLine = firstLine
	h.markDoneIfReady()
	h.mu.Unlock()
	return nil
}

func (h *handler) NewPacketConnection(ctx context.Context, conn N.PacketConn, metadata M.Metadata) error {
	defer conn.Close()
	packet := buf.NewPacket()
	defer packet.Release()
	destination, err := conn.ReadPacket(packet)
	if err != nil {
		return err
	}
	payload := append([]byte(nil), packet.Bytes()...)
	if _, err := bufio.WritePacketBuffer(conn, buf.As(payload), destination); err != nil {
		return err
	}
	h.mu.Lock()
	h.result.UDPDestination = socksaddrJSON(metadata.Destination)
	h.result.UDPPayload = string(payload)
	h.result.PacketReplyCount++
	h.markDoneIfReady()
	h.mu.Unlock()
	return nil
}

func main() {
	portFile := os.Args[1]
	resultFile := os.Args[2]
	protocol := os.Args[3]
	padding := os.Args[4] == "true"

	listener, err := net.Listen("tcp4", "127.0.0.1:0")
	if err != nil {
		panic(err)
	}
	defer listener.Close()
	port := listener.Addr().(*net.TCPAddr).Port
	if err := os.WriteFile(portFile, []byte(fmt.Sprint(port)), 0o644); err != nil {
		panic(err)
	}

	h := &handler{done: make(chan struct{})}
	h.result.Protocol = protocol
	h.result.Padding = padding
	service, err := mux.NewService(mux.ServiceOptions{
		NewStreamContext: func(ctx context.Context, conn net.Conn) context.Context { return ctx },
		Logger:           noopLogger{},
		Handler:          h,
	})
	if err != nil {
		panic(err)
	}

	raw, err := listener.Accept()
	if err != nil {
		panic(err)
	}
	raw.SetDeadline(time.Now().Add(20 * time.Second))
	target, err := readSocksAddr(raw)
	if err != nil {
		panic(err)
	}
	h.mu.Lock()
	h.result.MuxTarget = target
	h.mu.Unlock()

	go func() {
		if err := service.NewConnection(context.Background(), raw, M.Metadata{}); err != nil {
			fmt.Fprintln(os.Stderr, err)
		}
	}()

	select {
	case <-h.done:
	case <-time.After(20 * time.Second):
		panic("timed out waiting for sing-mux TCP and UDP probes")
	}

	h.mu.Lock()
	output, err := json.MarshalIndent(h.result, "", "  ")
	h.mu.Unlock()
	if err != nil {
		panic(err)
	}
	if err := os.WriteFile(resultFile, append(output, '\n'), 0o644); err != nil {
		panic(err)
	}
}

func readHTTPHeaders(reader io.Reader) ([]byte, error) {
	data := make([]byte, 0, 512)
	tmp := make([]byte, 1)
	for !strings.Contains(string(data), "\r\n\r\n") {
		if _, err := io.ReadFull(reader, tmp); err != nil {
			return nil, err
		}
		data = append(data, tmp[0])
		if len(data) > 65535 {
			return nil, fmt.Errorf("HTTP headers exceeded 64 KiB")
		}
	}
	return data, nil
}

func readSocksAddr(reader io.Reader) (map[string]any, error) {
	destination, err := M.SocksaddrSerializer.ReadAddrPort(reader)
	if err != nil {
		return nil, err
	}
	return socksaddrJSON(destination), nil
}

func socksaddrJSON(destination M.Socksaddr) map[string]any {
	if destination.IsFqdn() {
		return map[string]any{
			"host": destination.Fqdn,
			"port": destination.Port,
		}
	}
	return map[string]any{
		"host": destination.Addr.String(),
		"port": destination.Port,
	}
}

func (h *handler) NewPacket(ctx context.Context, conn N.PacketConn, buffer *buf.Buffer, metadata M.Metadata) error {
	return fmt.Errorf("unexpected packet handler path")
}
EOF

(
  cd "$go_dir"
  go mod tidy
)

(
  cd "$go_dir"
  go run . "$port_file" "$result_file" "$protocol" "$padding"
) >"$server_log" 2>&1 &
server_pid=$!

for _ in {1..600}; do
  if [[ -s "$port_file" ]]; then
    break
  fi
  if ! kill -0 "$server_pid" >/dev/null 2>&1; then
    echo "sing-mux peer exited before listening" >&2
    cat "$server_log" >&2 || true
    exit 1
  fi
  sleep 0.05
done

if [[ ! -s "$port_file" ]]; then
  echo "timed out waiting for sing-mux peer" >&2
  cat "$server_log" >&2 || true
  exit 1
fi
server_port="$(cat "$port_file")"

cat >"$payload" <<EOF
{
  "name": "burrow sing-mux interop",
  "source": { "url": "selftest://sing-mux" },
  "detected_format": "uri-list",
  "selected_ordinal": 1,
  "selected_name": "local shadowsocks sing-mux $protocol padding $padding",
  "nodes": [
    {
      "ordinal": 1,
      "name": "local shadowsocks sing-mux $protocol padding $padding",
      "protocol": "shadowsocks",
      "server": "127.0.0.1",
      "port": $server_port,
      "warnings": [],
      "config": {
        "type": "shadowsocks",
        "cipher": "none",
        "password": "unused",
        "udp": $node_udp,
        "udp_over_tcp": false,
        "udp_over_tcp_version": 1,
        "client_fingerprint": null,
        "plugin": null,
        "smux": {
          "enabled": true,
          "protocol": "$protocol",
          "max_connections": 1,
          "min_streams": 0,
          "max_streams": 16,
          "padding": $padding,
          "statistic": false,
          "only_tcp": false,
          "brutal": null
        }
      }
    }
  ],
  "warnings": []
}
EOF

(
  cd "$tmpdir"
  BURROW_SOCKET_PATH="$socket" "$burrow_bin" daemon >"$daemon_log" 2>&1
) &
daemon_pid=$!

for _ in {1..100}; do
  if [[ -S "$socket" ]]; then
    break
  fi
  if ! kill -0 "$daemon_pid" >/dev/null 2>&1; then
    echo "Burrow daemon exited before creating socket" >&2
    cat "$daemon_log" >&2 || true
    exit 1
  fi
  sleep 0.05
done

if [[ ! -S "$socket" ]]; then
  echo "timed out waiting for Burrow daemon socket" >&2
  cat "$daemon_log" >&2 || true
  exit 1
fi

(
  cd "$tmpdir"
  BURROW_SOCKET_PATH="$socket" "$burrow_bin" network-add 1 3 "$payload"
  BURROW_SOCKET_PATH="$socket" "$burrow_bin" tunnel-config
  BURROW_SOCKET_PATH="$socket" "$burrow_bin" proxy-tcp-probe \
    --dns-name selftest.invalid \
    --remote 1.1.1.1:80 \
    --timeout-ms 8000
  BURROW_SOCKET_PATH="$socket" "$burrow_bin" proxy-udp-echo \
    --message burrow-singmux-udp-selftest \
    --timeout-ms 8000 \
    9.9.9.9:5353
)

wait "$server_pid"
server_pid=""

echo
echo "Mihomo sing-mux peer observed:"
cat "$result_file"
echo

uv run python - "$result_file" "$protocol" "$padding" <<'PY'
from __future__ import annotations

import json
import sys

with open(sys.argv[1], "r", encoding="utf-8") as f:
    result = json.load(f)
expected_protocol = sys.argv[2]
expected_padding = sys.argv[3] == "true"


def expect(value: object, expected: object, label: str) -> None:
    if value != expected:
        raise AssertionError(f"{label}: expected {expected!r}, got {value!r}")


def expect_target(key: str, host: str, port: int) -> None:
    target = result.get(key)
    if not isinstance(target, dict):
        raise AssertionError(f"{key}: missing target object")
    expect(target.get("host"), host, f"{key}.host")
    expect(target.get("port"), port, f"{key}.port")


expect(result.get("protocol"), expected_protocol, "sing-mux protocol")
expect(result.get("padding", False), expected_padding, "sing-mux padding")
expect_target("mux_target", "sp.mux.sing-box.arpa", 444)
expect_target("tcp_destination", "selftest.invalid", 80)
expect(result.get("tcp_first_line"), "GET / HTTP/1.1", "tcp first line")
expect_target("udp_destination", "9.9.9.9", 5353)
expect(result.get("udp_payload"), "burrow-singmux-udp-selftest", "udp payload")
expect(result.get("packet_reply_count"), 1, "packet reply count")
padding_label = "padded" if expected_padding else "unpadded"
print(f"Mihomo sing-mux {expected_protocol} {padding_label} interop assertions passed")
PY