feat: add channel monitor safety helpers

Add default probe prompts and SSRF validation for channel monitor targets, with focused tests for blocked internal addresses.

service/channel_monitor_ssrf.go

@@ -0,0 +1,123 @@
+package service
+
+import (
+	"fmt"
+	"net"
+	"net/url"
+	"strconv"
+	"strings"
+)
+
+var channelMonitorBlockedIPNets = mustParseChannelMonitorCIDRs([]string{
+	"0.0.0.0/8",
+	"10.0.0.0/8",
+	"100.64.0.0/10",
+	"127.0.0.0/8",
+	"169.254.0.0/16",
+	"172.16.0.0/12",
+	"192.0.0.0/24",
+	"192.0.2.0/24",
+	"192.168.0.0/16",
+	"198.18.0.0/15",
+	"198.51.100.0/24",
+	"203.0.113.0/24",
+	"224.0.0.0/4",
+	"240.0.0.0/4",
+	"255.255.255.255/32",
+	"::/128",
+	"::1/128",
+	"64:ff9b::/96",
+	"100::/64",
+	"2001::/23",
+	"2001:db8::/32",
+	"fc00::/7",
+	"fe80::/10",
+	"ff00::/8",
+})
+
+func ValidateChannelMonitorTargetURL(target string) error {
+	target = strings.TrimSpace(target)
+	if target == "" {
+		return fmt.Errorf("empty URL")
+	}
+
+	parsedURL, err := url.Parse(target)
+	if err != nil {
+		return fmt.Errorf("invalid URL: %w", err)
+	}
+	if parsedURL.Scheme != "http" && parsedURL.Scheme != "https" {
+		return fmt.Errorf("only http and https allowed")
+	}
+
+	host := parsedURL.Hostname()
+	if host == "" {
+		return fmt.Errorf("empty host")
+	}
+	if strings.ContainsAny(host, " \t\r\n") {
+		return fmt.Errorf("invalid host: contains whitespace")
+	}
+	if port := parsedURL.Port(); port != "" {
+		portNumber, err := strconv.Atoi(port)
+		if err != nil || portNumber < 1 || portNumber > 65535 {
+			return fmt.Errorf("invalid port: %s", port)
+		}
+	}
+
+	ips, err := resolveChannelMonitorTargetHost(host)
+	if err != nil {
+		return err
+	}
+	for _, ip := range ips {
+		if isBlockedChannelMonitorTargetIP(ip) {
+			return fmt.Errorf("private/internal IP not allowed: %s resolves to %s", host, ip.String())
+		}
+	}
+	return nil
+}
+
+func resolveChannelMonitorTargetHost(host string) ([]net.IP, error) {
+	if ip := net.ParseIP(host); ip != nil {
+		return []net.IP{ip}, nil
+	}
+
+	ips, err := net.LookupIP(host)
+	if err != nil {
+		return nil, fmt.Errorf("DNS lookup failed for %s: %w", host, err)
+	}
+	if len(ips) == 0 {
+		return nil, fmt.Errorf("no IP resolved for host: %s", host)
+	}
+	return ips, nil
+}
+
+func isBlockedChannelMonitorTargetIP(ip net.IP) bool {
+	if ip == nil {
+		return true
+	}
+	if ip.IsUnspecified() || ip.IsLoopback() || ip.IsPrivate() || ip.IsLinkLocalUnicast() || ip.IsLinkLocalMulticast() || ip.IsInterfaceLocalMulticast() {
+		return true
+	}
+
+	checkIP := ip
+	if ipv4 := ip.To4(); ipv4 != nil {
+		checkIP = ipv4
+	}
+	for _, blockedNet := range channelMonitorBlockedIPNets {
+		if blockedNet.Contains(checkIP) || blockedNet.Contains(ip) {
+			return true
+		}
+	}
+	return false
+}
+
+func mustParseChannelMonitorCIDRs(cidrs []string) []net.IPNet {
+	nets := make([]net.IPNet, 0, len(cidrs))
+	for _, cidr := range cidrs {
+		_, parsedNet, err := net.ParseCIDR(cidr)
+		if err != nil {
+			panic(fmt.Sprintf("invalid channel monitor SSRF CIDR %s: %v", cidr, err))
+		}
+		nets = append(nets, *parsedNet)
+	}
+	return nets
+}

service/channel_monitor_ssrf_test.go

@@ -0,0 +1,87 @@
+package service
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/require"
+)
+
+func TestValidateChannelMonitorTargetURLRejectsInternalAndInvalidTargets(t *testing.T) {
+	t.Parallel()
+
+	testCases := []struct {
+		name          string
+		target        string
+		errorContains string
+	}{
+		{
+			name:          "reject loopback IPv4",
+			target:        "http://127.0.0.1:3000/v1/chat/completions",
+			errorContains: "private/internal IP not allowed",
+		},
+		{
+			name:          "reject localhost resolved IP",
+			target:        "http://localhost:3000/v1/chat/completions",
+			errorContains: "private/internal IP not allowed",
+		},
+		{
+			name:          "reject 10 private range",
+			target:        "http://10.0.0.1/v1/chat/completions",
+			errorContains: "private/internal IP not allowed",
+		},
+		{
+			name:          "reject 172.16 private range",
+			target:        "http://172.16.0.1/v1/chat/completions",
+			errorContains: "private/internal IP not allowed",
+		},
+		{
+			name:          "reject 172.31 private range",
+			target:        "https://172.31.255.255/v1/chat/completions",
+			errorContains: "private/internal IP not allowed",
+		},
+		{
+			name:          "reject 192.168 private range",
+			target:        "http://192.168.1.10/v1/chat/completions",
+			errorContains: "private/internal IP not allowed",
+		},
+		{
+			name:          "reject link-local IPv4",
+			target:        "http://169.254.169.254/latest/meta-data",
+			errorContains: "private/internal IP not allowed",
+		},
+		{
+			name:          "reject loopback IPv6",
+			target:        "http://[::1]:3000/v1/chat/completions",
+			errorContains: "private/internal IP not allowed",
+		},
+		{
+			name:          "reject unsupported scheme",
+			target:        "file:///etc/passwd",
+			errorContains: "only http and https allowed",
+		},
+		{
+			name:          "reject invalid URL",
+			target:        "http://[::1",
+			errorContains: "invalid URL",
+		},
+	}
+
+	for _, tc := range testCases {
+		tc := tc
+		t.Run(tc.name, func(t *testing.T) {
+			t.Parallel()
+
+			err := ValidateChannelMonitorTargetURL(tc.target)
+			require.Error(t, err)
+			require.Contains(t, err.Error(), tc.errorContains)
+		})
+	}
+}
+
+func TestValidateChannelMonitorTargetURLAllowsPublicHTTPSTarget(t *testing.T) {
+	t.Parallel()
+
+	err := ValidateChannelMonitorTargetURL("https://8.8.8.8/v1/chat/completions")
+
+	require.NoError(t, err)
+}

service/channel_monitor_templates.go

@@ -0,0 +1,70 @@
+package service
+
+import (
+	"math/rand"
+	"strings"
+
+	"github.com/QuantumNous/new-api/common"
+)
+
+var DefaultChannelMonitorRequestTemplates = []string{
+	"Explain what an API is in one sentence.",
+	"Write a Python function to check if a number is prime.",
+	"What is 15 percent of 200? Return only the number.",
+	"Translate 'Good morning' into Spanish.",
+	"What is the capital of Japan? Answer in one word.",
+	"Explain the difference between HTTP and HTTPS briefly.",
+	"Write a bash command to count lines in a file.",
+	"What does CPU stand for?",
+	"Convert 100 kilometers to miles. Return only the number.",
+	"Name three primary colors.",
+	"Write a SQL query to select all users from a users table.",
+	"What is the chemical symbol for water?",
+	"Explain gravity in one sentence.",
+	"What is 2 to the power of 10? Return only the number.",
+	"Explain recursion in one sentence.",
+	"Write a Linux command to list files in a directory.",
+	"What year did World War II end? Return only the year.",
+	"Explain what DNS does in one sentence.",
+	"What is the square root of 144? Return only the number.",
+	"Summarize photosynthesis in one sentence.",
+}
+
+func ParseChannelMonitorRequestTemplates(raw string) []string {
+	raw = strings.TrimSpace(raw)
+	if raw == "" {
+		return cloneDefaultChannelMonitorRequestTemplates()
+	}
+
+	var templates []string
+	if err := common.Unmarshal([]byte(raw), &templates); err != nil {
+		return cloneDefaultChannelMonitorRequestTemplates()
+	}
+
+	cleaned := make([]string, 0, len(templates))
+	for _, template := range templates {
+		template = strings.TrimSpace(template)
+		if template != "" {
+			cleaned = append(cleaned, template)
+		}
+	}
+	if len(cleaned) == 0 {
+		return cloneDefaultChannelMonitorRequestTemplates()
+	}
+	return cleaned
+}
+
+func SelectChannelMonitorRequestTemplate(raw string, rng *rand.Rand) string {
+	templates := ParseChannelMonitorRequestTemplates(raw)
+	if len(templates) == 0 {
+		return ""
+	}
+	if rng != nil {
+		return templates[rng.Intn(len(templates))]
+	}
+	return templates[rand.Intn(len(templates))]
+}
+
+func cloneDefaultChannelMonitorRequestTemplates() []string {
+	return append([]string(nil), DefaultChannelMonitorRequestTemplates...)
+}

service/channel_monitor_templates_test.go

@@ -0,0 +1,63 @@
+package service
+
+import (
+	"math/rand"
+	"testing"
+
+	"github.com/stretchr/testify/require"
+)
+
+func TestDefaultChannelMonitorRequestTemplatesHasTwentyShortEnglishPrompts(t *testing.T) {
+	t.Parallel()
+
+	require.Len(t, DefaultChannelMonitorRequestTemplates, 20)
+	for _, template := range DefaultChannelMonitorRequestTemplates {
+		require.NotEmpty(t, template)
+		require.LessOrEqual(t, len(template), 120)
+	}
+}
+
+func TestParseChannelMonitorRequestTemplatesUsesCustomTemplates(t *testing.T) {
+	t.Parallel()
+
+	templates := ParseChannelMonitorRequestTemplates(`["Say hello.","Solve 2+2."]`)
+
+	require.Equal(t, []string{"Say hello.", "Solve 2+2."}, templates)
+}
+
+func TestParseChannelMonitorRequestTemplatesFallsBackForEmptyInvalidOrBlankTemplates(t *testing.T) {
+	t.Parallel()
+
+	testCases := []struct {
+		name string
+		raw  string
+	}{
+		{name: "empty", raw: ""},
+		{name: "blank", raw: "   "},
+		{name: "invalid json", raw: "{bad json"},
+		{name: "wrong json type", raw: `{"prompt":"hello"}`},
+		{name: "only empty strings", raw: `["", "   "]`},
+	}
+
+	for _, tc := range testCases {
+		tc := tc
+		t.Run(tc.name, func(t *testing.T) {
+			t.Parallel()
+
+			templates := ParseChannelMonitorRequestTemplates(tc.raw)
+
+			require.Equal(t, DefaultChannelMonitorRequestTemplates, templates)
+		})
+	}
+}
+
+func TestSelectChannelMonitorRequestTemplateUsesDeterministicRandWhenProvided(t *testing.T) {
+	t.Parallel()
+
+	raw := `["first","second","third"]`
+	first := SelectChannelMonitorRequestTemplate(raw, rand.New(rand.NewSource(42)))
+	second := SelectChannelMonitorRequestTemplate(raw, rand.New(rand.NewSource(42)))
+
+	require.Equal(t, first, second)
+	require.Contains(t, []string{"first", "second", "third"}, first)
+}