Du kannst nicht mehr als 25 Themen auswählen Themen müssen entweder mit einem Buchstaben oder einer Ziffer beginnen. Sie können Bindestriche („-“) enthalten und bis zu 35 Zeichen lang sein.
 
 
 

246 Zeilen
7.7 KiB

  1. //go:build ignore
  2. package main
  3. import (
  4. "crypto/tls"
  5. "encoding/json"
  6. "fmt"
  7. "io"
  8. "net/http"
  9. "os"
  10. "strings"
  11. "time"
  12. )
  13. var (
  14. baseURL = "http://localhost:3001"
  15. apiKey string
  16. client = &http.Client{Timeout: 30 * time.Second}
  17. )
  18. func init() {
  19. // Skip TLS verification if needed
  20. client.Transport = &http.Transport{
  21. TLSClientConfig: &tls.Config{InsecureSkipVerify: true},
  22. }
  23. }
  24. func main() {
  25. // Use test token created earlier
  26. apiKey = "sk-3bJC6cCSSFuTqOxaZQ8lL6kOfYAl3pt0uuwWMiU40F6KkxOS"
  27. if len(os.Args) > 1 {
  28. apiKey = os.Args[1]
  29. }
  30. if len(os.Args) > 2 {
  31. baseURL = os.Args[2]
  32. }
  33. pass, fail := 0, 0
  34. fmt.Println("==========================================")
  35. fmt.Println("E2E: Endpoint Format Restriction")
  36. fmt.Println("==========================================")
  37. // Test 1: Feature OFF - Claude via /v1/chat/completions should pass
  38. fmt.Println("\n--- Test 1: Feature OFF, Claude via /v1/chat/completions ---")
  39. code, body := doRequest("POST", "/v1/chat/completions", `{"model":"anthropic/claude-sonnet-4-20250514","messages":[{"role":"user","content":"hi"}],"max_tokens":5}`)
  40. fmt.Printf(" HTTP: %d\n", code)
  41. if code != 400 || !strings.Contains(body, "endpoint_format_restricted") {
  42. fmt.Println(" PASS (not blocked)")
  43. pass++
  44. } else {
  45. fmt.Println(" FAIL: should not be blocked when feature is off")
  46. fail++
  47. }
  48. // Enable feature via admin API
  49. fmt.Println("\n--- Enabling feature via admin login ---")
  50. session := adminLogin()
  51. if session == "" {
  52. fmt.Println(" WARNING: Could not login as admin, trying direct DB approach")
  53. } else {
  54. enableFeature(session)
  55. }
  56. // Test 2: Feature ON - Claude via /v1/chat/completions should be BLOCKED
  57. fmt.Println("\n--- Test 2: Feature ON, Claude via /v1/chat/completions (should BLOCK) ---")
  58. code, body = doRequest("POST", "/v1/chat/completions", `{"model":"anthropic/claude-sonnet-4-20250514","messages":[{"role":"user","content":"hi"}],"max_tokens":5}`)
  59. fmt.Printf(" HTTP: %d\n", code)
  60. fmt.Printf(" Body: %s\n", truncate(body, 300))
  61. if code == 400 && strings.Contains(body, "endpoint_format_restricted") {
  62. fmt.Println(" PASS (correctly blocked)")
  63. pass++
  64. } else {
  65. fmt.Println(" FAIL: expected 400 with endpoint_format_restricted")
  66. fail++
  67. }
  68. // Test 3: Feature ON - Claude via /v1/messages should PASS
  69. fmt.Println("\n--- Test 3: Feature ON, Claude via /v1/messages (should PASS) ---")
  70. code, body = doRequestWithHeaders("POST", "/v1/messages",
  71. `{"model":"anthropic/claude-sonnet-4-20250514","messages":[{"role":"user","content":"hi"}],"max_tokens":5}`,
  72. map[string]string{"anthropic-version": "2023-06-01"})
  73. fmt.Printf(" HTTP: %d\n", code)
  74. if code != 400 || !strings.Contains(body, "endpoint_format_restricted") {
  75. fmt.Println(" PASS (correctly allowed)")
  76. pass++
  77. } else {
  78. fmt.Println(" FAIL: should not be blocked for matching format")
  79. fail++
  80. }
  81. // Test 4: Feature ON - DeepSeek via /v1/chat/completions should PASS
  82. fmt.Println("\n--- Test 4: Feature ON, DeepSeek via /v1/chat/completions (should PASS) ---")
  83. code, body = doRequest("POST", "/v1/chat/completions", `{"model":"deepseek-ai/DeepSeek-V3-0324","messages":[{"role":"user","content":"hi"}],"max_tokens":5}`)
  84. fmt.Printf(" HTTP: %d\n", code)
  85. if code != 400 || !strings.Contains(body, "endpoint_format_restricted") {
  86. fmt.Println(" PASS (correctly allowed)")
  87. pass++
  88. } else {
  89. fmt.Println(" FAIL: should not be blocked for openai channel")
  90. fail++
  91. }
  92. // Test 5: Check error message content
  93. fmt.Println("\n--- Test 5: Error message contains native endpoint recommendation ---")
  94. if strings.Contains(body, "/v1/messages") || strings.Contains(body, "Anthropic") || strings.Contains(body, "endpoint_format_restricted") {
  95. // Re-test to get the error body
  96. code2, body2 := doRequest("POST", "/v1/chat/completions", `{"model":"anthropic/claude-sonnet-4-20250514","messages":[{"role":"user","content":"hi"}],"max_tokens":5}`)
  97. _ = code2
  98. if strings.Contains(body2, "/v1/messages") && strings.Contains(body2, "endpoint_format_restricted") {
  99. fmt.Println(" PASS (error recommends /v1/messages)")
  100. pass++
  101. } else {
  102. fmt.Printf(" Body: %s\n", truncate(body2, 300))
  103. fmt.Println(" FAIL: error message missing native endpoint")
  104. fail++
  105. }
  106. } else {
  107. fmt.Println(" SKIP (previous test didn't produce blocking error)")
  108. }
  109. // Disable feature
  110. if session != "" {
  111. fmt.Println("\n--- Disabling feature ---")
  112. disableFeature(session)
  113. }
  114. fmt.Printf("\n==========================================\n")
  115. fmt.Printf("Results: %d passed, %d failed\n", pass, fail)
  116. fmt.Printf("==========================================\n")
  117. if fail > 0 {
  118. os.Exit(1)
  119. }
  120. }
  121. func doRequest(method, path, body string) (int, string) {
  122. return doRequestWithHeaders(method, path, body, nil)
  123. }
  124. func doRequestWithHeaders(method, path, body string, headers map[string]string) (int, string) {
  125. var req *http.Request
  126. var err error
  127. if body != "" {
  128. req, err = http.NewRequest(method, baseURL+path, strings.NewReader(body))
  129. } else {
  130. req, err = http.NewRequest(method, baseURL+path, nil)
  131. }
  132. if err != nil {
  133. return 0, err.Error()
  134. }
  135. req.Header.Set("Authorization", "Bearer "+apiKey)
  136. req.Header.Set("Content-Type", "application/json")
  137. for k, v := range headers {
  138. req.Header.Set(k, v)
  139. }
  140. resp, err := client.Do(req)
  141. if err != nil {
  142. return 0, err.Error()
  143. }
  144. defer resp.Body.Close()
  145. respBody, _ := io.ReadAll(resp.Body)
  146. return resp.StatusCode, string(respBody)
  147. }
  148. func adminLogin() string {
  149. // Try common passwords
  150. passwords := []string{"123456", "password", "admin", "root", "root123", "admin123"}
  151. for _, pw := range passwords {
  152. code, body := doRawRequest("POST", baseURL+"/api/user/login",
  153. fmt.Sprintf(`{"username":"root","password":"%s"}`, pw))
  154. if code == 200 {
  155. var result struct {
  156. Success bool `json:"success"`
  157. Data string `json:"data"`
  158. }
  159. json.Unmarshal([]byte(body), &result)
  160. if result.Success {
  161. fmt.Printf(" Logged in with password: %s\n", pw)
  162. return result.Data
  163. }
  164. }
  165. }
  166. fmt.Println(" Could not login with any common password")
  167. return ""
  168. }
  169. func doRawRequest(method, url, body string) (int, string) {
  170. req, err := http.NewRequest(method, url, strings.NewReader(body))
  171. if err != nil {
  172. return 0, err.Error()
  173. }
  174. req.Header.Set("Content-Type", "application/json")
  175. resp, err := client.Do(req)
  176. if err != nil {
  177. return 0, err.Error()
  178. }
  179. defer resp.Body.Close()
  180. respBody, _ := io.ReadAll(resp.Body)
  181. return resp.StatusCode, string(respBody)
  182. }
  183. func enableFeature(session string) {
  184. code, body := doRawRequest("PUT", baseURL+"/api/option/",
  185. `{"key":"global.endpoint_format_restriction_enabled","value":"true"}`)
  186. req, _ := http.NewRequest("PUT", baseURL+"/api/option/", strings.NewReader(
  187. `{"key":"global.endpoint_format_restriction_enabled","value":"true"}`))
  188. req.Header.Set("Content-Type", "application/json")
  189. req.Header.Set("Authorization", "Bearer "+session)
  190. resp, err := client.Do(req)
  191. if err != nil {
  192. fmt.Printf(" Error: %v\n", err)
  193. return
  194. }
  195. defer resp.Body.Close()
  196. respBody, _ := io.ReadAll(resp.Body)
  197. fmt.Printf(" Enable response: %d %s\n", resp.StatusCode, truncate(string(respBody), 100))
  198. _ = code
  199. _ = body
  200. time.Sleep(1 * time.Second)
  201. }
  202. func disableFeature(session string) {
  203. req, _ := http.NewRequest("PUT", baseURL+"/api/option/", strings.NewReader(
  204. `{"key":"global.endpoint_format_restriction_enabled","value":"false"}`))
  205. req.Header.Set("Content-Type", "application/json")
  206. req.Header.Set("Authorization", "Bearer "+session)
  207. resp, err := client.Do(req)
  208. if err != nil {
  209. fmt.Printf(" Error: %v\n", err)
  210. return
  211. }
  212. defer resp.Body.Close()
  213. respBody, _ := io.ReadAll(resp.Body)
  214. fmt.Printf(" Disable response: %d %s\n", resp.StatusCode, truncate(string(respBody), 100))
  215. }
  216. func truncate(s string, n int) string {
  217. if len(s) <= n {
  218. return s
  219. }
  220. return s[:n] + "..."
  221. }