|
- //go:build ignore
-
- package main
-
- import (
- "crypto/tls"
- "encoding/json"
- "fmt"
- "io"
- "net/http"
- "os"
- "strings"
- "time"
- )
-
- var (
- baseURL = "http://localhost:3001"
- apiKey string
- client = &http.Client{Timeout: 30 * time.Second}
- )
-
- func init() {
- // Skip TLS verification if needed
- client.Transport = &http.Transport{
- TLSClientConfig: &tls.Config{InsecureSkipVerify: true},
- }
- }
-
- func main() {
- // Use test token created earlier
- apiKey = "sk-3bJC6cCSSFuTqOxaZQ8lL6kOfYAl3pt0uuwWMiU40F6KkxOS"
- if len(os.Args) > 1 {
- apiKey = os.Args[1]
- }
- if len(os.Args) > 2 {
- baseURL = os.Args[2]
- }
-
- pass, fail := 0, 0
-
- fmt.Println("==========================================")
- fmt.Println("E2E: Endpoint Format Restriction")
- fmt.Println("==========================================")
-
- // Test 1: Feature OFF - Claude via /v1/chat/completions should pass
- fmt.Println("\n--- Test 1: Feature OFF, Claude via /v1/chat/completions ---")
- code, body := doRequest("POST", "/v1/chat/completions", `{"model":"anthropic/claude-sonnet-4-20250514","messages":[{"role":"user","content":"hi"}],"max_tokens":5}`)
- fmt.Printf(" HTTP: %d\n", code)
- if code != 400 || !strings.Contains(body, "endpoint_format_restricted") {
- fmt.Println(" PASS (not blocked)")
- pass++
- } else {
- fmt.Println(" FAIL: should not be blocked when feature is off")
- fail++
- }
-
- // Enable feature via admin API
- fmt.Println("\n--- Enabling feature via admin login ---")
- session := adminLogin()
- if session == "" {
- fmt.Println(" WARNING: Could not login as admin, trying direct DB approach")
- } else {
- enableFeature(session)
- }
-
- // Test 2: Feature ON - Claude via /v1/chat/completions should be BLOCKED
- fmt.Println("\n--- Test 2: Feature ON, Claude via /v1/chat/completions (should BLOCK) ---")
- code, body = doRequest("POST", "/v1/chat/completions", `{"model":"anthropic/claude-sonnet-4-20250514","messages":[{"role":"user","content":"hi"}],"max_tokens":5}`)
- fmt.Printf(" HTTP: %d\n", code)
- fmt.Printf(" Body: %s\n", truncate(body, 300))
- if code == 400 && strings.Contains(body, "endpoint_format_restricted") {
- fmt.Println(" PASS (correctly blocked)")
- pass++
- } else {
- fmt.Println(" FAIL: expected 400 with endpoint_format_restricted")
- fail++
- }
-
- // Test 3: Feature ON - Claude via /v1/messages should PASS
- fmt.Println("\n--- Test 3: Feature ON, Claude via /v1/messages (should PASS) ---")
- code, body = doRequestWithHeaders("POST", "/v1/messages",
- `{"model":"anthropic/claude-sonnet-4-20250514","messages":[{"role":"user","content":"hi"}],"max_tokens":5}`,
- map[string]string{"anthropic-version": "2023-06-01"})
- fmt.Printf(" HTTP: %d\n", code)
- if code != 400 || !strings.Contains(body, "endpoint_format_restricted") {
- fmt.Println(" PASS (correctly allowed)")
- pass++
- } else {
- fmt.Println(" FAIL: should not be blocked for matching format")
- fail++
- }
-
- // Test 4: Feature ON - DeepSeek via /v1/chat/completions should PASS
- fmt.Println("\n--- Test 4: Feature ON, DeepSeek via /v1/chat/completions (should PASS) ---")
- code, body = doRequest("POST", "/v1/chat/completions", `{"model":"deepseek-ai/DeepSeek-V3-0324","messages":[{"role":"user","content":"hi"}],"max_tokens":5}`)
- fmt.Printf(" HTTP: %d\n", code)
- if code != 400 || !strings.Contains(body, "endpoint_format_restricted") {
- fmt.Println(" PASS (correctly allowed)")
- pass++
- } else {
- fmt.Println(" FAIL: should not be blocked for openai channel")
- fail++
- }
-
- // Test 5: Check error message content
- fmt.Println("\n--- Test 5: Error message contains native endpoint recommendation ---")
- if strings.Contains(body, "/v1/messages") || strings.Contains(body, "Anthropic") || strings.Contains(body, "endpoint_format_restricted") {
- // Re-test to get the error body
- code2, body2 := doRequest("POST", "/v1/chat/completions", `{"model":"anthropic/claude-sonnet-4-20250514","messages":[{"role":"user","content":"hi"}],"max_tokens":5}`)
- _ = code2
- if strings.Contains(body2, "/v1/messages") && strings.Contains(body2, "endpoint_format_restricted") {
- fmt.Println(" PASS (error recommends /v1/messages)")
- pass++
- } else {
- fmt.Printf(" Body: %s\n", truncate(body2, 300))
- fmt.Println(" FAIL: error message missing native endpoint")
- fail++
- }
- } else {
- fmt.Println(" SKIP (previous test didn't produce blocking error)")
- }
-
- // Disable feature
- if session != "" {
- fmt.Println("\n--- Disabling feature ---")
- disableFeature(session)
- }
-
- fmt.Printf("\n==========================================\n")
- fmt.Printf("Results: %d passed, %d failed\n", pass, fail)
- fmt.Printf("==========================================\n")
-
- if fail > 0 {
- os.Exit(1)
- }
- }
-
- func doRequest(method, path, body string) (int, string) {
- return doRequestWithHeaders(method, path, body, nil)
- }
-
- func doRequestWithHeaders(method, path, body string, headers map[string]string) (int, string) {
- var req *http.Request
- var err error
- if body != "" {
- req, err = http.NewRequest(method, baseURL+path, strings.NewReader(body))
- } else {
- req, err = http.NewRequest(method, baseURL+path, nil)
- }
- if err != nil {
- return 0, err.Error()
- }
- req.Header.Set("Authorization", "Bearer "+apiKey)
- req.Header.Set("Content-Type", "application/json")
- for k, v := range headers {
- req.Header.Set(k, v)
- }
-
- resp, err := client.Do(req)
- if err != nil {
- return 0, err.Error()
- }
- defer resp.Body.Close()
- respBody, _ := io.ReadAll(resp.Body)
- return resp.StatusCode, string(respBody)
- }
-
- func adminLogin() string {
- // Try common passwords
- passwords := []string{"123456", "password", "admin", "root", "root123", "admin123"}
- for _, pw := range passwords {
- code, body := doRawRequest("POST", baseURL+"/api/user/login",
- fmt.Sprintf(`{"username":"root","password":"%s"}`, pw))
- if code == 200 {
- var result struct {
- Success bool `json:"success"`
- Data string `json:"data"`
- }
- json.Unmarshal([]byte(body), &result)
- if result.Success {
- fmt.Printf(" Logged in with password: %s\n", pw)
- return result.Data
- }
- }
- }
- fmt.Println(" Could not login with any common password")
- return ""
- }
-
- func doRawRequest(method, url, body string) (int, string) {
- req, err := http.NewRequest(method, url, strings.NewReader(body))
- if err != nil {
- return 0, err.Error()
- }
- req.Header.Set("Content-Type", "application/json")
- resp, err := client.Do(req)
- if err != nil {
- return 0, err.Error()
- }
- defer resp.Body.Close()
- respBody, _ := io.ReadAll(resp.Body)
- return resp.StatusCode, string(respBody)
- }
-
- func enableFeature(session string) {
- code, body := doRawRequest("PUT", baseURL+"/api/option/",
- `{"key":"global.endpoint_format_restriction_enabled","value":"true"}`)
- req, _ := http.NewRequest("PUT", baseURL+"/api/option/", strings.NewReader(
- `{"key":"global.endpoint_format_restriction_enabled","value":"true"}`))
- req.Header.Set("Content-Type", "application/json")
- req.Header.Set("Authorization", "Bearer "+session)
- resp, err := client.Do(req)
- if err != nil {
- fmt.Printf(" Error: %v\n", err)
- return
- }
- defer resp.Body.Close()
- respBody, _ := io.ReadAll(resp.Body)
- fmt.Printf(" Enable response: %d %s\n", resp.StatusCode, truncate(string(respBody), 100))
- _ = code
- _ = body
- time.Sleep(1 * time.Second)
- }
-
- func disableFeature(session string) {
- req, _ := http.NewRequest("PUT", baseURL+"/api/option/", strings.NewReader(
- `{"key":"global.endpoint_format_restriction_enabled","value":"false"}`))
- req.Header.Set("Content-Type", "application/json")
- req.Header.Set("Authorization", "Bearer "+session)
- resp, err := client.Do(req)
- if err != nil {
- fmt.Printf(" Error: %v\n", err)
- return
- }
- defer resp.Body.Close()
- respBody, _ := io.ReadAll(resp.Body)
- fmt.Printf(" Disable response: %d %s\n", resp.StatusCode, truncate(string(respBody), 100))
- }
-
- func truncate(s string, n int) string {
- if len(s) <= n {
- return s
- }
- return s[:n] + "..."
- }
|