65 lines
2.1 KiB
Go
65 lines
2.1 KiB
Go
|
package middleware
|
||
|
|
|
||
|
|
import (
|
||
|
|
"context"
|
||
|
|
"net/http"
|
||
|
|
"net/http/httptest"
|
||
|
|
"testing"
|
||
|
|
|
||
|
|
"github.com/Tencent/WeKnora/internal/accesspolicy"
|
||
|
|
"github.com/Tencent/WeKnora/internal/types"
|
||
|
|
"github.com/gin-gonic/gin"
|
||
|
|
"github.com/stretchr/testify/require"
|
||
|
|
)
|
||
|
|
|
||
|
|
func TestRequireKnowledgeExportFailsClosedWithoutResolvedKnowledgeBase(t *testing.T) {
|
||
|
|
gin.SetMode(gin.TestMode)
|
||
|
|
for _, tc := range []struct {
|
||
|
|
name string
|
||
|
|
access any
|
||
|
|
}{
|
||
|
|
{name: "absent"},
|
||
|
|
{name: "nil access", access: (*KBAccess)(nil)},
|
||
|
|
{name: "missing KB", access: &KBAccess{}},
|
||
|
|
{name: "invalid type", access: "kb-id"},
|
||
|
|
} {
|
||
|
|
t.Run(tc.name, func(t *testing.T) {
|
||
|
|
r := gin.New()
|
||
|
|
r.Use(ErrorHandler())
|
||
|
|
r.Use(func(c *gin.Context) {
|
||
|
|
ctx := accesspolicy.WithSnapshot(c.Request.Context(), &accesspolicy.Snapshot{
|
||
|
|
Publications: map[string]bool{"knowledge:kb-id": true},
|
||
|
|
})
|
||
|
|
c.Request = c.Request.WithContext(ctx)
|
||
|
|
if tc.access != nil {
|
||
|
|
c.Set(KBAccessContextKey, tc.access)
|
||
|
|
}
|
||
|
|
c.Next()
|
||
|
|
})
|
||
|
|
r.GET("/export", RequireKnowledgeExport(), func(c *gin.Context) {
|
||
|
|
t.Fatal("export handler ran without an authorized KB")
|
||
|
|
})
|
||
|
|
rec := httptest.NewRecorder()
|
||
|
|
r.ServeHTTP(rec, httptest.NewRequest(http.MethodGet, "/export", nil))
|
||
|
|
require.Equal(t, http.StatusServiceUnavailable, rec.Code)
|
||
|
|
})
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
func TestRequireKnowledgeExportPreservesSystemAdministrator(t *testing.T) {
|
||
|
|
gin.SetMode(gin.TestMode)
|
||
|
|
r := gin.New()
|
||
|
|
r.Use(func(c *gin.Context) {
|
||
|
|
ctx := context.WithValue(c.Request.Context(), types.SystemAdminContextKey, true)
|
||
|
|
ctx = context.WithValue(ctx, types.TenantRoleContextKey, types.TenantRoleViewer)
|
||
|
|
ctx = accesspolicy.WithSnapshot(ctx, &accesspolicy.Snapshot{Publications: map[string]bool{"knowledge:kb-id": true}})
|
||
|
|
c.Request = c.Request.WithContext(ctx)
|
||
|
|
c.Set(KBAccessContextKey, &KBAccess{KnowledgeBase: &types.KnowledgeBase{ID: "kb-id"}})
|
||
|
|
c.Next()
|
||
|
|
})
|
||
|
|
r.GET("/export", RequireKnowledgeExport(), func(c *gin.Context) { c.Status(http.StatusNoContent) })
|
||
|
|
rec := httptest.NewRecorder()
|
||
|
|
r.ServeHTTP(rec, httptest.NewRequest(http.MethodGet, "/export", nil))
|
||
|
|
require.Equal(t, http.StatusNoContent, rec.Code)
|
||
|
|
}
|