feat(captain): secure custom tool auth config

This commit is contained in:
2026-06-07 16:32:13 +08:00
parent 8f9adecd04
commit f45fbfd5f0
4 changed files with 135 additions and 11 deletions
@@ -65,6 +65,12 @@ func (s *CaptainCustomToolCRUDTestSuite) SetupSuite() {
// Unified router: :id = account_id, :tool_id = tool_id
s.router = gin.New()
s.router.RedirectTrailingSlash = false
s.router.Use(func(c *gin.Context) {
if role := c.GetHeader("X-Test-Role"); role != "" {
c.Set("role", role)
}
c.Next()
})
accGroup := s.router.Group("/api/v1/accounts/:account_id")
{
ctGroup := accGroup.Group("/captain/custom_tools")
@@ -97,6 +103,10 @@ func (s *CaptainCustomToolCRUDTestSuite) accountPath() string {
// helper: send HTTP request via unified router
func (s *CaptainCustomToolCRUDTestSuite) makeRequest(method, path string, body interface{}) *httptest.ResponseRecorder {
return s.makeRequestWithRole(method, path, body, "")
}
func (s *CaptainCustomToolCRUDTestSuite) makeRequestWithRole(method, path string, body interface{}, role string) *httptest.ResponseRecorder {
var bodyBytes []byte
if body != nil {
bodyBytes, _ = json.Marshal(body)
@@ -106,6 +116,9 @@ func (s *CaptainCustomToolCRUDTestSuite) makeRequest(method, path string, body i
if body != nil {
req.Header.Set("Content-Type", "application/json")
}
if role != "" {
req.Header.Set("X-Test-Role", role)
}
s.router.ServeHTTP(w, req)
return w
}
@@ -165,6 +178,58 @@ func (s *CaptainCustomToolCRUDTestSuite) TestCreate_默认GET方法() {
assert.Equal(s.T(), "GET", resp["http_method"]) // default HTTP method
}
func (s *CaptainCustomToolCRUDTestSuite) TestPayload_AuthConfigOnlyForAdministrators() {
body := map[string]interface{}{
"custom_tool": map[string]interface{}{
"title": "Bearer tool",
"endpoint_url": "https://example.com/bearer",
"auth_type": "bearer",
"auth_config": map[string]interface{}{"token": "secret-token"},
},
}
w := s.makeRequest("POST", s.accountPath()+"/captain/custom_tools/", body)
assert.Equal(s.T(), http.StatusOK, w.Code)
var created map[string]interface{}
s.Require().NoError(json.Unmarshal(w.Body.Bytes(), &created))
assert.NotContains(s.T(), created, "auth_config")
toolID := strconv.FormatFloat(created["id"].(float64), 'f', -1, 64)
w = s.makeRequestWithRole("GET", s.accountPath()+"/captain/custom_tools/"+toolID, nil, "agent")
assert.Equal(s.T(), http.StatusOK, w.Code)
var agentResp map[string]interface{}
s.Require().NoError(json.Unmarshal(w.Body.Bytes(), &agentResp))
assert.NotContains(s.T(), agentResp, "auth_config")
w = s.makeRequestWithRole("GET", s.accountPath()+"/captain/custom_tools/"+toolID, nil, "administrator")
assert.Equal(s.T(), http.StatusOK, w.Code)
var adminResp map[string]interface{}
s.Require().NoError(json.Unmarshal(w.Body.Bytes(), &adminResp))
assert.Equal(s.T(), map[string]interface{}{"token": "secret-token"}, adminResp["auth_config"])
}
func (s *CaptainCustomToolCRUDTestSuite) TestMutations_AgentRoleForbidden() {
toolID := s.createToolAndGetID("Agent forbidden", "agent-forbidden", "https://example.com/agent")
w := s.makeRequestWithRole("POST", s.accountPath()+"/captain/custom_tools/", map[string]interface{}{
"custom_tool": map[string]interface{}{
"title": "Agent create",
"endpoint_url": "https://example.com/create",
},
}, "agent")
assert.Equal(s.T(), http.StatusForbidden, w.Code)
w = s.makeRequestWithRole("PUT", s.accountPath()+"/captain/custom_tools/"+toolID, map[string]interface{}{"title": "blocked"}, "agent")
assert.Equal(s.T(), http.StatusForbidden, w.Code)
w = s.makeRequestWithRole("DELETE", s.accountPath()+"/captain/custom_tools/"+toolID, nil, "agent")
assert.Equal(s.T(), http.StatusForbidden, w.Code)
var resp map[string]interface{}
s.Require().NoError(json.Unmarshal(w.Body.Bytes(), &resp))
assert.Equal(s.T(), "You are not authorized to do this action", resp["error"])
}
func (s *CaptainCustomToolCRUDTestSuite) TestCreate_省略Slug时使用ChatwootCustom前缀和下划线() {
body := map[string]interface{}{
"custom_tool": map[string]interface{}{
@@ -33,6 +33,9 @@ func (h *CaptainCustomToolHandler) Create(c *gin.Context) {
if !h.ensureCustomToolsEnabled(c, accountID) {
return
}
if !h.ensureCustomToolAdmin(c) {
return
}
var req service.CreateCustomToolRequest
if err := bindNestedJSONPayload(c, "custom_tool", &req); err != nil {
@@ -54,7 +57,7 @@ func (h *CaptainCustomToolHandler) Create(c *gin.Context) {
return
}
c.JSON(http.StatusOK, captainCustomToolPayload(tool))
c.JSON(http.StatusOK, captainCustomToolPayload(c, tool))
}
// Get retrieves a custom tool by ID.
@@ -81,7 +84,7 @@ func (h *CaptainCustomToolHandler) Get(c *gin.Context) {
return
}
c.JSON(http.StatusOK, captainCustomToolPayload(tool))
c.JSON(http.StatusOK, captainCustomToolPayload(c, tool))
}
// Update updates an existing custom tool.
@@ -95,6 +98,9 @@ func (h *CaptainCustomToolHandler) Update(c *gin.Context) {
if !h.ensureCustomToolsEnabled(c, accountID) {
return
}
if !h.ensureCustomToolAdmin(c) {
return
}
id, err := parseUintAnyParam(c, "tool_id", "id")
if err != nil {
response.AbortWithStatusError(c, http.StatusBadRequest, response.ErrBadRequest, "invalid id")
@@ -117,7 +123,7 @@ func (h *CaptainCustomToolHandler) Update(c *gin.Context) {
return
}
c.JSON(http.StatusOK, captainCustomToolPayload(tool))
c.JSON(http.StatusOK, captainCustomToolPayload(c, tool))
}
// Delete deletes a custom tool.
@@ -131,6 +137,9 @@ func (h *CaptainCustomToolHandler) Delete(c *gin.Context) {
if !h.ensureCustomToolsEnabled(c, accountID) {
return
}
if !h.ensureCustomToolAdmin(c) {
return
}
id, err := parseUintAnyParam(c, "tool_id", "id")
if err != nil {
response.AbortWithStatusError(c, http.StatusBadRequest, response.ErrBadRequest, "invalid id")
@@ -167,7 +176,7 @@ func (h *CaptainCustomToolHandler) List(c *gin.Context) {
payload := make([]gin.H, 0, len(tools))
for i := range tools {
payload = append(payload, captainCustomToolPayload(&tools[i]))
payload = append(payload, captainCustomToolPayload(c, &tools[i]))
}
c.JSON(http.StatusOK, gin.H{"payload": payload, "meta": gin.H{"total_count": count, "page": 1}})
}
@@ -208,6 +217,9 @@ func (h *CaptainCustomToolHandler) TestTool(c *gin.Context) {
if !h.ensureCustomToolsEnabled(c, accountID) {
return
}
if !h.ensureCustomToolAdmin(c) {
return
}
var req service.TestToolRequest
if err := bindNestedJSONPayload(c, "custom_tool", &req); err != nil {
@@ -233,6 +245,17 @@ func (h *CaptainCustomToolHandler) ensureCustomToolsEnabled(c *gin.Context, acco
return false
}
func (h *CaptainCustomToolHandler) ensureCustomToolAdmin(c *gin.Context) bool {
if role, exists := c.Get("role"); exists {
if role == "administrator" || role == "super_admin" {
return true
}
c.JSON(http.StatusForbidden, gin.H{"error": "You are not authorized to do this action"})
return false
}
return true
}
func renderCaptainCustomToolValidationError(c *gin.Context, err error) bool {
var validationErr *service.CaptainCustomToolValidationError
if !errors.As(err, &validationErr) {
@@ -245,8 +268,8 @@ func renderCaptainCustomToolValidationError(c *gin.Context, err error) bool {
return true
}
func captainCustomToolPayload(tool *model.CaptainCustomTool) gin.H {
return gin.H{
func captainCustomToolPayload(c *gin.Context, tool *model.CaptainCustomTool) gin.H {
payload := gin.H{
"id": tool.ID,
"slug": tool.Slug,
"title": tool.Title,
@@ -256,11 +279,19 @@ func captainCustomToolPayload(tool *model.CaptainCustomTool) gin.H {
"request_template": tool.RequestTemplate,
"response_template": tool.ResponseTemplate,
"auth_type": tool.AuthType,
"auth_config": rawJSONValue(tool.AuthConfig),
"param_schema": rawJSONValue(tool.ParamSchema),
"enabled": tool.Enabled,
"account_id": tool.AccountID,
"created_at": tool.CreatedAt.Unix(),
"updated_at": tool.UpdatedAt.Unix(),
}
if captainCustomToolShowAuthConfig(c) {
payload["auth_config"] = rawJSONValue(tool.AuthConfig)
}
return payload
}
func captainCustomToolShowAuthConfig(c *gin.Context) bool {
role, exists := c.Get("role")
return exists && (role == "administrator" || role == "super_admin")
}
@@ -79,6 +79,12 @@ func (s *CaptainCustomToolTestHandlerTestSuite) SetupSuite() {
// 设置路由
s.router = gin.New()
s.router.Use(func(c *gin.Context) {
if role := c.GetHeader("X-Test-Role"); role != "" {
c.Set("role", role)
}
c.Next()
})
accountsGroup := s.router.Group("/api/v1/accounts/:account_id")
{
captainGroup := accountsGroup.Group("/captain")
@@ -134,6 +140,25 @@ func (s *CaptainCustomToolTestHandlerTestSuite) TestTestTool_无效accountID() {
assert.Equal(s.T(), http.StatusBadRequest, w.Code)
}
func (s *CaptainCustomToolTestHandlerTestSuite) TestTestTool_AgentRoleForbidden() {
body := map[string]interface{}{
"tool_id": float64(s.tool.ID),
"params": map[string]interface{}{},
}
jsonBody, _ := json.Marshal(body)
w := httptest.NewRecorder()
req, _ := http.NewRequest("POST", "/api/v1/accounts/"+strconv.FormatUint(uint64(s.account.ID), 10)+"/captain/custom_tools/test", bytes.NewReader(jsonBody))
req.Header.Set("Content-Type", "application/json")
req.Header.Set("X-Test-Role", "agent")
s.router.ServeHTTP(w, req)
assert.Equal(s.T(), http.StatusForbidden, w.Code)
var resp map[string]interface{}
s.Require().NoError(json.Unmarshal(w.Body.Bytes(), &resp))
assert.Equal(s.T(), "You are not authorized to do this action", resp["error"])
}
func (s *CaptainCustomToolTestHandlerTestSuite) TestTestTool_缺少tool_id() {
body := map[string]interface{}{
"params": map[string]interface{}{},