diff --git a/service/channel_monitor_ssrf.go b/service/channel_monitor_ssrf.go new file mode 100644 index 00000000..066386ae --- /dev/null +++ b/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 +} diff --git a/service/channel_monitor_ssrf_test.go b/service/channel_monitor_ssrf_test.go new file mode 100644 index 00000000..cf68117c --- /dev/null +++ b/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) +} diff --git a/service/channel_monitor_templates.go b/service/channel_monitor_templates.go new file mode 100644 index 00000000..d4e9982a --- /dev/null +++ b/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...) +} diff --git a/service/channel_monitor_templates_test.go b/service/channel_monitor_templates_test.go new file mode 100644 index 00000000..7ae0f563 --- /dev/null +++ b/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) +}