chore: import upstream snapshot with attribution
This commit is contained in:
@@ -0,0 +1,121 @@
|
||||
package auth
|
||||
|
||||
import (
|
||||
"sync"
|
||||
"time"
|
||||
|
||||
"pentagi/pkg/server/models"
|
||||
|
||||
"github.com/jinzhu/gorm"
|
||||
)
|
||||
|
||||
// tokenCacheEntry represents a cached token status entry
|
||||
type tokenCacheEntry struct {
|
||||
status models.TokenStatus
|
||||
privileges []string
|
||||
notFound bool // negative caching
|
||||
expiresAt time.Time
|
||||
}
|
||||
|
||||
// TokenCache provides caching for token status lookups
|
||||
type TokenCache struct {
|
||||
cache sync.Map
|
||||
ttl time.Duration
|
||||
db *gorm.DB
|
||||
}
|
||||
|
||||
// NewTokenCache creates a new token cache instance
|
||||
func NewTokenCache(db *gorm.DB) *TokenCache {
|
||||
return &TokenCache{
|
||||
ttl: 5 * time.Minute,
|
||||
db: db,
|
||||
}
|
||||
}
|
||||
|
||||
// SetTTL sets the TTL for the token cache
|
||||
func (tc *TokenCache) SetTTL(ttl time.Duration) {
|
||||
tc.ttl = ttl
|
||||
}
|
||||
|
||||
// GetStatus retrieves token status and privileges from cache or database
|
||||
func (tc *TokenCache) GetStatus(tokenID string) (models.TokenStatus, []string, error) {
|
||||
// check cache first
|
||||
if entry, ok := tc.cache.Load(tokenID); ok {
|
||||
cached := entry.(tokenCacheEntry)
|
||||
if time.Now().Before(cached.expiresAt) {
|
||||
// return cached "not found" error
|
||||
if cached.notFound {
|
||||
return "", nil, gorm.ErrRecordNotFound
|
||||
}
|
||||
return cached.status, cached.privileges, nil
|
||||
}
|
||||
// cache entry expired, remove it
|
||||
tc.cache.Delete(tokenID)
|
||||
}
|
||||
|
||||
// load from database with role privileges
|
||||
var token models.APIToken
|
||||
if err := tc.db.Where("token_id = ? AND deleted_at IS NULL", tokenID).First(&token).Error; err != nil {
|
||||
if gorm.IsRecordNotFoundError(err) {
|
||||
// cache negative result (token not found)
|
||||
tc.cache.Store(tokenID, tokenCacheEntry{
|
||||
notFound: true,
|
||||
expiresAt: time.Now().Add(tc.ttl),
|
||||
})
|
||||
return "", nil, gorm.ErrRecordNotFound
|
||||
}
|
||||
return "", nil, err
|
||||
}
|
||||
|
||||
// load privileges for the token's role
|
||||
var privileges []models.Privilege
|
||||
if err := tc.db.Where("role_id = ?", token.RoleID).Find(&privileges).Error; err != nil {
|
||||
return "", nil, err
|
||||
}
|
||||
|
||||
// extract privilege names
|
||||
privNames := make([]string, len(privileges))
|
||||
for i, priv := range privileges {
|
||||
privNames[i] = priv.Name
|
||||
}
|
||||
|
||||
// always add automation privilege for API tokens
|
||||
privNames = append(privNames, PrivilegeAutomation)
|
||||
|
||||
// update cache with positive result
|
||||
tc.cache.Store(tokenID, tokenCacheEntry{
|
||||
status: token.Status,
|
||||
privileges: privNames,
|
||||
notFound: false,
|
||||
expiresAt: time.Now().Add(tc.ttl),
|
||||
})
|
||||
|
||||
return token.Status, privNames, nil
|
||||
}
|
||||
|
||||
// Invalidate removes a specific token from cache
|
||||
func (tc *TokenCache) Invalidate(tokenID string) {
|
||||
tc.cache.Delete(tokenID)
|
||||
}
|
||||
|
||||
// InvalidateUser removes all tokens for a specific user from cache
|
||||
func (tc *TokenCache) InvalidateUser(userID uint64) {
|
||||
// load all tokens for this user
|
||||
var tokens []models.APIToken
|
||||
if err := tc.db.Where("user_id = ? AND deleted_at IS NULL", userID).Find(&tokens).Error; err != nil {
|
||||
return
|
||||
}
|
||||
|
||||
// invalidate each token in cache
|
||||
for _, token := range tokens {
|
||||
tc.cache.Delete(token.TokenID)
|
||||
}
|
||||
}
|
||||
|
||||
// InvalidateAll clears the entire cache
|
||||
func (tc *TokenCache) InvalidateAll() {
|
||||
tc.cache.Range(func(key, value any) bool {
|
||||
tc.cache.Delete(key)
|
||||
return true
|
||||
})
|
||||
}
|
||||
@@ -0,0 +1,322 @@
|
||||
package auth_test
|
||||
|
||||
import (
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"pentagi/pkg/server/auth"
|
||||
"pentagi/pkg/server/models"
|
||||
|
||||
"github.com/jinzhu/gorm"
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
func TestTokenCache_GetStatus(t *testing.T) {
|
||||
db := setupTestDB(t)
|
||||
defer db.Close()
|
||||
|
||||
cache := auth.NewTokenCache(db)
|
||||
tokenID := "testtoken1"
|
||||
|
||||
// Insert test token
|
||||
token := models.APIToken{
|
||||
TokenID: tokenID,
|
||||
UserID: 1,
|
||||
RoleID: 2,
|
||||
TTL: 3600,
|
||||
Status: models.TokenStatusActive,
|
||||
}
|
||||
err := db.Create(&token).Error
|
||||
require.NoError(t, err)
|
||||
|
||||
// Test: Get status (should hit database)
|
||||
status, privileges, err := cache.GetStatus(tokenID)
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, models.TokenStatusActive, status)
|
||||
assert.NotEmpty(t, privileges)
|
||||
assert.Contains(t, privileges, auth.PrivilegeAutomation)
|
||||
assert.Contains(t, privileges, "flows.create")
|
||||
assert.Contains(t, privileges, "settings.tokens.view")
|
||||
|
||||
// Test: Get status again (should hit cache)
|
||||
status, privileges, err = cache.GetStatus(tokenID)
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, models.TokenStatusActive, status)
|
||||
assert.NotEmpty(t, privileges)
|
||||
assert.Contains(t, privileges, auth.PrivilegeAutomation)
|
||||
|
||||
// Test: Non-existent token
|
||||
_, _, err = cache.GetStatus("nonexistent")
|
||||
assert.Error(t, err)
|
||||
assert.Equal(t, gorm.ErrRecordNotFound, err)
|
||||
}
|
||||
|
||||
func TestTokenCache_Invalidate(t *testing.T) {
|
||||
db := setupTestDB(t)
|
||||
defer db.Close()
|
||||
|
||||
cache := auth.NewTokenCache(db)
|
||||
tokenID := "testtoken2"
|
||||
|
||||
// Insert test token
|
||||
token := models.APIToken{
|
||||
TokenID: tokenID,
|
||||
UserID: 1,
|
||||
RoleID: 2,
|
||||
TTL: 3600,
|
||||
Status: models.TokenStatusActive,
|
||||
}
|
||||
err := db.Create(&token).Error
|
||||
require.NoError(t, err)
|
||||
|
||||
// Get status to populate cache
|
||||
status, privileges, err := cache.GetStatus(tokenID)
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, models.TokenStatusActive, status)
|
||||
assert.NotEmpty(t, privileges)
|
||||
|
||||
// Update token in database
|
||||
db.Model(&token).Update("status", models.TokenStatusRevoked)
|
||||
|
||||
// Status should still be active (from cache)
|
||||
status, privileges, err = cache.GetStatus(tokenID)
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, models.TokenStatusActive, status)
|
||||
assert.NotEmpty(t, privileges)
|
||||
|
||||
// Invalidate cache
|
||||
cache.Invalidate(tokenID)
|
||||
|
||||
// Status should now be revoked (from database)
|
||||
status, privileges, err = cache.GetStatus(tokenID)
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, models.TokenStatusRevoked, status)
|
||||
assert.NotEmpty(t, privileges)
|
||||
}
|
||||
|
||||
func TestTokenCache_InvalidateUser(t *testing.T) {
|
||||
db := setupTestDB(t)
|
||||
defer db.Close()
|
||||
|
||||
cache := auth.NewTokenCache(db)
|
||||
userID := uint64(1)
|
||||
|
||||
// Insert multiple tokens for user
|
||||
tokens := []models.APIToken{
|
||||
{
|
||||
TokenID: "token1",
|
||||
UserID: userID,
|
||||
RoleID: 2,
|
||||
TTL: 3600,
|
||||
Status: models.TokenStatusActive,
|
||||
},
|
||||
{
|
||||
TokenID: "token2",
|
||||
UserID: userID,
|
||||
RoleID: 2,
|
||||
TTL: 3600,
|
||||
Status: models.TokenStatusActive,
|
||||
},
|
||||
}
|
||||
|
||||
for _, token := range tokens {
|
||||
err := db.Create(&token).Error
|
||||
require.NoError(t, err)
|
||||
}
|
||||
|
||||
// Populate cache
|
||||
for _, token := range tokens {
|
||||
_, _, err := cache.GetStatus(token.TokenID)
|
||||
require.NoError(t, err)
|
||||
}
|
||||
|
||||
// Update tokens in database
|
||||
db.Model(&models.APIToken{}).Where("user_id = ?", userID).Update("status", models.TokenStatusRevoked)
|
||||
|
||||
// Invalidate all user tokens
|
||||
cache.InvalidateUser(userID)
|
||||
|
||||
// All tokens should now show revoked status
|
||||
for _, token := range tokens {
|
||||
status, privileges, err := cache.GetStatus(token.TokenID)
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, models.TokenStatusRevoked, status)
|
||||
assert.NotEmpty(t, privileges)
|
||||
}
|
||||
}
|
||||
|
||||
func TestTokenCache_Expiration(t *testing.T) {
|
||||
db := setupTestDB(t)
|
||||
defer db.Close()
|
||||
|
||||
// Create cache with very short TTL for testing
|
||||
cache := auth.NewTokenCache(db)
|
||||
cache.SetTTL(300 * time.Millisecond)
|
||||
|
||||
tokenID := "testtoken3"
|
||||
|
||||
// Insert test token
|
||||
token := models.APIToken{
|
||||
TokenID: tokenID,
|
||||
UserID: 1,
|
||||
RoleID: 2,
|
||||
TTL: 3600,
|
||||
Status: models.TokenStatusActive,
|
||||
}
|
||||
err := db.Create(&token).Error
|
||||
require.NoError(t, err)
|
||||
|
||||
// Get status to populate cache
|
||||
status, privileges, err := cache.GetStatus(tokenID)
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, models.TokenStatusActive, status)
|
||||
assert.NotEmpty(t, privileges)
|
||||
|
||||
// Update token in database
|
||||
db.Model(&token).Update("status", models.TokenStatusRevoked)
|
||||
|
||||
// Wait for cache to expire
|
||||
time.Sleep(500 * time.Millisecond)
|
||||
|
||||
// Status should now be revoked (cache expired, reading from DB)
|
||||
status, privileges, err = cache.GetStatus(tokenID)
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, models.TokenStatusRevoked, status)
|
||||
assert.NotEmpty(t, privileges)
|
||||
}
|
||||
|
||||
func TestTokenCache_PrivilegesByRole(t *testing.T) {
|
||||
db := setupTestDB(t)
|
||||
defer db.Close()
|
||||
|
||||
cache := auth.NewTokenCache(db)
|
||||
|
||||
// Test Admin token (role_id = 1)
|
||||
adminTokenID := "admin_token"
|
||||
adminToken := models.APIToken{
|
||||
TokenID: adminTokenID,
|
||||
UserID: 1,
|
||||
RoleID: 1,
|
||||
TTL: 3600,
|
||||
Status: models.TokenStatusActive,
|
||||
}
|
||||
err := db.Create(&adminToken).Error
|
||||
require.NoError(t, err)
|
||||
|
||||
status, adminPrivs, err := cache.GetStatus(adminTokenID)
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, models.TokenStatusActive, status)
|
||||
assert.NotEmpty(t, adminPrivs)
|
||||
assert.Contains(t, adminPrivs, auth.PrivilegeAutomation)
|
||||
assert.Contains(t, adminPrivs, "users.create")
|
||||
assert.Contains(t, adminPrivs, "users.delete")
|
||||
assert.Contains(t, adminPrivs, "settings.tokens.admin")
|
||||
|
||||
// Test User token (role_id = 2)
|
||||
userTokenID := "user_token"
|
||||
userToken := models.APIToken{
|
||||
TokenID: userTokenID,
|
||||
UserID: 2,
|
||||
RoleID: 2,
|
||||
TTL: 3600,
|
||||
Status: models.TokenStatusActive,
|
||||
}
|
||||
err = db.Create(&userToken).Error
|
||||
require.NoError(t, err)
|
||||
|
||||
status, userPrivs, err := cache.GetStatus(userTokenID)
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, models.TokenStatusActive, status)
|
||||
assert.NotEmpty(t, userPrivs)
|
||||
assert.Contains(t, userPrivs, auth.PrivilegeAutomation)
|
||||
assert.Contains(t, userPrivs, "flows.create")
|
||||
assert.Contains(t, userPrivs, "settings.tokens.view")
|
||||
|
||||
// User should NOT have admin privileges
|
||||
assert.NotContains(t, userPrivs, "users.create")
|
||||
assert.NotContains(t, userPrivs, "users.delete")
|
||||
assert.NotContains(t, userPrivs, "settings.tokens.admin")
|
||||
|
||||
// Admin should have more privileges than User
|
||||
assert.Greater(t, len(adminPrivs), len(userPrivs))
|
||||
}
|
||||
|
||||
func TestTokenCache_NegativeCaching(t *testing.T) {
|
||||
db := setupTestDB(t)
|
||||
defer db.Close()
|
||||
|
||||
cache := auth.NewTokenCache(db)
|
||||
nonExistentTokenID := "nonexistent"
|
||||
|
||||
// First call - should hit database and cache the "not found"
|
||||
_, _, err := cache.GetStatus(nonExistentTokenID)
|
||||
require.Error(t, err)
|
||||
assert.Equal(t, gorm.ErrRecordNotFound, err)
|
||||
|
||||
// Second call - should return from cache without hitting DB
|
||||
// We can verify this by checking error is still the same
|
||||
_, _, err = cache.GetStatus(nonExistentTokenID)
|
||||
require.Error(t, err)
|
||||
assert.Equal(t, gorm.ErrRecordNotFound, err, "Should return cached not found error")
|
||||
|
||||
// Now create the token in DB
|
||||
token := models.APIToken{
|
||||
TokenID: nonExistentTokenID,
|
||||
UserID: 1,
|
||||
RoleID: 2,
|
||||
TTL: 3600,
|
||||
Status: models.TokenStatusActive,
|
||||
}
|
||||
err = db.Create(&token).Error
|
||||
require.NoError(t, err)
|
||||
|
||||
// Should still return cached "not found" until invalidated
|
||||
_, _, err = cache.GetStatus(nonExistentTokenID)
|
||||
require.Error(t, err)
|
||||
assert.Equal(t, gorm.ErrRecordNotFound, err, "Should still return cached not found")
|
||||
|
||||
// Invalidate cache
|
||||
cache.Invalidate(nonExistentTokenID)
|
||||
|
||||
// Now should find the token
|
||||
status, privileges, err := cache.GetStatus(nonExistentTokenID)
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, models.TokenStatusActive, status)
|
||||
assert.NotEmpty(t, privileges)
|
||||
}
|
||||
|
||||
func TestTokenCache_NegativeCachingExpiration(t *testing.T) {
|
||||
db := setupTestDB(t)
|
||||
defer db.Close()
|
||||
|
||||
cache := auth.NewTokenCache(db)
|
||||
cache.SetTTL(300 * time.Millisecond)
|
||||
|
||||
nonExistentTokenID := "temp_nonexistent"
|
||||
|
||||
// First call - cache the "not found"
|
||||
_, _, err := cache.GetStatus(nonExistentTokenID)
|
||||
require.Error(t, err)
|
||||
assert.Equal(t, gorm.ErrRecordNotFound, err)
|
||||
|
||||
// Create token in DB
|
||||
token := models.APIToken{
|
||||
TokenID: nonExistentTokenID,
|
||||
UserID: 1,
|
||||
RoleID: 2,
|
||||
TTL: 3600,
|
||||
Status: models.TokenStatusActive,
|
||||
}
|
||||
err = db.Create(&token).Error
|
||||
require.NoError(t, err)
|
||||
|
||||
// Wait for cache to expire
|
||||
time.Sleep(500 * time.Millisecond)
|
||||
|
||||
// Now should find the token (cache expired)
|
||||
status, privileges, err := cache.GetStatus(nonExistentTokenID)
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, models.TokenStatusActive, status)
|
||||
assert.NotEmpty(t, privileges)
|
||||
}
|
||||
@@ -0,0 +1,28 @@
|
||||
package auth
|
||||
|
||||
import (
|
||||
"crypto/rand"
|
||||
"fmt"
|
||||
"math/big"
|
||||
)
|
||||
|
||||
const (
|
||||
Base62Chars = "0123456789ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz"
|
||||
TokenIDLength = 10
|
||||
)
|
||||
|
||||
// GenerateTokenID generates a random base62 string of specified length
|
||||
func GenerateTokenID() (string, error) {
|
||||
b := make([]byte, TokenIDLength)
|
||||
maxIdx := big.NewInt(int64(len(Base62Chars)))
|
||||
|
||||
for i := range b {
|
||||
idx, err := rand.Int(rand.Reader, maxIdx)
|
||||
if err != nil {
|
||||
return "", fmt.Errorf("error generating token ID: %w", err)
|
||||
}
|
||||
b[i] = Base62Chars[idx.Int64()]
|
||||
}
|
||||
|
||||
return string(b), nil
|
||||
}
|
||||
@@ -0,0 +1,49 @@
|
||||
package auth_test
|
||||
|
||||
import (
|
||||
"testing"
|
||||
|
||||
"pentagi/pkg/server/auth"
|
||||
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
func TestGenerateTokenID(t *testing.T) {
|
||||
// Test basic generation
|
||||
tokenID, err := auth.GenerateTokenID()
|
||||
require.NoError(t, err)
|
||||
assert.Len(t, tokenID, auth.TokenIDLength, "Token ID should have correct length")
|
||||
|
||||
// Test that all characters are from base62 charset
|
||||
for _, char := range tokenID {
|
||||
assert.Contains(t, auth.Base62Chars, string(char), "Token ID should only contain base62 characters")
|
||||
}
|
||||
|
||||
// Test uniqueness (generate multiple tokens and check they're different)
|
||||
tokens := make(map[string]bool)
|
||||
for i := 0; i < 100; i++ {
|
||||
token, err := auth.GenerateTokenID()
|
||||
require.NoError(t, err)
|
||||
assert.Len(t, token, auth.TokenIDLength)
|
||||
assert.False(t, tokens[token], "Generated tokens should be unique")
|
||||
tokens[token] = true
|
||||
}
|
||||
}
|
||||
|
||||
func TestGenerateTokenIDFormat(t *testing.T) {
|
||||
// Test that token IDs match expected format
|
||||
tokenID, err := auth.GenerateTokenID()
|
||||
require.NoError(t, err)
|
||||
|
||||
// Should be exactly 10 characters
|
||||
assert.Equal(t, 10, len(tokenID))
|
||||
|
||||
// Should only contain alphanumeric characters
|
||||
for _, char := range tokenID {
|
||||
isValid := (char >= '0' && char <= '9') ||
|
||||
(char >= 'A' && char <= 'Z') ||
|
||||
(char >= 'a' && char <= 'z')
|
||||
assert.True(t, isValid, "Character %c should be alphanumeric", char)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,62 @@
|
||||
package auth
|
||||
|
||||
import (
|
||||
"errors"
|
||||
"fmt"
|
||||
"time"
|
||||
|
||||
"pentagi/pkg/server/models"
|
||||
|
||||
"github.com/golang-jwt/jwt/v5"
|
||||
)
|
||||
|
||||
func MakeAPIToken(globalSalt string, claims jwt.Claims) (string, error) {
|
||||
token := jwt.NewWithClaims(jwt.SigningMethodHS256, claims)
|
||||
tokenString, err := token.SignedString(MakeJWTSigningKey(globalSalt))
|
||||
if err != nil {
|
||||
return "", fmt.Errorf("failed to sign token: %w", err)
|
||||
}
|
||||
|
||||
return tokenString, nil
|
||||
}
|
||||
|
||||
func MakeAPITokenClaims(tokenID, uhash string, uid, rid, ttl uint64) jwt.Claims {
|
||||
now := time.Now()
|
||||
return models.APITokenClaims{
|
||||
TokenID: tokenID,
|
||||
RID: rid,
|
||||
UID: uid,
|
||||
UHASH: uhash,
|
||||
RegisteredClaims: jwt.RegisteredClaims{
|
||||
ExpiresAt: jwt.NewNumericDate(now.Add(time.Duration(ttl) * time.Second)),
|
||||
IssuedAt: jwt.NewNumericDate(now),
|
||||
Subject: "api_token",
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
func ValidateAPIToken(tokenString, globalSalt string) (*models.APITokenClaims, error) {
|
||||
var claims models.APITokenClaims
|
||||
token, err := jwt.ParseWithClaims(tokenString, &claims, func(token *jwt.Token) (any, error) {
|
||||
// verify signing algorithm to prevent "alg: none"
|
||||
if _, ok := token.Method.(*jwt.SigningMethodHMAC); !ok {
|
||||
return nil, fmt.Errorf("unexpected signing method: %v", token.Header["alg"])
|
||||
}
|
||||
return MakeJWTSigningKey(globalSalt), nil
|
||||
})
|
||||
if err != nil {
|
||||
if errors.Is(err, jwt.ErrTokenMalformed) {
|
||||
return nil, fmt.Errorf("token is malformed")
|
||||
} else if errors.Is(err, jwt.ErrTokenExpired) || errors.Is(err, jwt.ErrTokenNotValidYet) {
|
||||
return nil, fmt.Errorf("token is either expired or not active yet")
|
||||
} else {
|
||||
return nil, fmt.Errorf("token invalid: %w", err)
|
||||
}
|
||||
}
|
||||
|
||||
if !token.Valid {
|
||||
return nil, fmt.Errorf("token is invalid")
|
||||
}
|
||||
|
||||
return &claims, nil
|
||||
}
|
||||
@@ -0,0 +1,479 @@
|
||||
package auth_test
|
||||
|
||||
import (
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"pentagi/pkg/server/auth"
|
||||
"pentagi/pkg/server/models"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/golang-jwt/jwt/v5"
|
||||
"github.com/jinzhu/gorm"
|
||||
_ "github.com/jinzhu/gorm/dialects/sqlite"
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
// setupTestDB creates an in-memory SQLite database for testing
|
||||
func setupTestDB(t *testing.T) *gorm.DB {
|
||||
t.Helper()
|
||||
db, err := gorm.Open("sqlite3", ":memory:")
|
||||
require.NoError(t, err)
|
||||
|
||||
// Create roles table
|
||||
result := db.Exec(`
|
||||
CREATE TABLE roles (
|
||||
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
||||
name TEXT NOT NULL UNIQUE
|
||||
)
|
||||
`)
|
||||
require.NoError(t, result.Error, "Failed to create roles table")
|
||||
|
||||
// Create privileges table
|
||||
result = db.Exec(`
|
||||
CREATE TABLE privileges (
|
||||
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
||||
role_id INTEGER NOT NULL,
|
||||
name TEXT NOT NULL,
|
||||
UNIQUE(role_id, name)
|
||||
)
|
||||
`)
|
||||
require.NoError(t, result.Error, "Failed to create privileges table")
|
||||
|
||||
// Create api_tokens table for testing
|
||||
result = db.Exec(`
|
||||
CREATE TABLE api_tokens (
|
||||
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
||||
token_id TEXT NOT NULL UNIQUE,
|
||||
user_id INTEGER NOT NULL,
|
||||
role_id INTEGER NOT NULL,
|
||||
name TEXT,
|
||||
ttl INTEGER NOT NULL,
|
||||
status TEXT NOT NULL DEFAULT 'active',
|
||||
created_at DATETIME DEFAULT CURRENT_TIMESTAMP,
|
||||
updated_at DATETIME DEFAULT CURRENT_TIMESTAMP,
|
||||
deleted_at DATETIME
|
||||
)
|
||||
`)
|
||||
require.NoError(t, result.Error, "Failed to create api_tokens table")
|
||||
|
||||
// Insert test roles
|
||||
db.Exec("INSERT INTO roles (id, name) VALUES (1, 'Admin'), (2, 'User')")
|
||||
|
||||
// Insert test privileges for Admin role
|
||||
db.Exec(`INSERT INTO privileges (role_id, name) VALUES
|
||||
(1, 'users.create'),
|
||||
(1, 'users.delete'),
|
||||
(1, 'users.edit'),
|
||||
(1, 'users.view'),
|
||||
(1, 'roles.view'),
|
||||
(1, 'flows.admin'),
|
||||
(1, 'flows.create'),
|
||||
(1, 'flows.delete'),
|
||||
(1, 'flows.edit'),
|
||||
(1, 'flows.view'),
|
||||
(1, 'settings.tokens.create'),
|
||||
(1, 'settings.tokens.view'),
|
||||
(1, 'settings.tokens.edit'),
|
||||
(1, 'settings.tokens.delete'),
|
||||
(1, 'settings.tokens.admin')`)
|
||||
|
||||
// Insert test privileges for User role
|
||||
db.Exec(`INSERT INTO privileges (role_id, name) VALUES
|
||||
(2, 'roles.view'),
|
||||
(2, 'flows.create'),
|
||||
(2, 'flows.delete'),
|
||||
(2, 'flows.edit'),
|
||||
(2, 'flows.view'),
|
||||
(2, 'settings.tokens.create'),
|
||||
(2, 'settings.tokens.view'),
|
||||
(2, 'settings.tokens.edit'),
|
||||
(2, 'settings.tokens.delete')`)
|
||||
|
||||
// Create users table
|
||||
db.Exec(`
|
||||
CREATE TABLE users (
|
||||
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
||||
hash TEXT NOT NULL UNIQUE,
|
||||
type TEXT NOT NULL DEFAULT 'local',
|
||||
mail TEXT NOT NULL UNIQUE,
|
||||
name TEXT NOT NULL DEFAULT '',
|
||||
status TEXT NOT NULL DEFAULT 'active',
|
||||
role_id INTEGER NOT NULL DEFAULT 2,
|
||||
password TEXT,
|
||||
password_change_required BOOLEAN NOT NULL DEFAULT false,
|
||||
provider TEXT,
|
||||
created_at DATETIME DEFAULT CURRENT_TIMESTAMP,
|
||||
deleted_at DATETIME
|
||||
)
|
||||
`)
|
||||
|
||||
// Insert test users
|
||||
db.Exec("INSERT INTO users (id, hash, mail, name, status, role_id) VALUES (1, 'testhash', 'user1@test.com', 'User 1', 'active', 2)")
|
||||
db.Exec("INSERT INTO users (id, hash, mail, name, status, role_id) VALUES (2, 'testhash2', 'user2@test.com', 'User 2', 'active', 2)")
|
||||
|
||||
// Create user_preferences table
|
||||
db.Exec(`
|
||||
CREATE TABLE user_preferences (
|
||||
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
||||
user_id INTEGER NOT NULL UNIQUE,
|
||||
preferences TEXT NOT NULL DEFAULT '{"favoriteFlows": []}',
|
||||
created_at DATETIME DEFAULT CURRENT_TIMESTAMP,
|
||||
updated_at DATETIME DEFAULT CURRENT_TIMESTAMP,
|
||||
FOREIGN KEY (user_id) REFERENCES users(id) ON DELETE CASCADE
|
||||
)
|
||||
`)
|
||||
|
||||
// Insert preferences for test users
|
||||
db.Exec("INSERT INTO user_preferences (user_id, preferences) VALUES (1, '{\"favoriteFlows\": []}')")
|
||||
db.Exec("INSERT INTO user_preferences (user_id, preferences) VALUES (2, '{\"favoriteFlows\": []}')")
|
||||
|
||||
time.Sleep(200 * time.Millisecond) // wait for database to be ready
|
||||
|
||||
return db
|
||||
}
|
||||
|
||||
func TestValidateAPIToken(t *testing.T) {
|
||||
globalSalt := "test_salt"
|
||||
|
||||
testCases := []struct {
|
||||
name string
|
||||
setup func() string
|
||||
expectError bool
|
||||
errorMsg string
|
||||
}{
|
||||
{
|
||||
name: "valid token",
|
||||
setup: func() string {
|
||||
claims := models.APITokenClaims{
|
||||
TokenID: "abc123xyz9",
|
||||
RID: 2,
|
||||
UID: 1,
|
||||
UHASH: "testhash",
|
||||
RegisteredClaims: jwt.RegisteredClaims{
|
||||
ExpiresAt: jwt.NewNumericDate(time.Now().Add(1 * time.Hour)),
|
||||
IssuedAt: jwt.NewNumericDate(time.Now()),
|
||||
Subject: "api_token",
|
||||
},
|
||||
}
|
||||
token := jwt.NewWithClaims(jwt.SigningMethodHS256, claims)
|
||||
tokenString, _ := token.SignedString(auth.MakeJWTSigningKey(globalSalt))
|
||||
return tokenString
|
||||
},
|
||||
expectError: false,
|
||||
},
|
||||
{
|
||||
name: "expired token",
|
||||
setup: func() string {
|
||||
claims := models.APITokenClaims{
|
||||
TokenID: "abc123xyz9",
|
||||
RID: 2,
|
||||
UID: 1,
|
||||
UHASH: "testhash",
|
||||
RegisteredClaims: jwt.RegisteredClaims{
|
||||
ExpiresAt: jwt.NewNumericDate(time.Now().Add(-1 * time.Hour)),
|
||||
IssuedAt: jwt.NewNumericDate(time.Now().Add(-2 * time.Hour)),
|
||||
Subject: "api_token",
|
||||
},
|
||||
}
|
||||
token := jwt.NewWithClaims(jwt.SigningMethodHS256, claims)
|
||||
tokenString, _ := token.SignedString(auth.MakeJWTSigningKey(globalSalt))
|
||||
return tokenString
|
||||
},
|
||||
expectError: true,
|
||||
errorMsg: "expired",
|
||||
},
|
||||
{
|
||||
name: "invalid signature",
|
||||
setup: func() string {
|
||||
claims := models.APITokenClaims{
|
||||
TokenID: "abc123xyz9",
|
||||
RID: 2,
|
||||
UID: 1,
|
||||
UHASH: "testhash",
|
||||
RegisteredClaims: jwt.RegisteredClaims{
|
||||
ExpiresAt: jwt.NewNumericDate(time.Now().Add(1 * time.Hour)),
|
||||
IssuedAt: jwt.NewNumericDate(time.Now()),
|
||||
Subject: "api_token",
|
||||
},
|
||||
}
|
||||
token := jwt.NewWithClaims(jwt.SigningMethodHS256, claims)
|
||||
tokenString, _ := token.SignedString([]byte("wrong_key"))
|
||||
return tokenString
|
||||
},
|
||||
expectError: true,
|
||||
errorMsg: "invalid",
|
||||
},
|
||||
{
|
||||
name: "malformed token",
|
||||
setup: func() string {
|
||||
return "not.a.valid.jwt.token"
|
||||
},
|
||||
expectError: true,
|
||||
errorMsg: "malformed",
|
||||
},
|
||||
}
|
||||
|
||||
for _, tc := range testCases {
|
||||
t.Run(tc.name, func(t *testing.T) {
|
||||
tokenString := tc.setup()
|
||||
claims, err := auth.ValidateAPIToken(tokenString, globalSalt)
|
||||
|
||||
if tc.expectError {
|
||||
assert.Error(t, err)
|
||||
if tc.errorMsg != "" {
|
||||
assert.Contains(t, err.Error(), tc.errorMsg)
|
||||
}
|
||||
assert.Nil(t, claims)
|
||||
} else {
|
||||
assert.NoError(t, err)
|
||||
assert.NotNil(t, claims)
|
||||
assert.Equal(t, "abc123xyz9", claims.TokenID)
|
||||
assert.Equal(t, uint64(1), claims.UID)
|
||||
assert.Equal(t, uint64(2), claims.RID)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestAPITokenAuthentication_CacheExpiration(t *testing.T) {
|
||||
db := setupTestDB(t)
|
||||
defer db.Close()
|
||||
|
||||
// Create cache with short TTL for testing
|
||||
tokenCache := auth.NewTokenCache(db)
|
||||
tokenCache.SetTTL(100 * time.Millisecond)
|
||||
userCache := auth.NewUserCache(db)
|
||||
authMiddleware := auth.NewAuthMiddleware("/base/url", "test", tokenCache, userCache)
|
||||
|
||||
// Create active token
|
||||
tokenID, err := auth.GenerateTokenID()
|
||||
require.NoError(t, err)
|
||||
apiToken := models.APIToken{
|
||||
TokenID: tokenID,
|
||||
UserID: 1,
|
||||
RoleID: 2,
|
||||
TTL: 3600,
|
||||
Status: models.TokenStatusActive,
|
||||
}
|
||||
err = db.Create(&apiToken).Error
|
||||
require.NoError(t, err)
|
||||
|
||||
// Create JWT
|
||||
claims := models.APITokenClaims{
|
||||
TokenID: tokenID,
|
||||
RID: 2,
|
||||
UID: 1,
|
||||
UHASH: "testhash",
|
||||
RegisteredClaims: jwt.RegisteredClaims{
|
||||
ExpiresAt: jwt.NewNumericDate(time.Now().Add(1 * time.Hour)),
|
||||
IssuedAt: jwt.NewNumericDate(time.Now()),
|
||||
Subject: "api_token",
|
||||
},
|
||||
}
|
||||
token := jwt.NewWithClaims(jwt.SigningMethodHS256, claims)
|
||||
tokenString, err := token.SignedString(auth.MakeJWTSigningKey("test"))
|
||||
require.NoError(t, err)
|
||||
|
||||
server := newTestServer(t, "/test", db, authMiddleware.AuthTokenRequired)
|
||||
defer server.Close()
|
||||
|
||||
// First call: should work (status active, cached)
|
||||
assert.True(t, server.CallAndGetStatus(t, "Bearer "+tokenString))
|
||||
|
||||
// Revoke token in DB
|
||||
db.Model(&apiToken).Update("status", models.TokenStatusRevoked)
|
||||
|
||||
// Second call: should still work (cache not expired)
|
||||
assert.True(t, server.CallAndGetStatus(t, "Bearer "+tokenString))
|
||||
|
||||
// Wait for cache to expire
|
||||
time.Sleep(150 * time.Millisecond)
|
||||
|
||||
// Third call: should fail (cache expired, reads from DB)
|
||||
assert.False(t, server.CallAndGetStatus(t, "Bearer "+tokenString))
|
||||
}
|
||||
|
||||
func TestAPITokenAuthentication_DefaultSalt(t *testing.T) {
|
||||
db := setupTestDB(t)
|
||||
defer db.Close()
|
||||
|
||||
testCases := []struct {
|
||||
name string
|
||||
globalSalt string
|
||||
shouldSkip bool
|
||||
}{
|
||||
{
|
||||
name: "default salt 'salt'",
|
||||
globalSalt: "salt",
|
||||
shouldSkip: true,
|
||||
},
|
||||
{
|
||||
name: "empty salt",
|
||||
globalSalt: "",
|
||||
shouldSkip: true,
|
||||
},
|
||||
{
|
||||
name: "custom salt",
|
||||
globalSalt: "custom_secure_salt",
|
||||
shouldSkip: false,
|
||||
},
|
||||
}
|
||||
|
||||
for _, tc := range testCases {
|
||||
t.Run(tc.name, func(t *testing.T) {
|
||||
tokenCache := auth.NewTokenCache(db)
|
||||
userCache := auth.NewUserCache(db)
|
||||
authMiddleware := auth.NewAuthMiddleware("/base/url", tc.globalSalt, tokenCache, userCache)
|
||||
|
||||
// Create a token (even with default salt, for testing)
|
||||
tokenID, err := auth.GenerateTokenID()
|
||||
require.NoError(t, err)
|
||||
claims := models.APITokenClaims{
|
||||
TokenID: tokenID,
|
||||
RID: 2,
|
||||
UID: 1,
|
||||
UHASH: "testhash",
|
||||
RegisteredClaims: jwt.RegisteredClaims{
|
||||
ExpiresAt: jwt.NewNumericDate(time.Now().Add(1 * time.Hour)),
|
||||
IssuedAt: jwt.NewNumericDate(time.Now()),
|
||||
Subject: "api_token",
|
||||
},
|
||||
}
|
||||
token := jwt.NewWithClaims(jwt.SigningMethodHS256, claims)
|
||||
tokenString, _ := token.SignedString(auth.MakeJWTSigningKey(tc.globalSalt))
|
||||
|
||||
server := newTestServer(t, "/test", db, authMiddleware.AuthTokenRequired)
|
||||
defer server.Close()
|
||||
|
||||
// With default salt, token validation should be skipped
|
||||
result := server.CallAndGetStatus(t, "Bearer "+tokenString)
|
||||
|
||||
if tc.shouldSkip {
|
||||
// Should skip token auth and try cookie (which will fail)
|
||||
assert.False(t, result)
|
||||
} else {
|
||||
// With custom salt but no DB record, should fail with "not found"
|
||||
assert.False(t, result)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestAPITokenAuthentication_SoftDelete(t *testing.T) {
|
||||
db := setupTestDB(t)
|
||||
defer db.Close()
|
||||
|
||||
tokenCache := auth.NewTokenCache(db)
|
||||
userCache := auth.NewUserCache(db)
|
||||
authMiddleware := auth.NewAuthMiddleware("/base/url", "test", tokenCache, userCache)
|
||||
|
||||
// Create token
|
||||
tokenID, err := auth.GenerateTokenID()
|
||||
require.NoError(t, err)
|
||||
apiToken := models.APIToken{
|
||||
TokenID: tokenID,
|
||||
UserID: 1,
|
||||
RoleID: 2,
|
||||
TTL: 3600,
|
||||
Status: models.TokenStatusActive,
|
||||
}
|
||||
err = db.Create(&apiToken).Error
|
||||
require.NoError(t, err)
|
||||
|
||||
// Create JWT
|
||||
claims := models.APITokenClaims{
|
||||
TokenID: tokenID,
|
||||
RID: 2,
|
||||
UID: 1,
|
||||
UHASH: "testhash",
|
||||
RegisteredClaims: jwt.RegisteredClaims{
|
||||
ExpiresAt: jwt.NewNumericDate(time.Now().Add(1 * time.Hour)),
|
||||
IssuedAt: jwt.NewNumericDate(time.Now()),
|
||||
Subject: "api_token",
|
||||
},
|
||||
}
|
||||
token := jwt.NewWithClaims(jwt.SigningMethodHS256, claims)
|
||||
tokenString, err := token.SignedString(auth.MakeJWTSigningKey("test"))
|
||||
require.NoError(t, err)
|
||||
|
||||
server := newTestServer(t, "/test", db, authMiddleware.AuthTokenRequired)
|
||||
defer server.Close()
|
||||
|
||||
// Should work initially
|
||||
assert.True(t, server.CallAndGetStatus(t, "Bearer "+tokenString))
|
||||
|
||||
// Soft delete
|
||||
now := time.Now()
|
||||
db.Model(&apiToken).Update("deleted_at", now)
|
||||
tokenCache.Invalidate(tokenID)
|
||||
|
||||
// Should fail after soft delete
|
||||
assert.False(t, server.CallAndGetStatus(t, "Bearer "+tokenString))
|
||||
}
|
||||
|
||||
func TestAPITokenAuthentication_AlgNoneAttack(t *testing.T) {
|
||||
db := setupTestDB(t)
|
||||
defer db.Close()
|
||||
|
||||
tokenCache := auth.NewTokenCache(db)
|
||||
userCache := auth.NewUserCache(db)
|
||||
authMiddleware := auth.NewAuthMiddleware("/base/url", "test", tokenCache, userCache)
|
||||
|
||||
tokenID, err := auth.GenerateTokenID()
|
||||
require.NoError(t, err)
|
||||
|
||||
// Create token with "none" algorithm (security attack)
|
||||
claims := models.APITokenClaims{
|
||||
TokenID: tokenID,
|
||||
RID: 2,
|
||||
UID: 1,
|
||||
UHASH: "testhash",
|
||||
RegisteredClaims: jwt.RegisteredClaims{
|
||||
ExpiresAt: jwt.NewNumericDate(time.Now().Add(1 * time.Hour)),
|
||||
IssuedAt: jwt.NewNumericDate(time.Now()),
|
||||
Subject: "api_token",
|
||||
},
|
||||
}
|
||||
|
||||
// Try to use "none" algorithm
|
||||
token := jwt.NewWithClaims(jwt.SigningMethodNone, claims)
|
||||
tokenString, err := token.SignedString(jwt.UnsafeAllowNoneSignatureType)
|
||||
require.NoError(t, err)
|
||||
|
||||
server := newTestServer(t, "/test", db, authMiddleware.AuthTokenRequired)
|
||||
defer server.Close()
|
||||
|
||||
// Should reject "none" algorithm
|
||||
assert.False(t, server.CallAndGetStatus(t, "Bearer "+tokenString))
|
||||
}
|
||||
|
||||
func TestAPITokenAuthentication_LegacyProtoToken(t *testing.T) {
|
||||
db := setupTestDB(t)
|
||||
defer db.Close()
|
||||
|
||||
tokenCache := auth.NewTokenCache(db)
|
||||
userCache := auth.NewUserCache(db)
|
||||
authMiddleware := auth.NewAuthMiddleware("/base/url", "test", tokenCache, userCache)
|
||||
|
||||
server := newTestServer(t, "/test", db, authMiddleware.AuthTokenRequired)
|
||||
defer server.Close()
|
||||
|
||||
// Authorize with cookie to get legacy proto token
|
||||
server.Authorize(t, []string{auth.PrivilegeAutomation})
|
||||
legacyToken := server.GetToken(t)
|
||||
require.NotEmpty(t, legacyToken)
|
||||
|
||||
// Unauthorize cookie
|
||||
server.Unauthorize(t)
|
||||
|
||||
// Legacy proto token should still work (fallback mechanism)
|
||||
server.SetSessionCheckFunc(func(t *testing.T, c *gin.Context) {
|
||||
t.Helper()
|
||||
assert.Equal(t, uint64(1), c.GetUint64("uid"))
|
||||
assert.Equal(t, "automation", c.GetString("cpt"))
|
||||
})
|
||||
|
||||
assert.True(t, server.CallAndGetStatus(t, "Bearer "+legacyToken))
|
||||
}
|
||||
@@ -0,0 +1,237 @@
|
||||
package auth
|
||||
|
||||
import (
|
||||
"errors"
|
||||
"fmt"
|
||||
"slices"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"pentagi/pkg/server/models"
|
||||
"pentagi/pkg/server/rdb"
|
||||
"pentagi/pkg/server/response"
|
||||
|
||||
"github.com/gin-contrib/sessions"
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/jinzhu/gorm"
|
||||
)
|
||||
|
||||
type authResult int
|
||||
|
||||
const (
|
||||
authResultOk authResult = iota
|
||||
authResultSkip
|
||||
authResultFail
|
||||
authResultAbort
|
||||
)
|
||||
|
||||
type AuthMiddleware struct {
|
||||
globalSalt string
|
||||
tokenCache *TokenCache
|
||||
userCache *UserCache
|
||||
}
|
||||
|
||||
func NewAuthMiddleware(baseURL, globalSalt string, tokenCache *TokenCache, userCache *UserCache) *AuthMiddleware {
|
||||
return &AuthMiddleware{
|
||||
globalSalt: globalSalt,
|
||||
tokenCache: tokenCache,
|
||||
userCache: userCache,
|
||||
}
|
||||
}
|
||||
|
||||
func (p *AuthMiddleware) AuthUserRequired(c *gin.Context) {
|
||||
p.tryAuth(c, true, p.tryUserCookieAuthentication)
|
||||
}
|
||||
|
||||
func (p *AuthMiddleware) AuthTokenRequired(c *gin.Context) {
|
||||
p.tryAuth(c, true, p.tryProtoTokenAuthentication, p.tryUserCookieAuthentication)
|
||||
}
|
||||
|
||||
func (p *AuthMiddleware) TryAuth(c *gin.Context) {
|
||||
p.tryAuth(c, false, p.tryProtoTokenAuthentication, p.tryUserCookieAuthentication)
|
||||
}
|
||||
|
||||
func (p *AuthMiddleware) tryAuth(
|
||||
c *gin.Context,
|
||||
withFail bool,
|
||||
authMethods ...func(c *gin.Context) (authResult, error),
|
||||
) {
|
||||
if c.IsAborted() {
|
||||
return
|
||||
}
|
||||
|
||||
result := authResultSkip
|
||||
var authErr error
|
||||
for _, authMethod := range authMethods {
|
||||
result, authErr = authMethod(c)
|
||||
if c.IsAborted() || result == authResultAbort {
|
||||
return
|
||||
}
|
||||
if result != authResultSkip {
|
||||
break
|
||||
}
|
||||
}
|
||||
|
||||
if withFail && result != authResultOk {
|
||||
response.Error(c, response.ErrAuthRequired, authErr)
|
||||
return
|
||||
}
|
||||
c.Next()
|
||||
}
|
||||
|
||||
func (p *AuthMiddleware) tryUserCookieAuthentication(c *gin.Context) (authResult, error) {
|
||||
sessionObject, exists := c.Get(sessions.DefaultKey)
|
||||
if !exists {
|
||||
return authResultSkip, errors.New("can't find session object")
|
||||
}
|
||||
|
||||
session, ok := sessionObject.(sessions.Session)
|
||||
if !ok {
|
||||
return authResultFail, errors.New("not a session object")
|
||||
}
|
||||
|
||||
uid := session.Get("uid")
|
||||
uhash := session.Get("uhash")
|
||||
rid := session.Get("rid")
|
||||
prm := session.Get("prm")
|
||||
exp := session.Get("exp")
|
||||
gtm := session.Get("gtm")
|
||||
tid := session.Get("tid")
|
||||
uname := session.Get("uname")
|
||||
|
||||
for _, attr := range []any{uid, rid, prm, exp, gtm, uname, uhash, tid} {
|
||||
if attr == nil {
|
||||
return authResultFail, errors.New("cookie claim invalid")
|
||||
}
|
||||
}
|
||||
|
||||
prms, ok := prm.([]string)
|
||||
if !ok {
|
||||
return authResultFail, errors.New("no permissions granted")
|
||||
}
|
||||
|
||||
// Verify session expiration
|
||||
expVal, ok := exp.(int64)
|
||||
if !ok {
|
||||
return authResultFail, errors.New("token claim invalid")
|
||||
}
|
||||
if time.Now().Unix() > expVal {
|
||||
return authResultFail, errors.New("session expired")
|
||||
}
|
||||
|
||||
// Verify user hash matches database
|
||||
userID := uid.(uint64)
|
||||
sessionHash := uhash.(string)
|
||||
|
||||
dbHash, userStatus, err := p.userCache.GetUserHash(userID)
|
||||
if err != nil {
|
||||
if errors.Is(err, gorm.ErrRecordNotFound) {
|
||||
return authResultFail, errors.New("user has been deleted")
|
||||
}
|
||||
return authResultFail, fmt.Errorf("error checking user status: %w", err)
|
||||
}
|
||||
|
||||
switch userStatus {
|
||||
case models.UserStatusBlocked:
|
||||
return authResultFail, errors.New("user has been blocked")
|
||||
case models.UserStatusCreated:
|
||||
return authResultFail, errors.New("user is not ready")
|
||||
case models.UserStatusActive:
|
||||
}
|
||||
|
||||
if dbHash != sessionHash {
|
||||
return authResultFail, errors.New("user hash mismatch - session invalid for this installation")
|
||||
}
|
||||
|
||||
c.Set("prm", prms)
|
||||
c.Set("uid", userID)
|
||||
c.Set("uhash", sessionHash)
|
||||
c.Set("rid", rid.(uint64))
|
||||
c.Set("exp", exp.(int64))
|
||||
c.Set("gtm", gtm.(int64))
|
||||
c.Set("tid", tid.(string))
|
||||
c.Set("uname", uname.(string))
|
||||
|
||||
if slices.Contains(prms, PrivilegeAutomation) {
|
||||
c.Set("cpt", "automation")
|
||||
}
|
||||
|
||||
return authResultOk, nil
|
||||
}
|
||||
|
||||
const PrivilegeAutomation = "pentagi.automation"
|
||||
|
||||
func (p *AuthMiddleware) tryProtoTokenAuthentication(c *gin.Context) (authResult, error) {
|
||||
authHeader := c.Request.Header.Get("Authorization")
|
||||
if authHeader == "" {
|
||||
return authResultSkip, errors.New("token required")
|
||||
}
|
||||
|
||||
if !strings.HasPrefix(authHeader, "Bearer ") {
|
||||
return authResultSkip, errors.New("bearer scheme must be used")
|
||||
}
|
||||
token := authHeader[7:]
|
||||
if token == "" {
|
||||
return authResultSkip, errors.New("token can't be empty")
|
||||
}
|
||||
|
||||
// skip validation if using default salt (for backward compatibility)
|
||||
if p.globalSalt == "" || p.globalSalt == "salt" {
|
||||
return authResultSkip, errors.New("token validation disabled with default salt")
|
||||
}
|
||||
|
||||
// try to validate as API token first (new format with JWT signing key)
|
||||
apiClaims, apiErr := ValidateAPIToken(token, p.globalSalt)
|
||||
if apiErr != nil {
|
||||
return authResultFail, errors.New("token is invalid")
|
||||
}
|
||||
|
||||
// check token status and get privileges through cache
|
||||
status, privileges, err := p.tokenCache.GetStatus(apiClaims.TokenID)
|
||||
if err != nil {
|
||||
if errors.Is(err, gorm.ErrRecordNotFound) {
|
||||
return authResultFail, errors.New("token not found in database")
|
||||
}
|
||||
return authResultFail, fmt.Errorf("error checking token status: %w", err)
|
||||
}
|
||||
if status != models.TokenStatusActive {
|
||||
return authResultFail, errors.New("token has been revoked")
|
||||
}
|
||||
|
||||
// Verify user hash matches database
|
||||
dbHash, userStatus, err := p.userCache.GetUserHash(apiClaims.UID)
|
||||
if err != nil {
|
||||
if errors.Is(err, gorm.ErrRecordNotFound) {
|
||||
return authResultFail, errors.New("user has been deleted")
|
||||
}
|
||||
return authResultFail, fmt.Errorf("error checking user status: %w", err)
|
||||
}
|
||||
|
||||
if userStatus == models.UserStatusBlocked {
|
||||
return authResultFail, errors.New("user has been blocked")
|
||||
}
|
||||
|
||||
if dbHash != apiClaims.UHASH {
|
||||
return authResultFail, errors.New("user hash mismatch - token invalid for this installation")
|
||||
}
|
||||
|
||||
// generate UUID from user hash (fallback to empty string if hash is invalid)
|
||||
uuid, err := rdb.MakeUuidStrFromHash(apiClaims.UHASH)
|
||||
if err != nil {
|
||||
// Use empty UUID for invalid hashes (e.g., in tests)
|
||||
uuid = ""
|
||||
}
|
||||
|
||||
// set session fields similar to regular login
|
||||
c.Set("uid", apiClaims.UID)
|
||||
c.Set("uhash", apiClaims.UHASH)
|
||||
c.Set("rid", apiClaims.RID)
|
||||
c.Set("tid", models.UserTypeAPI.String())
|
||||
c.Set("prm", privileges)
|
||||
c.Set("gtm", time.Now().Unix())
|
||||
c.Set("exp", apiClaims.ExpiresAt.Unix())
|
||||
c.Set("uuid", uuid)
|
||||
c.Set("cpt", "automation")
|
||||
|
||||
return authResultOk, nil
|
||||
}
|
||||
@@ -0,0 +1,337 @@
|
||||
package auth_test
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"io"
|
||||
"math/rand"
|
||||
"net/http"
|
||||
"net/http/cookiejar"
|
||||
"net/http/httptest"
|
||||
"net/url"
|
||||
"strconv"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"pentagi/pkg/server/auth"
|
||||
"pentagi/pkg/server/models"
|
||||
|
||||
"github.com/gin-contrib/sessions"
|
||||
"github.com/gin-contrib/sessions/cookie"
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/jinzhu/gorm"
|
||||
_ "github.com/jinzhu/gorm/dialects/sqlite"
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
func TestAuthTokenProtoRequiredAuthWithCookie(t *testing.T) {
|
||||
db := setupTestDB(t)
|
||||
defer db.Close()
|
||||
|
||||
tokenCache := auth.NewTokenCache(db)
|
||||
userCache := auth.NewUserCache(db)
|
||||
authMiddleware := auth.NewAuthMiddleware("/base/url", "test", tokenCache, userCache)
|
||||
|
||||
t.Run("test URL", func(t *testing.T) {
|
||||
server := newTestServer(t, "/test", db, authMiddleware.AuthTokenRequired)
|
||||
defer server.Close()
|
||||
|
||||
assert.False(t, server.CallAndGetStatus(t))
|
||||
|
||||
server.SetSessionCheckFunc(func(t *testing.T, c *gin.Context) {
|
||||
t.Helper()
|
||||
assert.Equal(t, "", c.GetString("cpt"))
|
||||
})
|
||||
|
||||
server.Authorize(t, []string{})
|
||||
assert.True(t, server.CallAndGetStatus(t))
|
||||
|
||||
server.Authorize(t, []string{"wrong.permission"})
|
||||
assert.True(t, server.CallAndGetStatus(t))
|
||||
|
||||
server.SetSessionCheckFunc(func(t *testing.T, c *gin.Context) {
|
||||
t.Helper()
|
||||
assert.Equal(t, "automation", c.GetString("cpt"))
|
||||
})
|
||||
|
||||
server.Authorize(t, []string{auth.PrivilegeAutomation})
|
||||
assert.True(t, server.CallAndGetStatus(t))
|
||||
|
||||
server.Authorize(t, []string{"wrong.permission", auth.PrivilegeAutomation})
|
||||
assert.True(t, server.CallAndGetStatus(t))
|
||||
})
|
||||
}
|
||||
|
||||
func TestAuthTokenProtoRequiredAuthWithToken(t *testing.T) {
|
||||
db := setupTestDB(t)
|
||||
defer db.Close()
|
||||
|
||||
tokenCache := auth.NewTokenCache(db)
|
||||
userCache := auth.NewUserCache(db)
|
||||
authMiddleware := auth.NewAuthMiddleware("/base/url", "test", tokenCache, userCache)
|
||||
|
||||
server := newTestServer(t, "/test", db, authMiddleware.AuthTokenRequired)
|
||||
defer server.Close()
|
||||
|
||||
server.Authorize(t, []string{auth.PrivilegeAutomation})
|
||||
token := server.GetToken(t)
|
||||
require.NotEmpty(t, token)
|
||||
|
||||
server.Unauthorize(t)
|
||||
assert.False(t, server.CallAndGetStatus(t))
|
||||
|
||||
assert.False(t, server.CallAndGetStatus(t, token))
|
||||
assert.False(t, server.CallAndGetStatus(t, "not a bearer "+token))
|
||||
assert.False(t, server.CallAndGetStatus(t, "Bearer"+token))
|
||||
assert.False(t, server.CallAndGetStatus(t, "Bearer not_a_token"))
|
||||
|
||||
server.SetSessionCheckFunc(func(t *testing.T, c *gin.Context) {
|
||||
t.Helper()
|
||||
assert.Equal(t, uint64(1), c.GetUint64("uid"))
|
||||
assert.Equal(t, uint64(2), c.GetUint64("rid"))
|
||||
assert.NotNil(t, c.GetStringSlice("prm"))
|
||||
|
||||
// gtm and exp should now be set for API tokens
|
||||
gtm := c.GetInt64("gtm")
|
||||
assert.Greater(t, gtm, int64(0), "GTM should be set")
|
||||
|
||||
exp := c.GetInt64("exp")
|
||||
assert.Greater(t, exp, gtm, "EXP should be greater than GTM")
|
||||
|
||||
// uuid will be empty for invalid hash (test uses "123" which is not valid MD5)
|
||||
assert.NotNil(t, c.GetString("uuid"))
|
||||
|
||||
assert.Equal(t, "automation", c.GetString("cpt"))
|
||||
assert.Empty(t, c.GetString("uname"))
|
||||
})
|
||||
|
||||
assert.True(t, server.CallAndGetStatus(t, "Bearer "+token))
|
||||
}
|
||||
|
||||
func TestAuthRequiredAuthWithCookie(t *testing.T) {
|
||||
db := setupTestDB(t)
|
||||
defer db.Close()
|
||||
|
||||
tokenCache := auth.NewTokenCache(db)
|
||||
userCache := auth.NewUserCache(db)
|
||||
authMiddleware := auth.NewAuthMiddleware("/base/url", "test", tokenCache, userCache)
|
||||
|
||||
server := newTestServer(t, "/test", db, authMiddleware.AuthUserRequired)
|
||||
defer server.Close()
|
||||
|
||||
server.SetSessionCheckFunc(func(t *testing.T, c *gin.Context) {
|
||||
t.Helper()
|
||||
assert.Equal(t, uint64(1), c.GetUint64("uid"))
|
||||
assert.Equal(t, uint64(2), c.GetUint64("rid"))
|
||||
assert.NotNil(t, c.GetStringSlice("prm"))
|
||||
assert.NotNil(t, c.GetInt64("gtm"))
|
||||
assert.NotNil(t, c.GetInt64("exp"))
|
||||
assert.Empty(t, c.GetString("uuid"))
|
||||
assert.Equal(t, "User 1", c.GetString("uname"))
|
||||
})
|
||||
|
||||
assert.False(t, server.CallAndGetStatus(t))
|
||||
|
||||
server.Authorize(t, []string{"some.permission"})
|
||||
assert.True(t, server.CallAndGetStatus(t))
|
||||
}
|
||||
|
||||
type testServer struct {
|
||||
testEndpoint string
|
||||
client *http.Client
|
||||
calls map[string]struct{}
|
||||
sessionCheckFunc func(t *testing.T, c *gin.Context)
|
||||
db *gorm.DB
|
||||
*httptest.Server
|
||||
}
|
||||
|
||||
func newTestServer(t *testing.T, testEndpoint string, db *gorm.DB, middlewares ...gin.HandlerFunc) *testServer {
|
||||
t.Helper()
|
||||
|
||||
server := &testServer{
|
||||
db: db,
|
||||
}
|
||||
|
||||
router := gin.New()
|
||||
globalSalt := "test"
|
||||
cookieStore := cookie.NewStore(auth.MakeCookieStoreKey(globalSalt)...)
|
||||
router.Use(sessions.Sessions("auth", cookieStore))
|
||||
|
||||
server.calls = map[string]struct{}{}
|
||||
|
||||
if testEndpoint == "" {
|
||||
testEndpoint = "/test"
|
||||
}
|
||||
server.testEndpoint = testEndpoint
|
||||
|
||||
router.GET("/auth", func(c *gin.Context) {
|
||||
t.Helper()
|
||||
privs, _ := c.GetQueryArray("privileges")
|
||||
expString, ok := c.GetQuery("expiration")
|
||||
assert.True(t, ok)
|
||||
exp, err := strconv.Atoi(expString)
|
||||
assert.NoError(t, err)
|
||||
setTestSession(t, c, privs, exp)
|
||||
})
|
||||
|
||||
authRoutes := router.Group("")
|
||||
for _, middleware := range middlewares {
|
||||
authRoutes.Use(middleware)
|
||||
}
|
||||
|
||||
authRoutes.GET(server.testEndpoint, func(c *gin.Context) {
|
||||
t.Helper()
|
||||
|
||||
id, _ := c.GetQuery("id")
|
||||
require.NotEmpty(t, id)
|
||||
|
||||
if server.sessionCheckFunc != nil {
|
||||
server.sessionCheckFunc(t, c)
|
||||
}
|
||||
server.calls[id] = struct{}{}
|
||||
})
|
||||
|
||||
authRoutes.GET("/auth_token", func(c *gin.Context) {
|
||||
t.Helper()
|
||||
|
||||
tokenID, err := auth.GenerateTokenID()
|
||||
require.NoError(t, err)
|
||||
uhash := "testhash"
|
||||
uid := uint64(1)
|
||||
rid := uint64(2)
|
||||
ttl := uint64(3600)
|
||||
claims := auth.MakeAPITokenClaims(tokenID, uhash, uid, rid, ttl)
|
||||
token, err := auth.MakeAPIToken(globalSalt, claims)
|
||||
require.NoError(t, err)
|
||||
|
||||
db.Create(&models.APIToken{
|
||||
TokenID: tokenID,
|
||||
UserID: uid,
|
||||
RoleID: rid,
|
||||
TTL: ttl,
|
||||
Status: models.TokenStatusActive,
|
||||
})
|
||||
|
||||
c.Writer.Write([]byte(token))
|
||||
})
|
||||
|
||||
server.Server = httptest.NewServer(router)
|
||||
server.client = server.Client()
|
||||
jar, err := cookiejar.New(nil)
|
||||
require.NoError(t, err)
|
||||
server.client.Jar = jar
|
||||
|
||||
return server
|
||||
}
|
||||
|
||||
func (s *testServer) Authorize(t *testing.T, privileges []string) {
|
||||
t.Helper()
|
||||
request, err := http.NewRequest(http.MethodGet, s.URL+"/auth", nil)
|
||||
require.NoError(t, err)
|
||||
query := url.Values{}
|
||||
for _, p := range privileges {
|
||||
query.Add("privileges", p)
|
||||
}
|
||||
query.Add("expiration", strconv.Itoa(5*60))
|
||||
request.URL.RawQuery = query.Encode()
|
||||
|
||||
resp, err := s.client.Do(request)
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, http.StatusOK, resp.StatusCode)
|
||||
}
|
||||
|
||||
func (s *testServer) GetToken(t *testing.T) string {
|
||||
t.Helper()
|
||||
request, err := http.NewRequest(http.MethodGet, s.URL+"/auth_token", nil)
|
||||
require.NoError(t, err)
|
||||
|
||||
resp, err := s.client.Do(request)
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, http.StatusOK, resp.StatusCode)
|
||||
token, err := io.ReadAll(resp.Body)
|
||||
require.NoError(t, err)
|
||||
return string(token)
|
||||
}
|
||||
|
||||
func (s *testServer) SetSessionCheckFunc(f func(t *testing.T, c *gin.Context)) {
|
||||
s.sessionCheckFunc = f
|
||||
}
|
||||
|
||||
func (s *testServer) Unauthorize(t *testing.T) {
|
||||
t.Helper()
|
||||
request, err := http.NewRequest(http.MethodGet, s.URL+"/auth", nil)
|
||||
require.NoError(t, err)
|
||||
query := url.Values{}
|
||||
query.Add("expiration", strconv.Itoa(-1))
|
||||
request.URL.RawQuery = query.Encode()
|
||||
|
||||
resp, err := s.client.Do(request)
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, http.StatusOK, resp.StatusCode)
|
||||
}
|
||||
|
||||
func (s *testServer) TestCall(t *testing.T, token ...string) (string, bool) {
|
||||
t.Helper()
|
||||
id := strconv.Itoa(rand.Int())
|
||||
|
||||
request, err := http.NewRequest(http.MethodGet, s.URL+s.testEndpoint+"?id="+id, nil)
|
||||
require.NoError(t, err)
|
||||
if len(token) == 1 {
|
||||
request.Header.Add("Authorization", token[0])
|
||||
}
|
||||
|
||||
resp, err := s.client.Do(request)
|
||||
require.NoError(t, err)
|
||||
|
||||
assert.True(t, resp.StatusCode == http.StatusOK ||
|
||||
resp.StatusCode == http.StatusForbidden)
|
||||
|
||||
return id, resp.StatusCode == http.StatusOK
|
||||
}
|
||||
|
||||
func (s *testServer) TestCallWithData(t *testing.T, data string) (string, bool) {
|
||||
t.Helper()
|
||||
id := strconv.Itoa(rand.Int())
|
||||
|
||||
request, err := http.NewRequest(http.MethodGet, s.URL+s.testEndpoint+"?id="+id, bytes.NewBufferString(data))
|
||||
require.NoError(t, err)
|
||||
|
||||
resp, err := s.client.Do(request)
|
||||
require.NoError(t, err)
|
||||
|
||||
assert.True(t, resp.StatusCode == http.StatusOK ||
|
||||
resp.StatusCode == http.StatusForbidden)
|
||||
|
||||
return id, resp.StatusCode == http.StatusOK
|
||||
}
|
||||
|
||||
func (s *testServer) Called(id string) bool {
|
||||
_, ok := s.calls[id]
|
||||
return ok
|
||||
}
|
||||
|
||||
func (s *testServer) CallAndGetStatus(t *testing.T, token ...string) bool {
|
||||
t.Helper()
|
||||
id, ok := s.TestCall(t, token...)
|
||||
assert.Equal(t, ok, s.Called(id))
|
||||
return ok
|
||||
}
|
||||
|
||||
func setTestSession(t *testing.T, c *gin.Context, privileges []string, expires int) {
|
||||
t.Helper()
|
||||
session := sessions.Default(c)
|
||||
session.Set("uid", uint64(1))
|
||||
session.Set("uhash", "testhash")
|
||||
session.Set("rid", uint64(2))
|
||||
session.Set("tid", models.UserTypeLocal.String())
|
||||
session.Set("prm", privileges)
|
||||
session.Set("gtm", time.Now().Unix())
|
||||
session.Set("exp", time.Now().Add(time.Duration(expires)*time.Second).Unix())
|
||||
session.Set("uuid", "uuid1")
|
||||
session.Set("uname", "User 1")
|
||||
session.Options(sessions.Options{
|
||||
HttpOnly: true,
|
||||
MaxAge: expires,
|
||||
})
|
||||
require.NoError(t, session.Save())
|
||||
}
|
||||
@@ -0,0 +1,890 @@
|
||||
package auth_test
|
||||
|
||||
import (
|
||||
"sync"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"pentagi/pkg/server/auth"
|
||||
"pentagi/pkg/server/models"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/golang-jwt/jwt/v5"
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
// TestEndToEndAPITokenFlow tests complete flow from creation to usage
|
||||
func TestEndToEndAPITokenFlow(t *testing.T) {
|
||||
db := setupTestDB(t)
|
||||
defer db.Close()
|
||||
|
||||
tokenCache := auth.NewTokenCache(db)
|
||||
userCache := auth.NewUserCache(db)
|
||||
authMiddleware := auth.NewAuthMiddleware("/base/url", "test_salt", tokenCache, userCache)
|
||||
|
||||
testCases := []struct {
|
||||
name string
|
||||
tokenID string
|
||||
status models.TokenStatus
|
||||
shouldPass bool
|
||||
errorContains string
|
||||
}{
|
||||
{
|
||||
name: "active token authenticates successfully",
|
||||
tokenID: "active123",
|
||||
status: models.TokenStatusActive,
|
||||
shouldPass: true,
|
||||
},
|
||||
{
|
||||
name: "revoked token is rejected",
|
||||
tokenID: "revoked456",
|
||||
status: models.TokenStatusRevoked,
|
||||
shouldPass: false,
|
||||
errorContains: "revoked",
|
||||
},
|
||||
}
|
||||
|
||||
for _, tc := range testCases {
|
||||
t.Run(tc.name, func(t *testing.T) {
|
||||
// Create token in database
|
||||
apiToken := models.APIToken{
|
||||
TokenID: tc.tokenID,
|
||||
UserID: 1,
|
||||
RoleID: 2,
|
||||
TTL: 3600,
|
||||
Status: tc.status,
|
||||
}
|
||||
err := db.Create(&apiToken).Error
|
||||
require.NoError(t, err)
|
||||
|
||||
// Create JWT token
|
||||
claims := models.APITokenClaims{
|
||||
TokenID: tc.tokenID,
|
||||
RID: 2,
|
||||
UID: 1,
|
||||
UHASH: "testhash",
|
||||
RegisteredClaims: jwt.RegisteredClaims{
|
||||
ExpiresAt: jwt.NewNumericDate(time.Now().Add(1 * time.Hour)),
|
||||
IssuedAt: jwt.NewNumericDate(time.Now()),
|
||||
Subject: "api_token",
|
||||
},
|
||||
}
|
||||
token := jwt.NewWithClaims(jwt.SigningMethodHS256, claims)
|
||||
tokenString, err := token.SignedString(auth.MakeJWTSigningKey("test_salt"))
|
||||
require.NoError(t, err)
|
||||
|
||||
// Test authentication
|
||||
server := newTestServer(t, "/protected", db, authMiddleware.AuthTokenRequired)
|
||||
defer server.Close()
|
||||
|
||||
success := server.CallAndGetStatus(t, "Bearer "+tokenString)
|
||||
assert.Equal(t, tc.shouldPass, success)
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
// TestAPIToken_RoleIsolation verifies that token inherits creator's role
|
||||
func TestAPIToken_RoleIsolation(t *testing.T) {
|
||||
testCases := []struct {
|
||||
name string
|
||||
creatorRole uint64
|
||||
tokenRole uint64
|
||||
expectMatch bool
|
||||
}{
|
||||
{
|
||||
name: "user creates token with user role",
|
||||
creatorRole: 2,
|
||||
tokenRole: 2,
|
||||
expectMatch: true,
|
||||
},
|
||||
{
|
||||
name: "admin creates token with admin role",
|
||||
creatorRole: 1,
|
||||
tokenRole: 1,
|
||||
expectMatch: true,
|
||||
},
|
||||
}
|
||||
|
||||
for _, tc := range testCases {
|
||||
t.Run(tc.name, func(t *testing.T) {
|
||||
tokenID, err := auth.GenerateTokenID()
|
||||
require.NoError(t, err)
|
||||
|
||||
// Create JWT with specific role
|
||||
claims := models.APITokenClaims{
|
||||
TokenID: tokenID,
|
||||
RID: tc.tokenRole,
|
||||
UID: 1,
|
||||
UHASH: "testhash",
|
||||
RegisteredClaims: jwt.RegisteredClaims{
|
||||
ExpiresAt: jwt.NewNumericDate(time.Now().Add(1 * time.Hour)),
|
||||
IssuedAt: jwt.NewNumericDate(time.Now()),
|
||||
Subject: "api_token",
|
||||
},
|
||||
}
|
||||
|
||||
token := jwt.NewWithClaims(jwt.SigningMethodHS256, claims)
|
||||
tokenString, err := token.SignedString(auth.MakeJWTSigningKey("test"))
|
||||
require.NoError(t, err)
|
||||
|
||||
// Validate and check role
|
||||
validated, err := auth.ValidateAPIToken(tokenString, "test")
|
||||
require.NoError(t, err)
|
||||
|
||||
if tc.expectMatch {
|
||||
assert.Equal(t, tc.tokenRole, validated.RID)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
// TestAPIToken_SignatureVerification tests various signature attacks
|
||||
func TestAPIToken_SignatureVerification(t *testing.T) {
|
||||
correctSalt := "correct_salt"
|
||||
wrongSalt := "wrong_salt"
|
||||
|
||||
testCases := []struct {
|
||||
name string
|
||||
signSalt string
|
||||
verifySalt string
|
||||
expectValid bool
|
||||
errorContains string
|
||||
}{
|
||||
{
|
||||
name: "matching salt - valid",
|
||||
signSalt: correctSalt,
|
||||
verifySalt: correctSalt,
|
||||
expectValid: true,
|
||||
},
|
||||
{
|
||||
name: "mismatched salt - invalid",
|
||||
signSalt: correctSalt,
|
||||
verifySalt: wrongSalt,
|
||||
expectValid: false,
|
||||
errorContains: "invalid",
|
||||
},
|
||||
}
|
||||
|
||||
for _, tc := range testCases {
|
||||
t.Run(tc.name, func(t *testing.T) {
|
||||
tokenID, err := auth.GenerateTokenID()
|
||||
require.NoError(t, err)
|
||||
|
||||
claims := models.APITokenClaims{
|
||||
TokenID: tokenID,
|
||||
RID: 2,
|
||||
UID: 1,
|
||||
UHASH: "testhash",
|
||||
RegisteredClaims: jwt.RegisteredClaims{
|
||||
ExpiresAt: jwt.NewNumericDate(time.Now().Add(1 * time.Hour)),
|
||||
IssuedAt: jwt.NewNumericDate(time.Now()),
|
||||
Subject: "api_token",
|
||||
},
|
||||
}
|
||||
|
||||
token := jwt.NewWithClaims(jwt.SigningMethodHS256, claims)
|
||||
tokenString, err := token.SignedString(auth.MakeJWTSigningKey(tc.signSalt))
|
||||
require.NoError(t, err)
|
||||
|
||||
validated, err := auth.ValidateAPIToken(tokenString, tc.verifySalt)
|
||||
|
||||
if tc.expectValid {
|
||||
assert.NoError(t, err)
|
||||
assert.NotNil(t, validated)
|
||||
} else {
|
||||
assert.Error(t, err)
|
||||
if tc.errorContains != "" {
|
||||
assert.Contains(t, err.Error(), tc.errorContains)
|
||||
}
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
// TestAPIToken_CacheInvalidation verifies cache invalidation scenarios
|
||||
func TestAPIToken_CacheInvalidation(t *testing.T) {
|
||||
db := setupTestDB(t)
|
||||
defer db.Close()
|
||||
|
||||
tokenCache := auth.NewTokenCache(db)
|
||||
|
||||
// Create token
|
||||
tokenID, err := auth.GenerateTokenID()
|
||||
require.NoError(t, err)
|
||||
apiToken := models.APIToken{
|
||||
TokenID: tokenID,
|
||||
UserID: 1,
|
||||
RoleID: 2,
|
||||
TTL: 3600,
|
||||
Status: models.TokenStatusActive,
|
||||
}
|
||||
err = db.Create(&apiToken).Error
|
||||
require.NoError(t, err)
|
||||
|
||||
// Load into cache
|
||||
status1, _, err := tokenCache.GetStatus(tokenID)
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, models.TokenStatusActive, status1)
|
||||
|
||||
// Update in DB
|
||||
db.Model(&apiToken).Update("status", models.TokenStatusRevoked)
|
||||
|
||||
// Should still return active from cache
|
||||
status2, _, err := tokenCache.GetStatus(tokenID)
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, models.TokenStatusActive, status2, "Cache should return stale value")
|
||||
|
||||
// Invalidate cache
|
||||
tokenCache.Invalidate(tokenID)
|
||||
|
||||
// Should now return revoked from DB
|
||||
status3, _, err := tokenCache.GetStatus(tokenID)
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, models.TokenStatusRevoked, status3, "Cache should be refreshed from DB")
|
||||
}
|
||||
|
||||
// TestAPIToken_ConcurrentAccess tests thread-safety of cache
|
||||
func TestAPIToken_ConcurrentAccess(t *testing.T) {
|
||||
db := setupTestDB(t)
|
||||
defer db.Close()
|
||||
|
||||
tokenCache := auth.NewTokenCache(db)
|
||||
|
||||
// Create multiple tokens
|
||||
tokenIDs := make([]string, 10)
|
||||
for i := range 10 {
|
||||
tokenID, err := auth.GenerateTokenID()
|
||||
require.NoError(t, err)
|
||||
tokenIDs[i] = tokenID
|
||||
apiToken := models.APIToken{
|
||||
TokenID: tokenID,
|
||||
UserID: 1,
|
||||
RoleID: 2,
|
||||
TTL: 3600,
|
||||
Status: models.TokenStatusActive,
|
||||
}
|
||||
err = db.Create(&apiToken).Error
|
||||
require.NoError(t, err)
|
||||
}
|
||||
|
||||
// Verify tokens were created
|
||||
var count int
|
||||
db.Model(&models.APIToken{}).Where("deleted_at IS NULL").Count(&count)
|
||||
require.Equal(t, 10, count)
|
||||
|
||||
// Warm up cache
|
||||
for i := range 10 {
|
||||
status, _, err := tokenCache.GetStatus(tokenIDs[i])
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, models.TokenStatusActive, status)
|
||||
}
|
||||
|
||||
// Concurrent cache access using channels for error reporting
|
||||
type testResult struct {
|
||||
success bool
|
||||
err error
|
||||
}
|
||||
results := make(chan testResult, 10)
|
||||
|
||||
var wg sync.WaitGroup
|
||||
wg.Add(10)
|
||||
for i := range 10 {
|
||||
go func(tokenID string) {
|
||||
defer wg.Done()
|
||||
for range 100 {
|
||||
status, _, err := tokenCache.GetStatus(tokenID)
|
||||
if err != nil {
|
||||
results <- testResult{success: false, err: err}
|
||||
return
|
||||
}
|
||||
if status != models.TokenStatusActive {
|
||||
results <- testResult{success: false, err: assert.AnError}
|
||||
return
|
||||
}
|
||||
}
|
||||
results <- testResult{success: true, err: nil}
|
||||
}(tokenIDs[i])
|
||||
}
|
||||
|
||||
wg.Wait()
|
||||
close(results)
|
||||
|
||||
// Wait and check all results
|
||||
for result := range results {
|
||||
assert.NoError(t, result.err)
|
||||
assert.True(t, result.success, "Goroutine should complete successfully")
|
||||
}
|
||||
}
|
||||
|
||||
// TestAPIToken_JSONStructure verifies JWT payload structure
|
||||
func TestAPIToken_JSONStructure(t *testing.T) {
|
||||
tokenID, err := auth.GenerateTokenID()
|
||||
require.NoError(t, err)
|
||||
|
||||
claims := models.APITokenClaims{
|
||||
TokenID: tokenID,
|
||||
RID: 2,
|
||||
UID: 1,
|
||||
UHASH: "testhash",
|
||||
RegisteredClaims: jwt.RegisteredClaims{
|
||||
ExpiresAt: jwt.NewNumericDate(time.Now().Add(1 * time.Hour)),
|
||||
IssuedAt: jwt.NewNumericDate(time.Now()),
|
||||
Subject: "api_token",
|
||||
},
|
||||
}
|
||||
|
||||
token := jwt.NewWithClaims(jwt.SigningMethodHS256, claims)
|
||||
tokenString, err := token.SignedString(auth.MakeJWTSigningKey("test"))
|
||||
require.NoError(t, err)
|
||||
|
||||
// Parse and verify all fields
|
||||
parsed, err := auth.ValidateAPIToken(tokenString, "test")
|
||||
require.NoError(t, err)
|
||||
|
||||
assert.Equal(t, tokenID, parsed.TokenID, "TokenID should match")
|
||||
assert.Equal(t, uint64(2), parsed.RID, "RID should match")
|
||||
assert.Equal(t, uint64(1), parsed.UID, "UID should match")
|
||||
assert.Equal(t, "testhash", parsed.UHASH, "UHASH should match")
|
||||
assert.Equal(t, "api_token", parsed.Subject, "Subject should match")
|
||||
assert.NotNil(t, parsed.ExpiresAt, "ExpiresAt should be set")
|
||||
assert.NotNil(t, parsed.IssuedAt, "IssuedAt should be set")
|
||||
}
|
||||
|
||||
// TestAPIToken_Expiration verifies TTL enforcement
|
||||
func TestAPIToken_Expiration(t *testing.T) {
|
||||
testCases := []struct {
|
||||
name string
|
||||
ttl time.Duration
|
||||
expectValid bool
|
||||
}{
|
||||
{
|
||||
name: "future expiration - valid",
|
||||
ttl: 1 * time.Hour,
|
||||
expectValid: true,
|
||||
},
|
||||
{
|
||||
name: "past expiration - invalid",
|
||||
ttl: -1 * time.Hour,
|
||||
expectValid: false,
|
||||
},
|
||||
{
|
||||
name: "just expired - invalid",
|
||||
ttl: -1 * time.Second,
|
||||
expectValid: false,
|
||||
},
|
||||
}
|
||||
|
||||
for _, tc := range testCases {
|
||||
t.Run(tc.name, func(t *testing.T) {
|
||||
tokenID, err := auth.GenerateTokenID()
|
||||
require.NoError(t, err)
|
||||
|
||||
claims := models.APITokenClaims{
|
||||
TokenID: tokenID,
|
||||
RID: 2,
|
||||
UID: 1,
|
||||
UHASH: "testhash",
|
||||
RegisteredClaims: jwt.RegisteredClaims{
|
||||
ExpiresAt: jwt.NewNumericDate(time.Now().Add(tc.ttl)),
|
||||
IssuedAt: jwt.NewNumericDate(time.Now().Add(-1 * time.Hour)),
|
||||
Subject: "api_token",
|
||||
},
|
||||
}
|
||||
|
||||
token := jwt.NewWithClaims(jwt.SigningMethodHS256, claims)
|
||||
tokenString, err := token.SignedString(auth.MakeJWTSigningKey("test"))
|
||||
require.NoError(t, err)
|
||||
|
||||
validated, err := auth.ValidateAPIToken(tokenString, "test")
|
||||
|
||||
if tc.expectValid {
|
||||
assert.NoError(t, err)
|
||||
assert.NotNil(t, validated)
|
||||
} else {
|
||||
assert.Error(t, err)
|
||||
assert.Contains(t, err.Error(), "expired")
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
// TestDualAuthentication verifies both cookie and token auth work together
|
||||
func TestDualAuthentication(t *testing.T) {
|
||||
db := setupTestDB(t)
|
||||
defer db.Close()
|
||||
|
||||
tokenCache := auth.NewTokenCache(db)
|
||||
userCache := auth.NewUserCache(db)
|
||||
authMiddleware := auth.NewAuthMiddleware("/base/url", "test", tokenCache, userCache)
|
||||
|
||||
server := newTestServer(t, "/test", db, authMiddleware.AuthTokenRequired)
|
||||
defer server.Close()
|
||||
|
||||
// Test 1: Cookie authentication
|
||||
server.Authorize(t, []string{auth.PrivilegeAutomation})
|
||||
assert.True(t, server.CallAndGetStatus(t), "Cookie auth should work")
|
||||
|
||||
// Test 2: Create and use API token
|
||||
tokenID, err := auth.GenerateTokenID()
|
||||
require.NoError(t, err)
|
||||
apiToken := models.APIToken{
|
||||
TokenID: tokenID,
|
||||
UserID: 1,
|
||||
RoleID: 2,
|
||||
TTL: 3600,
|
||||
Status: models.TokenStatusActive,
|
||||
}
|
||||
err = db.Create(&apiToken).Error
|
||||
require.NoError(t, err)
|
||||
|
||||
claims := models.APITokenClaims{
|
||||
TokenID: tokenID,
|
||||
RID: 2,
|
||||
UID: 1,
|
||||
UHASH: "testhash",
|
||||
RegisteredClaims: jwt.RegisteredClaims{
|
||||
ExpiresAt: jwt.NewNumericDate(time.Now().Add(1 * time.Hour)),
|
||||
IssuedAt: jwt.NewNumericDate(time.Now()),
|
||||
Subject: "api_token",
|
||||
},
|
||||
}
|
||||
token := jwt.NewWithClaims(jwt.SigningMethodHS256, claims)
|
||||
tokenString, _ := token.SignedString(auth.MakeJWTSigningKey("test"))
|
||||
|
||||
// Unauthorize cookie
|
||||
server.Unauthorize(t)
|
||||
|
||||
// Test 3: Token authentication should work
|
||||
assert.True(t, server.CallAndGetStatus(t, "Bearer "+tokenString), "Token auth should work")
|
||||
|
||||
// Test 4: Both should work simultaneously
|
||||
server.Authorize(t, []string{auth.PrivilegeAutomation})
|
||||
assert.True(t, server.CallAndGetStatus(t, "Bearer "+tokenString), "Both auth methods should work")
|
||||
}
|
||||
|
||||
// TestSecurityAudit_ClaimsInJWT verifies all security-critical data is in JWT
|
||||
func TestSecurityAudit_ClaimsInJWT(t *testing.T) {
|
||||
// Create token in DB with certain values
|
||||
tokenID, err := auth.GenerateTokenID()
|
||||
require.NoError(t, err)
|
||||
dbToken := models.APIToken{
|
||||
TokenID: tokenID,
|
||||
UserID: 1,
|
||||
RoleID: 2, // User role in DB
|
||||
}
|
||||
|
||||
// Create JWT with different role (simulating compromise scenario)
|
||||
jwtClaims := models.APITokenClaims{
|
||||
TokenID: tokenID,
|
||||
RID: 1, // Admin role in JWT (different from DB!)
|
||||
UID: 1,
|
||||
UHASH: "testhash",
|
||||
RegisteredClaims: jwt.RegisteredClaims{
|
||||
ExpiresAt: jwt.NewNumericDate(time.Now().Add(1 * time.Hour)),
|
||||
IssuedAt: jwt.NewNumericDate(time.Now()),
|
||||
Subject: "api_token",
|
||||
},
|
||||
}
|
||||
|
||||
token := jwt.NewWithClaims(jwt.SigningMethodHS256, jwtClaims)
|
||||
tokenString, _ := token.SignedString(auth.MakeJWTSigningKey("test"))
|
||||
|
||||
// Validate token
|
||||
validated, err := auth.ValidateAPIToken(tokenString, "test")
|
||||
require.NoError(t, err)
|
||||
|
||||
// We trust JWT claims, not DB values
|
||||
assert.Equal(t, uint64(1), validated.RID, "Should use role from JWT, not DB")
|
||||
assert.NotEqual(t, dbToken.RoleID, validated.RID, "JWT role differs from DB role")
|
||||
assert.Equal(t, dbToken.UserID, validated.UID)
|
||||
assert.Equal(t, dbToken.TokenID, validated.TokenID)
|
||||
|
||||
// This is CORRECT behavior: DB only stores metadata for management
|
||||
// Actual authorization data comes from signed JWT
|
||||
}
|
||||
|
||||
// TestSecurityAudit_TokenIDUniqueness verifies token ID collision resistance
|
||||
func TestSecurityAudit_TokenIDUniqueness(t *testing.T) {
|
||||
iterations := 10000
|
||||
tokens := make(map[string]bool, iterations)
|
||||
|
||||
for i := 0; i < iterations; i++ {
|
||||
tokenID, err := auth.GenerateTokenID()
|
||||
require.NoError(t, err)
|
||||
|
||||
// Check format
|
||||
assert.Len(t, tokenID, 10)
|
||||
|
||||
// Check uniqueness
|
||||
if tokens[tokenID] {
|
||||
t.Fatalf("Duplicate token ID generated: %s", tokenID)
|
||||
}
|
||||
tokens[tokenID] = true
|
||||
}
|
||||
|
||||
t.Logf("Generated %d unique token IDs without collision", iterations)
|
||||
}
|
||||
|
||||
// TestSecurityAudit_SaltIsolation verifies JWT and Cookie keys are different
|
||||
func TestSecurityAudit_SaltIsolation(t *testing.T) {
|
||||
salts := []string{"salt1", "salt2", "production_salt"}
|
||||
|
||||
for _, salt := range salts {
|
||||
t.Run("salt="+salt, func(t *testing.T) {
|
||||
jwtKey := auth.MakeJWTSigningKey(salt)
|
||||
cookieKeys := auth.MakeCookieStoreKey(salt)
|
||||
|
||||
// JWT key must be different from both cookie keys
|
||||
assert.NotEqual(t, jwtKey, cookieKeys[0], "JWT key must differ from cookie auth key")
|
||||
assert.NotEqual(t, jwtKey, cookieKeys[1], "JWT key must differ from cookie encryption key")
|
||||
|
||||
// Verify key lengths
|
||||
assert.Len(t, jwtKey, 32, "JWT key must be 32 bytes")
|
||||
assert.Len(t, cookieKeys[0], 64, "Cookie auth key must be 64 bytes")
|
||||
assert.Len(t, cookieKeys[1], 32, "Cookie encryption key must be 32 bytes")
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
// TestAPIToken_ContextSetup verifies correct context values are set
|
||||
func TestAPIToken_ContextSetup(t *testing.T) {
|
||||
db := setupTestDB(t)
|
||||
defer db.Close()
|
||||
|
||||
tokenCache := auth.NewTokenCache(db)
|
||||
userCache := auth.NewUserCache(db)
|
||||
authMiddleware := auth.NewAuthMiddleware("/base/url", "test", tokenCache, userCache)
|
||||
|
||||
// Create token
|
||||
tokenID, err := auth.GenerateTokenID()
|
||||
require.NoError(t, err)
|
||||
apiToken := models.APIToken{
|
||||
TokenID: tokenID,
|
||||
UserID: 5,
|
||||
RoleID: 3,
|
||||
TTL: 3600,
|
||||
Status: models.TokenStatusActive,
|
||||
}
|
||||
err = db.Create(&apiToken).Error
|
||||
require.NoError(t, err)
|
||||
|
||||
user := models.User{
|
||||
ID: 5,
|
||||
Hash: "user5hash",
|
||||
Mail: "user5@example.com",
|
||||
Name: "User 5",
|
||||
Status: models.UserStatusActive,
|
||||
RoleID: 2,
|
||||
}
|
||||
err = db.Create(&user).Error
|
||||
require.NoError(t, err)
|
||||
|
||||
claims := models.APITokenClaims{
|
||||
TokenID: tokenID,
|
||||
RID: 2,
|
||||
UID: 5,
|
||||
UHASH: "user5hash",
|
||||
RegisteredClaims: jwt.RegisteredClaims{
|
||||
ExpiresAt: jwt.NewNumericDate(time.Now().Add(1 * time.Hour)),
|
||||
IssuedAt: jwt.NewNumericDate(time.Now()),
|
||||
Subject: "api_token",
|
||||
},
|
||||
}
|
||||
token := jwt.NewWithClaims(jwt.SigningMethodHS256, claims)
|
||||
tokenString, _ := token.SignedString(auth.MakeJWTSigningKey("test"))
|
||||
|
||||
server := newTestServer(t, "/test", db, authMiddleware.AuthTokenRequired)
|
||||
defer server.Close()
|
||||
|
||||
server.SetSessionCheckFunc(func(t *testing.T, c *gin.Context) {
|
||||
t.Helper()
|
||||
|
||||
// Verify all context values are set correctly
|
||||
assert.Equal(t, uint64(5), c.GetUint64("uid"), "UID from JWT")
|
||||
assert.Equal(t, uint64(2), c.GetUint64("rid"), "RID from JWT")
|
||||
assert.Equal(t, "user5hash", c.GetString("uhash"), "UHASH from JWT")
|
||||
assert.Equal(t, "automation", c.GetString("cpt"), "CPT from JWT")
|
||||
assert.Equal(t, "api", c.GetString("tid"), "TID should be 'api' for API tokens")
|
||||
|
||||
prms := c.GetStringSlice("prm")
|
||||
assert.Contains(t, prms, auth.PrivilegeAutomation, "Should have automation privilege")
|
||||
|
||||
// Verify session timing fields
|
||||
gtm := c.GetInt64("gtm")
|
||||
assert.Greater(t, gtm, int64(0), "GTM (generation time) should be set")
|
||||
|
||||
exp := c.GetInt64("exp")
|
||||
assert.Greater(t, exp, gtm, "EXP (expiration time) should be greater than GTM")
|
||||
|
||||
// UUID might be empty if hash is invalid (which is expected in tests)
|
||||
uuid := c.GetString("uuid")
|
||||
assert.NotNil(t, uuid, "UUID should be set (even if empty)")
|
||||
})
|
||||
|
||||
assert.True(t, server.CallAndGetStatus(t, "Bearer "+tokenString))
|
||||
}
|
||||
|
||||
// TestUserHashValidation_CookieAuth tests uhash validation with cookie authentication
|
||||
func TestUserHashValidation_CookieAuth(t *testing.T) {
|
||||
db := setupTestDB(t)
|
||||
defer db.Close()
|
||||
|
||||
tokenCache := auth.NewTokenCache(db)
|
||||
userCache := auth.NewUserCache(db)
|
||||
authMiddleware := auth.NewAuthMiddleware("/base/url", "test", tokenCache, userCache)
|
||||
|
||||
server := newTestServer(t, "/test", db, authMiddleware.AuthUserRequired)
|
||||
defer server.Close()
|
||||
|
||||
// Create test user with ID=1 and hash="123" to match session
|
||||
var count int
|
||||
db.Model(&models.User{}).Where("id = ?", 1).Count(&count)
|
||||
testUser := models.User{
|
||||
ID: 1,
|
||||
Hash: "123",
|
||||
Mail: "test_user@example.com",
|
||||
Name: "Test User",
|
||||
Status: models.UserStatusActive,
|
||||
RoleID: 2,
|
||||
}
|
||||
if count == 0 {
|
||||
err := db.Create(&testUser).Error
|
||||
require.NoError(t, err)
|
||||
} else {
|
||||
db.First(&testUser, 1)
|
||||
}
|
||||
|
||||
t.Run("correct uhash succeeds", func(t *testing.T) {
|
||||
server.Authorize(t, []string{"test.permission"})
|
||||
assert.True(t, server.CallAndGetStatus(t))
|
||||
})
|
||||
|
||||
t.Run("modified uhash in database fails", func(t *testing.T) {
|
||||
// Update user hash in database
|
||||
db.Model(&testUser).Where("id = ?", 1).Update("hash", "modified_hash")
|
||||
userCache.Invalidate(1)
|
||||
|
||||
// Try to authenticate with old session (has hash="123")
|
||||
assert.False(t, server.CallAndGetStatus(t))
|
||||
})
|
||||
|
||||
t.Run("blocked user fails", func(t *testing.T) {
|
||||
// Restore original hash
|
||||
db.Model(&testUser).Where("id = ?", 1).Update("hash", "123")
|
||||
// Block user
|
||||
db.Model(&testUser).Where("id = ?", 1).Update("status", models.UserStatusBlocked)
|
||||
userCache.Invalidate(1)
|
||||
|
||||
assert.False(t, server.CallAndGetStatus(t))
|
||||
})
|
||||
|
||||
t.Run("deleted user fails", func(t *testing.T) {
|
||||
// Undelete and unblock first
|
||||
db.Model(&models.User{}).Unscoped().Where("id = ?", 1).Update("deleted_at", nil)
|
||||
db.Model(&testUser).Where("id = ?", 1).Update("status", models.UserStatusActive)
|
||||
|
||||
// Delete user
|
||||
db.Delete(&testUser, 1)
|
||||
userCache.Invalidate(1)
|
||||
|
||||
assert.False(t, server.CallAndGetStatus(t))
|
||||
})
|
||||
}
|
||||
|
||||
// TestUserHashValidation_TokenAuth tests uhash validation with token authentication
|
||||
func TestUserHashValidation_TokenAuth(t *testing.T) {
|
||||
db := setupTestDB(t)
|
||||
defer db.Close()
|
||||
|
||||
tokenCache := auth.NewTokenCache(db)
|
||||
userCache := auth.NewUserCache(db)
|
||||
authMiddleware := auth.NewAuthMiddleware("/base/url", "test_salt", tokenCache, userCache)
|
||||
|
||||
server := newTestServer(t, "/protected", db, authMiddleware.AuthTokenRequired)
|
||||
defer server.Close()
|
||||
|
||||
// Create test user
|
||||
testUser := models.User{
|
||||
ID: 200,
|
||||
Hash: "token_test_hash",
|
||||
Mail: "token_user@example.com",
|
||||
Name: "Token Test User",
|
||||
Status: models.UserStatusActive,
|
||||
RoleID: 2,
|
||||
}
|
||||
err := db.Create(&testUser).Error
|
||||
require.NoError(t, err)
|
||||
|
||||
// Create API token
|
||||
tokenID, err := auth.GenerateTokenID()
|
||||
require.NoError(t, err)
|
||||
apiToken := models.APIToken{
|
||||
TokenID: tokenID,
|
||||
UserID: 200,
|
||||
RoleID: 2,
|
||||
TTL: 3600,
|
||||
Status: models.TokenStatusActive,
|
||||
}
|
||||
err = db.Create(&apiToken).Error
|
||||
require.NoError(t, err)
|
||||
|
||||
// Create JWT token with correct hash
|
||||
claims := models.APITokenClaims{
|
||||
TokenID: tokenID,
|
||||
RID: 2,
|
||||
UID: 200,
|
||||
UHASH: "token_test_hash",
|
||||
RegisteredClaims: jwt.RegisteredClaims{
|
||||
ExpiresAt: jwt.NewNumericDate(time.Now().Add(1 * time.Hour)),
|
||||
IssuedAt: jwt.NewNumericDate(time.Now()),
|
||||
Subject: "api_token",
|
||||
},
|
||||
}
|
||||
token := jwt.NewWithClaims(jwt.SigningMethodHS256, claims)
|
||||
tokenString, err := token.SignedString(auth.MakeJWTSigningKey("test_salt"))
|
||||
require.NoError(t, err)
|
||||
|
||||
t.Run("correct uhash succeeds", func(t *testing.T) {
|
||||
success := server.CallAndGetStatus(t, "Bearer "+tokenString)
|
||||
assert.True(t, success)
|
||||
})
|
||||
|
||||
t.Run("modified uhash in database fails", func(t *testing.T) {
|
||||
// Update user hash in database
|
||||
db.Model(&testUser).Update("hash", "different_hash")
|
||||
userCache.Invalidate(200)
|
||||
|
||||
// Try to authenticate with token (has original hash)
|
||||
success := server.CallAndGetStatus(t, "Bearer "+tokenString)
|
||||
assert.False(t, success)
|
||||
})
|
||||
|
||||
t.Run("blocked user fails", func(t *testing.T) {
|
||||
// Restore original hash
|
||||
db.Model(&testUser).Update("hash", "token_test_hash")
|
||||
// Block user
|
||||
db.Model(&testUser).Update("status", models.UserStatusBlocked)
|
||||
userCache.Invalidate(200)
|
||||
|
||||
success := server.CallAndGetStatus(t, "Bearer "+tokenString)
|
||||
assert.False(t, success)
|
||||
})
|
||||
|
||||
t.Run("deleted user fails", func(t *testing.T) {
|
||||
// Unblock and restore for clean state
|
||||
db.Model(&models.User{}).Unscoped().Where("id = ?", 200).Update("deleted_at", nil)
|
||||
db.Model(&testUser).Update("status", models.UserStatusActive)
|
||||
|
||||
// Delete user
|
||||
db.Delete(&testUser)
|
||||
userCache.Invalidate(200)
|
||||
|
||||
success := server.CallAndGetStatus(t, "Bearer "+tokenString)
|
||||
assert.False(t, success)
|
||||
})
|
||||
}
|
||||
|
||||
// TestUserHashValidation_CrossInstallation simulates different installations
|
||||
func TestUserHashValidation_CrossInstallation(t *testing.T) {
|
||||
db := setupTestDB(t)
|
||||
defer db.Close()
|
||||
|
||||
tokenCache := auth.NewTokenCache(db)
|
||||
userCache := auth.NewUserCache(db)
|
||||
authMiddleware := auth.NewAuthMiddleware("/base/url", "test_salt", tokenCache, userCache)
|
||||
|
||||
server := newTestServer(t, "/protected", db, authMiddleware.AuthTokenRequired)
|
||||
defer server.Close()
|
||||
|
||||
// Simulate Installation A
|
||||
userInstallationA := models.User{
|
||||
ID: 300,
|
||||
Hash: "installation_a_hash",
|
||||
Mail: "cross@example.com",
|
||||
Name: "Cross Installation User",
|
||||
Status: models.UserStatusActive,
|
||||
RoleID: 2,
|
||||
}
|
||||
err := db.Create(&userInstallationA).Error
|
||||
require.NoError(t, err)
|
||||
|
||||
// Create API token for Installation A
|
||||
tokenID, err := auth.GenerateTokenID()
|
||||
require.NoError(t, err)
|
||||
apiToken := models.APIToken{
|
||||
TokenID: tokenID,
|
||||
UserID: 300,
|
||||
RoleID: 2,
|
||||
TTL: 3600,
|
||||
Status: models.TokenStatusActive,
|
||||
}
|
||||
err = db.Create(&apiToken).Error
|
||||
require.NoError(t, err)
|
||||
|
||||
// Create JWT token with Installation A hash
|
||||
claimsA := models.APITokenClaims{
|
||||
TokenID: tokenID,
|
||||
RID: 2,
|
||||
UID: 300,
|
||||
UHASH: "installation_a_hash",
|
||||
RegisteredClaims: jwt.RegisteredClaims{
|
||||
ExpiresAt: jwt.NewNumericDate(time.Now().Add(1 * time.Hour)),
|
||||
IssuedAt: jwt.NewNumericDate(time.Now()),
|
||||
Subject: "api_token",
|
||||
},
|
||||
}
|
||||
tokenA := jwt.NewWithClaims(jwt.SigningMethodHS256, claimsA)
|
||||
tokenStringA, err := tokenA.SignedString(auth.MakeJWTSigningKey("test_salt"))
|
||||
require.NoError(t, err)
|
||||
|
||||
t.Run("token works on Installation A", func(t *testing.T) {
|
||||
success := server.CallAndGetStatus(t, "Bearer "+tokenStringA)
|
||||
assert.True(t, success)
|
||||
})
|
||||
|
||||
t.Run("token from Installation A fails on Installation B", func(t *testing.T) {
|
||||
// Simulate Installation B - user has different hash
|
||||
db.Model(&userInstallationA).Update("hash", "installation_b_hash")
|
||||
userCache.Invalidate(300)
|
||||
|
||||
// Try to use token from Installation A (has installation_a_hash)
|
||||
success := server.CallAndGetStatus(t, "Bearer "+tokenStringA)
|
||||
assert.False(t, success, "Token from Installation A should not work on Installation B")
|
||||
})
|
||||
|
||||
t.Run("new token from Installation B works", func(t *testing.T) {
|
||||
// Create new token for Installation B
|
||||
tokenIDB, err := auth.GenerateTokenID()
|
||||
require.NoError(t, err)
|
||||
apiTokenB := models.APIToken{
|
||||
TokenID: tokenIDB,
|
||||
UserID: 300,
|
||||
RoleID: 2,
|
||||
TTL: 3600,
|
||||
Status: models.TokenStatusActive,
|
||||
}
|
||||
err = db.Create(&apiTokenB).Error
|
||||
require.NoError(t, err)
|
||||
|
||||
// Create JWT token with Installation B hash
|
||||
claimsB := models.APITokenClaims{
|
||||
TokenID: tokenIDB,
|
||||
RID: 2,
|
||||
UID: 300,
|
||||
UHASH: "installation_b_hash", // correct hash for Installation B
|
||||
RegisteredClaims: jwt.RegisteredClaims{
|
||||
ExpiresAt: jwt.NewNumericDate(time.Now().Add(1 * time.Hour)),
|
||||
IssuedAt: jwt.NewNumericDate(time.Now()),
|
||||
Subject: "api_token",
|
||||
},
|
||||
}
|
||||
tokenB := jwt.NewWithClaims(jwt.SigningMethodHS256, claimsB)
|
||||
tokenStringB, err := tokenB.SignedString(auth.MakeJWTSigningKey("test_salt"))
|
||||
require.NoError(t, err)
|
||||
|
||||
// Token from Installation B should work
|
||||
success := server.CallAndGetStatus(t, "Bearer "+tokenStringB)
|
||||
assert.True(t, success, "New token from Installation B should work")
|
||||
})
|
||||
}
|
||||
@@ -0,0 +1,46 @@
|
||||
package auth
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"slices"
|
||||
|
||||
"pentagi/pkg/server/response"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
)
|
||||
|
||||
func getPrms(c *gin.Context) ([]string, error) {
|
||||
prms := c.GetStringSlice("prm")
|
||||
if len(prms) == 0 {
|
||||
return nil, fmt.Errorf("privileges are not set")
|
||||
}
|
||||
return prms, nil
|
||||
}
|
||||
|
||||
func PrivilegesRequired(privs ...string) gin.HandlerFunc {
|
||||
return func(c *gin.Context) {
|
||||
if c.IsAborted() {
|
||||
return
|
||||
}
|
||||
|
||||
prms, err := getPrms(c)
|
||||
if err != nil {
|
||||
response.Error(c, response.ErrPrivilegesRequired, err)
|
||||
c.Abort()
|
||||
return
|
||||
}
|
||||
|
||||
for _, priv := range privs {
|
||||
if !LookupPerm(prms, priv) {
|
||||
response.Error(c, response.ErrPrivilegesRequired, fmt.Errorf("'%s' is not set", priv))
|
||||
c.Abort()
|
||||
return
|
||||
}
|
||||
}
|
||||
c.Next()
|
||||
}
|
||||
}
|
||||
|
||||
func LookupPerm(prm []string, perm string) bool {
|
||||
return slices.Contains(prm, perm)
|
||||
}
|
||||
@@ -0,0 +1,36 @@
|
||||
package auth_test
|
||||
|
||||
import (
|
||||
"pentagi/pkg/server/auth"
|
||||
"testing"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/stretchr/testify/assert"
|
||||
)
|
||||
|
||||
func TestPrivilegesRequired(t *testing.T) {
|
||||
db := setupTestDB(t)
|
||||
defer db.Close()
|
||||
|
||||
tokenCache := auth.NewTokenCache(db)
|
||||
userCache := auth.NewUserCache(db)
|
||||
authMiddleware := auth.NewAuthMiddleware("/base/url", "test", tokenCache, userCache)
|
||||
server := newTestServer(t, "/test", db, authMiddleware.AuthTokenRequired, auth.PrivilegesRequired("priv1", "priv2"))
|
||||
defer server.Close()
|
||||
|
||||
server.SetSessionCheckFunc(func(t *testing.T, c *gin.Context) {
|
||||
t.Helper()
|
||||
assert.Equal(t, uint64(1), c.GetUint64("uid"))
|
||||
})
|
||||
|
||||
assert.False(t, server.CallAndGetStatus(t))
|
||||
|
||||
server.Authorize(t, []string{"some.permission"})
|
||||
assert.False(t, server.CallAndGetStatus(t))
|
||||
|
||||
server.Authorize(t, []string{"priv1"})
|
||||
assert.False(t, server.CallAndGetStatus(t))
|
||||
|
||||
server.Authorize(t, []string{"priv1", "priv2"})
|
||||
assert.True(t, server.CallAndGetStatus(t))
|
||||
}
|
||||
@@ -0,0 +1,71 @@
|
||||
package auth
|
||||
|
||||
import (
|
||||
"crypto/sha512"
|
||||
"strings"
|
||||
"sync"
|
||||
|
||||
"golang.org/x/crypto/pbkdf2"
|
||||
)
|
||||
|
||||
var (
|
||||
cookieStoreKeys sync.Map // cache of cookie keys per salt
|
||||
jwtSigningKeys sync.Map // cache of JWT signing keys per salt
|
||||
)
|
||||
|
||||
const (
|
||||
pbkdf2Iterations = 210000 // OWASP 2023 recommendation
|
||||
jwtKeyLength = 32 // 256 bits for HS256
|
||||
authKeyLength = 64 // 512 bits for cookie auth key
|
||||
encKeyLength = 32 // 256 bits for cookie encryption key
|
||||
)
|
||||
|
||||
// MakeCookieStoreKey is function to generate auth and encryption keys for cookie store
|
||||
func MakeCookieStoreKey(globalSalt string) [][]byte {
|
||||
// Check cache for existing keys
|
||||
if cached, ok := cookieStoreKeys.Load(globalSalt); ok {
|
||||
return cached.([][]byte)
|
||||
}
|
||||
|
||||
// Generate new keys for this salt using PBKDF2
|
||||
password := []byte(strings.Join([]string{
|
||||
"a8d0abae36f749588f4393e6fc292690",
|
||||
globalSalt,
|
||||
"7c9be62adec5076970fa946e78f256e2",
|
||||
}, "|"))
|
||||
|
||||
// Auth key (64 bytes) - using salt variant 1
|
||||
authSalt := []byte("pentagi.cookie.auth|" + globalSalt)
|
||||
authKey := pbkdf2.Key(password, authSalt, pbkdf2Iterations, authKeyLength, sha512.New)
|
||||
|
||||
// Encryption key (32 bytes) - using salt variant 2
|
||||
encSalt := []byte("pentagi.cookie.enc|" + globalSalt)
|
||||
encKey := pbkdf2.Key(password, encSalt, pbkdf2Iterations, encKeyLength, sha512.New)
|
||||
|
||||
newKeys := [][]byte{authKey, encKey}
|
||||
|
||||
// Store in cache (LoadOrStore handles concurrent access)
|
||||
actual, _ := cookieStoreKeys.LoadOrStore(globalSalt, newKeys)
|
||||
return actual.([][]byte)
|
||||
}
|
||||
|
||||
// MakeJWTSigningKey is function to generate signing key for JWT tokens
|
||||
func MakeJWTSigningKey(globalSalt string) []byte {
|
||||
// Check cache for existing key
|
||||
if cached, ok := jwtSigningKeys.Load(globalSalt); ok {
|
||||
return cached.([]byte)
|
||||
}
|
||||
|
||||
// Generate new key for this salt using PBKDF2
|
||||
password := []byte(strings.Join([]string{
|
||||
"4c1e9cb77df7f9a58fcc5f52d40af685",
|
||||
globalSalt,
|
||||
"09784e190148d13d48885aa47cf8a297",
|
||||
}, "|"))
|
||||
salt := []byte("pentagi.jwt.signing|" + globalSalt)
|
||||
newKey := pbkdf2.Key(password, salt, pbkdf2Iterations, jwtKeyLength, sha512.New)
|
||||
|
||||
// Store in cache (LoadOrStore handles concurrent access)
|
||||
actual, _ := jwtSigningKeys.LoadOrStore(globalSalt, newKey)
|
||||
return actual.([]byte)
|
||||
}
|
||||
@@ -0,0 +1,61 @@
|
||||
package auth_test
|
||||
|
||||
import (
|
||||
"pentagi/pkg/server/auth"
|
||||
"testing"
|
||||
|
||||
"github.com/stretchr/testify/assert"
|
||||
)
|
||||
|
||||
func TestMakeJWTSigningKey(t *testing.T) {
|
||||
salt1 := "test_salt_1"
|
||||
salt2 := "test_salt_2"
|
||||
|
||||
// Test that key is generated
|
||||
key1 := auth.MakeJWTSigningKey(salt1)
|
||||
assert.NotNil(t, key1)
|
||||
assert.Len(t, key1, 32, "JWT signing key should be 32 bytes (256 bits)")
|
||||
|
||||
// Test that same salt produces same key (cached)
|
||||
key1Again := auth.MakeJWTSigningKey(salt1)
|
||||
assert.Equal(t, key1, key1Again, "Same salt should produce same key from cache")
|
||||
|
||||
// Test that different salts produce different keys
|
||||
key2 := auth.MakeJWTSigningKey(salt2)
|
||||
assert.NotEqual(t, key1, key2, "Different salts should produce different keys")
|
||||
assert.Len(t, key2, 32, "JWT signing key should be 32 bytes (256 bits)")
|
||||
|
||||
// Verify consistency for salt2
|
||||
key2Again := auth.MakeJWTSigningKey(salt2)
|
||||
assert.Equal(t, key2, key2Again, "Same salt should produce same key from cache")
|
||||
}
|
||||
|
||||
func TestMakeCookieStoreKey(t *testing.T) {
|
||||
salt := "test_salt"
|
||||
|
||||
// Test that keys are generated
|
||||
keys := auth.MakeCookieStoreKey(salt)
|
||||
assert.NotNil(t, keys)
|
||||
assert.Len(t, keys, 2, "Should return auth and encryption keys")
|
||||
|
||||
// Test that auth key is 64 bytes (SHA512)
|
||||
assert.Len(t, keys[0], 64, "Auth key should be 64 bytes")
|
||||
|
||||
// Test that encryption key is 32 bytes (SHA256)
|
||||
assert.Len(t, keys[1], 32, "Encryption key should be 32 bytes")
|
||||
|
||||
// Test consistency
|
||||
keysAgain := auth.MakeCookieStoreKey(salt)
|
||||
assert.Equal(t, keys, keysAgain, "Same salt should produce same keys")
|
||||
}
|
||||
|
||||
func TestMakeJWTSigningKeyDifferentFromCookieKey(t *testing.T) {
|
||||
salt := "test_salt"
|
||||
|
||||
jwtKey := auth.MakeJWTSigningKey(salt)
|
||||
cookieKeys := auth.MakeCookieStoreKey(salt)
|
||||
|
||||
// JWT signing key should be different from both cookie keys
|
||||
assert.NotEqual(t, jwtKey, cookieKeys[0], "JWT key should differ from cookie auth key")
|
||||
assert.NotEqual(t, jwtKey, cookieKeys[1], "JWT key should differ from cookie encryption key")
|
||||
}
|
||||
@@ -0,0 +1,92 @@
|
||||
package auth
|
||||
|
||||
import (
|
||||
"sync"
|
||||
"time"
|
||||
|
||||
"pentagi/pkg/server/models"
|
||||
|
||||
"github.com/jinzhu/gorm"
|
||||
)
|
||||
|
||||
// userCacheEntry represents a cached user status entry
|
||||
type userCacheEntry struct {
|
||||
hash string
|
||||
status models.UserStatus
|
||||
notFound bool // negative caching
|
||||
expiresAt time.Time
|
||||
}
|
||||
|
||||
// UserCache provides caching for user hash lookups
|
||||
type UserCache struct {
|
||||
cache sync.Map
|
||||
ttl time.Duration
|
||||
db *gorm.DB
|
||||
}
|
||||
|
||||
// NewUserCache creates a new user cache instance
|
||||
func NewUserCache(db *gorm.DB) *UserCache {
|
||||
return &UserCache{
|
||||
ttl: 5 * time.Minute,
|
||||
db: db,
|
||||
}
|
||||
}
|
||||
|
||||
// SetTTL sets the TTL for the user cache
|
||||
func (uc *UserCache) SetTTL(ttl time.Duration) {
|
||||
uc.ttl = ttl
|
||||
}
|
||||
|
||||
// GetUserHash retrieves user hash and status from cache or database
|
||||
func (uc *UserCache) GetUserHash(userID uint64) (string, models.UserStatus, error) {
|
||||
// check cache first
|
||||
if entry, ok := uc.cache.Load(userID); ok {
|
||||
cached := entry.(userCacheEntry)
|
||||
if time.Now().Before(cached.expiresAt) {
|
||||
// return cached "not found" error
|
||||
if cached.notFound {
|
||||
return "", "", gorm.ErrRecordNotFound
|
||||
}
|
||||
return cached.hash, cached.status, nil
|
||||
}
|
||||
// cache entry expired, remove it
|
||||
uc.cache.Delete(userID)
|
||||
}
|
||||
|
||||
// load from database
|
||||
var user models.User
|
||||
if err := uc.db.Where("id = ?", userID).First(&user).Error; err != nil {
|
||||
if gorm.IsRecordNotFoundError(err) {
|
||||
// cache negative result (user not found)
|
||||
uc.cache.Store(userID, userCacheEntry{
|
||||
notFound: true,
|
||||
expiresAt: time.Now().Add(uc.ttl),
|
||||
})
|
||||
return "", "", gorm.ErrRecordNotFound
|
||||
}
|
||||
return "", "", err
|
||||
}
|
||||
|
||||
// update cache with positive result
|
||||
uc.cache.Store(userID, userCacheEntry{
|
||||
hash: user.Hash,
|
||||
status: user.Status,
|
||||
notFound: false,
|
||||
expiresAt: time.Now().Add(uc.ttl),
|
||||
})
|
||||
|
||||
return user.Hash, user.Status, nil
|
||||
}
|
||||
|
||||
// Invalidate removes a specific user from cache
|
||||
func (uc *UserCache) Invalidate(userID uint64) {
|
||||
uc.cache.Delete(userID)
|
||||
}
|
||||
|
||||
// InvalidateAll clears the entire cache
|
||||
func (uc *UserCache) InvalidateAll() {
|
||||
uc.cache.Range(func(key, value any) bool {
|
||||
uc.cache.Delete(key)
|
||||
return true
|
||||
})
|
||||
}
|
||||
@@ -0,0 +1,443 @@
|
||||
package auth_test
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"sync"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"pentagi/pkg/server/auth"
|
||||
"pentagi/pkg/server/models"
|
||||
|
||||
"github.com/jinzhu/gorm"
|
||||
_ "github.com/jinzhu/gorm/dialects/sqlite"
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
func setupUserTestDB(t *testing.T) *gorm.DB {
|
||||
t.Helper()
|
||||
db, err := gorm.Open("sqlite3", ":memory:")
|
||||
require.NoError(t, err)
|
||||
|
||||
// Create users table
|
||||
db.Exec(`
|
||||
CREATE TABLE users (
|
||||
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
||||
hash TEXT NOT NULL UNIQUE,
|
||||
type TEXT NOT NULL DEFAULT 'local',
|
||||
mail TEXT NOT NULL UNIQUE,
|
||||
name TEXT NOT NULL DEFAULT '',
|
||||
status TEXT NOT NULL DEFAULT 'active',
|
||||
role_id INTEGER NOT NULL DEFAULT 2,
|
||||
password TEXT,
|
||||
password_change_required BOOLEAN NOT NULL DEFAULT false,
|
||||
provider TEXT,
|
||||
created_at DATETIME DEFAULT CURRENT_TIMESTAMP,
|
||||
deleted_at DATETIME
|
||||
)
|
||||
`)
|
||||
|
||||
// Create user_preferences table
|
||||
db.Exec(`
|
||||
CREATE TABLE user_preferences (
|
||||
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
||||
user_id INTEGER NOT NULL UNIQUE,
|
||||
preferences TEXT NOT NULL DEFAULT '{"favoriteFlows": []}',
|
||||
created_at DATETIME DEFAULT CURRENT_TIMESTAMP,
|
||||
updated_at DATETIME DEFAULT CURRENT_TIMESTAMP,
|
||||
FOREIGN KEY (user_id) REFERENCES users(id) ON DELETE CASCADE
|
||||
)
|
||||
`)
|
||||
|
||||
time.Sleep(200 * time.Millisecond) // wait for database to be ready
|
||||
|
||||
return db
|
||||
}
|
||||
|
||||
func TestUserCache_GetUserHash(t *testing.T) {
|
||||
db := setupUserTestDB(t)
|
||||
defer db.Close()
|
||||
|
||||
cache := auth.NewUserCache(db)
|
||||
|
||||
// Insert test user
|
||||
user := models.User{
|
||||
ID: 1,
|
||||
Hash: "test_hash_123",
|
||||
Mail: "test@example.com",
|
||||
Name: "Test User",
|
||||
Status: models.UserStatusActive,
|
||||
RoleID: 2,
|
||||
}
|
||||
err := db.Create(&user).Error
|
||||
require.NoError(t, err)
|
||||
|
||||
// Test: Get user hash (should hit database)
|
||||
hash, status, err := cache.GetUserHash(1)
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, "test_hash_123", hash)
|
||||
assert.Equal(t, models.UserStatusActive, status)
|
||||
|
||||
// Test: Get user hash again (should hit cache)
|
||||
hash, status, err = cache.GetUserHash(1)
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, "test_hash_123", hash)
|
||||
assert.Equal(t, models.UserStatusActive, status)
|
||||
|
||||
// Test: Non-existent user
|
||||
_, _, err = cache.GetUserHash(999)
|
||||
assert.Error(t, err)
|
||||
assert.Equal(t, gorm.ErrRecordNotFound, err)
|
||||
}
|
||||
|
||||
func TestUserCache_Invalidate(t *testing.T) {
|
||||
db := setupUserTestDB(t)
|
||||
defer db.Close()
|
||||
|
||||
cache := auth.NewUserCache(db)
|
||||
|
||||
// Insert test user
|
||||
user := models.User{
|
||||
ID: 1,
|
||||
Hash: "test_hash_456",
|
||||
Mail: "test2@example.com",
|
||||
Name: "Test User 2",
|
||||
Status: models.UserStatusActive,
|
||||
RoleID: 2,
|
||||
}
|
||||
err := db.Create(&user).Error
|
||||
require.NoError(t, err)
|
||||
|
||||
// Get hash to populate cache
|
||||
hash, status, err := cache.GetUserHash(1)
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, "test_hash_456", hash)
|
||||
assert.Equal(t, models.UserStatusActive, status)
|
||||
|
||||
// Update user in database
|
||||
db.Model(&user).Update("status", models.UserStatusBlocked)
|
||||
|
||||
// Status should still be active (from cache)
|
||||
hash, status, err = cache.GetUserHash(1)
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, "test_hash_456", hash)
|
||||
assert.Equal(t, models.UserStatusActive, status)
|
||||
|
||||
// Invalidate cache
|
||||
cache.Invalidate(1)
|
||||
|
||||
// Status should now be blocked (from database)
|
||||
hash, status, err = cache.GetUserHash(1)
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, "test_hash_456", hash)
|
||||
assert.Equal(t, models.UserStatusBlocked, status)
|
||||
}
|
||||
|
||||
func TestUserCache_Expiration(t *testing.T) {
|
||||
db := setupUserTestDB(t)
|
||||
defer db.Close()
|
||||
|
||||
// Create cache with very short TTL for testing
|
||||
cache := auth.NewUserCache(db)
|
||||
cache.SetTTL(300 * time.Millisecond)
|
||||
|
||||
// Insert test user
|
||||
user := models.User{
|
||||
ID: 1,
|
||||
Hash: "test_hash_789",
|
||||
Mail: "test3@example.com",
|
||||
Name: "Test User 3",
|
||||
Status: models.UserStatusActive,
|
||||
RoleID: 2,
|
||||
}
|
||||
err := db.Create(&user).Error
|
||||
require.NoError(t, err)
|
||||
|
||||
// Get hash to populate cache
|
||||
hash, status, err := cache.GetUserHash(1)
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, "test_hash_789", hash)
|
||||
assert.Equal(t, models.UserStatusActive, status)
|
||||
|
||||
// Update user in database
|
||||
db.Model(&user).Update("status", models.UserStatusBlocked)
|
||||
|
||||
// Wait for cache to expire
|
||||
time.Sleep(500 * time.Millisecond)
|
||||
|
||||
// Status should now be blocked (cache expired, reading from DB)
|
||||
hash, status, err = cache.GetUserHash(1)
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, "test_hash_789", hash)
|
||||
assert.Equal(t, models.UserStatusBlocked, status)
|
||||
}
|
||||
|
||||
func TestUserCache_UserStatuses(t *testing.T) {
|
||||
db := setupUserTestDB(t)
|
||||
defer db.Close()
|
||||
|
||||
cache := auth.NewUserCache(db)
|
||||
|
||||
testCases := []struct {
|
||||
name string
|
||||
userStatus models.UserStatus
|
||||
expectedStatus models.UserStatus
|
||||
}{
|
||||
{
|
||||
name: "active user",
|
||||
userStatus: models.UserStatusActive,
|
||||
expectedStatus: models.UserStatusActive,
|
||||
},
|
||||
{
|
||||
name: "blocked user",
|
||||
userStatus: models.UserStatusBlocked,
|
||||
expectedStatus: models.UserStatusBlocked,
|
||||
},
|
||||
{
|
||||
name: "created user",
|
||||
userStatus: models.UserStatusCreated,
|
||||
expectedStatus: models.UserStatusCreated,
|
||||
},
|
||||
}
|
||||
|
||||
for i, tc := range testCases {
|
||||
t.Run(tc.name, func(t *testing.T) {
|
||||
user := models.User{
|
||||
ID: uint64(i + 1),
|
||||
Hash: "hash_" + tc.name,
|
||||
Mail: tc.name + "@example.com",
|
||||
Name: tc.name,
|
||||
Status: tc.userStatus,
|
||||
RoleID: 2,
|
||||
}
|
||||
err := db.Create(&user).Error
|
||||
require.NoError(t, err)
|
||||
|
||||
hash, status, err := cache.GetUserHash(user.ID)
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, user.Hash, hash)
|
||||
assert.Equal(t, tc.expectedStatus, status)
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestUserCache_DeletedUser(t *testing.T) {
|
||||
db := setupUserTestDB(t)
|
||||
defer db.Close()
|
||||
|
||||
cache := auth.NewUserCache(db)
|
||||
|
||||
// Insert test user
|
||||
user := models.User{
|
||||
ID: 1,
|
||||
Hash: "deleted_hash",
|
||||
Mail: "deleted@example.com",
|
||||
Name: "Deleted User",
|
||||
Status: models.UserStatusActive,
|
||||
RoleID: 2,
|
||||
}
|
||||
err := db.Create(&user).Error
|
||||
require.NoError(t, err)
|
||||
|
||||
// Get hash to populate cache
|
||||
hash, status, err := cache.GetUserHash(1)
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, "deleted_hash", hash)
|
||||
assert.Equal(t, models.UserStatusActive, status)
|
||||
|
||||
// Soft delete user
|
||||
db.Delete(&user)
|
||||
|
||||
// Invalidate cache
|
||||
cache.Invalidate(1)
|
||||
|
||||
// Should return error for deleted user
|
||||
_, _, err = cache.GetUserHash(1)
|
||||
assert.Error(t, err)
|
||||
assert.Equal(t, gorm.ErrRecordNotFound, err)
|
||||
}
|
||||
|
||||
func TestUserCache_ConcurrentAccess(t *testing.T) {
|
||||
db := setupUserTestDB(t)
|
||||
defer db.Close()
|
||||
|
||||
cache := auth.NewUserCache(db)
|
||||
|
||||
// Insert test users
|
||||
for i := 1; i <= 10; i++ {
|
||||
user := models.User{
|
||||
ID: uint64(i),
|
||||
Hash: fmt.Sprintf("concurrent_hash_%d", i),
|
||||
Mail: fmt.Sprintf("concurrent%d@example.com", i),
|
||||
Name: "Concurrent User",
|
||||
Status: models.UserStatusActive,
|
||||
RoleID: 2,
|
||||
}
|
||||
err := db.Create(&user).Error
|
||||
require.NoError(t, err)
|
||||
}
|
||||
|
||||
// warm up cache
|
||||
for i := range 10 {
|
||||
_, _, err := cache.GetUserHash(uint64(i%10 + 1))
|
||||
require.NoError(t, err)
|
||||
}
|
||||
|
||||
var wg sync.WaitGroup
|
||||
errors := make(chan error, 100)
|
||||
|
||||
// Concurrent reads
|
||||
for i := range 10 {
|
||||
wg.Add(1)
|
||||
go func(userID uint64) {
|
||||
defer wg.Done()
|
||||
for range 10 {
|
||||
_, _, err := cache.GetUserHash(userID)
|
||||
if err != nil {
|
||||
errors <- err
|
||||
}
|
||||
}
|
||||
}(uint64(i%10 + 1))
|
||||
}
|
||||
|
||||
// Concurrent invalidations
|
||||
for i := range 5 {
|
||||
wg.Add(1)
|
||||
go func(userID uint64) {
|
||||
defer wg.Done()
|
||||
for range 5 {
|
||||
cache.Invalidate(userID)
|
||||
time.Sleep(10 * time.Millisecond)
|
||||
}
|
||||
}(uint64(i%10 + 1))
|
||||
}
|
||||
|
||||
wg.Wait()
|
||||
close(errors)
|
||||
|
||||
// Check for errors
|
||||
for err := range errors {
|
||||
t.Errorf("Concurrent access error: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestUserCache_InvalidateAll(t *testing.T) {
|
||||
db := setupUserTestDB(t)
|
||||
defer db.Close()
|
||||
|
||||
cache := auth.NewUserCache(db)
|
||||
|
||||
// Insert multiple users
|
||||
for i := 1; i <= 5; i++ {
|
||||
user := models.User{
|
||||
ID: uint64(i),
|
||||
Hash: fmt.Sprintf("invalidate_all_%d", i),
|
||||
Mail: fmt.Sprintf("all%d@example.com", i),
|
||||
Name: fmt.Sprintf("User %d", i),
|
||||
Status: models.UserStatusActive,
|
||||
RoleID: 2,
|
||||
}
|
||||
err := db.Create(&user).Error
|
||||
require.NoError(t, err)
|
||||
}
|
||||
|
||||
// Populate cache
|
||||
for i := 1; i <= 5; i++ {
|
||||
_, _, err := cache.GetUserHash(uint64(i))
|
||||
require.NoError(t, err)
|
||||
}
|
||||
|
||||
// Update all users in database
|
||||
db.Model(&models.User{}).Where("id > 0").Update("status", models.UserStatusBlocked)
|
||||
|
||||
// Invalidate all
|
||||
cache.InvalidateAll()
|
||||
|
||||
// All users should now show blocked status
|
||||
for i := 1; i <= 5; i++ {
|
||||
_, status, err := cache.GetUserHash(uint64(i))
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, models.UserStatusBlocked, status)
|
||||
}
|
||||
}
|
||||
|
||||
func TestUserCache_NegativeCaching(t *testing.T) {
|
||||
db := setupUserTestDB(t)
|
||||
defer db.Close()
|
||||
|
||||
cache := auth.NewUserCache(db)
|
||||
nonExistentUserID := uint64(9999)
|
||||
|
||||
// First call - should hit database and cache the "not found"
|
||||
_, _, err := cache.GetUserHash(nonExistentUserID)
|
||||
require.Error(t, err)
|
||||
assert.Equal(t, gorm.ErrRecordNotFound, err)
|
||||
|
||||
// Second call - should return from cache without hitting DB
|
||||
_, _, err = cache.GetUserHash(nonExistentUserID)
|
||||
require.Error(t, err)
|
||||
assert.Equal(t, gorm.ErrRecordNotFound, err, "Should return cached not found error")
|
||||
|
||||
// Now create the user in DB
|
||||
user := models.User{
|
||||
ID: nonExistentUserID,
|
||||
Hash: "new_user_hash",
|
||||
Mail: "new@example.com",
|
||||
Name: "New User",
|
||||
Status: models.UserStatusActive,
|
||||
RoleID: 2,
|
||||
}
|
||||
err = db.Create(&user).Error
|
||||
require.NoError(t, err)
|
||||
|
||||
// Should still return cached "not found" until invalidated
|
||||
_, _, err = cache.GetUserHash(nonExistentUserID)
|
||||
require.Error(t, err)
|
||||
assert.Equal(t, gorm.ErrRecordNotFound, err, "Should still return cached not found")
|
||||
|
||||
// Invalidate cache
|
||||
cache.Invalidate(nonExistentUserID)
|
||||
|
||||
// Now should find the user
|
||||
hash, status, err := cache.GetUserHash(nonExistentUserID)
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, "new_user_hash", hash)
|
||||
assert.Equal(t, models.UserStatusActive, status)
|
||||
}
|
||||
|
||||
func TestUserCache_NegativeCachingExpiration(t *testing.T) {
|
||||
db := setupUserTestDB(t)
|
||||
defer db.Close()
|
||||
|
||||
cache := auth.NewUserCache(db)
|
||||
cache.SetTTL(300 * time.Millisecond)
|
||||
|
||||
nonExistentUserID := uint64(8888)
|
||||
|
||||
// First call - cache the "not found"
|
||||
_, _, err := cache.GetUserHash(nonExistentUserID)
|
||||
require.Error(t, err)
|
||||
assert.Equal(t, gorm.ErrRecordNotFound, err)
|
||||
|
||||
// Create user in DB
|
||||
user := models.User{
|
||||
ID: nonExistentUserID,
|
||||
Hash: "temp_user_hash",
|
||||
Mail: "temp@example.com",
|
||||
Name: "Temp User",
|
||||
Status: models.UserStatusActive,
|
||||
RoleID: 2,
|
||||
}
|
||||
err = db.Create(&user).Error
|
||||
require.NoError(t, err)
|
||||
|
||||
// Wait for cache to expire
|
||||
time.Sleep(500 * time.Millisecond)
|
||||
|
||||
// Now should find the user (cache expired)
|
||||
hash, status, err := cache.GetUserHash(nonExistentUserID)
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, "temp_user_hash", hash)
|
||||
assert.Equal(t, models.UserStatusActive, status)
|
||||
}
|
||||
@@ -0,0 +1,61 @@
|
||||
package context
|
||||
|
||||
import (
|
||||
"github.com/gin-contrib/sessions"
|
||||
"github.com/gin-gonic/gin"
|
||||
)
|
||||
|
||||
// GetInt64 is function to get some int64 value from gin context
|
||||
func GetInt64(c *gin.Context, key string) (int64, bool) {
|
||||
if iv, ok := c.Get(key); !ok {
|
||||
return 0, false
|
||||
} else if v, ok := iv.(int64); !ok {
|
||||
return 0, false
|
||||
} else {
|
||||
return v, true
|
||||
}
|
||||
}
|
||||
|
||||
// GetUint64 is function to get some uint64 value from gin context
|
||||
func GetUint64(c *gin.Context, key string) (uint64, bool) {
|
||||
if iv, ok := c.Get(key); !ok {
|
||||
return 0, false
|
||||
} else if v, ok := iv.(uint64); !ok {
|
||||
return 0, false
|
||||
} else {
|
||||
return v, true
|
||||
}
|
||||
}
|
||||
|
||||
// GetString is function to get some string value from gin context
|
||||
func GetString(c *gin.Context, key string) (string, bool) {
|
||||
if iv, ok := c.Get(key); !ok {
|
||||
return "", false
|
||||
} else if v, ok := iv.(string); !ok {
|
||||
return "", false
|
||||
} else {
|
||||
return v, true
|
||||
}
|
||||
}
|
||||
|
||||
// GetStringArray is function to get some string array value from gin context
|
||||
func GetStringArray(c *gin.Context, key string) ([]string, bool) {
|
||||
if iv, ok := c.Get(key); !ok {
|
||||
return []string{}, false
|
||||
} else if v, ok := iv.([]string); !ok {
|
||||
return []string{}, false
|
||||
} else {
|
||||
return v, true
|
||||
}
|
||||
}
|
||||
|
||||
func GetStringFromSession(c *gin.Context, key string) (string, bool) {
|
||||
session := sessions.Default(c)
|
||||
if iv := session.Get(key); iv == nil {
|
||||
return "", false
|
||||
} else if v, ok := iv.(string); !ok {
|
||||
return "", false
|
||||
} else {
|
||||
return v, true
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,510 @@
|
||||
package context
|
||||
|
||||
import (
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"testing"
|
||||
|
||||
"github.com/gin-contrib/sessions"
|
||||
"github.com/gin-contrib/sessions/cookie"
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/stretchr/testify/assert"
|
||||
)
|
||||
|
||||
func init() {
|
||||
gin.SetMode(gin.TestMode)
|
||||
}
|
||||
|
||||
func TestGetInt64(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
tests := []struct {
|
||||
name string
|
||||
setup func(c *gin.Context)
|
||||
key string
|
||||
wantVal int64
|
||||
wantOK bool
|
||||
}{
|
||||
{
|
||||
name: "found",
|
||||
setup: func(c *gin.Context) { c.Set("id", int64(42)) },
|
||||
key: "id",
|
||||
wantVal: 42,
|
||||
wantOK: true,
|
||||
},
|
||||
{
|
||||
name: "missing",
|
||||
setup: func(c *gin.Context) {},
|
||||
key: "id",
|
||||
wantVal: 0,
|
||||
wantOK: false,
|
||||
},
|
||||
{
|
||||
name: "wrong type string",
|
||||
setup: func(c *gin.Context) { c.Set("id", "not-an-int") },
|
||||
key: "id",
|
||||
wantVal: 0,
|
||||
wantOK: false,
|
||||
},
|
||||
{
|
||||
name: "wrong type uint64",
|
||||
setup: func(c *gin.Context) { c.Set("id", uint64(99)) },
|
||||
key: "id",
|
||||
wantVal: 0,
|
||||
wantOK: false,
|
||||
},
|
||||
{
|
||||
name: "zero value",
|
||||
setup: func(c *gin.Context) { c.Set("id", int64(0)) },
|
||||
key: "id",
|
||||
wantVal: 0,
|
||||
wantOK: true,
|
||||
},
|
||||
{
|
||||
name: "negative value",
|
||||
setup: func(c *gin.Context) { c.Set("id", int64(-100)) },
|
||||
key: "id",
|
||||
wantVal: -100,
|
||||
wantOK: true,
|
||||
},
|
||||
{
|
||||
name: "max int64",
|
||||
setup: func(c *gin.Context) { c.Set("id", int64(9223372036854775807)) },
|
||||
key: "id",
|
||||
wantVal: 9223372036854775807,
|
||||
wantOK: true,
|
||||
},
|
||||
{
|
||||
name: "different key",
|
||||
setup: func(c *gin.Context) { c.Set("other", int64(123)) },
|
||||
key: "id",
|
||||
wantVal: 0,
|
||||
wantOK: false,
|
||||
},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
c, _ := gin.CreateTestContext(httptest.NewRecorder())
|
||||
tt.setup(c)
|
||||
val, ok := GetInt64(c, tt.key)
|
||||
assert.Equal(t, tt.wantOK, ok)
|
||||
assert.Equal(t, tt.wantVal, val)
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestGetUint64(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
tests := []struct {
|
||||
name string
|
||||
setup func(c *gin.Context)
|
||||
key string
|
||||
wantVal uint64
|
||||
wantOK bool
|
||||
}{
|
||||
{
|
||||
name: "found",
|
||||
setup: func(c *gin.Context) { c.Set("uid", uint64(99)) },
|
||||
key: "uid",
|
||||
wantVal: 99,
|
||||
wantOK: true,
|
||||
},
|
||||
{
|
||||
name: "missing",
|
||||
setup: func(c *gin.Context) {},
|
||||
key: "uid",
|
||||
wantVal: 0,
|
||||
wantOK: false,
|
||||
},
|
||||
{
|
||||
name: "wrong type int64",
|
||||
setup: func(c *gin.Context) { c.Set("uid", int64(99)) },
|
||||
key: "uid",
|
||||
wantVal: 0,
|
||||
wantOK: false,
|
||||
},
|
||||
{
|
||||
name: "wrong type string",
|
||||
setup: func(c *gin.Context) { c.Set("uid", "99") },
|
||||
key: "uid",
|
||||
wantVal: 0,
|
||||
wantOK: false,
|
||||
},
|
||||
{
|
||||
name: "zero value",
|
||||
setup: func(c *gin.Context) { c.Set("uid", uint64(0)) },
|
||||
key: "uid",
|
||||
wantVal: 0,
|
||||
wantOK: true,
|
||||
},
|
||||
{
|
||||
name: "large value",
|
||||
setup: func(c *gin.Context) { c.Set("uid", uint64(18446744073709551615)) },
|
||||
key: "uid",
|
||||
wantVal: 18446744073709551615,
|
||||
wantOK: true,
|
||||
},
|
||||
{
|
||||
name: "different key",
|
||||
setup: func(c *gin.Context) { c.Set("other", uint64(456)) },
|
||||
key: "uid",
|
||||
wantVal: 0,
|
||||
wantOK: false,
|
||||
},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
c, _ := gin.CreateTestContext(httptest.NewRecorder())
|
||||
tt.setup(c)
|
||||
val, ok := GetUint64(c, tt.key)
|
||||
assert.Equal(t, tt.wantOK, ok)
|
||||
assert.Equal(t, tt.wantVal, val)
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestGetString(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
tests := []struct {
|
||||
name string
|
||||
setup func(c *gin.Context)
|
||||
key string
|
||||
wantVal string
|
||||
wantOK bool
|
||||
}{
|
||||
{
|
||||
name: "found",
|
||||
setup: func(c *gin.Context) { c.Set("name", "alice") },
|
||||
key: "name",
|
||||
wantVal: "alice",
|
||||
wantOK: true,
|
||||
},
|
||||
{
|
||||
name: "missing",
|
||||
setup: func(c *gin.Context) {},
|
||||
key: "name",
|
||||
wantVal: "",
|
||||
wantOK: false,
|
||||
},
|
||||
{
|
||||
name: "wrong type int",
|
||||
setup: func(c *gin.Context) { c.Set("name", 123) },
|
||||
key: "name",
|
||||
wantVal: "",
|
||||
wantOK: false,
|
||||
},
|
||||
{
|
||||
name: "wrong type bool",
|
||||
setup: func(c *gin.Context) { c.Set("name", true) },
|
||||
key: "name",
|
||||
wantVal: "",
|
||||
wantOK: false,
|
||||
},
|
||||
{
|
||||
name: "empty string",
|
||||
setup: func(c *gin.Context) { c.Set("name", "") },
|
||||
key: "name",
|
||||
wantVal: "",
|
||||
wantOK: true,
|
||||
},
|
||||
{
|
||||
name: "long string",
|
||||
setup: func(c *gin.Context) { c.Set("name", "very-long-string-with-special-chars-@#$%") },
|
||||
key: "name",
|
||||
wantVal: "very-long-string-with-special-chars-@#$%",
|
||||
wantOK: true,
|
||||
},
|
||||
{
|
||||
name: "different key",
|
||||
setup: func(c *gin.Context) { c.Set("other", "value") },
|
||||
key: "name",
|
||||
wantVal: "",
|
||||
wantOK: false,
|
||||
},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
c, _ := gin.CreateTestContext(httptest.NewRecorder())
|
||||
tt.setup(c)
|
||||
val, ok := GetString(c, tt.key)
|
||||
assert.Equal(t, tt.wantOK, ok)
|
||||
assert.Equal(t, tt.wantVal, val)
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestGetStringArray(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
tests := []struct {
|
||||
name string
|
||||
setup func(c *gin.Context)
|
||||
key string
|
||||
wantVal []string
|
||||
wantOK bool
|
||||
}{
|
||||
{
|
||||
name: "found",
|
||||
setup: func(c *gin.Context) { c.Set("perms", []string{"read", "write"}) },
|
||||
key: "perms",
|
||||
wantVal: []string{"read", "write"},
|
||||
wantOK: true,
|
||||
},
|
||||
{
|
||||
name: "missing",
|
||||
setup: func(c *gin.Context) {},
|
||||
key: "perms",
|
||||
wantVal: []string{},
|
||||
wantOK: false,
|
||||
},
|
||||
{
|
||||
name: "wrong type string",
|
||||
setup: func(c *gin.Context) { c.Set("perms", "not-a-slice") },
|
||||
key: "perms",
|
||||
wantVal: []string{},
|
||||
wantOK: false,
|
||||
},
|
||||
{
|
||||
name: "wrong type int slice",
|
||||
setup: func(c *gin.Context) { c.Set("perms", []int{1, 2, 3}) },
|
||||
key: "perms",
|
||||
wantVal: []string{},
|
||||
wantOK: false,
|
||||
},
|
||||
{
|
||||
name: "empty array",
|
||||
setup: func(c *gin.Context) { c.Set("perms", []string{}) },
|
||||
key: "perms",
|
||||
wantVal: []string{},
|
||||
wantOK: true,
|
||||
},
|
||||
{
|
||||
name: "nil array",
|
||||
setup: func(c *gin.Context) { c.Set("perms", []string(nil)) },
|
||||
key: "perms",
|
||||
wantVal: nil,
|
||||
wantOK: true,
|
||||
},
|
||||
{
|
||||
name: "single element",
|
||||
setup: func(c *gin.Context) { c.Set("perms", []string{"admin"}) },
|
||||
key: "perms",
|
||||
wantVal: []string{"admin"},
|
||||
wantOK: true,
|
||||
},
|
||||
{
|
||||
name: "many elements",
|
||||
setup: func(c *gin.Context) { c.Set("perms", []string{"a", "b", "c", "d", "e"}) },
|
||||
key: "perms",
|
||||
wantVal: []string{"a", "b", "c", "d", "e"},
|
||||
wantOK: true,
|
||||
},
|
||||
{
|
||||
name: "different key",
|
||||
setup: func(c *gin.Context) { c.Set("other", []string{"val"}) },
|
||||
key: "perms",
|
||||
wantVal: []string{},
|
||||
wantOK: false,
|
||||
},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
c, _ := gin.CreateTestContext(httptest.NewRecorder())
|
||||
tt.setup(c)
|
||||
val, ok := GetStringArray(c, tt.key)
|
||||
assert.Equal(t, tt.wantOK, ok)
|
||||
assert.Equal(t, tt.wantVal, val)
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestGetStringFromSession(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
tests := []struct {
|
||||
name string
|
||||
setup func(session sessions.Session)
|
||||
key string
|
||||
wantVal string
|
||||
wantOK bool
|
||||
}{
|
||||
{
|
||||
name: "found",
|
||||
setup: func(s sessions.Session) {
|
||||
s.Set("token", "abc123")
|
||||
_ = s.Save()
|
||||
},
|
||||
key: "token",
|
||||
wantVal: "abc123",
|
||||
wantOK: true,
|
||||
},
|
||||
{
|
||||
name: "missing",
|
||||
setup: func(s sessions.Session) {},
|
||||
key: "token",
|
||||
wantVal: "",
|
||||
wantOK: false,
|
||||
},
|
||||
{
|
||||
name: "wrong type int",
|
||||
setup: func(s sessions.Session) {
|
||||
s.Set("token", 999)
|
||||
_ = s.Save()
|
||||
},
|
||||
key: "token",
|
||||
wantVal: "",
|
||||
wantOK: false,
|
||||
},
|
||||
{
|
||||
name: "wrong type bool",
|
||||
setup: func(s sessions.Session) {
|
||||
s.Set("token", true)
|
||||
_ = s.Save()
|
||||
},
|
||||
key: "token",
|
||||
wantVal: "",
|
||||
wantOK: false,
|
||||
},
|
||||
{
|
||||
name: "empty string",
|
||||
setup: func(s sessions.Session) {
|
||||
s.Set("token", "")
|
||||
_ = s.Save()
|
||||
},
|
||||
key: "token",
|
||||
wantVal: "",
|
||||
wantOK: true,
|
||||
},
|
||||
{
|
||||
name: "different key",
|
||||
setup: func(s sessions.Session) {
|
||||
s.Set("other", "value")
|
||||
_ = s.Save()
|
||||
},
|
||||
key: "token",
|
||||
wantVal: "",
|
||||
wantOK: false,
|
||||
},
|
||||
{
|
||||
name: "multiple values in session",
|
||||
setup: func(s sessions.Session) {
|
||||
s.Set("token", "abc123")
|
||||
s.Set("user", "alice")
|
||||
s.Set("role", "admin")
|
||||
_ = s.Save()
|
||||
},
|
||||
key: "token",
|
||||
wantVal: "abc123",
|
||||
wantOK: true,
|
||||
},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
store := cookie.NewStore([]byte("test-secret"))
|
||||
router := gin.New()
|
||||
router.Use(sessions.Sessions("test", store))
|
||||
|
||||
var val string
|
||||
var ok bool
|
||||
|
||||
router.GET("/test", func(c *gin.Context) {
|
||||
session := sessions.Default(c)
|
||||
tt.setup(session)
|
||||
val, ok = GetStringFromSession(c, tt.key)
|
||||
})
|
||||
|
||||
w := httptest.NewRecorder()
|
||||
req := httptest.NewRequest(http.MethodGet, "/test", nil)
|
||||
router.ServeHTTP(w, req)
|
||||
|
||||
assert.Equal(t, tt.wantOK, ok)
|
||||
assert.Equal(t, tt.wantVal, val)
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestMultipleValuesInContext(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
c, _ := gin.CreateTestContext(httptest.NewRecorder())
|
||||
|
||||
c.Set("int64_val", int64(123))
|
||||
c.Set("uint64_val", uint64(456))
|
||||
c.Set("string_val", "test")
|
||||
c.Set("array_val", []string{"a", "b"})
|
||||
|
||||
// Verify all values are independently accessible
|
||||
intVal, ok := GetInt64(c, "int64_val")
|
||||
assert.True(t, ok)
|
||||
assert.Equal(t, int64(123), intVal)
|
||||
|
||||
uintVal, ok := GetUint64(c, "uint64_val")
|
||||
assert.True(t, ok)
|
||||
assert.Equal(t, uint64(456), uintVal)
|
||||
|
||||
strVal, ok := GetString(c, "string_val")
|
||||
assert.True(t, ok)
|
||||
assert.Equal(t, "test", strVal)
|
||||
|
||||
arrVal, ok := GetStringArray(c, "array_val")
|
||||
assert.True(t, ok)
|
||||
assert.Equal(t, []string{"a", "b"}, arrVal)
|
||||
}
|
||||
|
||||
func TestContextOverwrite(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
c, _ := gin.CreateTestContext(httptest.NewRecorder())
|
||||
|
||||
// Set initial value
|
||||
c.Set("key", "original")
|
||||
val, ok := GetString(c, "key")
|
||||
assert.True(t, ok)
|
||||
assert.Equal(t, "original", val)
|
||||
|
||||
// Overwrite with new value
|
||||
c.Set("key", "updated")
|
||||
val, ok = GetString(c, "key")
|
||||
assert.True(t, ok)
|
||||
assert.Equal(t, "updated", val)
|
||||
}
|
||||
|
||||
func TestContextTypeChange(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
c, _ := gin.CreateTestContext(httptest.NewRecorder())
|
||||
|
||||
// Set as string
|
||||
c.Set("value", "123")
|
||||
strVal, ok := GetString(c, "value")
|
||||
assert.True(t, ok)
|
||||
assert.Equal(t, "123", strVal)
|
||||
|
||||
// Try to get as int64 - should fail
|
||||
intVal, ok := GetInt64(c, "value")
|
||||
assert.False(t, ok)
|
||||
assert.Equal(t, int64(0), intVal)
|
||||
|
||||
// Overwrite with int64
|
||||
c.Set("value", int64(123))
|
||||
intVal, ok = GetInt64(c, "value")
|
||||
assert.True(t, ok)
|
||||
assert.Equal(t, int64(123), intVal)
|
||||
}
|
||||
File diff suppressed because it is too large
Load Diff
File diff suppressed because it is too large
Load Diff
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,99 @@
|
||||
package logger
|
||||
|
||||
import (
|
||||
"context"
|
||||
"time"
|
||||
|
||||
obs "pentagi/pkg/observability"
|
||||
|
||||
"github.com/99designs/gqlgen/graphql"
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/sirupsen/logrus"
|
||||
)
|
||||
|
||||
// FromContext is function to get logrus Entry with context
|
||||
func FromContext(c *gin.Context) *logrus.Entry {
|
||||
return logrus.WithContext(c.Request.Context())
|
||||
}
|
||||
|
||||
func TraceEnabled() bool {
|
||||
return logrus.IsLevelEnabled(logrus.TraceLevel)
|
||||
}
|
||||
|
||||
func WithGinLogger(service string) gin.HandlerFunc {
|
||||
return func(c *gin.Context) {
|
||||
start := time.Now()
|
||||
uri := c.Request.URL.Path
|
||||
raw := c.Request.URL.RawQuery
|
||||
if raw != "" {
|
||||
uri = uri + "?" + raw
|
||||
}
|
||||
|
||||
entry := logrus.WithFields(logrus.Fields{
|
||||
"component": "api",
|
||||
"net_peer_ip": c.ClientIP(),
|
||||
"http_uri": uri,
|
||||
"http_path": c.Request.URL.Path,
|
||||
"http_host_name": c.Request.Host,
|
||||
"http_method": c.Request.Method,
|
||||
})
|
||||
if c.FullPath() == "" {
|
||||
entry = entry.WithField("request", "proxy handled")
|
||||
} else {
|
||||
entry = entry.WithField("request", "api handled")
|
||||
}
|
||||
|
||||
// serve the request to the next middleware
|
||||
c.Next()
|
||||
|
||||
if len(c.Errors) > 0 {
|
||||
entry = entry.WithField("gin.errors", c.Errors.String())
|
||||
}
|
||||
|
||||
entry = entry.WithFields(logrus.Fields{
|
||||
"duration": time.Since(start).String(),
|
||||
"http_status_code": c.Writer.Status(),
|
||||
"http_resp_size": c.Writer.Size(),
|
||||
}).WithContext(c.Request.Context())
|
||||
if c.Writer.Status() >= 400 {
|
||||
entry.Error("http request handled error")
|
||||
} else {
|
||||
entry.Debug("http request handled success")
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func WithGqlLogger(service string) func(ctx context.Context, next graphql.ResponseHandler) *graphql.Response {
|
||||
return func(ctx context.Context, next graphql.ResponseHandler) *graphql.Response {
|
||||
ctx, span := obs.Observer.NewSpan(ctx, obs.SpanKindServer, "graphql.handler")
|
||||
defer span.End()
|
||||
|
||||
start := time.Now()
|
||||
entry := logrus.WithContext(ctx).WithField("component", service)
|
||||
|
||||
res := next(ctx)
|
||||
|
||||
op := graphql.GetOperationContext(ctx)
|
||||
if op != nil && op.Operation != nil {
|
||||
entry = entry.WithFields(logrus.Fields{
|
||||
"operation_name": op.OperationName,
|
||||
"operation_type": op.Operation.Operation,
|
||||
})
|
||||
}
|
||||
|
||||
entry = entry.WithField("duration", time.Since(start).String())
|
||||
|
||||
if res == nil {
|
||||
return res
|
||||
}
|
||||
|
||||
if len(res.Errors) > 0 {
|
||||
entry = entry.WithField("gql.errors", res.Errors.Error())
|
||||
entry.Error("graphql request handled with errors")
|
||||
} else {
|
||||
entry.Debug("graphql request handled success")
|
||||
}
|
||||
|
||||
return res
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,36 @@
|
||||
package router
|
||||
|
||||
import (
|
||||
"pentagi/pkg/server/models"
|
||||
"pentagi/pkg/server/response"
|
||||
|
||||
"github.com/gin-contrib/sessions"
|
||||
"github.com/gin-gonic/gin"
|
||||
)
|
||||
|
||||
func localUserRequired() gin.HandlerFunc {
|
||||
return func(c *gin.Context) {
|
||||
if c.IsAborted() {
|
||||
return
|
||||
}
|
||||
|
||||
session := sessions.Default(c)
|
||||
tid, ok := session.Get("tid").(string)
|
||||
|
||||
if !ok || tid != models.UserTypeLocal.String() {
|
||||
response.Error(c, response.ErrLocalUserRequired, nil)
|
||||
return
|
||||
}
|
||||
|
||||
c.Next()
|
||||
}
|
||||
}
|
||||
|
||||
func noCacheMiddleware() gin.HandlerFunc {
|
||||
return func(c *gin.Context) {
|
||||
c.Header("Cache-Control", "no-cache, no-store, must-revalidate") // HTTP 1.1
|
||||
c.Header("Pragma", "no-cache") // HTTP 1.0
|
||||
c.Header("Expires", "0") // prevents caching at the proxy server
|
||||
c.Next()
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,38 @@
|
||||
package models
|
||||
|
||||
import (
|
||||
"time"
|
||||
|
||||
"github.com/jinzhu/gorm"
|
||||
)
|
||||
|
||||
// Agentlog is model to contain agent task and result information
|
||||
// nolint:lll
|
||||
type Agentlog struct {
|
||||
ID uint64 `form:"id" json:"id" validate:"min=0,numeric" gorm:"type:BIGINT;NOT NULL;PRIMARY_KEY;AUTO_INCREMENT"`
|
||||
Initiator MsgchainType `json:"initiator" validate:"valid,required" gorm:"type:MSGCHAIN_TYPE;NOT NULL"`
|
||||
Executor MsgchainType `json:"executor" validate:"valid,required" gorm:"type:MSGCHAIN_TYPE;NOT NULL"`
|
||||
Task string `json:"task" validate:"required" gorm:"type:TEXT;NOT NULL"`
|
||||
Result string `json:"result" validate:"omitempty" gorm:"type:TEXT;NOT NULL;default:''"`
|
||||
FlowID uint64 `form:"flow_id" json:"flow_id" validate:"min=0,numeric,required" gorm:"type:BIGINT;NOT NULL"`
|
||||
TaskID *uint64 `form:"task_id,omitempty" json:"task_id,omitempty" validate:"omitnil,min=0" gorm:"type:BIGINT;NOT NULL"`
|
||||
SubtaskID *uint64 `form:"subtask_id,omitempty" json:"subtask_id,omitempty" validate:"omitnil,min=0" gorm:"type:BIGINT;NOT NULL"`
|
||||
CreatedAt time.Time `form:"created_at,omitempty" json:"created_at,omitempty" validate:"omitempty" gorm:"type:TIMESTAMPTZ;default:CURRENT_TIMESTAMP"`
|
||||
}
|
||||
|
||||
// TableName returns the table name string to guaranty use correct table
|
||||
func (ml *Agentlog) TableName() string {
|
||||
return "agentlogs"
|
||||
}
|
||||
|
||||
// Valid is function to control input/output data
|
||||
func (ml Agentlog) Valid() error {
|
||||
return validate.Struct(ml)
|
||||
}
|
||||
|
||||
// Validate is function to use callback to control input/output data
|
||||
func (ml Agentlog) Validate(db *gorm.DB) {
|
||||
if err := ml.Valid(); err != nil {
|
||||
db.AddError(err)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,546 @@
|
||||
package models
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"time"
|
||||
|
||||
"github.com/jinzhu/gorm"
|
||||
)
|
||||
|
||||
// UsageStatsPeriod represents time period enum for analytics
|
||||
type UsageStatsPeriod string
|
||||
|
||||
const (
|
||||
UsageStatsPeriodWeek UsageStatsPeriod = "week"
|
||||
UsageStatsPeriodMonth UsageStatsPeriod = "month"
|
||||
UsageStatsPeriodQuarter UsageStatsPeriod = "quarter"
|
||||
)
|
||||
|
||||
func (p UsageStatsPeriod) String() string {
|
||||
return string(p)
|
||||
}
|
||||
|
||||
// Valid is function to control input/output data
|
||||
func (p UsageStatsPeriod) Valid() error {
|
||||
switch p {
|
||||
case UsageStatsPeriodWeek,
|
||||
UsageStatsPeriodMonth,
|
||||
UsageStatsPeriodQuarter:
|
||||
return nil
|
||||
default:
|
||||
return fmt.Errorf("invalid UsageStatsPeriod: %s", p)
|
||||
}
|
||||
}
|
||||
|
||||
// Validate is function to use callback to control input/output data
|
||||
func (p UsageStatsPeriod) Validate(db *gorm.DB) {
|
||||
if err := p.Valid(); err != nil {
|
||||
db.AddError(err)
|
||||
}
|
||||
}
|
||||
|
||||
// ==================== Basic Statistics Structures ====================
|
||||
|
||||
// UsageStats represents token usage statistics
|
||||
// nolint:lll
|
||||
type UsageStats struct {
|
||||
TotalUsageIn int `json:"total_usage_in" validate:"min=0"`
|
||||
TotalUsageOut int `json:"total_usage_out" validate:"min=0"`
|
||||
TotalUsageCacheIn int `json:"total_usage_cache_in" validate:"min=0"`
|
||||
TotalUsageCacheOut int `json:"total_usage_cache_out" validate:"min=0"`
|
||||
TotalUsageCostIn float64 `json:"total_usage_cost_in" validate:"min=0"`
|
||||
TotalUsageCostOut float64 `json:"total_usage_cost_out" validate:"min=0"`
|
||||
}
|
||||
|
||||
// Valid is function to control input/output data
|
||||
func (u UsageStats) Valid() error {
|
||||
return validate.Struct(u)
|
||||
}
|
||||
|
||||
// Validate is function to use callback to control input/output data
|
||||
func (u UsageStats) Validate(db *gorm.DB) {
|
||||
if err := u.Valid(); err != nil {
|
||||
db.AddError(err)
|
||||
}
|
||||
}
|
||||
|
||||
// ToolcallsStats represents toolcalls statistics
|
||||
// nolint:lll
|
||||
type ToolcallsStats struct {
|
||||
TotalCount int `json:"total_count" validate:"min=0"`
|
||||
TotalDurationSeconds float64 `json:"total_duration_seconds" validate:"min=0"`
|
||||
}
|
||||
|
||||
// Valid is function to control input/output data
|
||||
func (t ToolcallsStats) Valid() error {
|
||||
return validate.Struct(t)
|
||||
}
|
||||
|
||||
// Validate is function to use callback to control input/output data
|
||||
func (t ToolcallsStats) Validate(db *gorm.DB) {
|
||||
if err := t.Valid(); err != nil {
|
||||
db.AddError(err)
|
||||
}
|
||||
}
|
||||
|
||||
// FlowsStats represents flows/tasks/subtasks counts
|
||||
// nolint:lll
|
||||
type FlowsStats struct {
|
||||
TotalFlowsCount int `json:"total_flows_count" validate:"min=0"`
|
||||
TotalTasksCount int `json:"total_tasks_count" validate:"min=0"`
|
||||
TotalSubtasksCount int `json:"total_subtasks_count" validate:"min=0"`
|
||||
TotalAssistantsCount int `json:"total_assistants_count" validate:"min=0"`
|
||||
}
|
||||
|
||||
// Valid is function to control input/output data
|
||||
func (f FlowsStats) Valid() error {
|
||||
return validate.Struct(f)
|
||||
}
|
||||
|
||||
// Validate is function to use callback to control input/output data
|
||||
func (f FlowsStats) Validate(db *gorm.DB) {
|
||||
if err := f.Valid(); err != nil {
|
||||
db.AddError(err)
|
||||
}
|
||||
}
|
||||
|
||||
// FlowStats represents single flow statistics
|
||||
// nolint:lll
|
||||
type FlowStats struct {
|
||||
TotalTasksCount int `json:"total_tasks_count" validate:"min=0"`
|
||||
TotalSubtasksCount int `json:"total_subtasks_count" validate:"min=0"`
|
||||
TotalAssistantsCount int `json:"total_assistants_count" validate:"min=0"`
|
||||
}
|
||||
|
||||
// Valid is function to control input/output data
|
||||
func (f FlowStats) Valid() error {
|
||||
return validate.Struct(f)
|
||||
}
|
||||
|
||||
// Validate is function to use callback to control input/output data
|
||||
func (f FlowStats) Validate(db *gorm.DB) {
|
||||
if err := f.Valid(); err != nil {
|
||||
db.AddError(err)
|
||||
}
|
||||
}
|
||||
|
||||
// ==================== Time-series Statistics ====================
|
||||
|
||||
// DailyUsageStats for time-series usage data
|
||||
// nolint:lll
|
||||
type DailyUsageStats struct {
|
||||
Date time.Time `json:"date" validate:"required"`
|
||||
Stats *UsageStats `json:"stats" validate:"required"`
|
||||
}
|
||||
|
||||
// Valid is function to control input/output data
|
||||
func (d DailyUsageStats) Valid() error {
|
||||
if err := validate.Struct(d); err != nil {
|
||||
return err
|
||||
}
|
||||
if d.Stats != nil {
|
||||
return d.Stats.Valid()
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// Validate is function to use callback to control input/output data
|
||||
func (d DailyUsageStats) Validate(db *gorm.DB) {
|
||||
if err := d.Valid(); err != nil {
|
||||
db.AddError(err)
|
||||
}
|
||||
}
|
||||
|
||||
// DailyToolcallsStats for time-series toolcalls data
|
||||
// nolint:lll
|
||||
type DailyToolcallsStats struct {
|
||||
Date time.Time `json:"date" validate:"required"`
|
||||
Stats *ToolcallsStats `json:"stats" validate:"required"`
|
||||
}
|
||||
|
||||
// Valid is function to control input/output data
|
||||
func (d DailyToolcallsStats) Valid() error {
|
||||
if err := validate.Struct(d); err != nil {
|
||||
return err
|
||||
}
|
||||
if d.Stats != nil {
|
||||
return d.Stats.Valid()
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// Validate is function to use callback to control input/output data
|
||||
func (d DailyToolcallsStats) Validate(db *gorm.DB) {
|
||||
if err := d.Valid(); err != nil {
|
||||
db.AddError(err)
|
||||
}
|
||||
}
|
||||
|
||||
// DailyFlowsStats for time-series flows data
|
||||
// nolint:lll
|
||||
type DailyFlowsStats struct {
|
||||
Date time.Time `json:"date" validate:"required"`
|
||||
Stats *FlowsStats `json:"stats" validate:"required"`
|
||||
}
|
||||
|
||||
// Valid is function to control input/output data
|
||||
func (d DailyFlowsStats) Valid() error {
|
||||
if err := validate.Struct(d); err != nil {
|
||||
return err
|
||||
}
|
||||
if d.Stats != nil {
|
||||
return d.Stats.Valid()
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// Validate is function to use callback to control input/output data
|
||||
func (d DailyFlowsStats) Validate(db *gorm.DB) {
|
||||
if err := d.Valid(); err != nil {
|
||||
db.AddError(err)
|
||||
}
|
||||
}
|
||||
|
||||
// ==================== Grouped Statistics ====================
|
||||
|
||||
// ProviderUsageStats for provider-specific usage statistics
|
||||
// nolint:lll
|
||||
type ProviderUsageStats struct {
|
||||
Provider string `json:"provider" validate:"required"`
|
||||
Stats *UsageStats `json:"stats" validate:"required"`
|
||||
}
|
||||
|
||||
// Valid is function to control input/output data
|
||||
func (p ProviderUsageStats) Valid() error {
|
||||
if err := validate.Struct(p); err != nil {
|
||||
return err
|
||||
}
|
||||
if p.Stats != nil {
|
||||
return p.Stats.Valid()
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// Validate is function to use callback to control input/output data
|
||||
func (p ProviderUsageStats) Validate(db *gorm.DB) {
|
||||
if err := p.Valid(); err != nil {
|
||||
db.AddError(err)
|
||||
}
|
||||
}
|
||||
|
||||
// ModelUsageStats for model-specific usage statistics
|
||||
// nolint:lll
|
||||
type ModelUsageStats struct {
|
||||
Model string `json:"model" validate:"required"`
|
||||
Provider string `json:"provider" validate:"required"`
|
||||
Stats *UsageStats `json:"stats" validate:"required"`
|
||||
}
|
||||
|
||||
// Valid is function to control input/output data
|
||||
func (m ModelUsageStats) Valid() error {
|
||||
if err := validate.Struct(m); err != nil {
|
||||
return err
|
||||
}
|
||||
if m.Stats != nil {
|
||||
return m.Stats.Valid()
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// Validate is function to use callback to control input/output data
|
||||
func (m ModelUsageStats) Validate(db *gorm.DB) {
|
||||
if err := m.Valid(); err != nil {
|
||||
db.AddError(err)
|
||||
}
|
||||
}
|
||||
|
||||
// AgentTypeUsageStats for agent type usage statistics
|
||||
// nolint:lll
|
||||
type AgentTypeUsageStats struct {
|
||||
AgentType MsgchainType `json:"agent_type" validate:"valid,required"`
|
||||
Stats *UsageStats `json:"stats" validate:"required"`
|
||||
}
|
||||
|
||||
// Valid is function to control input/output data
|
||||
func (a AgentTypeUsageStats) Valid() error {
|
||||
if err := validate.Struct(a); err != nil {
|
||||
return err
|
||||
}
|
||||
if a.Stats != nil {
|
||||
return a.Stats.Valid()
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// Validate is function to use callback to control input/output data
|
||||
func (a AgentTypeUsageStats) Validate(db *gorm.DB) {
|
||||
if err := a.Valid(); err != nil {
|
||||
db.AddError(err)
|
||||
}
|
||||
}
|
||||
|
||||
// FunctionToolcallsStats for function-specific toolcalls statistics
|
||||
// nolint:lll
|
||||
type FunctionToolcallsStats struct {
|
||||
FunctionName string `json:"function_name" validate:"required"`
|
||||
IsAgent bool `json:"is_agent"`
|
||||
TotalCount int `json:"total_count" validate:"min=0"`
|
||||
TotalDurationSeconds float64 `json:"total_duration_seconds" validate:"min=0"`
|
||||
AvgDurationSeconds float64 `json:"avg_duration_seconds" validate:"min=0"`
|
||||
}
|
||||
|
||||
// Valid is function to control input/output data
|
||||
func (f FunctionToolcallsStats) Valid() error {
|
||||
return validate.Struct(f)
|
||||
}
|
||||
|
||||
// Validate is function to use callback to control input/output data
|
||||
func (f FunctionToolcallsStats) Validate(db *gorm.DB) {
|
||||
if err := f.Valid(); err != nil {
|
||||
db.AddError(err)
|
||||
}
|
||||
}
|
||||
|
||||
// ==================== Execution Statistics ====================
|
||||
|
||||
// SubtaskExecutionStats represents execution statistics for a subtask
|
||||
// nolint:lll
|
||||
type SubtaskExecutionStats struct {
|
||||
SubtaskID int64 `json:"subtask_id" validate:"min=0"`
|
||||
SubtaskTitle string `json:"subtask_title" validate:"required"`
|
||||
TotalDurationSeconds float64 `json:"total_duration_seconds" validate:"min=0"`
|
||||
TotalToolcallsCount int `json:"total_toolcalls_count" validate:"min=0"`
|
||||
}
|
||||
|
||||
// Valid is function to control input/output data
|
||||
func (s SubtaskExecutionStats) Valid() error {
|
||||
return validate.Struct(s)
|
||||
}
|
||||
|
||||
// Validate is function to use callback to control input/output data
|
||||
func (s SubtaskExecutionStats) Validate(db *gorm.DB) {
|
||||
if err := s.Valid(); err != nil {
|
||||
db.AddError(err)
|
||||
}
|
||||
}
|
||||
|
||||
// TaskExecutionStats represents execution statistics for a task
|
||||
// nolint:lll
|
||||
type TaskExecutionStats struct {
|
||||
TaskID int64 `json:"task_id" validate:"min=0"`
|
||||
TaskTitle string `json:"task_title" validate:"required"`
|
||||
TotalDurationSeconds float64 `json:"total_duration_seconds" validate:"min=0"`
|
||||
TotalToolcallsCount int `json:"total_toolcalls_count" validate:"min=0"`
|
||||
Subtasks []SubtaskExecutionStats `json:"subtasks" validate:"omitempty"`
|
||||
}
|
||||
|
||||
// Valid is function to control input/output data
|
||||
func (t TaskExecutionStats) Valid() error {
|
||||
if err := validate.Struct(t); err != nil {
|
||||
return err
|
||||
}
|
||||
for i := range t.Subtasks {
|
||||
if err := t.Subtasks[i].Valid(); err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// Validate is function to use callback to control input/output data
|
||||
func (t TaskExecutionStats) Validate(db *gorm.DB) {
|
||||
if err := t.Valid(); err != nil {
|
||||
db.AddError(err)
|
||||
}
|
||||
}
|
||||
|
||||
// FlowExecutionStats represents execution statistics for a flow
|
||||
// nolint:lll
|
||||
type FlowExecutionStats struct {
|
||||
FlowID int64 `json:"flow_id" validate:"min=0"`
|
||||
FlowTitle string `json:"flow_title" validate:"required"`
|
||||
TotalDurationSeconds float64 `json:"total_duration_seconds" validate:"min=0"`
|
||||
TotalToolcallsCount int `json:"total_toolcalls_count" validate:"min=0"`
|
||||
TotalAssistantsCount int `json:"total_assistants_count" validate:"min=0"`
|
||||
Tasks []TaskExecutionStats `json:"tasks" validate:"omitempty"`
|
||||
}
|
||||
|
||||
// Valid is function to control input/output data
|
||||
func (f FlowExecutionStats) Valid() error {
|
||||
if err := validate.Struct(f); err != nil {
|
||||
return err
|
||||
}
|
||||
for i := range f.Tasks {
|
||||
if err := f.Tasks[i].Valid(); err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// Validate is function to use callback to control input/output data
|
||||
func (f FlowExecutionStats) Validate(db *gorm.DB) {
|
||||
if err := f.Valid(); err != nil {
|
||||
db.AddError(err)
|
||||
}
|
||||
}
|
||||
|
||||
// ==================== Aggregated Response Models ====================
|
||||
|
||||
// SystemUsageResponse represents system-wide analytics response
|
||||
// nolint:lll
|
||||
type SystemUsageResponse struct {
|
||||
UsageStatsTotal *UsageStats `json:"usage_stats_total" validate:"required"`
|
||||
ToolcallsStatsTotal *ToolcallsStats `json:"toolcalls_stats_total" validate:"required"`
|
||||
FlowsStatsTotal *FlowsStats `json:"flows_stats_total" validate:"required"`
|
||||
UsageStatsByProvider []ProviderUsageStats `json:"usage_stats_by_provider" validate:"omitempty"`
|
||||
UsageStatsByModel []ModelUsageStats `json:"usage_stats_by_model" validate:"omitempty"`
|
||||
UsageStatsByAgentType []AgentTypeUsageStats `json:"usage_stats_by_agent_type" validate:"omitempty"`
|
||||
ToolcallsStatsByFunction []FunctionToolcallsStats `json:"toolcalls_stats_by_function" validate:"omitempty"`
|
||||
}
|
||||
|
||||
// Valid is function to control input/output data
|
||||
func (s SystemUsageResponse) Valid() error {
|
||||
if err := validate.Struct(s); err != nil {
|
||||
return err
|
||||
}
|
||||
if s.UsageStatsTotal != nil {
|
||||
if err := s.UsageStatsTotal.Valid(); err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
if s.ToolcallsStatsTotal != nil {
|
||||
if err := s.ToolcallsStatsTotal.Valid(); err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
if s.FlowsStatsTotal != nil {
|
||||
if err := s.FlowsStatsTotal.Valid(); err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
for i := range s.UsageStatsByProvider {
|
||||
if err := s.UsageStatsByProvider[i].Valid(); err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
for i := range s.UsageStatsByModel {
|
||||
if err := s.UsageStatsByModel[i].Valid(); err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
for i := range s.UsageStatsByAgentType {
|
||||
if err := s.UsageStatsByAgentType[i].Valid(); err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
for i := range s.ToolcallsStatsByFunction {
|
||||
if err := s.ToolcallsStatsByFunction[i].Valid(); err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// Validate is function to use callback to control input/output data
|
||||
func (s SystemUsageResponse) Validate(db *gorm.DB) {
|
||||
if err := s.Valid(); err != nil {
|
||||
db.AddError(err)
|
||||
}
|
||||
}
|
||||
|
||||
// PeriodUsageResponse represents period-based analytics response
|
||||
// nolint:lll
|
||||
type PeriodUsageResponse struct {
|
||||
Period string `json:"period" validate:"required"`
|
||||
UsageStatsByPeriod []DailyUsageStats `json:"usage_stats_by_period" validate:"omitempty"`
|
||||
ToolcallsStatsByPeriod []DailyToolcallsStats `json:"toolcalls_stats_by_period" validate:"omitempty"`
|
||||
FlowsStatsByPeriod []DailyFlowsStats `json:"flows_stats_by_period" validate:"omitempty"`
|
||||
FlowsExecutionStatsByPeriod []FlowExecutionStats `json:"flows_execution_stats_by_period" validate:"omitempty"`
|
||||
}
|
||||
|
||||
// Valid is function to control input/output data
|
||||
func (p PeriodUsageResponse) Valid() error {
|
||||
if err := validate.Struct(p); err != nil {
|
||||
return err
|
||||
}
|
||||
for i := range p.UsageStatsByPeriod {
|
||||
if err := p.UsageStatsByPeriod[i].Valid(); err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
for i := range p.ToolcallsStatsByPeriod {
|
||||
if err := p.ToolcallsStatsByPeriod[i].Valid(); err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
for i := range p.FlowsStatsByPeriod {
|
||||
if err := p.FlowsStatsByPeriod[i].Valid(); err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
for i := range p.FlowsExecutionStatsByPeriod {
|
||||
if err := p.FlowsExecutionStatsByPeriod[i].Valid(); err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// Validate is function to use callback to control input/output data
|
||||
func (p PeriodUsageResponse) Validate(db *gorm.DB) {
|
||||
if err := p.Valid(); err != nil {
|
||||
db.AddError(err)
|
||||
}
|
||||
}
|
||||
|
||||
// FlowUsageResponse represents flow-specific analytics response
|
||||
// nolint:lll
|
||||
type FlowUsageResponse struct {
|
||||
FlowID int64 `json:"flow_id" validate:"min=0"`
|
||||
UsageStatsByFlow *UsageStats `json:"usage_stats_by_flow" validate:"required"`
|
||||
UsageStatsByAgentTypeForFlow []AgentTypeUsageStats `json:"usage_stats_by_agent_type_for_flow" validate:"omitempty"`
|
||||
ToolcallsStatsByFlow *ToolcallsStats `json:"toolcalls_stats_by_flow" validate:"required"`
|
||||
ToolcallsStatsByFunctionForFlow []FunctionToolcallsStats `json:"toolcalls_stats_by_function_for_flow" validate:"omitempty"`
|
||||
FlowStatsByFlow *FlowStats `json:"flow_stats_by_flow" validate:"required"`
|
||||
}
|
||||
|
||||
// Valid is function to control input/output data
|
||||
func (f FlowUsageResponse) Valid() error {
|
||||
if err := validate.Struct(f); err != nil {
|
||||
return err
|
||||
}
|
||||
if f.UsageStatsByFlow != nil {
|
||||
if err := f.UsageStatsByFlow.Valid(); err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
for i := range f.UsageStatsByAgentTypeForFlow {
|
||||
if err := f.UsageStatsByAgentTypeForFlow[i].Valid(); err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
if f.ToolcallsStatsByFlow != nil {
|
||||
if err := f.ToolcallsStatsByFlow.Valid(); err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
for i := range f.ToolcallsStatsByFunctionForFlow {
|
||||
if err := f.ToolcallsStatsByFunctionForFlow[i].Valid(); err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
if f.FlowStatsByFlow != nil {
|
||||
if err := f.FlowStatsByFlow.Valid(); err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// Validate is function to use callback to control input/output data
|
||||
func (f FlowUsageResponse) Validate(db *gorm.DB) {
|
||||
if err := f.Valid(); err != nil {
|
||||
db.AddError(err)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,133 @@
|
||||
package models
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"time"
|
||||
|
||||
"github.com/golang-jwt/jwt/v5"
|
||||
"github.com/jinzhu/gorm"
|
||||
)
|
||||
|
||||
// TokenStatus represents the status of an API token
|
||||
type TokenStatus string
|
||||
|
||||
const (
|
||||
TokenStatusActive TokenStatus = "active"
|
||||
TokenStatusRevoked TokenStatus = "revoked"
|
||||
TokenStatusExpired TokenStatus = "expired"
|
||||
)
|
||||
|
||||
func (s TokenStatus) String() string {
|
||||
return string(s)
|
||||
}
|
||||
|
||||
// Valid is function to control input/output data
|
||||
func (s TokenStatus) Valid() error {
|
||||
switch s {
|
||||
case TokenStatusActive, TokenStatusRevoked, TokenStatusExpired:
|
||||
return nil
|
||||
default:
|
||||
return fmt.Errorf("invalid TokenStatus: %s", s)
|
||||
}
|
||||
}
|
||||
|
||||
// Validate is function to use callback to control input/output data
|
||||
func (s TokenStatus) Validate(db *gorm.DB) {
|
||||
if err := s.Valid(); err != nil {
|
||||
db.AddError(err)
|
||||
}
|
||||
}
|
||||
|
||||
// APIToken is model to contain API token metadata
|
||||
// nolint:lll
|
||||
type APIToken struct {
|
||||
ID uint64 `form:"id" json:"id" validate:"min=0,numeric" gorm:"type:BIGINT;NOT NULL;PRIMARY_KEY;AUTO_INCREMENT"`
|
||||
TokenID string `form:"token_id" json:"token_id" validate:"required,len=10" gorm:"type:TEXT;NOT NULL;UNIQUE_INDEX"`
|
||||
UserID uint64 `form:"user_id" json:"user_id" validate:"min=0,numeric" gorm:"type:BIGINT;NOT NULL"`
|
||||
RoleID uint64 `form:"role_id" json:"role_id" validate:"min=0,numeric" gorm:"type:BIGINT;NOT NULL"`
|
||||
Name *string `form:"name,omitempty" json:"name,omitempty" validate:"omitempty,max=100" gorm:"type:TEXT"`
|
||||
TTL uint64 `form:"ttl" json:"ttl" validate:"required,min=60,max=94608000" gorm:"type:BIGINT;NOT NULL"`
|
||||
Status TokenStatus `form:"status" json:"status" validate:"valid,required" gorm:"type:TOKEN_STATUS;NOT NULL;default:'active'"`
|
||||
CreatedAt time.Time `form:"created_at" json:"created_at" validate:"required" gorm:"type:TIMESTAMPTZ;NOT NULL;default:CURRENT_TIMESTAMP"`
|
||||
UpdatedAt time.Time `form:"updated_at" json:"updated_at" validate:"required" gorm:"type:TIMESTAMPTZ;NOT NULL;default:CURRENT_TIMESTAMP"`
|
||||
DeletedAt *time.Time `form:"deleted_at,omitempty" json:"deleted_at,omitempty" validate:"omitempty" sql:"index" gorm:"type:TIMESTAMPTZ"`
|
||||
}
|
||||
|
||||
// TableName returns the table name string to guaranty use correct table
|
||||
func (at *APIToken) TableName() string {
|
||||
return "api_tokens"
|
||||
}
|
||||
|
||||
// Valid is function to control input/output data
|
||||
func (at APIToken) Valid() error {
|
||||
if err := at.Status.Valid(); err != nil {
|
||||
return err
|
||||
}
|
||||
return validate.Struct(at)
|
||||
}
|
||||
|
||||
// Validate is function to use callback to control input/output data
|
||||
func (at APIToken) Validate(db *gorm.DB) {
|
||||
if err := at.Valid(); err != nil {
|
||||
db.AddError(err)
|
||||
}
|
||||
}
|
||||
|
||||
// APITokenWithSecret is model to contain API token with the JWT token string (returned only on creation)
|
||||
// nolint:lll
|
||||
type APITokenWithSecret struct {
|
||||
APIToken `form:"" json:""`
|
||||
Token string `form:"token" json:"token" validate:"required,jwt" gorm:"-"`
|
||||
}
|
||||
|
||||
// Valid is function to control input/output data
|
||||
func (ats APITokenWithSecret) Valid() error {
|
||||
if err := ats.APIToken.Valid(); err != nil {
|
||||
return err
|
||||
}
|
||||
return validate.Struct(ats)
|
||||
}
|
||||
|
||||
// CreateAPITokenRequest is model to contain request data for creating an API token
|
||||
// nolint:lll
|
||||
type CreateAPITokenRequest struct {
|
||||
Name *string `form:"name,omitempty" json:"name,omitempty" validate:"omitempty,max=100"`
|
||||
TTL uint64 `form:"ttl" json:"ttl" validate:"required,min=60,max=94608000"` // from 1 minute to 3 years
|
||||
}
|
||||
|
||||
// Valid is function to control input/output data
|
||||
func (catr CreateAPITokenRequest) Valid() error {
|
||||
return validate.Struct(catr)
|
||||
}
|
||||
|
||||
// UpdateAPITokenRequest is model to contain request data for updating an API token
|
||||
// nolint:lll
|
||||
type UpdateAPITokenRequest struct {
|
||||
Name *string `form:"name,omitempty" json:"name,omitempty" validate:"omitempty,max=100"`
|
||||
Status TokenStatus `form:"status,omitempty" json:"status,omitempty" validate:"omitempty,valid"`
|
||||
}
|
||||
|
||||
// Valid is function to control input/output data
|
||||
func (uatr UpdateAPITokenRequest) Valid() error {
|
||||
if uatr.Status != "" {
|
||||
if err := uatr.Status.Valid(); err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
return validate.Struct(uatr)
|
||||
}
|
||||
|
||||
// APITokenClaims is model to contain JWT claims for API tokens
|
||||
// nolint:lll
|
||||
type APITokenClaims struct {
|
||||
TokenID string `json:"tid" validate:"required,len=10"`
|
||||
RID uint64 `json:"rid" validate:"min=0,max=10000"`
|
||||
UID uint64 `json:"uid" validate:"min=0,max=10000"`
|
||||
UHASH string `json:"uhash" validate:"required"`
|
||||
jwt.RegisteredClaims
|
||||
}
|
||||
|
||||
// Valid is function to control input/output data
|
||||
func (atc APITokenClaims) Valid() error {
|
||||
return validate.Struct(atc)
|
||||
}
|
||||
@@ -0,0 +1,255 @@
|
||||
package models
|
||||
|
||||
import (
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/stretchr/testify/assert"
|
||||
)
|
||||
|
||||
func TestTokenStatusValid(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
tests := []struct {
|
||||
name string
|
||||
status TokenStatus
|
||||
wantErr bool
|
||||
}{
|
||||
{"valid active", TokenStatusActive, false},
|
||||
{"valid revoked", TokenStatusRevoked, false},
|
||||
{"valid expired", TokenStatusExpired, false},
|
||||
{"invalid empty", TokenStatus(""), true},
|
||||
{"invalid unknown", TokenStatus("unknown"), true},
|
||||
{"invalid deleted", TokenStatus("deleted"), true},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
tt := tt
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
t.Parallel()
|
||||
err := tt.status.Valid()
|
||||
if tt.wantErr {
|
||||
assert.Error(t, err)
|
||||
assert.Contains(t, err.Error(), "invalid TokenStatus")
|
||||
} else {
|
||||
assert.NoError(t, err)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestTokenStatusString(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
assert.Equal(t, "active", TokenStatusActive.String())
|
||||
assert.Equal(t, "revoked", TokenStatusRevoked.String())
|
||||
assert.Equal(t, "expired", TokenStatusExpired.String())
|
||||
}
|
||||
|
||||
func TestAPITokenValid(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
now := time.Now()
|
||||
validToken := APIToken{
|
||||
TokenID: "abcdefghij",
|
||||
UserID: 1,
|
||||
RoleID: 1,
|
||||
TTL: 3600,
|
||||
Status: TokenStatusActive,
|
||||
CreatedAt: now,
|
||||
UpdatedAt: now,
|
||||
}
|
||||
|
||||
t.Run("valid token", func(t *testing.T) {
|
||||
t.Parallel()
|
||||
assert.NoError(t, validToken.Valid())
|
||||
})
|
||||
|
||||
t.Run("invalid status", func(t *testing.T) {
|
||||
t.Parallel()
|
||||
token := validToken
|
||||
token.Status = TokenStatus("invalid")
|
||||
assert.Error(t, token.Valid())
|
||||
})
|
||||
|
||||
t.Run("token id wrong length", func(t *testing.T) {
|
||||
t.Parallel()
|
||||
token := validToken
|
||||
token.TokenID = "short"
|
||||
assert.Error(t, token.Valid())
|
||||
})
|
||||
|
||||
t.Run("ttl too small", func(t *testing.T) {
|
||||
t.Parallel()
|
||||
token := validToken
|
||||
token.TTL = 10
|
||||
assert.Error(t, token.Valid())
|
||||
})
|
||||
|
||||
t.Run("ttl too large", func(t *testing.T) {
|
||||
t.Parallel()
|
||||
token := validToken
|
||||
token.TTL = 94608001
|
||||
assert.Error(t, token.Valid())
|
||||
})
|
||||
|
||||
t.Run("ttl minimum boundary", func(t *testing.T) {
|
||||
t.Parallel()
|
||||
token := validToken
|
||||
token.TTL = 60
|
||||
assert.NoError(t, token.Valid())
|
||||
})
|
||||
|
||||
t.Run("ttl maximum boundary", func(t *testing.T) {
|
||||
t.Parallel()
|
||||
token := validToken
|
||||
token.TTL = 94608000
|
||||
assert.NoError(t, token.Valid())
|
||||
})
|
||||
}
|
||||
|
||||
func TestAPITokenTableName(t *testing.T) {
|
||||
t.Parallel()
|
||||
at := &APIToken{}
|
||||
assert.Equal(t, "api_tokens", at.TableName())
|
||||
}
|
||||
|
||||
func TestCreateAPITokenRequestValid(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
t.Run("valid request", func(t *testing.T) {
|
||||
t.Parallel()
|
||||
req := CreateAPITokenRequest{TTL: 3600}
|
||||
assert.NoError(t, req.Valid())
|
||||
})
|
||||
|
||||
t.Run("valid request with name", func(t *testing.T) {
|
||||
t.Parallel()
|
||||
name := "my-token"
|
||||
req := CreateAPITokenRequest{Name: &name, TTL: 3600}
|
||||
assert.NoError(t, req.Valid())
|
||||
})
|
||||
|
||||
t.Run("ttl too small", func(t *testing.T) {
|
||||
t.Parallel()
|
||||
req := CreateAPITokenRequest{TTL: 30}
|
||||
assert.Error(t, req.Valid())
|
||||
})
|
||||
|
||||
t.Run("zero ttl", func(t *testing.T) {
|
||||
t.Parallel()
|
||||
req := CreateAPITokenRequest{TTL: 0}
|
||||
assert.Error(t, req.Valid())
|
||||
})
|
||||
}
|
||||
|
||||
func TestUpdateAPITokenRequestValid(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
t.Run("valid update with status", func(t *testing.T) {
|
||||
t.Parallel()
|
||||
req := UpdateAPITokenRequest{Status: TokenStatusRevoked}
|
||||
assert.NoError(t, req.Valid())
|
||||
})
|
||||
|
||||
t.Run("valid empty update", func(t *testing.T) {
|
||||
t.Parallel()
|
||||
req := UpdateAPITokenRequest{}
|
||||
assert.NoError(t, req.Valid())
|
||||
})
|
||||
|
||||
t.Run("invalid status", func(t *testing.T) {
|
||||
t.Parallel()
|
||||
req := UpdateAPITokenRequest{Status: TokenStatus("invalid")}
|
||||
assert.Error(t, req.Valid())
|
||||
})
|
||||
}
|
||||
|
||||
func TestAPITokenWithSecretValid(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
now := time.Now()
|
||||
validJWT := "eyJhbGciOiJSUzI1NiIsInR5cCI6IkpXVCJ9.eyJzdWIiOiIxMjM0NTY3ODkwIn0.POstGetfAytaZS82wHcjoTyoqhMyxXiWdR7Nn7A29DNSl0EiXLdwJ6xC6AfgZWF1bOsS_TuYI3OG85AmiExREkrS6tDfTQ2B3WXlrr-wp5AokiRbz3_oB4OxG-W9KcEEbDRcZc0nH3L7LzYptiy1PtAylQGxHTWZXtGz4ht0bAecBgmpdgXMguEIcoqPJ1n3pIWk_dUZegpqx0Lka21H6XxUTxiy8OcaarA8zdnPUnV6AmNP3ecFawIFYdvJB_cm-GvpCSbr8G8y_Mllj8f4x9nBH8pQux89_6gUY618iYv7tuPWBFfEbLxtF2pZS6YC1aSfLQxaOoaBSTNRg"
|
||||
|
||||
baseToken := APIToken{
|
||||
TokenID: "abcdefghij",
|
||||
UserID: 1,
|
||||
RoleID: 1,
|
||||
TTL: 3600,
|
||||
Status: TokenStatusActive,
|
||||
CreatedAt: now,
|
||||
UpdatedAt: now,
|
||||
}
|
||||
|
||||
t.Run("valid token with secret", func(t *testing.T) {
|
||||
t.Parallel()
|
||||
ats := APITokenWithSecret{
|
||||
APIToken: baseToken,
|
||||
Token: validJWT,
|
||||
}
|
||||
assert.NoError(t, ats.Valid())
|
||||
})
|
||||
|
||||
t.Run("invalid embedded api token", func(t *testing.T) {
|
||||
t.Parallel()
|
||||
badToken := baseToken
|
||||
badToken.Status = TokenStatus("invalid")
|
||||
ats := APITokenWithSecret{
|
||||
APIToken: badToken,
|
||||
Token: validJWT,
|
||||
}
|
||||
assert.Error(t, ats.Valid())
|
||||
})
|
||||
|
||||
t.Run("invalid jwt token string", func(t *testing.T) {
|
||||
t.Parallel()
|
||||
ats := APITokenWithSecret{
|
||||
APIToken: baseToken,
|
||||
Token: "not-a-jwt",
|
||||
}
|
||||
assert.Error(t, ats.Valid())
|
||||
})
|
||||
}
|
||||
|
||||
func TestAPITokenClaimsValid(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
t.Run("valid claims", func(t *testing.T) {
|
||||
t.Parallel()
|
||||
claims := APITokenClaims{
|
||||
TokenID: "abcdefghij",
|
||||
RID: 1,
|
||||
UID: 1,
|
||||
UHASH: "somehash",
|
||||
}
|
||||
assert.NoError(t, claims.Valid())
|
||||
})
|
||||
|
||||
t.Run("missing token id", func(t *testing.T) {
|
||||
t.Parallel()
|
||||
claims := APITokenClaims{
|
||||
TokenID: "",
|
||||
UHASH: "somehash",
|
||||
}
|
||||
assert.Error(t, claims.Valid())
|
||||
})
|
||||
|
||||
t.Run("missing uhash", func(t *testing.T) {
|
||||
t.Parallel()
|
||||
claims := APITokenClaims{
|
||||
TokenID: "abcdefghij",
|
||||
UHASH: "",
|
||||
}
|
||||
assert.Error(t, claims.Valid())
|
||||
})
|
||||
|
||||
t.Run("uid too large", func(t *testing.T) {
|
||||
t.Parallel()
|
||||
claims := APITokenClaims{
|
||||
TokenID: "abcdefghij",
|
||||
UID: 10001,
|
||||
UHASH: "somehash",
|
||||
}
|
||||
assert.Error(t, claims.Valid())
|
||||
})
|
||||
}
|
||||
@@ -0,0 +1,38 @@
|
||||
package models
|
||||
|
||||
import (
|
||||
"time"
|
||||
|
||||
"github.com/jinzhu/gorm"
|
||||
)
|
||||
|
||||
// Assistantlog is model to contain log record information from agents about their actions
|
||||
// nolint:lll
|
||||
type Assistantlog struct {
|
||||
ID uint64 `form:"id" json:"id" validate:"min=0,numeric" gorm:"type:BIGINT;NOT NULL;PRIMARY_KEY;AUTO_INCREMENT"`
|
||||
Type MsglogType `form:"type" json:"type" validate:"valid,required" gorm:"type:MSGLOG_TYPE;NOT NULL"`
|
||||
Message string `form:"message" json:"message" validate:"omitempty" gorm:"type:TEXT;NOT NULL"`
|
||||
Thinking string `form:"thinking" json:"thinking" validate:"omitempty" gorm:"type:TEXT;NULL"`
|
||||
Result string `form:"result" json:"result" validate:"omitempty" gorm:"type:TEXT;NOT NULL;default:''"`
|
||||
ResultFormat MsglogResultFormat `form:"result_format" json:"result_format" validate:"valid,required" gorm:"type:MSGLOG_RESULT_FORMAT;NOT NULL;default:plain"`
|
||||
FlowID uint64 `form:"flow_id" json:"flow_id" validate:"min=0,numeric,required" gorm:"type:BIGINT;NOT NULL"`
|
||||
AssistantID uint64 `form:"assistant_id" json:"assistant_id" validate:"min=0,numeric,required" gorm:"type:BIGINT;NOT NULL"`
|
||||
CreatedAt time.Time `form:"created_at,omitempty" json:"created_at,omitempty" validate:"omitempty" gorm:"type:TIMESTAMPTZ;default:CURRENT_TIMESTAMP"`
|
||||
}
|
||||
|
||||
// TableName returns the table name string to guaranty use correct table
|
||||
func (al *Assistantlog) TableName() string {
|
||||
return "assistantlogs"
|
||||
}
|
||||
|
||||
// Valid is function to control input/output data
|
||||
func (al Assistantlog) Valid() error {
|
||||
return validate.Struct(al)
|
||||
}
|
||||
|
||||
// Validate is function to use callback to control input/output data
|
||||
func (al Assistantlog) Validate(db *gorm.DB) {
|
||||
if err := al.Valid(); err != nil {
|
||||
db.AddError(err)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,133 @@
|
||||
package models
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"pentagi/pkg/tools"
|
||||
"time"
|
||||
|
||||
"github.com/jinzhu/gorm"
|
||||
)
|
||||
|
||||
type AssistantStatus string
|
||||
|
||||
const (
|
||||
AssistantStatusCreated AssistantStatus = "created"
|
||||
AssistantStatusRunning AssistantStatus = "running"
|
||||
AssistantStatusWaiting AssistantStatus = "waiting"
|
||||
AssistantStatusFinished AssistantStatus = "finished"
|
||||
AssistantStatusFailed AssistantStatus = "failed"
|
||||
)
|
||||
|
||||
func (s AssistantStatus) String() string {
|
||||
return string(s)
|
||||
}
|
||||
|
||||
// Valid is function to control input/output data
|
||||
func (s AssistantStatus) Valid() error {
|
||||
switch s {
|
||||
case AssistantStatusCreated,
|
||||
AssistantStatusRunning,
|
||||
AssistantStatusWaiting,
|
||||
AssistantStatusFinished,
|
||||
AssistantStatusFailed:
|
||||
return nil
|
||||
default:
|
||||
return fmt.Errorf("invalid AssistantStatus: %s", s)
|
||||
}
|
||||
}
|
||||
|
||||
// Validate is function to use callback to control input/output data
|
||||
func (s AssistantStatus) Validate(db *gorm.DB) {
|
||||
if err := s.Valid(); err != nil {
|
||||
db.AddError(err)
|
||||
}
|
||||
}
|
||||
|
||||
// Assistant is model to contain assistant information
|
||||
// nolint:lll
|
||||
type Assistant struct {
|
||||
ID uint64 `form:"id" json:"id" validate:"min=0,numeric" gorm:"type:BIGINT;NOT NULL;PRIMARY_KEY;AUTO_INCREMENT"`
|
||||
Status AssistantStatus `form:"status" json:"status" validate:"valid,required" gorm:"type:ASSISTANT_STATUS;NOT NULL;default:'created'"`
|
||||
Title string `form:"title" json:"title" validate:"required" gorm:"type:TEXT;NOT NULL;default:'untitled'"`
|
||||
Model string `form:"model" json:"model" validate:"max=70,required" gorm:"type:TEXT;NOT NULL"`
|
||||
ModelProviderName string `form:"model_provider_name" json:"model_provider_name" validate:"max=70,required" gorm:"type:TEXT;NOT NULL"`
|
||||
ModelProviderType ProviderType `form:"model_provider_type" json:"model_provider_type" validate:"valid,required" gorm:"type:PROVIDER_TYPE;NOT NULL"`
|
||||
Language string `form:"language" json:"language" validate:"max=70,required" gorm:"type:TEXT;NOT NULL"`
|
||||
Functions *tools.Functions `form:"functions,omitempty" json:"functions,omitempty" validate:"omitempty,valid" gorm:"type:JSON;NOT NULL;default:'{}'"`
|
||||
FlowID uint64 `form:"flow_id" json:"flow_id" validate:"min=0,numeric,required" gorm:"type:BIGINT;NOT NULL"`
|
||||
MsgchainID *uint64 `form:"msgchain_id" json:"msgchain_id" validate:"min=0,numeric" gorm:"type:BIGINT;NOT NULL"`
|
||||
TraceID *string `form:"trace_id" json:"trace_id" validate:"max=70,required" gorm:"type:TEXT;NOT NULL"`
|
||||
ToolCallIDTemplate string `form:"tool_call_id_template" json:"tool_call_id_template" validate:"max=70,required" gorm:"type:TEXT;NOT NULL"`
|
||||
UseAgents bool `form:"use_agents" json:"use_agents" validate:"omitempty" gorm:"type:BOOLEAN;NOT NULL;default:false"`
|
||||
CreatedAt time.Time `form:"created_at,omitempty" json:"created_at,omitempty" validate:"omitempty" gorm:"type:TIMESTAMPTZ;default:CURRENT_TIMESTAMP"`
|
||||
UpdatedAt time.Time `form:"updated_at,omitempty" json:"updated_at,omitempty" validate:"omitempty" gorm:"type:TIMESTAMPTZ;default:CURRENT_TIMESTAMP"`
|
||||
DeletedAt *time.Time `form:"deleted_at,omitempty" json:"deleted_at,omitempty" validate:"omitempty" sql:"index" gorm:"type:TIMESTAMPTZ"`
|
||||
}
|
||||
|
||||
// TableName returns the table name string to guaranty use correct table
|
||||
func (a *Assistant) TableName() string {
|
||||
return "assistants"
|
||||
}
|
||||
|
||||
// Valid is function to control input/output data
|
||||
func (a Assistant) Valid() error {
|
||||
return validate.Struct(a)
|
||||
}
|
||||
|
||||
// Validate is function to use callback to control input/output data
|
||||
func (a Assistant) Validate(db *gorm.DB) {
|
||||
if err := a.Valid(); err != nil {
|
||||
db.AddError(err)
|
||||
}
|
||||
}
|
||||
|
||||
// CreateAssistant is model to contain assistant creation paylaod
|
||||
// nolint:lll
|
||||
type CreateAssistant struct {
|
||||
Input string `form:"input" json:"input" validate:"required" example:"user input for running assistant"`
|
||||
Provider string `form:"provider" json:"provider" validate:"required" example:"openai"`
|
||||
UseAgents bool `form:"use_agents" json:"use_agents" validate:"omitempty" example:"true"`
|
||||
Functions *tools.Functions `form:"functions,omitempty" json:"functions,omitempty" validate:"omitempty,valid"`
|
||||
ResourceIDs []uint64 `form:"resource_ids,omitempty" json:"resource_ids,omitempty" validate:"omitempty" swaggertype:"array,integer"`
|
||||
}
|
||||
|
||||
// Valid is function to control input/output data
|
||||
func (ca CreateAssistant) Valid() error {
|
||||
return validate.Struct(ca)
|
||||
}
|
||||
|
||||
// PatchAssistant is model to contain assistant patching paylaod
|
||||
// nolint:lll
|
||||
type PatchAssistant struct {
|
||||
Action string `form:"action" json:"action" validate:"required,oneof=stop input" enums:"stop,input" default:"stop"`
|
||||
Input *string `form:"input,omitempty" json:"input,omitempty" validate:"required_if=Action input" example:"user input for waiting assistant"`
|
||||
UseAgents bool `form:"use_agents" json:"use_agents" validate:"omitempty" example:"true"`
|
||||
ResourceIDs []uint64 `form:"resource_ids,omitempty" json:"resource_ids,omitempty" validate:"omitempty" swaggertype:"array,integer"`
|
||||
}
|
||||
|
||||
// Valid is function to control input/output data
|
||||
func (pa PatchAssistant) Valid() error {
|
||||
return validate.Struct(pa)
|
||||
}
|
||||
|
||||
// AssistantFlow is model to contain assistant information linked with flow
|
||||
// nolint:lll
|
||||
type AssistantFlow struct {
|
||||
Flow Flow `form:"flow,omitempty" json:"flow,omitempty" gorm:"association_autoupdate:false;association_autocreate:false"`
|
||||
Assistant `form:"" json:""`
|
||||
}
|
||||
|
||||
// Valid is function to control input/output data
|
||||
func (af AssistantFlow) Valid() error {
|
||||
if err := af.Flow.Valid(); err != nil {
|
||||
return err
|
||||
}
|
||||
return af.Assistant.Valid()
|
||||
}
|
||||
|
||||
// Validate is function to use callback to control input/output data
|
||||
func (af AssistantFlow) Validate(db *gorm.DB) {
|
||||
if err := af.Valid(); err != nil {
|
||||
db.AddError(err)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,211 @@
|
||||
package models
|
||||
|
||||
import (
|
||||
"testing"
|
||||
|
||||
"github.com/stretchr/testify/assert"
|
||||
)
|
||||
|
||||
func TestAssistantStatusValid(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
tests := []struct {
|
||||
name string
|
||||
status AssistantStatus
|
||||
wantErr bool
|
||||
}{
|
||||
{"valid created", AssistantStatusCreated, false},
|
||||
{"valid running", AssistantStatusRunning, false},
|
||||
{"valid waiting", AssistantStatusWaiting, false},
|
||||
{"valid finished", AssistantStatusFinished, false},
|
||||
{"valid failed", AssistantStatusFailed, false},
|
||||
{"invalid empty", AssistantStatus(""), true},
|
||||
{"invalid unknown", AssistantStatus("unknown"), true},
|
||||
{"invalid paused", AssistantStatus("paused"), true},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
tt := tt
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
t.Parallel()
|
||||
err := tt.status.Valid()
|
||||
if tt.wantErr {
|
||||
assert.Error(t, err)
|
||||
assert.Contains(t, err.Error(), "invalid AssistantStatus")
|
||||
} else {
|
||||
assert.NoError(t, err)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestAssistantStatusString(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
assert.Equal(t, "created", AssistantStatusCreated.String())
|
||||
assert.Equal(t, "running", AssistantStatusRunning.String())
|
||||
assert.Equal(t, "finished", AssistantStatusFinished.String())
|
||||
}
|
||||
|
||||
func TestAssistantValid(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
traceID := "trace-123"
|
||||
msgchainID := uint64(1)
|
||||
validAssistant := Assistant{
|
||||
Status: AssistantStatusCreated,
|
||||
Title: "test assistant",
|
||||
Model: "gpt-4",
|
||||
ModelProviderName: "openai",
|
||||
ModelProviderType: ProviderType("openai"),
|
||||
Language: "en",
|
||||
ToolCallIDTemplate: "call_{id}",
|
||||
TraceID: &traceID,
|
||||
FlowID: 1,
|
||||
MsgchainID: &msgchainID,
|
||||
}
|
||||
|
||||
t.Run("valid assistant", func(t *testing.T) {
|
||||
t.Parallel()
|
||||
assert.NoError(t, validAssistant.Valid())
|
||||
})
|
||||
|
||||
t.Run("invalid status", func(t *testing.T) {
|
||||
t.Parallel()
|
||||
a := validAssistant
|
||||
a.Status = AssistantStatus("invalid")
|
||||
assert.Error(t, a.Valid())
|
||||
})
|
||||
|
||||
t.Run("missing title", func(t *testing.T) {
|
||||
t.Parallel()
|
||||
a := validAssistant
|
||||
a.Title = ""
|
||||
assert.Error(t, a.Valid())
|
||||
})
|
||||
|
||||
t.Run("missing model", func(t *testing.T) {
|
||||
t.Parallel()
|
||||
a := validAssistant
|
||||
a.Model = ""
|
||||
assert.Error(t, a.Valid())
|
||||
})
|
||||
}
|
||||
|
||||
func TestAssistantTableName(t *testing.T) {
|
||||
t.Parallel()
|
||||
a := &Assistant{}
|
||||
assert.Equal(t, "assistants", a.TableName())
|
||||
}
|
||||
|
||||
func TestCreateAssistantValid(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
t.Run("valid create assistant", func(t *testing.T) {
|
||||
t.Parallel()
|
||||
ca := CreateAssistant{Input: "hello", Provider: "openai"}
|
||||
assert.NoError(t, ca.Valid())
|
||||
})
|
||||
|
||||
t.Run("missing input", func(t *testing.T) {
|
||||
t.Parallel()
|
||||
ca := CreateAssistant{Input: "", Provider: "openai"}
|
||||
assert.Error(t, ca.Valid())
|
||||
})
|
||||
|
||||
t.Run("missing provider", func(t *testing.T) {
|
||||
t.Parallel()
|
||||
ca := CreateAssistant{Input: "hello", Provider: ""}
|
||||
assert.Error(t, ca.Valid())
|
||||
})
|
||||
}
|
||||
|
||||
func TestPatchAssistantValid(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
t.Run("valid stop action", func(t *testing.T) {
|
||||
t.Parallel()
|
||||
pa := PatchAssistant{Action: "stop"}
|
||||
assert.NoError(t, pa.Valid())
|
||||
})
|
||||
|
||||
t.Run("valid input action with input", func(t *testing.T) {
|
||||
t.Parallel()
|
||||
input := "user response"
|
||||
pa := PatchAssistant{Action: "input", Input: &input}
|
||||
assert.NoError(t, pa.Valid())
|
||||
})
|
||||
|
||||
t.Run("invalid action", func(t *testing.T) {
|
||||
t.Parallel()
|
||||
pa := PatchAssistant{Action: "restart"}
|
||||
assert.Error(t, pa.Valid())
|
||||
})
|
||||
|
||||
t.Run("input action without input", func(t *testing.T) {
|
||||
t.Parallel()
|
||||
pa := PatchAssistant{Action: "input"}
|
||||
assert.Error(t, pa.Valid())
|
||||
})
|
||||
}
|
||||
|
||||
func TestAssistantFlowValid(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
traceID := "trace-123"
|
||||
mcID := uint64(1)
|
||||
validAssistant := Assistant{
|
||||
Status: AssistantStatusCreated,
|
||||
Title: "test",
|
||||
Model: "gpt-4",
|
||||
ModelProviderName: "openai",
|
||||
ModelProviderType: ProviderType("openai"),
|
||||
Language: "en",
|
||||
ToolCallIDTemplate: "call_{id}",
|
||||
TraceID: &traceID,
|
||||
FlowID: 1,
|
||||
MsgchainID: &mcID,
|
||||
}
|
||||
validFlow := Flow{
|
||||
Status: FlowStatusCreated,
|
||||
Title: "flow",
|
||||
Model: "gpt-4",
|
||||
ModelProviderName: "openai",
|
||||
ModelProviderType: ProviderType("openai"),
|
||||
Language: "en",
|
||||
ToolCallIDTemplate: "call_{id}",
|
||||
TraceID: &traceID,
|
||||
UserID: 1,
|
||||
}
|
||||
|
||||
t.Run("valid assistant flow", func(t *testing.T) {
|
||||
t.Parallel()
|
||||
af := AssistantFlow{
|
||||
Flow: validFlow,
|
||||
Assistant: validAssistant,
|
||||
}
|
||||
assert.NoError(t, af.Valid())
|
||||
})
|
||||
|
||||
t.Run("invalid flow", func(t *testing.T) {
|
||||
t.Parallel()
|
||||
badFlow := validFlow
|
||||
badFlow.Title = ""
|
||||
af := AssistantFlow{
|
||||
Flow: badFlow,
|
||||
Assistant: validAssistant,
|
||||
}
|
||||
assert.Error(t, af.Valid())
|
||||
})
|
||||
|
||||
t.Run("invalid assistant", func(t *testing.T) {
|
||||
t.Parallel()
|
||||
badAssistant := validAssistant
|
||||
badAssistant.Title = ""
|
||||
af := AssistantFlow{
|
||||
Flow: validFlow,
|
||||
Assistant: badAssistant,
|
||||
}
|
||||
assert.Error(t, af.Valid())
|
||||
})
|
||||
}
|
||||
@@ -0,0 +1,103 @@
|
||||
package models
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"time"
|
||||
|
||||
"github.com/jinzhu/gorm"
|
||||
)
|
||||
|
||||
type ContainerStatus string
|
||||
|
||||
const (
|
||||
ContainerStatusStarting ContainerStatus = "starting"
|
||||
ContainerStatusRunning ContainerStatus = "running"
|
||||
ContainerStatusStopped ContainerStatus = "stopped"
|
||||
ContainerStatusDeleted ContainerStatus = "deleted"
|
||||
ContainerStatusFailed ContainerStatus = "failed"
|
||||
)
|
||||
|
||||
func (s ContainerStatus) String() string {
|
||||
return string(s)
|
||||
}
|
||||
|
||||
// Valid is function to control input/output data
|
||||
func (s ContainerStatus) Valid() error {
|
||||
switch s {
|
||||
case ContainerStatusStarting,
|
||||
ContainerStatusRunning,
|
||||
ContainerStatusStopped,
|
||||
ContainerStatusDeleted,
|
||||
ContainerStatusFailed:
|
||||
return nil
|
||||
default:
|
||||
return fmt.Errorf("invalid ContainerStatus: %s", s)
|
||||
}
|
||||
}
|
||||
|
||||
// Validate is function to use callback to control input/output data
|
||||
func (s ContainerStatus) Validate(db *gorm.DB) {
|
||||
if err := s.Valid(); err != nil {
|
||||
db.AddError(err)
|
||||
}
|
||||
}
|
||||
|
||||
type ContainerType string
|
||||
|
||||
const (
|
||||
ContainerTypePrimary ContainerType = "primary"
|
||||
ContainerTypeSecondary ContainerType = "secondary"
|
||||
)
|
||||
|
||||
func (t ContainerType) String() string {
|
||||
return string(t)
|
||||
}
|
||||
|
||||
// Valid is function to control input/output data
|
||||
func (t ContainerType) Valid() error {
|
||||
switch t {
|
||||
case ContainerTypePrimary, ContainerTypeSecondary:
|
||||
return nil
|
||||
default:
|
||||
return fmt.Errorf("invalid ContainerType: %s", t)
|
||||
}
|
||||
}
|
||||
|
||||
// Validate is function to use callback to control input/output data
|
||||
func (t ContainerType) Validate(db *gorm.DB) {
|
||||
if err := t.Valid(); err != nil {
|
||||
db.AddError(err)
|
||||
}
|
||||
}
|
||||
|
||||
// Container is model to contain container information
|
||||
// nolint:lll
|
||||
type Container struct {
|
||||
ID uint64 `form:"id" json:"id" validate:"min=0,numeric" gorm:"type:BIGINT;NOT NULL;PRIMARY_KEY;AUTO_INCREMENT"`
|
||||
Type ContainerType `form:"type" json:"type" validate:"valid,required" gorm:"type:CONTAINER_TYPE;NOT NULL;default:'primary'"`
|
||||
Name string `form:"name" json:"name" validate:"required" gorm:"type:TEXT;NOT NULL;default:MD5(RANDOM()::text)"`
|
||||
Image string `form:"image" json:"image" validate:"required" gorm:"type:TEXT;NOT NULL"`
|
||||
Status ContainerStatus `form:"status" json:"status" validate:"valid,required" gorm:"type:CONTAINER_STATUS;NOT NULL;default:'starting'"`
|
||||
LocalID string `form:"local_id" json:"local_id" validate:"required" gorm:"type:TEXT;NOT NULL"`
|
||||
LocalDir string `form:"local_dir" json:"local_dir" validate:"required" gorm:"type:TEXT;NOT NULL"`
|
||||
FlowID uint64 `form:"flow_id" json:"flow_id" validate:"min=0,numeric,required" gorm:"type:BIGINT;NOT NULL"`
|
||||
CreatedAt time.Time `form:"created_at,omitempty" json:"created_at,omitempty" validate:"omitempty" gorm:"type:TIMESTAMPTZ;default:CURRENT_TIMESTAMP"`
|
||||
UpdatedAt time.Time `form:"updated_at,omitempty" json:"updated_at,omitempty" validate:"omitempty" gorm:"type:TIMESTAMPTZ;default:CURRENT_TIMESTAMP"`
|
||||
}
|
||||
|
||||
// TableName returns the table name string to guaranty use correct table
|
||||
func (c *Container) TableName() string {
|
||||
return "containers"
|
||||
}
|
||||
|
||||
// Valid is function to control input/output data
|
||||
func (c Container) Valid() error {
|
||||
return validate.Struct(c)
|
||||
}
|
||||
|
||||
// Validate is function to use callback to control input/output data
|
||||
func (c Container) Validate(db *gorm.DB) {
|
||||
if err := c.Valid(); err != nil {
|
||||
db.AddError(err)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,79 @@
|
||||
package models
|
||||
|
||||
import "time"
|
||||
|
||||
// FlowFile represents a single entry in the flow's local file cache.
|
||||
type FlowFile struct {
|
||||
ID string `json:"id"`
|
||||
Name string `json:"name"`
|
||||
Path string `json:"path"` // relative: "uploads/<x>" or resources/<x> or container/<x>"
|
||||
Size int64 `json:"size"`
|
||||
IsDir bool `json:"is_dir"`
|
||||
ModifiedAt time.Time `json:"modified_at"`
|
||||
}
|
||||
|
||||
// FlowFiles is the list response for flow file operations.
|
||||
type FlowFiles struct {
|
||||
Files []FlowFile `json:"files"`
|
||||
Total uint64 `json:"total"`
|
||||
}
|
||||
|
||||
// ContainerFile represents a single entry in the running container's filesystem.
|
||||
type ContainerFile struct {
|
||||
ID string `json:"id"`
|
||||
Name string `json:"name"`
|
||||
Path string `json:"path"`
|
||||
Size int64 `json:"size"`
|
||||
IsDir bool `json:"is_dir"`
|
||||
ModifiedAt time.Time `json:"modified_at"`
|
||||
}
|
||||
|
||||
// ContainerFiles is the list response for container directory listing.
|
||||
type ContainerFiles struct {
|
||||
Path string `json:"path"`
|
||||
Files []ContainerFile `json:"files"`
|
||||
Total uint64 `json:"total"`
|
||||
}
|
||||
|
||||
// PullFlowFilesRequest is the request body for pulling files from a container.
|
||||
type PullFlowFilesRequest struct {
|
||||
// Path is a single container path (maintained for backward compatibility).
|
||||
// Combined with Paths if both are provided.
|
||||
Path string `json:"path"`
|
||||
// Paths is a list of container paths to pull. Combined with Path if provided.
|
||||
Paths []string `json:"paths"`
|
||||
// Force overwrites local cache entries that already exist.
|
||||
Force bool `json:"force"`
|
||||
}
|
||||
|
||||
// AddResourcesRequest is the request body for copying user resources into a flow.
|
||||
type AddResourcesRequest struct {
|
||||
// IDs is the list of user resource IDs to copy into the flow resources directory.
|
||||
IDs []uint64 `json:"ids" binding:"required,min=1"`
|
||||
// Force overwrites files that already exist in the flow resources directory.
|
||||
Force bool `json:"force"`
|
||||
}
|
||||
|
||||
// AddResourceFromFlowRequest is the request body for promoting one or more flow
|
||||
// files / directories into the user's global resource store.
|
||||
//
|
||||
// At least one of Source or Sources must be non-empty after deduplication.
|
||||
//
|
||||
// - Single source (Source XOR one entry in Sources):
|
||||
// Destination is the exact target path (file) or root directory (dir tree).
|
||||
// - Multiple sources:
|
||||
// Destination is treated as a base directory; each source's base name is
|
||||
// appended automatically, e.g. sources ["uploads/a.txt", "container/b.txt"]
|
||||
// with destination "results" → "results/a.txt" and "results/b.txt".
|
||||
type AddResourceFromFlowRequest struct {
|
||||
// Source is a single relative path within the flow cache (kept for backward
|
||||
// compatibility). Combined with Sources when both are provided.
|
||||
Source string `json:"source"`
|
||||
// Sources is a list of relative paths within the flow cache. Combined with
|
||||
// Source when both are provided; duplicate entries are silently removed.
|
||||
Sources []string `json:"sources"`
|
||||
// Destination is the virtual path prefix in the user's resource tree.
|
||||
Destination string `json:"destination" binding:"required"`
|
||||
// Force overwrites existing resources at the target paths when true.
|
||||
Force bool `json:"force"`
|
||||
}
|
||||
@@ -0,0 +1,168 @@
|
||||
package models
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"time"
|
||||
|
||||
"pentagi/pkg/tools"
|
||||
|
||||
"github.com/jinzhu/gorm"
|
||||
)
|
||||
|
||||
type FlowStatus string
|
||||
|
||||
const (
|
||||
FlowStatusCreated FlowStatus = "created"
|
||||
FlowStatusRunning FlowStatus = "running"
|
||||
FlowStatusWaiting FlowStatus = "waiting"
|
||||
FlowStatusFinished FlowStatus = "finished"
|
||||
FlowStatusFailed FlowStatus = "failed"
|
||||
)
|
||||
|
||||
func (s FlowStatus) String() string {
|
||||
return string(s)
|
||||
}
|
||||
|
||||
// Valid is function to control input/output data
|
||||
func (s FlowStatus) Valid() error {
|
||||
switch s {
|
||||
case FlowStatusCreated,
|
||||
FlowStatusRunning,
|
||||
FlowStatusWaiting,
|
||||
FlowStatusFinished,
|
||||
FlowStatusFailed:
|
||||
return nil
|
||||
default:
|
||||
return fmt.Errorf("invalid FlowStatus: %s", s)
|
||||
}
|
||||
}
|
||||
|
||||
// Validate is function to use callback to control input/output data
|
||||
func (s FlowStatus) Validate(db *gorm.DB) {
|
||||
if err := s.Valid(); err != nil {
|
||||
db.AddError(err)
|
||||
}
|
||||
}
|
||||
|
||||
// Flow is model to contain flow information
|
||||
// nolint:lll
|
||||
type Flow struct {
|
||||
ID uint64 `form:"id" json:"id" validate:"min=0,numeric" gorm:"type:BIGINT;NOT NULL;PRIMARY_KEY;AUTO_INCREMENT"`
|
||||
Status FlowStatus `form:"status" json:"status" validate:"valid,required" gorm:"type:FLOW_STATUS;NOT NULL;default:'created'"`
|
||||
Title string `form:"title" json:"title" validate:"required" gorm:"type:TEXT;NOT NULL;default:'untitled'"`
|
||||
Model string `form:"model" json:"model" validate:"max=70,required" gorm:"type:TEXT;NOT NULL"`
|
||||
ModelProviderName string `form:"model_provider_name" json:"model_provider_name" validate:"max=70,required" gorm:"type:TEXT;NOT NULL"`
|
||||
ModelProviderType ProviderType `form:"model_provider_type" json:"model_provider_type" validate:"valid,required" gorm:"type:PROVIDER_TYPE;NOT NULL"`
|
||||
Language string `form:"language" json:"language" validate:"max=70,required" gorm:"type:TEXT;NOT NULL"`
|
||||
Functions *tools.Functions `form:"functions,omitempty" json:"functions,omitempty" validate:"omitempty,valid" gorm:"type:JSON;NOT NULL;default:'{}'"`
|
||||
ToolCallIDTemplate string `form:"tool_call_id_template" json:"tool_call_id_template" validate:"max=70,required" gorm:"type:TEXT;NOT NULL"`
|
||||
TraceID *string `form:"trace_id" json:"trace_id" validate:"max=70,required" gorm:"type:TEXT;NOT NULL"`
|
||||
UserID uint64 `form:"user_id" json:"user_id" validate:"min=0,numeric,required" gorm:"type:BIGINT;NOT NULL"`
|
||||
CreatedAt time.Time `form:"created_at,omitempty" json:"created_at,omitempty" validate:"omitempty" gorm:"type:TIMESTAMPTZ;default:CURRENT_TIMESTAMP"`
|
||||
UpdatedAt time.Time `form:"updated_at,omitempty" json:"updated_at,omitempty" validate:"omitempty" gorm:"type:TIMESTAMPTZ;default:CURRENT_TIMESTAMP"`
|
||||
DeletedAt *time.Time `form:"deleted_at,omitempty" json:"deleted_at,omitempty" validate:"omitempty" sql:"index" gorm:"type:TIMESTAMPTZ"`
|
||||
}
|
||||
|
||||
// TableName returns the table name string to guaranty use correct table
|
||||
func (f *Flow) TableName() string {
|
||||
return "flows"
|
||||
}
|
||||
|
||||
// Valid is function to control input/output data
|
||||
func (f Flow) Valid() error {
|
||||
return validate.Struct(f)
|
||||
}
|
||||
|
||||
// Validate is function to use callback to control input/output data
|
||||
func (f Flow) Validate(db *gorm.DB) {
|
||||
if err := f.Valid(); err != nil {
|
||||
db.AddError(err)
|
||||
}
|
||||
}
|
||||
|
||||
// CreateFlow is model to contain flow creation paylaod
|
||||
// nolint:lll
|
||||
type CreateFlow struct {
|
||||
Input string `form:"input" json:"input" validate:"required" example:"user input for first task in the flow"`
|
||||
Provider string `form:"provider" json:"provider" validate:"required" example:"openai"`
|
||||
Functions *tools.Functions `form:"functions,omitempty" json:"functions,omitempty" validate:"omitempty,valid"`
|
||||
ResourceIDs []uint64 `form:"resource_ids,omitempty" json:"resource_ids,omitempty" validate:"omitempty" swaggertype:"array,integer"`
|
||||
}
|
||||
|
||||
// Valid is function to control input/output data
|
||||
func (cf CreateFlow) Valid() error {
|
||||
return validate.Struct(cf)
|
||||
}
|
||||
|
||||
// PatchFlow is model to contain flow patching paylaod
|
||||
// nolint:lll
|
||||
type PatchFlow struct {
|
||||
Action string `form:"action" json:"action" validate:"required,oneof=stop finish input rename" enums:"stop,finish,input,rename" default:"stop"`
|
||||
Input *string `form:"input,omitempty" json:"input,omitempty" validate:"required_if=Action input" example:"user input for waiting flow"`
|
||||
Provider *string `form:"provider,omitempty" json:"provider,omitempty" validate:"omitempty"`
|
||||
Name *string `form:"name,omitempty" json:"name,omitempty" validate:"required_if=Action rename" example:"new flow name"`
|
||||
ResourceIDs []uint64 `form:"resource_ids,omitempty" json:"resource_ids,omitempty" validate:"omitempty" swaggertype:"array,integer"`
|
||||
}
|
||||
|
||||
// Valid is function to control input/output data
|
||||
func (pf PatchFlow) Valid() error {
|
||||
return validate.Struct(pf)
|
||||
}
|
||||
|
||||
// FlowTasksSubtasks is model to contain flow, linded tasks and linked subtasks information
|
||||
// nolint:lll
|
||||
type FlowTasksSubtasks struct {
|
||||
Tasks []TaskSubtasks `form:"tasks" json:"tasks" validate:"required" gorm:"foreignkey:FlowID;association_autoupdate:false;association_autocreate:false"`
|
||||
Flow `form:"" json:""`
|
||||
}
|
||||
|
||||
// TableName returns the table name string to guaranty use correct table
|
||||
func (fts *FlowTasksSubtasks) TableName() string {
|
||||
return "flows"
|
||||
}
|
||||
|
||||
// Valid is function to control input/output data
|
||||
func (fts FlowTasksSubtasks) Valid() error {
|
||||
for i := range fts.Tasks {
|
||||
if err := fts.Tasks[i].Valid(); err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
return fts.Flow.Valid()
|
||||
}
|
||||
|
||||
// Validate is function to use callback to control input/output data
|
||||
func (fts FlowTasksSubtasks) Validate(db *gorm.DB) {
|
||||
if err := fts.Valid(); err != nil {
|
||||
db.AddError(err)
|
||||
}
|
||||
}
|
||||
|
||||
// FlowContainers is model to contain flow and linked containers information
|
||||
// nolint:lll
|
||||
type FlowContainers struct {
|
||||
Containers []Container `form:"containers" json:"containers" validate:"required" gorm:"foreignkey:FlowID;association_autoupdate:false;association_autocreate:false"`
|
||||
Flow `form:"" json:""`
|
||||
}
|
||||
|
||||
// TableName returns the table name string to guaranty use correct table
|
||||
func (fc *FlowContainers) TableName() string {
|
||||
return "flows"
|
||||
}
|
||||
|
||||
// Valid is function to control input/output data
|
||||
func (fc FlowContainers) Valid() error {
|
||||
for i := range fc.Containers {
|
||||
if err := fc.Containers[i].Valid(); err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
return fc.Flow.Valid()
|
||||
}
|
||||
|
||||
// Validate is function to use callback to control input/output data
|
||||
func (fc FlowContainers) Validate(db *gorm.DB) {
|
||||
if err := fc.Valid(); err != nil {
|
||||
db.AddError(err)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,711 @@
|
||||
package models
|
||||
|
||||
import (
|
||||
"testing"
|
||||
|
||||
"github.com/stretchr/testify/assert"
|
||||
)
|
||||
|
||||
func TestFlowStatusValid(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
tests := []struct {
|
||||
name string
|
||||
status FlowStatus
|
||||
wantErr bool
|
||||
}{
|
||||
{"valid created", FlowStatusCreated, false},
|
||||
{"valid running", FlowStatusRunning, false},
|
||||
{"valid waiting", FlowStatusWaiting, false},
|
||||
{"valid finished", FlowStatusFinished, false},
|
||||
{"valid failed", FlowStatusFailed, false},
|
||||
{"invalid empty", FlowStatus(""), true},
|
||||
{"invalid unknown", FlowStatus("unknown"), true},
|
||||
{"invalid paused", FlowStatus("paused"), true},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
tt := tt
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
t.Parallel()
|
||||
err := tt.status.Valid()
|
||||
if tt.wantErr {
|
||||
assert.Error(t, err)
|
||||
assert.Contains(t, err.Error(), "invalid FlowStatus")
|
||||
} else {
|
||||
assert.NoError(t, err)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestFlowStatusString(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
assert.Equal(t, "created", FlowStatusCreated.String())
|
||||
assert.Equal(t, "running", FlowStatusRunning.String())
|
||||
assert.Equal(t, "finished", FlowStatusFinished.String())
|
||||
}
|
||||
|
||||
func TestFlowTableName(t *testing.T) {
|
||||
t.Parallel()
|
||||
f := &Flow{}
|
||||
assert.Equal(t, "flows", f.TableName())
|
||||
}
|
||||
|
||||
func TestCreateFlowValid(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
t.Run("valid create flow", func(t *testing.T) {
|
||||
t.Parallel()
|
||||
cf := CreateFlow{Input: "scan target", Provider: "openai"}
|
||||
assert.NoError(t, cf.Valid())
|
||||
})
|
||||
|
||||
t.Run("missing input", func(t *testing.T) {
|
||||
t.Parallel()
|
||||
cf := CreateFlow{Input: "", Provider: "openai"}
|
||||
assert.Error(t, cf.Valid())
|
||||
})
|
||||
|
||||
t.Run("missing provider", func(t *testing.T) {
|
||||
t.Parallel()
|
||||
cf := CreateFlow{Input: "scan target", Provider: ""}
|
||||
assert.Error(t, cf.Valid())
|
||||
})
|
||||
}
|
||||
|
||||
func TestPatchFlowValid(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
t.Run("valid stop action", func(t *testing.T) {
|
||||
t.Parallel()
|
||||
pf := PatchFlow{Action: "stop"}
|
||||
assert.NoError(t, pf.Valid())
|
||||
})
|
||||
|
||||
t.Run("valid finish action", func(t *testing.T) {
|
||||
t.Parallel()
|
||||
pf := PatchFlow{Action: "finish"}
|
||||
assert.NoError(t, pf.Valid())
|
||||
})
|
||||
|
||||
t.Run("valid input action with input", func(t *testing.T) {
|
||||
t.Parallel()
|
||||
input := "new input"
|
||||
pf := PatchFlow{Action: "input", Input: &input}
|
||||
assert.NoError(t, pf.Valid())
|
||||
})
|
||||
|
||||
t.Run("valid rename action with name", func(t *testing.T) {
|
||||
t.Parallel()
|
||||
name := "new name"
|
||||
pf := PatchFlow{Action: "rename", Name: &name}
|
||||
assert.NoError(t, pf.Valid())
|
||||
})
|
||||
|
||||
t.Run("invalid action", func(t *testing.T) {
|
||||
t.Parallel()
|
||||
pf := PatchFlow{Action: "invalid"}
|
||||
assert.Error(t, pf.Valid())
|
||||
})
|
||||
|
||||
t.Run("empty action", func(t *testing.T) {
|
||||
t.Parallel()
|
||||
pf := PatchFlow{Action: ""}
|
||||
assert.Error(t, pf.Valid())
|
||||
})
|
||||
|
||||
t.Run("input action without input", func(t *testing.T) {
|
||||
t.Parallel()
|
||||
pf := PatchFlow{Action: "input"}
|
||||
assert.Error(t, pf.Valid())
|
||||
})
|
||||
|
||||
t.Run("rename action without name", func(t *testing.T) {
|
||||
t.Parallel()
|
||||
pf := PatchFlow{Action: "rename"}
|
||||
assert.Error(t, pf.Valid())
|
||||
})
|
||||
}
|
||||
|
||||
func TestTaskStatusValid(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
tests := []struct {
|
||||
name string
|
||||
status TaskStatus
|
||||
wantErr bool
|
||||
}{
|
||||
{"valid created", TaskStatusCreated, false},
|
||||
{"valid running", TaskStatusRunning, false},
|
||||
{"valid waiting", TaskStatusWaiting, false},
|
||||
{"valid finished", TaskStatusFinished, false},
|
||||
{"valid failed", TaskStatusFailed, false},
|
||||
{"invalid empty", TaskStatus(""), true},
|
||||
{"invalid unknown", TaskStatus("cancelled"), true},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
tt := tt
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
t.Parallel()
|
||||
err := tt.status.Valid()
|
||||
if tt.wantErr {
|
||||
assert.Error(t, err)
|
||||
assert.Contains(t, err.Error(), "invalid TaskStatus")
|
||||
} else {
|
||||
assert.NoError(t, err)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestTaskTableName(t *testing.T) {
|
||||
t.Parallel()
|
||||
tk := &Task{}
|
||||
assert.Equal(t, "tasks", tk.TableName())
|
||||
}
|
||||
|
||||
func TestSubtaskStatusValid(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
tests := []struct {
|
||||
name string
|
||||
status SubtaskStatus
|
||||
wantErr bool
|
||||
}{
|
||||
{"valid created", SubtaskStatusCreated, false},
|
||||
{"valid running", SubtaskStatusRunning, false},
|
||||
{"valid waiting", SubtaskStatusWaiting, false},
|
||||
{"valid finished", SubtaskStatusFinished, false},
|
||||
{"valid failed", SubtaskStatusFailed, false},
|
||||
{"invalid empty", SubtaskStatus(""), true},
|
||||
{"invalid unknown", SubtaskStatus("pending"), true},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
tt := tt
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
t.Parallel()
|
||||
err := tt.status.Valid()
|
||||
if tt.wantErr {
|
||||
assert.Error(t, err)
|
||||
assert.Contains(t, err.Error(), "invalid SubtaskStatus")
|
||||
} else {
|
||||
assert.NoError(t, err)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestSubtaskTableName(t *testing.T) {
|
||||
t.Parallel()
|
||||
s := &Subtask{}
|
||||
assert.Equal(t, "subtasks", s.TableName())
|
||||
}
|
||||
|
||||
func TestContainerStatusValid(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
tests := []struct {
|
||||
name string
|
||||
status ContainerStatus
|
||||
wantErr bool
|
||||
}{
|
||||
{"valid starting", ContainerStatusStarting, false},
|
||||
{"valid running", ContainerStatusRunning, false},
|
||||
{"valid stopped", ContainerStatusStopped, false},
|
||||
{"valid deleted", ContainerStatusDeleted, false},
|
||||
{"valid failed", ContainerStatusFailed, false},
|
||||
{"invalid empty", ContainerStatus(""), true},
|
||||
{"invalid unknown", ContainerStatus("restarting"), true},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
tt := tt
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
t.Parallel()
|
||||
err := tt.status.Valid()
|
||||
if tt.wantErr {
|
||||
assert.Error(t, err)
|
||||
assert.Contains(t, err.Error(), "invalid ContainerStatus")
|
||||
} else {
|
||||
assert.NoError(t, err)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestContainerTypeValid(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
tests := []struct {
|
||||
name string
|
||||
ct ContainerType
|
||||
wantErr bool
|
||||
}{
|
||||
{"valid primary", ContainerTypePrimary, false},
|
||||
{"valid secondary", ContainerTypeSecondary, false},
|
||||
{"invalid empty", ContainerType(""), true},
|
||||
{"invalid tertiary", ContainerType("tertiary"), true},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
tt := tt
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
t.Parallel()
|
||||
err := tt.ct.Valid()
|
||||
if tt.wantErr {
|
||||
assert.Error(t, err)
|
||||
assert.Contains(t, err.Error(), "invalid ContainerType")
|
||||
} else {
|
||||
assert.NoError(t, err)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestContainerTableName(t *testing.T) {
|
||||
t.Parallel()
|
||||
c := &Container{}
|
||||
assert.Equal(t, "containers", c.TableName())
|
||||
}
|
||||
|
||||
func TestRoleValid(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
t.Run("valid role", func(t *testing.T) {
|
||||
t.Parallel()
|
||||
r := Role{Name: "admin"}
|
||||
assert.NoError(t, r.Valid())
|
||||
})
|
||||
|
||||
t.Run("empty name", func(t *testing.T) {
|
||||
t.Parallel()
|
||||
r := Role{Name: ""}
|
||||
assert.Error(t, r.Valid())
|
||||
})
|
||||
}
|
||||
|
||||
func TestRoleTableName(t *testing.T) {
|
||||
t.Parallel()
|
||||
r := &Role{}
|
||||
assert.Equal(t, "roles", r.TableName())
|
||||
}
|
||||
|
||||
func TestPrivilegeValid(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
t.Run("valid privilege", func(t *testing.T) {
|
||||
t.Parallel()
|
||||
p := Privilege{Name: "read"}
|
||||
assert.NoError(t, p.Valid())
|
||||
})
|
||||
|
||||
t.Run("empty name", func(t *testing.T) {
|
||||
t.Parallel()
|
||||
p := Privilege{Name: ""}
|
||||
assert.Error(t, p.Valid())
|
||||
})
|
||||
}
|
||||
|
||||
func TestPrivilegeTableName(t *testing.T) {
|
||||
t.Parallel()
|
||||
p := &Privilege{}
|
||||
assert.Equal(t, "privileges", p.TableName())
|
||||
}
|
||||
|
||||
func TestRolePrivilegesValid(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
t.Run("valid role privileges", func(t *testing.T) {
|
||||
t.Parallel()
|
||||
rp := RolePrivileges{
|
||||
Privileges: []Privilege{{Name: "read"}, {Name: "write"}},
|
||||
Role: Role{Name: "admin"},
|
||||
}
|
||||
assert.NoError(t, rp.Valid())
|
||||
})
|
||||
|
||||
t.Run("invalid privilege", func(t *testing.T) {
|
||||
t.Parallel()
|
||||
rp := RolePrivileges{
|
||||
Privileges: []Privilege{{Name: ""}, {Name: "write"}},
|
||||
Role: Role{Name: "admin"},
|
||||
}
|
||||
assert.Error(t, rp.Valid())
|
||||
})
|
||||
|
||||
t.Run("invalid role", func(t *testing.T) {
|
||||
t.Parallel()
|
||||
rp := RolePrivileges{
|
||||
Privileges: []Privilege{{Name: "read"}},
|
||||
Role: Role{Name: ""},
|
||||
}
|
||||
assert.Error(t, rp.Valid())
|
||||
})
|
||||
}
|
||||
|
||||
func TestRolePrivilegesTableName(t *testing.T) {
|
||||
t.Parallel()
|
||||
rp := &RolePrivileges{}
|
||||
assert.Equal(t, "roles", rp.TableName())
|
||||
}
|
||||
|
||||
func TestFlowValid(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
traceID := "trace-123"
|
||||
validFlow := Flow{
|
||||
Status: FlowStatusCreated,
|
||||
Title: "test flow",
|
||||
Model: "gpt-4",
|
||||
ModelProviderName: "openai",
|
||||
ModelProviderType: ProviderType("openai"),
|
||||
Language: "en",
|
||||
ToolCallIDTemplate: "call_{id}",
|
||||
TraceID: &traceID,
|
||||
UserID: 1,
|
||||
}
|
||||
|
||||
t.Run("valid flow", func(t *testing.T) {
|
||||
t.Parallel()
|
||||
assert.NoError(t, validFlow.Valid())
|
||||
})
|
||||
|
||||
t.Run("missing title", func(t *testing.T) {
|
||||
t.Parallel()
|
||||
f := validFlow
|
||||
f.Title = ""
|
||||
assert.Error(t, f.Valid())
|
||||
})
|
||||
|
||||
t.Run("invalid status", func(t *testing.T) {
|
||||
t.Parallel()
|
||||
f := validFlow
|
||||
f.Status = FlowStatus("invalid")
|
||||
assert.Error(t, f.Valid())
|
||||
})
|
||||
|
||||
t.Run("invalid provider type", func(t *testing.T) {
|
||||
t.Parallel()
|
||||
f := validFlow
|
||||
f.ModelProviderType = ProviderType("invalid")
|
||||
assert.Error(t, f.Valid())
|
||||
})
|
||||
}
|
||||
|
||||
func TestFlowTasksSubtasksValid(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
traceID := "trace-123"
|
||||
validFlow := Flow{
|
||||
Status: FlowStatusCreated,
|
||||
Title: "test flow",
|
||||
Model: "gpt-4",
|
||||
ModelProviderName: "openai",
|
||||
ModelProviderType: ProviderType("openai"),
|
||||
Language: "en",
|
||||
ToolCallIDTemplate: "call_{id}",
|
||||
TraceID: &traceID,
|
||||
UserID: 1,
|
||||
}
|
||||
|
||||
validTask := Task{
|
||||
Status: TaskStatusCreated,
|
||||
Title: "task1",
|
||||
Input: "input1",
|
||||
FlowID: 1,
|
||||
}
|
||||
validSubtask := Subtask{
|
||||
Status: SubtaskStatusCreated,
|
||||
Title: "subtask1",
|
||||
Description: "desc1",
|
||||
TaskID: 1,
|
||||
}
|
||||
|
||||
t.Run("valid flow tasks subtasks", func(t *testing.T) {
|
||||
t.Parallel()
|
||||
fts := FlowTasksSubtasks{
|
||||
Tasks: []TaskSubtasks{{
|
||||
Subtasks: []Subtask{validSubtask},
|
||||
Task: validTask,
|
||||
}},
|
||||
Flow: validFlow,
|
||||
}
|
||||
assert.NoError(t, fts.Valid())
|
||||
})
|
||||
|
||||
t.Run("invalid flow", func(t *testing.T) {
|
||||
t.Parallel()
|
||||
badFlow := validFlow
|
||||
badFlow.Title = ""
|
||||
fts := FlowTasksSubtasks{
|
||||
Tasks: []TaskSubtasks{{
|
||||
Subtasks: []Subtask{validSubtask},
|
||||
Task: validTask,
|
||||
}},
|
||||
Flow: badFlow,
|
||||
}
|
||||
assert.Error(t, fts.Valid())
|
||||
})
|
||||
|
||||
t.Run("invalid nested subtask", func(t *testing.T) {
|
||||
t.Parallel()
|
||||
badSubtask := validSubtask
|
||||
badSubtask.Title = ""
|
||||
fts := FlowTasksSubtasks{
|
||||
Tasks: []TaskSubtasks{{
|
||||
Subtasks: []Subtask{badSubtask},
|
||||
Task: validTask,
|
||||
}},
|
||||
Flow: validFlow,
|
||||
}
|
||||
assert.Error(t, fts.Valid())
|
||||
})
|
||||
}
|
||||
|
||||
func TestFlowTasksSubtasksTableName(t *testing.T) {
|
||||
t.Parallel()
|
||||
fts := &FlowTasksSubtasks{}
|
||||
assert.Equal(t, "flows", fts.TableName())
|
||||
}
|
||||
|
||||
func TestFlowContainersValid(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
traceID := "trace-123"
|
||||
validFlow := Flow{
|
||||
Status: FlowStatusCreated,
|
||||
Title: "test flow",
|
||||
Model: "gpt-4",
|
||||
ModelProviderName: "openai",
|
||||
ModelProviderType: ProviderType("openai"),
|
||||
Language: "en",
|
||||
ToolCallIDTemplate: "call_{id}",
|
||||
TraceID: &traceID,
|
||||
UserID: 1,
|
||||
}
|
||||
validContainer := Container{
|
||||
Type: ContainerTypePrimary,
|
||||
Name: "test-container",
|
||||
Image: "alpine:latest",
|
||||
Status: ContainerStatusRunning,
|
||||
LocalID: "abc123",
|
||||
LocalDir: "/tmp/test",
|
||||
FlowID: 1,
|
||||
}
|
||||
|
||||
t.Run("valid flow containers", func(t *testing.T) {
|
||||
t.Parallel()
|
||||
fc := FlowContainers{
|
||||
Containers: []Container{validContainer},
|
||||
Flow: validFlow,
|
||||
}
|
||||
assert.NoError(t, fc.Valid())
|
||||
})
|
||||
|
||||
t.Run("invalid flow", func(t *testing.T) {
|
||||
t.Parallel()
|
||||
badFlow := validFlow
|
||||
badFlow.Title = ""
|
||||
fc := FlowContainers{
|
||||
Containers: []Container{validContainer},
|
||||
Flow: badFlow,
|
||||
}
|
||||
assert.Error(t, fc.Valid())
|
||||
})
|
||||
|
||||
t.Run("invalid container", func(t *testing.T) {
|
||||
t.Parallel()
|
||||
badContainer := validContainer
|
||||
badContainer.Name = ""
|
||||
fc := FlowContainers{
|
||||
Containers: []Container{badContainer},
|
||||
Flow: validFlow,
|
||||
}
|
||||
assert.Error(t, fc.Valid())
|
||||
})
|
||||
}
|
||||
|
||||
func TestFlowContainersTableName(t *testing.T) {
|
||||
t.Parallel()
|
||||
fc := &FlowContainers{}
|
||||
assert.Equal(t, "flows", fc.TableName())
|
||||
}
|
||||
|
||||
func TestTaskValid(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
validTask := Task{
|
||||
Status: TaskStatusCreated,
|
||||
Title: "test task",
|
||||
Input: "test input",
|
||||
FlowID: 1,
|
||||
}
|
||||
|
||||
t.Run("valid task", func(t *testing.T) {
|
||||
t.Parallel()
|
||||
assert.NoError(t, validTask.Valid())
|
||||
})
|
||||
|
||||
t.Run("missing title", func(t *testing.T) {
|
||||
t.Parallel()
|
||||
tk := validTask
|
||||
tk.Title = ""
|
||||
assert.Error(t, tk.Valid())
|
||||
})
|
||||
|
||||
t.Run("missing input", func(t *testing.T) {
|
||||
t.Parallel()
|
||||
tk := validTask
|
||||
tk.Input = ""
|
||||
assert.Error(t, tk.Valid())
|
||||
})
|
||||
|
||||
t.Run("invalid status", func(t *testing.T) {
|
||||
t.Parallel()
|
||||
tk := validTask
|
||||
tk.Status = TaskStatus("invalid")
|
||||
assert.Error(t, tk.Valid())
|
||||
})
|
||||
}
|
||||
|
||||
func TestTaskSubtasksValid(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
validTask := Task{
|
||||
Status: TaskStatusCreated,
|
||||
Title: "task1",
|
||||
Input: "input1",
|
||||
FlowID: 1,
|
||||
}
|
||||
validSubtask := Subtask{
|
||||
Status: SubtaskStatusCreated,
|
||||
Title: "subtask1",
|
||||
Description: "desc1",
|
||||
TaskID: 1,
|
||||
}
|
||||
|
||||
t.Run("valid task subtasks", func(t *testing.T) {
|
||||
t.Parallel()
|
||||
ts := TaskSubtasks{
|
||||
Subtasks: []Subtask{validSubtask},
|
||||
Task: validTask,
|
||||
}
|
||||
assert.NoError(t, ts.Valid())
|
||||
})
|
||||
|
||||
t.Run("invalid task", func(t *testing.T) {
|
||||
t.Parallel()
|
||||
badTask := validTask
|
||||
badTask.Title = ""
|
||||
ts := TaskSubtasks{
|
||||
Subtasks: []Subtask{validSubtask},
|
||||
Task: badTask,
|
||||
}
|
||||
assert.Error(t, ts.Valid())
|
||||
})
|
||||
|
||||
t.Run("invalid subtask", func(t *testing.T) {
|
||||
t.Parallel()
|
||||
badSubtask := validSubtask
|
||||
badSubtask.Title = ""
|
||||
ts := TaskSubtasks{
|
||||
Subtasks: []Subtask{badSubtask},
|
||||
Task: validTask,
|
||||
}
|
||||
assert.Error(t, ts.Valid())
|
||||
})
|
||||
}
|
||||
|
||||
func TestTaskSubtasksTableName(t *testing.T) {
|
||||
t.Parallel()
|
||||
ts := &TaskSubtasks{}
|
||||
assert.Equal(t, "tasks", ts.TableName())
|
||||
}
|
||||
|
||||
func TestSubtaskValid(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
validSubtask := Subtask{
|
||||
Status: SubtaskStatusCreated,
|
||||
Title: "subtask1",
|
||||
Description: "description1",
|
||||
TaskID: 1,
|
||||
}
|
||||
|
||||
t.Run("valid subtask", func(t *testing.T) {
|
||||
t.Parallel()
|
||||
assert.NoError(t, validSubtask.Valid())
|
||||
})
|
||||
|
||||
t.Run("missing title", func(t *testing.T) {
|
||||
t.Parallel()
|
||||
s := validSubtask
|
||||
s.Title = ""
|
||||
assert.Error(t, s.Valid())
|
||||
})
|
||||
|
||||
t.Run("missing description", func(t *testing.T) {
|
||||
t.Parallel()
|
||||
s := validSubtask
|
||||
s.Description = ""
|
||||
assert.Error(t, s.Valid())
|
||||
})
|
||||
|
||||
t.Run("invalid status", func(t *testing.T) {
|
||||
t.Parallel()
|
||||
s := validSubtask
|
||||
s.Status = SubtaskStatus("invalid")
|
||||
assert.Error(t, s.Valid())
|
||||
})
|
||||
}
|
||||
|
||||
func TestContainerValid(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
validContainer := Container{
|
||||
Type: ContainerTypePrimary,
|
||||
Name: "test-container",
|
||||
Image: "alpine:latest",
|
||||
Status: ContainerStatusRunning,
|
||||
LocalID: "abc123",
|
||||
LocalDir: "/tmp/test",
|
||||
FlowID: 1,
|
||||
}
|
||||
|
||||
t.Run("valid container", func(t *testing.T) {
|
||||
t.Parallel()
|
||||
assert.NoError(t, validContainer.Valid())
|
||||
})
|
||||
|
||||
t.Run("missing name", func(t *testing.T) {
|
||||
t.Parallel()
|
||||
c := validContainer
|
||||
c.Name = ""
|
||||
assert.Error(t, c.Valid())
|
||||
})
|
||||
|
||||
t.Run("missing image", func(t *testing.T) {
|
||||
t.Parallel()
|
||||
c := validContainer
|
||||
c.Image = ""
|
||||
assert.Error(t, c.Valid())
|
||||
})
|
||||
|
||||
t.Run("invalid status", func(t *testing.T) {
|
||||
t.Parallel()
|
||||
c := validContainer
|
||||
c.Status = ContainerStatus("invalid")
|
||||
assert.Error(t, c.Valid())
|
||||
})
|
||||
|
||||
t.Run("invalid type", func(t *testing.T) {
|
||||
t.Parallel()
|
||||
c := validContainer
|
||||
c.Type = ContainerType("invalid")
|
||||
assert.Error(t, c.Valid())
|
||||
})
|
||||
}
|
||||
@@ -0,0 +1,231 @@
|
||||
package models
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"reflect"
|
||||
"regexp"
|
||||
"strings"
|
||||
|
||||
"github.com/go-playground/validator/v10"
|
||||
"github.com/xeipuuv/gojsonschema"
|
||||
)
|
||||
|
||||
const (
|
||||
solidRegexString = "^[a-z0-9_\\-]+$"
|
||||
clDateRegexString = "^[0-9]{2}[.-][0-9]{2}[.-][0-9]{4}$"
|
||||
semverRegexString = "^[0-9]+\\.[0-9]+(\\.[0-9]+)?$"
|
||||
semverexRegexString = "^(v)?[0-9]+\\.[0-9]+(\\.[0-9]+)?(\\.[0-9]+)?(-[a-zA-Z0-9]+)?$"
|
||||
)
|
||||
|
||||
var (
|
||||
validate *validator.Validate
|
||||
)
|
||||
|
||||
func GetValidator() *validator.Validate {
|
||||
return validate
|
||||
}
|
||||
|
||||
// IValid is interface to control all models from user code
|
||||
type IValid interface {
|
||||
Valid() error
|
||||
}
|
||||
|
||||
func templateValidatorString(regexpString string) validator.Func {
|
||||
regexpValue := regexp.MustCompile(regexpString)
|
||||
return func(fl validator.FieldLevel) bool {
|
||||
field := fl.Field()
|
||||
matchString := func(str string) bool {
|
||||
if str == "" && fl.Param() == "omitempty" {
|
||||
return true
|
||||
}
|
||||
return regexpValue.MatchString(str)
|
||||
}
|
||||
|
||||
switch field.Kind() {
|
||||
case reflect.String:
|
||||
return matchString(fl.Field().String())
|
||||
case reflect.Slice, reflect.Array:
|
||||
for i := 0; i < field.Len(); i++ {
|
||||
if !matchString(field.Index(i).String()) {
|
||||
return false
|
||||
}
|
||||
}
|
||||
return true
|
||||
case reflect.Map:
|
||||
for _, k := range field.MapKeys() {
|
||||
if !matchString(field.MapIndex(k).String()) {
|
||||
return false
|
||||
}
|
||||
}
|
||||
return true
|
||||
default:
|
||||
return false
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func strongPasswordValidatorString() validator.Func {
|
||||
numberRegex := regexp.MustCompile("[0-9]")
|
||||
alphaLRegex := regexp.MustCompile("[a-z]")
|
||||
alphaURegex := regexp.MustCompile("[A-Z]")
|
||||
specRegex := regexp.MustCompile("[!@#$&*]")
|
||||
return func(fl validator.FieldLevel) bool {
|
||||
field := fl.Field()
|
||||
|
||||
switch field.Kind() {
|
||||
case reflect.String:
|
||||
password := fl.Field().String()
|
||||
return len(password) > 15 || (len(password) >= 8 &&
|
||||
numberRegex.MatchString(password) &&
|
||||
alphaLRegex.MatchString(password) &&
|
||||
alphaURegex.MatchString(password) &&
|
||||
specRegex.MatchString(password))
|
||||
default:
|
||||
return false
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func emailValidatorString() validator.Func {
|
||||
emailRegex := regexp.MustCompile(`^[a-z0-9._%+\-]+@[a-z0-9.\-]+\.[a-z]{2,4}$`)
|
||||
return func(fl validator.FieldLevel) bool {
|
||||
field := fl.Field()
|
||||
|
||||
switch field.Kind() {
|
||||
case reflect.String:
|
||||
email := fl.Field().String()
|
||||
if email == "admin" {
|
||||
return true
|
||||
}
|
||||
if err := validate.Var(email, "required,uuid"); err == nil {
|
||||
return true
|
||||
}
|
||||
return len(email) > 4 && emailRegex.MatchString(email)
|
||||
default:
|
||||
return false
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func oauthMinScope() validator.Func {
|
||||
scopeParts := []string{
|
||||
"openid",
|
||||
"email",
|
||||
}
|
||||
return func(fl validator.FieldLevel) bool {
|
||||
field := fl.Field()
|
||||
|
||||
switch field.Kind() {
|
||||
case reflect.String:
|
||||
scope := strings.ToLower(fl.Field().String())
|
||||
for _, part := range scopeParts {
|
||||
if !strings.Contains(scope, part) {
|
||||
return false
|
||||
}
|
||||
}
|
||||
return true
|
||||
default:
|
||||
return false
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func deepValidator() validator.Func {
|
||||
return func(fl validator.FieldLevel) bool {
|
||||
if iv, ok := fl.Field().Interface().(IValid); ok {
|
||||
if err := iv.Valid(); err != nil {
|
||||
return false
|
||||
}
|
||||
}
|
||||
|
||||
return true
|
||||
}
|
||||
}
|
||||
|
||||
func getMapKeys(kvmap interface{}) string {
|
||||
kl := []interface{}{}
|
||||
val := reflect.ValueOf(kvmap)
|
||||
if val.Kind() == reflect.Map {
|
||||
for _, e := range val.MapKeys() {
|
||||
v := val.MapIndex(e)
|
||||
kl = append(kl, v.Interface())
|
||||
}
|
||||
}
|
||||
kld, _ := json.Marshal(kl)
|
||||
return string(kld)
|
||||
}
|
||||
|
||||
func mismatchLenError(tag string, wants, current int) string {
|
||||
return fmt.Sprintf("%s wants len %d but current is %d", tag, wants, current)
|
||||
}
|
||||
|
||||
func keyIsNotExtistInMap(tag, key string, kvmap interface{}) string {
|
||||
return fmt.Sprintf("%s must present key %s in keys list %s", tag, key, getMapKeys(kvmap))
|
||||
}
|
||||
|
||||
func keyIsNotExtistInSlice(tag, key string, klist interface{}) string {
|
||||
kld, _ := json.Marshal(klist)
|
||||
return fmt.Sprintf("%s must present key %s in keys list %s", tag, key, string(kld))
|
||||
}
|
||||
|
||||
func keysAreNotExtistInSlice(tag, lkeys, rkeys interface{}) string {
|
||||
lkeysd, _ := json.Marshal(lkeys)
|
||||
rkeysd, _ := json.Marshal(rkeys)
|
||||
return fmt.Sprintf("%s must all keys present %s in keys list %s", tag, string(lkeysd), string(rkeysd))
|
||||
}
|
||||
|
||||
func contextError(tag string, id string, ctx interface{}) string {
|
||||
ctxd, _ := json.Marshal(ctx)
|
||||
return fmt.Sprintf("%s with %s ctx %s", tag, id, string(ctxd))
|
||||
}
|
||||
|
||||
func caughtValidationError(tag string, err error) string {
|
||||
return fmt.Sprintf("%s caught error %s", tag, err.Error())
|
||||
}
|
||||
|
||||
func caughtSchemaValidationError(tag string, errs []gojsonschema.ResultError) string {
|
||||
var arr []string
|
||||
for _, err := range errs {
|
||||
arr = append(arr, err.String())
|
||||
}
|
||||
errd, _ := json.Marshal(arr)
|
||||
return fmt.Sprintf("%s caught errors %s", tag, string(errd))
|
||||
}
|
||||
|
||||
func scanFromJSON(input interface{}, output interface{}) error {
|
||||
if v, ok := input.(string); ok {
|
||||
return json.Unmarshal([]byte(v), output)
|
||||
} else if v, ok := input.([]byte); ok {
|
||||
if err := json.Unmarshal(v, output); err != nil {
|
||||
return err
|
||||
}
|
||||
return nil
|
||||
}
|
||||
return fmt.Errorf("unsupported type of input value to scan")
|
||||
}
|
||||
|
||||
func init() {
|
||||
validate = validator.New()
|
||||
_ = validate.RegisterValidation("solid", templateValidatorString(solidRegexString))
|
||||
_ = validate.RegisterValidation("cldate", templateValidatorString(clDateRegexString))
|
||||
_ = validate.RegisterValidation("semver", templateValidatorString(semverRegexString))
|
||||
_ = validate.RegisterValidation("semverex", templateValidatorString(semverexRegexString))
|
||||
_ = validate.RegisterValidation("stpass", strongPasswordValidatorString())
|
||||
_ = validate.RegisterValidation("vmail", emailValidatorString())
|
||||
_ = validate.RegisterValidation("oauth_min_scope", oauthMinScope())
|
||||
_ = validate.RegisterValidation("valid", deepValidator())
|
||||
|
||||
// Check validation interface for all models
|
||||
_, _ = reflect.ValueOf(Login{}).Interface().(IValid)
|
||||
_, _ = reflect.ValueOf(AuthCallback{}).Interface().(IValid)
|
||||
|
||||
_, _ = reflect.ValueOf(User{}).Interface().(IValid)
|
||||
_, _ = reflect.ValueOf(Password{}).Interface().(IValid)
|
||||
|
||||
_, _ = reflect.ValueOf(Role{}).Interface().(IValid)
|
||||
_, _ = reflect.ValueOf(Prompt{}).Interface().(IValid)
|
||||
_, _ = reflect.ValueOf(Assistant{}).Interface().(IValid)
|
||||
_, _ = reflect.ValueOf(Flow{}).Interface().(IValid)
|
||||
_, _ = reflect.ValueOf(Provider{}).Interface().(IValid)
|
||||
}
|
||||
@@ -0,0 +1,250 @@
|
||||
package models
|
||||
|
||||
import (
|
||||
"testing"
|
||||
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
func TestGetValidator(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
v := GetValidator()
|
||||
require.NotNil(t, v)
|
||||
}
|
||||
|
||||
func TestStrongPasswordValidator(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
tests := []struct {
|
||||
name string
|
||||
pw string
|
||||
wantErr bool
|
||||
}{
|
||||
{"long password over 15 chars", "abcdefghijklmnop", false},
|
||||
{"exactly 16 chars", "abcdefghijklmnop", false},
|
||||
{"8 chars with all requirements", "Pass1!ab", false},
|
||||
{"short no special", "pass1", true},
|
||||
{"7 chars with requirements", "Pa1!abc", true},
|
||||
{"8 chars no number", "Pass!abc", true},
|
||||
{"8 chars no uppercase", "pass1!ab", true},
|
||||
{"8 chars no lowercase", "PASS1!AB", true},
|
||||
{"8 chars no special", "Pass1abc", true},
|
||||
{"empty password", "", true},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
tt := tt
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
t.Parallel()
|
||||
err := GetValidator().Var(tt.pw, "stpass")
|
||||
if tt.wantErr {
|
||||
assert.Error(t, err)
|
||||
} else {
|
||||
assert.NoError(t, err)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestEmailValidator(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
tests := []struct {
|
||||
name string
|
||||
email string
|
||||
wantErr bool
|
||||
}{
|
||||
{"valid email", "test@example.com", false},
|
||||
{"admin special case", "admin", false},
|
||||
{"email with subdomain", "user@mail.example.com", false},
|
||||
{"email too short", "a@b", true},
|
||||
{"empty email", "", true},
|
||||
{"no at sign", "testexample.com", true},
|
||||
{"no domain", "test@", true},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
tt := tt
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
t.Parallel()
|
||||
err := GetValidator().Var(tt.email, "vmail")
|
||||
if tt.wantErr {
|
||||
assert.Error(t, err)
|
||||
} else {
|
||||
assert.NoError(t, err)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestOAuthMinScopeValidator(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
tests := []struct {
|
||||
name string
|
||||
scope string
|
||||
wantErr bool
|
||||
}{
|
||||
{"valid openid email", "openid email", false},
|
||||
{"valid with extra scopes", "openid email profile", false},
|
||||
{"valid case insensitive", "OpenID Email", false},
|
||||
{"missing openid", "email profile", true},
|
||||
{"missing email", "openid profile", true},
|
||||
{"empty scope", "", true},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
tt := tt
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
t.Parallel()
|
||||
err := GetValidator().Var(tt.scope, "oauth_min_scope")
|
||||
if tt.wantErr {
|
||||
assert.Error(t, err)
|
||||
} else {
|
||||
assert.NoError(t, err)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestSolidValidator(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
tests := []struct {
|
||||
name string
|
||||
input string
|
||||
wantErr bool
|
||||
}{
|
||||
{"valid lowercase", "hello", false},
|
||||
{"valid with numbers", "abc123", false},
|
||||
{"valid with underscore", "my_name", false},
|
||||
{"valid with dash", "my-name", false},
|
||||
{"invalid uppercase", "Hello", true},
|
||||
{"invalid spaces", "hello world", true},
|
||||
{"invalid special chars", "hello@world", true},
|
||||
{"empty string", "", true},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
tt := tt
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
t.Parallel()
|
||||
err := GetValidator().Var(tt.input, "solid")
|
||||
if tt.wantErr {
|
||||
assert.Error(t, err)
|
||||
} else {
|
||||
assert.NoError(t, err)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestSemverValidator(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
tests := []struct {
|
||||
name string
|
||||
input string
|
||||
wantErr bool
|
||||
}{
|
||||
{"valid two part", "1.0", false},
|
||||
{"valid three part", "1.2.3", false},
|
||||
{"valid zero", "0.0.0", false},
|
||||
{"invalid prefix v", "v1.0.0", true},
|
||||
{"invalid single", "1", true},
|
||||
{"invalid text", "abc", true},
|
||||
{"empty", "", true},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
tt := tt
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
t.Parallel()
|
||||
err := GetValidator().Var(tt.input, "semver")
|
||||
if tt.wantErr {
|
||||
assert.Error(t, err)
|
||||
} else {
|
||||
assert.NoError(t, err)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestSemverExValidator(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
tests := []struct {
|
||||
name string
|
||||
input string
|
||||
wantErr bool
|
||||
}{
|
||||
{"valid two part", "1.0", false},
|
||||
{"valid three part", "1.2.3", false},
|
||||
{"valid with v prefix", "v1.0.0", false},
|
||||
{"valid with prerelease", "1.2.3-beta", false},
|
||||
{"valid four part", "1.2.3.4", false},
|
||||
{"invalid text only", "abc", true},
|
||||
{"empty", "", true},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
tt := tt
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
t.Parallel()
|
||||
err := GetValidator().Var(tt.input, "semverex")
|
||||
if tt.wantErr {
|
||||
assert.Error(t, err)
|
||||
} else {
|
||||
assert.NoError(t, err)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestScanFromJSON(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
type testStruct struct {
|
||||
Name string `json:"name"`
|
||||
}
|
||||
|
||||
t.Run("scan from string", func(t *testing.T) {
|
||||
t.Parallel()
|
||||
var out testStruct
|
||||
err := scanFromJSON(`{"name":"test"}`, &out)
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, "test", out.Name)
|
||||
})
|
||||
|
||||
t.Run("scan from bytes", func(t *testing.T) {
|
||||
t.Parallel()
|
||||
var out testStruct
|
||||
err := scanFromJSON([]byte(`{"name":"hello"}`), &out)
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, "hello", out.Name)
|
||||
})
|
||||
|
||||
t.Run("scan unsupported type", func(t *testing.T) {
|
||||
t.Parallel()
|
||||
var out testStruct
|
||||
err := scanFromJSON(12345, &out)
|
||||
assert.Error(t, err)
|
||||
assert.Contains(t, err.Error(), "unsupported type")
|
||||
})
|
||||
|
||||
t.Run("scan invalid json string", func(t *testing.T) {
|
||||
t.Parallel()
|
||||
var out testStruct
|
||||
err := scanFromJSON("not-json", &out)
|
||||
assert.Error(t, err)
|
||||
})
|
||||
|
||||
t.Run("scan invalid json bytes", func(t *testing.T) {
|
||||
t.Parallel()
|
||||
var out testStruct
|
||||
err := scanFromJSON([]byte("not-json"), &out)
|
||||
assert.Error(t, err)
|
||||
})
|
||||
}
|
||||
@@ -0,0 +1,476 @@
|
||||
package models
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
|
||||
gqlmodel "pentagi/pkg/graph/model"
|
||||
)
|
||||
|
||||
// knowledgeRawMeta mirrors the flat JSON stored in langchain_pg_embedding.cmetadata.
|
||||
// Used only within the models package for REST → gorm conversion.
|
||||
type knowledgeRawMeta struct {
|
||||
DocType string `json:"doc_type"`
|
||||
UserID int64 `json:"user_id"`
|
||||
FlowID *int64 `json:"flow_id"`
|
||||
TaskID *int64 `json:"task_id"`
|
||||
SubtaskID *int64 `json:"subtask_id"`
|
||||
Question string `json:"question"`
|
||||
Description string `json:"description"`
|
||||
GuideType string `json:"guide_type"`
|
||||
AnswerType string `json:"answer_type"`
|
||||
CodeLang string `json:"code_lang"`
|
||||
PartSize int `json:"part_size"`
|
||||
TotalSize int `json:"total_size"`
|
||||
Manual bool `json:"manual"`
|
||||
}
|
||||
|
||||
func knowledgeMetaFromJSON(raw string) knowledgeRawMeta {
|
||||
var m knowledgeRawMeta
|
||||
if raw != "" && raw != "{}" {
|
||||
_ = json.Unmarshal([]byte(raw), &m)
|
||||
}
|
||||
return m
|
||||
}
|
||||
|
||||
func knowledgeMetaToEntry(id, content string, m knowledgeRawMeta) KnowledgeDocEntry {
|
||||
entry := KnowledgeDocEntry{
|
||||
ID: id,
|
||||
DocType: KnowledgeDocType(m.DocType),
|
||||
Content: content,
|
||||
Question: m.Question,
|
||||
UserID: m.UserID,
|
||||
FlowID: m.FlowID,
|
||||
TaskID: m.TaskID,
|
||||
SubtaskID: m.SubtaskID,
|
||||
PartSize: m.PartSize,
|
||||
TotalSize: m.TotalSize,
|
||||
Manual: m.Manual,
|
||||
}
|
||||
if m.Description != "" {
|
||||
d := m.Description
|
||||
entry.Description = &d
|
||||
}
|
||||
if m.GuideType != "" {
|
||||
gt := KnowledgeGuideType(m.GuideType)
|
||||
entry.GuideType = >
|
||||
}
|
||||
if m.AnswerType != "" {
|
||||
at := KnowledgeAnswerType(m.AnswerType)
|
||||
entry.AnswerType = &at
|
||||
}
|
||||
if m.CodeLang != "" {
|
||||
cl := m.CodeLang
|
||||
entry.CodeLang = &cl
|
||||
}
|
||||
return entry
|
||||
}
|
||||
|
||||
// ==================== Enum types =============================================
|
||||
|
||||
// KnowledgeDocType mirrors model.KnowledgeDocType for the REST layer.
|
||||
type KnowledgeDocType string
|
||||
|
||||
const (
|
||||
KnowledgeDocTypeAnswer KnowledgeDocType = "answer"
|
||||
KnowledgeDocTypeGuide KnowledgeDocType = "guide"
|
||||
KnowledgeDocTypeCode KnowledgeDocType = "code"
|
||||
)
|
||||
|
||||
// Valid implements IValid.
|
||||
func (t KnowledgeDocType) Valid() error {
|
||||
switch t {
|
||||
case KnowledgeDocTypeAnswer, KnowledgeDocTypeGuide, KnowledgeDocTypeCode:
|
||||
return nil
|
||||
default:
|
||||
return fmt.Errorf("invalid KnowledgeDocType: %q (expected answer|guide|code)", t)
|
||||
}
|
||||
}
|
||||
|
||||
// KnowledgeGuideType mirrors model.KnowledgeGuideType for the REST layer.
|
||||
type KnowledgeGuideType string
|
||||
|
||||
const (
|
||||
KnowledgeGuideTypeInstall KnowledgeGuideType = "install"
|
||||
KnowledgeGuideTypeConfigure KnowledgeGuideType = "configure"
|
||||
KnowledgeGuideTypeUse KnowledgeGuideType = "use"
|
||||
KnowledgeGuideTypePentest KnowledgeGuideType = "pentest"
|
||||
KnowledgeGuideTypeDevelopment KnowledgeGuideType = "development"
|
||||
KnowledgeGuideTypeOther KnowledgeGuideType = "other"
|
||||
)
|
||||
|
||||
// Valid implements IValid.
|
||||
func (t KnowledgeGuideType) Valid() error {
|
||||
switch t {
|
||||
case KnowledgeGuideTypeInstall,
|
||||
KnowledgeGuideTypeConfigure,
|
||||
KnowledgeGuideTypeUse,
|
||||
KnowledgeGuideTypePentest,
|
||||
KnowledgeGuideTypeDevelopment,
|
||||
KnowledgeGuideTypeOther:
|
||||
return nil
|
||||
default:
|
||||
return fmt.Errorf("invalid KnowledgeGuideType: %q", t)
|
||||
}
|
||||
}
|
||||
|
||||
// KnowledgeAnswerType mirrors model.KnowledgeAnswerType for the REST layer.
|
||||
type KnowledgeAnswerType string
|
||||
|
||||
const (
|
||||
KnowledgeAnswerTypeGuide KnowledgeAnswerType = "guide"
|
||||
KnowledgeAnswerTypeVulnerability KnowledgeAnswerType = "vulnerability"
|
||||
KnowledgeAnswerTypeCode KnowledgeAnswerType = "code"
|
||||
KnowledgeAnswerTypeTool KnowledgeAnswerType = "tool"
|
||||
KnowledgeAnswerTypeOther KnowledgeAnswerType = "other"
|
||||
)
|
||||
|
||||
// Valid implements IValid.
|
||||
func (t KnowledgeAnswerType) Valid() error {
|
||||
switch t {
|
||||
case KnowledgeAnswerTypeGuide,
|
||||
KnowledgeAnswerTypeVulnerability,
|
||||
KnowledgeAnswerTypeCode,
|
||||
KnowledgeAnswerTypeTool,
|
||||
KnowledgeAnswerTypeOther:
|
||||
return nil
|
||||
default:
|
||||
return fmt.Errorf("invalid KnowledgeAnswerType: %q", t)
|
||||
}
|
||||
}
|
||||
|
||||
// ==================== Response models =========================================
|
||||
|
||||
// KnowledgeDocEntry is the REST response representation of a knowledge document.
|
||||
//
|
||||
//nolint:lll
|
||||
type KnowledgeDocEntry struct {
|
||||
ID string `json:"id"`
|
||||
DocType KnowledgeDocType `json:"doc_type"`
|
||||
Content string `json:"content"`
|
||||
Question string `json:"question"`
|
||||
Description *string `json:"description,omitempty"`
|
||||
UserID int64 `json:"user_id"`
|
||||
FlowID *int64 `json:"flow_id,omitempty"`
|
||||
TaskID *int64 `json:"task_id,omitempty"`
|
||||
SubtaskID *int64 `json:"subtask_id,omitempty"`
|
||||
GuideType *KnowledgeGuideType `json:"guide_type,omitempty"`
|
||||
AnswerType *KnowledgeAnswerType `json:"answer_type,omitempty"`
|
||||
CodeLang *string `json:"code_lang,omitempty"`
|
||||
PartSize int `json:"part_size"`
|
||||
TotalSize int `json:"total_size"`
|
||||
Manual bool `json:"manual"`
|
||||
}
|
||||
|
||||
// KnowledgeDocList is the REST list response (Total matches rdb.TableQuery return type).
|
||||
type KnowledgeDocList struct {
|
||||
Items []KnowledgeDocEntry `json:"items"`
|
||||
Total uint64 `json:"total"`
|
||||
}
|
||||
|
||||
// KnowledgeEmbeddingRow is a gorm scan target for langchain_pg_embedding rows.
|
||||
// Only the columns that the knowledge API exposes are selected (no embedding vector).
|
||||
type KnowledgeEmbeddingRow struct {
|
||||
ID string `gorm:"column:id"`
|
||||
Document string `gorm:"column:document"`
|
||||
Cmetadata string `gorm:"column:cmetadata"`
|
||||
}
|
||||
|
||||
// KnowledgeDocWithScore wraps a document with its semantic similarity score.
|
||||
type KnowledgeDocWithScore struct {
|
||||
Score float64 `json:"score"`
|
||||
Document KnowledgeDocEntry `json:"document"`
|
||||
}
|
||||
|
||||
// KnowledgeSearchResult is the REST search response.
|
||||
type KnowledgeSearchResult struct {
|
||||
Items []KnowledgeDocWithScore `json:"items"`
|
||||
Total int `json:"total"`
|
||||
}
|
||||
|
||||
// ==================== Query / request models =================================
|
||||
|
||||
// KnowledgeListQuery holds query-string parameters for the list endpoint.
|
||||
//
|
||||
//nolint:lll
|
||||
type KnowledgeListQuery struct {
|
||||
WithContent bool `form:"with_content" json:"with_content"`
|
||||
DocTypes []KnowledgeDocType `form:"doc_types[]" json:"doc_types,omitempty"`
|
||||
GuideTypes []KnowledgeGuideType `form:"guide_types[]" json:"guide_types,omitempty"`
|
||||
AnswerTypes []KnowledgeAnswerType `form:"answer_types[]" json:"answer_types,omitempty"`
|
||||
CodeLangs []string `form:"code_langs[]" json:"code_langs,omitempty"`
|
||||
FlowID *int64 `form:"flow_id" json:"flow_id,omitempty"`
|
||||
Manual *bool `form:"manual" json:"manual,omitempty"`
|
||||
}
|
||||
|
||||
// Valid validates all enum values supplied in the list query.
|
||||
func (q KnowledgeListQuery) Valid() error {
|
||||
for _, dt := range q.DocTypes {
|
||||
if err := dt.Valid(); err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
for _, gt := range q.GuideTypes {
|
||||
if err := gt.Valid(); err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
for _, at := range q.AnswerTypes {
|
||||
if err := at.Valid(); err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// CreateKnowledgeDocRequest is the POST body for creating a knowledge document.
|
||||
//
|
||||
//nolint:lll
|
||||
type CreateKnowledgeDocRequest struct {
|
||||
DocType KnowledgeDocType `json:"doc_type" validate:"required,valid"`
|
||||
Content string `json:"content" validate:"required,min=1,max=65536"`
|
||||
Question string `json:"question" validate:"required,min=1,max=2048"`
|
||||
Description *string `json:"description,omitempty" validate:"omitempty,max=1000"`
|
||||
GuideType *KnowledgeGuideType `json:"guide_type,omitempty" validate:"omitempty,valid"`
|
||||
AnswerType *KnowledgeAnswerType `json:"answer_type,omitempty" validate:"omitempty,valid"`
|
||||
CodeLang *string `json:"code_lang,omitempty" validate:"omitempty,max=100"`
|
||||
}
|
||||
|
||||
// Valid implements IValid.
|
||||
func (r CreateKnowledgeDocRequest) Valid() error {
|
||||
if err := r.DocType.Valid(); err != nil {
|
||||
return err
|
||||
}
|
||||
if r.GuideType != nil {
|
||||
if err := r.GuideType.Valid(); err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
if r.AnswerType != nil {
|
||||
if err := r.AnswerType.Valid(); err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
return validate.Struct(r)
|
||||
}
|
||||
|
||||
// UpdateKnowledgeDocRequest is the PUT body for updating a knowledge document.
|
||||
// content is mandatory (triggers re-embedding); all other fields are optional overrides.
|
||||
//
|
||||
//nolint:lll
|
||||
type UpdateKnowledgeDocRequest struct {
|
||||
Content string `json:"content" validate:"required,min=1,max=65536"`
|
||||
Question *string `json:"question,omitempty" validate:"omitempty,min=1,max=2048"`
|
||||
Description *string `json:"description,omitempty" validate:"omitempty,max=1000"`
|
||||
GuideType *KnowledgeGuideType `json:"guide_type,omitempty" validate:"omitempty,valid"`
|
||||
AnswerType *KnowledgeAnswerType `json:"answer_type,omitempty" validate:"omitempty,valid"`
|
||||
CodeLang *string `json:"code_lang,omitempty" validate:"omitempty,max=100"`
|
||||
}
|
||||
|
||||
// Valid implements IValid.
|
||||
func (r UpdateKnowledgeDocRequest) Valid() error {
|
||||
if r.GuideType != nil {
|
||||
if err := r.GuideType.Valid(); err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
if r.AnswerType != nil {
|
||||
if err := r.AnswerType.Valid(); err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
return validate.Struct(r)
|
||||
}
|
||||
|
||||
// KnowledgeSearchRequest is the POST body for semantic search.
|
||||
//
|
||||
//nolint:lll
|
||||
type KnowledgeSearchRequest struct {
|
||||
Query string `json:"query" validate:"required,min=1,max=2048"`
|
||||
Limit int `json:"limit,omitempty" validate:"omitempty,min=1,max=100"`
|
||||
DocTypes []KnowledgeDocType `json:"doc_types,omitempty" validate:"omitempty,dive,valid"`
|
||||
GuideTypes []KnowledgeGuideType `json:"guide_types,omitempty" validate:"omitempty,dive,valid"`
|
||||
AnswerTypes []KnowledgeAnswerType `json:"answer_types,omitempty" validate:"omitempty,dive,valid"`
|
||||
CodeLangs []string `json:"code_langs,omitempty" validate:"omitempty,dive,max=100"`
|
||||
FlowID *int64 `json:"flow_id,omitempty"`
|
||||
Manual *bool `json:"manual,omitempty"`
|
||||
}
|
||||
|
||||
// Valid implements IValid.
|
||||
func (r KnowledgeSearchRequest) Valid() error {
|
||||
for _, dt := range r.DocTypes {
|
||||
if err := dt.Valid(); err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
for _, gt := range r.GuideTypes {
|
||||
if err := gt.Valid(); err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
for _, at := range r.AnswerTypes {
|
||||
if err := at.Valid(); err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
return validate.Struct(r)
|
||||
}
|
||||
|
||||
// ==================== Converters (REST ↔ GraphQL model) ======================
|
||||
|
||||
// KnowledgeDocFromGQL converts a GraphQL KnowledgeDocument to a REST KnowledgeDocEntry.
|
||||
func KnowledgeDocFromGQL(doc *gqlmodel.KnowledgeDocument) KnowledgeDocEntry {
|
||||
entry := KnowledgeDocEntry{
|
||||
ID: doc.ID,
|
||||
DocType: KnowledgeDocType(doc.DocType),
|
||||
Content: doc.Content,
|
||||
Question: doc.Question,
|
||||
UserID: doc.UserID,
|
||||
FlowID: doc.FlowID,
|
||||
TaskID: doc.TaskID,
|
||||
SubtaskID: doc.SubtaskID,
|
||||
PartSize: doc.PartSize,
|
||||
TotalSize: doc.TotalSize,
|
||||
Manual: doc.Manual,
|
||||
}
|
||||
if doc.Description != nil {
|
||||
entry.Description = doc.Description
|
||||
}
|
||||
if doc.GuideType != nil {
|
||||
gt := KnowledgeGuideType(*doc.GuideType)
|
||||
entry.GuideType = >
|
||||
}
|
||||
if doc.AnswerType != nil {
|
||||
at := KnowledgeAnswerType(*doc.AnswerType)
|
||||
entry.AnswerType = &at
|
||||
}
|
||||
if doc.CodeLang != nil {
|
||||
entry.CodeLang = doc.CodeLang
|
||||
}
|
||||
return entry
|
||||
}
|
||||
|
||||
// KnowledgeDocListFromGQL converts a slice of GraphQL documents to a REST list response.
|
||||
func KnowledgeDocListFromGQL(docs []*gqlmodel.KnowledgeDocument) KnowledgeDocList {
|
||||
items := make([]KnowledgeDocEntry, 0, len(docs))
|
||||
for _, d := range docs {
|
||||
items = append(items, KnowledgeDocFromGQL(d))
|
||||
}
|
||||
return KnowledgeDocList{Items: items, Total: uint64(len(items))}
|
||||
}
|
||||
|
||||
// KnowledgeDocListFromRows converts raw gorm scan rows to a REST list response.
|
||||
func KnowledgeDocListFromRows(rows []KnowledgeEmbeddingRow, total uint64, withContent bool) KnowledgeDocList {
|
||||
items := make([]KnowledgeDocEntry, 0, len(rows))
|
||||
for _, r := range rows {
|
||||
items = append(items, KnowledgeDocEntryFromRow(r, withContent))
|
||||
}
|
||||
return KnowledgeDocList{Items: items, Total: total}
|
||||
}
|
||||
|
||||
// KnowledgeDocEntryFromRow converts a single gorm scan row to a REST KnowledgeDocEntry.
|
||||
func KnowledgeDocEntryFromRow(r KnowledgeEmbeddingRow, withContent bool) KnowledgeDocEntry {
|
||||
meta := knowledgeMetaFromJSON(r.Cmetadata)
|
||||
content := r.Document
|
||||
if !withContent {
|
||||
content = ""
|
||||
}
|
||||
return knowledgeMetaToEntry(r.ID, content, meta)
|
||||
}
|
||||
|
||||
// KnowledgeSearchResultFromGQL converts GraphQL search results to the REST response.
|
||||
func KnowledgeSearchResultFromGQL(results []*gqlmodel.KnowledgeDocumentWithScore) KnowledgeSearchResult {
|
||||
items := make([]KnowledgeDocWithScore, 0, len(results))
|
||||
for _, r := range results {
|
||||
items = append(items, KnowledgeDocWithScore{
|
||||
Score: r.Score,
|
||||
Document: KnowledgeDocFromGQL(r.Document),
|
||||
})
|
||||
}
|
||||
return KnowledgeSearchResult{Items: items, Total: len(items)}
|
||||
}
|
||||
|
||||
// CreateRequestToGQL converts a REST create request to the GraphQL input type.
|
||||
func (r CreateKnowledgeDocRequest) ToGQL() gqlmodel.CreateKnowledgeDocumentInput {
|
||||
input := gqlmodel.CreateKnowledgeDocumentInput{
|
||||
DocType: gqlmodel.KnowledgeDocType(r.DocType),
|
||||
Content: r.Content,
|
||||
Question: r.Question,
|
||||
Description: r.Description,
|
||||
CodeLang: r.CodeLang,
|
||||
}
|
||||
if r.GuideType != nil {
|
||||
gt := gqlmodel.KnowledgeGuideType(*r.GuideType)
|
||||
input.GuideType = >
|
||||
}
|
||||
if r.AnswerType != nil {
|
||||
at := gqlmodel.KnowledgeAnswerType(*r.AnswerType)
|
||||
input.AnswerType = &at
|
||||
}
|
||||
return input
|
||||
}
|
||||
|
||||
// UpdateRequestToGQL converts a REST update request to the GraphQL input type.
|
||||
func (r UpdateKnowledgeDocRequest) ToGQL() gqlmodel.UpdateKnowledgeDocumentInput {
|
||||
input := gqlmodel.UpdateKnowledgeDocumentInput{
|
||||
Content: r.Content,
|
||||
Question: r.Question,
|
||||
Description: r.Description,
|
||||
CodeLang: r.CodeLang,
|
||||
}
|
||||
if r.GuideType != nil {
|
||||
gt := gqlmodel.KnowledgeGuideType(*r.GuideType)
|
||||
input.GuideType = >
|
||||
}
|
||||
if r.AnswerType != nil {
|
||||
at := gqlmodel.KnowledgeAnswerType(*r.AnswerType)
|
||||
input.AnswerType = &at
|
||||
}
|
||||
return input
|
||||
}
|
||||
|
||||
// ListQueryToGQLFilter converts a REST list query to a GraphQL KnowledgeFilter.
|
||||
func (q KnowledgeListQuery) ToGQLFilter() *gqlmodel.KnowledgeFilter {
|
||||
f := &gqlmodel.KnowledgeFilter{
|
||||
FlowID: q.FlowID,
|
||||
Manual: q.Manual,
|
||||
}
|
||||
for _, dt := range q.DocTypes {
|
||||
f.DocTypes = append(f.DocTypes, gqlmodel.KnowledgeDocType(dt))
|
||||
}
|
||||
for _, gt := range q.GuideTypes {
|
||||
f.GuideTypes = append(f.GuideTypes, gqlmodel.KnowledgeGuideType(gt))
|
||||
}
|
||||
for _, at := range q.AnswerTypes {
|
||||
f.AnswerTypes = append(f.AnswerTypes, gqlmodel.KnowledgeAnswerType(at))
|
||||
}
|
||||
f.CodeLangs = append(f.CodeLangs, q.CodeLangs...)
|
||||
if f.FlowID == nil && f.Manual == nil && len(f.DocTypes) == 0 &&
|
||||
len(f.GuideTypes) == 0 && len(f.AnswerTypes) == 0 && len(f.CodeLangs) == 0 {
|
||||
return nil
|
||||
}
|
||||
return f
|
||||
}
|
||||
|
||||
// SearchRequestToGQLFilter converts a REST search request to a GraphQL KnowledgeFilter.
|
||||
func (r KnowledgeSearchRequest) ToGQLFilter() *gqlmodel.KnowledgeFilter {
|
||||
f := &gqlmodel.KnowledgeFilter{
|
||||
FlowID: r.FlowID,
|
||||
Manual: r.Manual,
|
||||
}
|
||||
for _, dt := range r.DocTypes {
|
||||
f.DocTypes = append(f.DocTypes, gqlmodel.KnowledgeDocType(dt))
|
||||
}
|
||||
for _, gt := range r.GuideTypes {
|
||||
f.GuideTypes = append(f.GuideTypes, gqlmodel.KnowledgeGuideType(gt))
|
||||
}
|
||||
for _, at := range r.AnswerTypes {
|
||||
f.AnswerTypes = append(f.AnswerTypes, gqlmodel.KnowledgeAnswerType(at))
|
||||
}
|
||||
f.CodeLangs = append(f.CodeLangs, r.CodeLangs...)
|
||||
if f.FlowID == nil && f.Manual == nil && len(f.DocTypes) == 0 &&
|
||||
len(f.GuideTypes) == 0 && len(f.AnswerTypes) == 0 && len(f.CodeLangs) == 0 {
|
||||
return nil
|
||||
}
|
||||
return f
|
||||
}
|
||||
@@ -0,0 +1,67 @@
|
||||
package models
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
|
||||
"github.com/jinzhu/gorm"
|
||||
)
|
||||
|
||||
type MsgchainType string
|
||||
|
||||
const (
|
||||
MsgchainTypePrimaryAgent MsgchainType = "primary_agent"
|
||||
MsgchainTypeReporter MsgchainType = "reporter"
|
||||
MsgchainTypeGenerator MsgchainType = "generator"
|
||||
MsgchainTypeRefiner MsgchainType = "refiner"
|
||||
MsgchainTypeReflector MsgchainType = "reflector"
|
||||
MsgchainTypeEnricher MsgchainType = "enricher"
|
||||
MsgchainTypeAdviser MsgchainType = "adviser"
|
||||
MsgchainTypeCoder MsgchainType = "coder"
|
||||
MsgchainTypeMemorist MsgchainType = "memorist"
|
||||
MsgchainTypeSearcher MsgchainType = "searcher"
|
||||
MsgchainTypeInstaller MsgchainType = "installer"
|
||||
MsgchainTypePentester MsgchainType = "pentester"
|
||||
MsgchainTypeSummarizer MsgchainType = "summarizer"
|
||||
MsgchainTypeToolCallFixer MsgchainType = "tool_call_fixer"
|
||||
MsgchainTypeAssistant MsgchainType = "assistant"
|
||||
)
|
||||
|
||||
func (e *MsgchainType) Scan(src interface{}) error {
|
||||
switch s := src.(type) {
|
||||
case []byte:
|
||||
*e = MsgchainType(s)
|
||||
case string:
|
||||
*e = MsgchainType(s)
|
||||
default:
|
||||
return fmt.Errorf("unsupported scan type for MsgchainType: %T", src)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func (s MsgchainType) String() string {
|
||||
return string(s)
|
||||
}
|
||||
|
||||
// Valid is function to control input/output data
|
||||
func (ml MsgchainType) Valid() error {
|
||||
switch ml {
|
||||
case MsgchainTypePrimaryAgent, MsgchainTypeReporter,
|
||||
MsgchainTypeGenerator, MsgchainTypeRefiner,
|
||||
MsgchainTypeReflector, MsgchainTypeEnricher,
|
||||
MsgchainTypeAdviser, MsgchainTypeCoder,
|
||||
MsgchainTypeMemorist, MsgchainTypeSearcher,
|
||||
MsgchainTypeInstaller, MsgchainTypePentester,
|
||||
MsgchainTypeSummarizer, MsgchainTypeToolCallFixer,
|
||||
MsgchainTypeAssistant:
|
||||
return nil
|
||||
default:
|
||||
return fmt.Errorf("invalid MsgchainType: %s", ml)
|
||||
}
|
||||
}
|
||||
|
||||
// Validate is function to use callback to control input/output data
|
||||
func (ml MsgchainType) Validate(db *gorm.DB) {
|
||||
if err := ml.Valid(); err != nil {
|
||||
db.AddError(err)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,111 @@
|
||||
package models
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"time"
|
||||
|
||||
"github.com/jinzhu/gorm"
|
||||
)
|
||||
|
||||
type MsglogType string
|
||||
|
||||
const (
|
||||
MsglogTypeAnswer MsglogType = "answer"
|
||||
MsglogTypeReport MsglogType = "report"
|
||||
MsglogTypeThoughts MsglogType = "thoughts"
|
||||
MsglogTypeBrowser MsglogType = "browser"
|
||||
MsglogTypeTerminal MsglogType = "terminal"
|
||||
MsglogTypeFile MsglogType = "file"
|
||||
MsglogTypeSearch MsglogType = "search"
|
||||
MsglogTypeAdvice MsglogType = "advice"
|
||||
MsglogTypeAsk MsglogType = "ask"
|
||||
MsglogTypeInput MsglogType = "input"
|
||||
MsglogTypeDone MsglogType = "done"
|
||||
)
|
||||
|
||||
func (s MsglogType) String() string {
|
||||
return string(s)
|
||||
}
|
||||
|
||||
// Valid is function to control input/output data
|
||||
func (s MsglogType) Valid() error {
|
||||
switch s {
|
||||
case MsglogTypeAnswer, MsglogTypeReport, MsglogTypeThoughts,
|
||||
MsglogTypeBrowser, MsglogTypeTerminal, MsglogTypeFile,
|
||||
MsglogTypeSearch, MsglogTypeAdvice, MsglogTypeAsk,
|
||||
MsglogTypeInput, MsglogTypeDone:
|
||||
return nil
|
||||
default:
|
||||
return fmt.Errorf("invalid MsglogType: %s", s)
|
||||
}
|
||||
}
|
||||
|
||||
// Validate is function to use callback to control input/output data
|
||||
func (s MsglogType) Validate(db *gorm.DB) {
|
||||
if err := s.Valid(); err != nil {
|
||||
db.AddError(err)
|
||||
}
|
||||
}
|
||||
|
||||
type MsglogResultFormat string
|
||||
|
||||
const (
|
||||
MsglogResultFormatPlain MsglogResultFormat = "plain"
|
||||
MsglogResultFormatMarkdown MsglogResultFormat = "markdown"
|
||||
MsglogResultFormatTerminal MsglogResultFormat = "terminal"
|
||||
)
|
||||
|
||||
func (s MsglogResultFormat) String() string {
|
||||
return string(s)
|
||||
}
|
||||
|
||||
// Valid is function to control input/output data
|
||||
func (s MsglogResultFormat) Valid() error {
|
||||
switch s {
|
||||
case MsglogResultFormatPlain,
|
||||
MsglogResultFormatMarkdown,
|
||||
MsglogResultFormatTerminal:
|
||||
return nil
|
||||
default:
|
||||
return fmt.Errorf("invalid MsglogResultFormat: %s", s)
|
||||
}
|
||||
}
|
||||
|
||||
// Validate is function to use callback to control input/output data
|
||||
func (s MsglogResultFormat) Validate(db *gorm.DB) {
|
||||
if err := s.Valid(); err != nil {
|
||||
db.AddError(err)
|
||||
}
|
||||
}
|
||||
|
||||
// Msglog is model to contain log record information from agents about their actions
|
||||
// nolint:lll
|
||||
type Msglog struct {
|
||||
ID uint64 `form:"id" json:"id" validate:"min=0,numeric" gorm:"type:BIGINT;NOT NULL;PRIMARY_KEY;AUTO_INCREMENT"`
|
||||
Type MsglogType `form:"type" json:"type" validate:"valid,required" gorm:"type:MSGLOG_TYPE;NOT NULL"`
|
||||
Message string `form:"message" json:"message" validate:"required" gorm:"type:TEXT;NOT NULL"`
|
||||
Thinking string `form:"thinking" json:"thinking" validate:"omitempty" gorm:"type:TEXT;NULL"`
|
||||
Result string `form:"result" json:"result" validate:"omitempty" gorm:"type:TEXT;NOT NULL;default:''"`
|
||||
ResultFormat MsglogResultFormat `form:"result_format" json:"result_format" validate:"valid,required" gorm:"type:MSGLOG_RESULT_FORMAT;NOT NULL;default:plain"`
|
||||
FlowID uint64 `form:"flow_id" json:"flow_id" validate:"min=0,numeric,required" gorm:"type:BIGINT;NOT NULL"`
|
||||
TaskID *uint64 `form:"task_id,omitempty" json:"task_id,omitempty" validate:"omitnil,min=0" gorm:"type:BIGINT;NOT NULL"`
|
||||
SubtaskID *uint64 `form:"subtask_id,omitempty" json:"subtask_id,omitempty" validate:"omitnil,min=0" gorm:"type:BIGINT;NOT NULL"`
|
||||
CreatedAt time.Time `form:"created_at,omitempty" json:"created_at,omitempty" validate:"omitempty" gorm:"type:TIMESTAMPTZ;default:CURRENT_TIMESTAMP"`
|
||||
}
|
||||
|
||||
// TableName returns the table name string to guaranty use correct table
|
||||
func (ml *Msglog) TableName() string {
|
||||
return "msglogs"
|
||||
}
|
||||
|
||||
// Valid is function to control input/output data
|
||||
func (ml Msglog) Valid() error {
|
||||
return validate.Struct(ml)
|
||||
}
|
||||
|
||||
// Validate is function to use callback to control input/output data
|
||||
func (ml Msglog) Validate(db *gorm.DB) {
|
||||
if err := ml.Valid(); err != nil {
|
||||
db.AddError(err)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,92 @@
|
||||
package models
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"time"
|
||||
|
||||
"pentagi/pkg/templates"
|
||||
|
||||
"github.com/jinzhu/gorm"
|
||||
)
|
||||
|
||||
// PromptType is an alias for templates.PromptType with validation methods for GORM
|
||||
type PromptType templates.PromptType
|
||||
|
||||
// String returns the string representation of PromptType
|
||||
func (s PromptType) String() string {
|
||||
return string(s)
|
||||
}
|
||||
|
||||
// Valid is function to control input/output data
|
||||
func (s PromptType) Valid() error {
|
||||
// Convert to templates.PromptType and validate against known constants
|
||||
templateType := templates.PromptType(s)
|
||||
switch templateType {
|
||||
case templates.PromptTypePrimaryAgent, templates.PromptTypeAssistant,
|
||||
templates.PromptTypePentester, templates.PromptTypeQuestionPentester,
|
||||
templates.PromptTypeCoder, templates.PromptTypeQuestionCoder,
|
||||
templates.PromptTypeInstaller, templates.PromptTypeQuestionInstaller,
|
||||
templates.PromptTypeSearcher, templates.PromptTypeQuestionSearcher,
|
||||
templates.PromptTypeMemorist, templates.PromptTypeQuestionMemorist,
|
||||
templates.PromptTypeAdviser, templates.PromptTypeQuestionAdviser,
|
||||
templates.PromptTypeGenerator, templates.PromptTypeSubtasksGenerator,
|
||||
templates.PromptTypeRefiner, templates.PromptTypeSubtasksRefiner,
|
||||
templates.PromptTypeReporter, templates.PromptTypeTaskReporter,
|
||||
templates.PromptTypeReflector, templates.PromptTypeQuestionReflector,
|
||||
templates.PromptTypeEnricher, templates.PromptTypeQuestionEnricher,
|
||||
templates.PromptTypeToolCallFixer, templates.PromptTypeInputToolCallFixer,
|
||||
templates.PromptTypeSummarizer, templates.PromptTypeImageChooser,
|
||||
templates.PromptTypeLanguageChooser, templates.PromptTypeFlowDescriptor,
|
||||
templates.PromptTypeTaskDescriptor, templates.PromptTypeExecutionLogs,
|
||||
templates.PromptTypeFullExecutionContext, templates.PromptTypeShortExecutionContext,
|
||||
templates.PromptTypeToolCallIDCollector, templates.PromptTypeToolCallIDDetector:
|
||||
return nil
|
||||
default:
|
||||
return fmt.Errorf("invalid PromptType: %s", s)
|
||||
}
|
||||
}
|
||||
|
||||
// Validate is function to use callback to control input/output data
|
||||
func (s PromptType) Validate(db *gorm.DB) {
|
||||
if err := s.Valid(); err != nil {
|
||||
db.AddError(err)
|
||||
}
|
||||
}
|
||||
|
||||
// Prompt is model to contain prompt information
|
||||
// nolint:lll
|
||||
type Prompt struct {
|
||||
ID uint64 `form:"id" json:"id" validate:"min=0,numeric" gorm:"type:BIGINT;NOT NULL;PRIMARY_KEY;AUTO_INCREMENT"`
|
||||
Type PromptType `form:"type" json:"type" validate:"valid,required" gorm:"type:PROMPT_TYPE;NOT NULL"`
|
||||
UserID uint64 `form:"user_id" json:"user_id" validate:"min=0,numeric" gorm:"type:BIGINT;NOT NULL"`
|
||||
Prompt string `form:"prompt" json:"prompt" validate:"required" gorm:"type:TEXT;NOT NULL"`
|
||||
CreatedAt time.Time `form:"created_at,omitempty" json:"created_at,omitempty" validate:"omitempty" gorm:"type:TIMESTAMPTZ;default:CURRENT_TIMESTAMP"`
|
||||
UpdatedAt time.Time `form:"updated_at,omitempty" json:"updated_at,omitempty" validate:"omitempty" gorm:"type:TIMESTAMPTZ;default:CURRENT_TIMESTAMP"`
|
||||
}
|
||||
|
||||
// TableName returns the table name string to guaranty use correct table
|
||||
func (p *Prompt) TableName() string {
|
||||
return "prompts"
|
||||
}
|
||||
|
||||
// Valid is function to control input/output data
|
||||
func (p Prompt) Valid() error {
|
||||
return validate.Struct(p)
|
||||
}
|
||||
|
||||
// Validate is function to use callback to control input/output data
|
||||
func (p Prompt) Validate(db *gorm.DB) {
|
||||
if err := p.Valid(); err != nil {
|
||||
db.AddError(err)
|
||||
}
|
||||
}
|
||||
|
||||
// PatchPrompt is model to contain prompt patching paylaod
|
||||
type PatchPrompt struct {
|
||||
Prompt string `form:"prompt" json:"prompt" validate:"required"`
|
||||
}
|
||||
|
||||
// Valid is function to control input/output data
|
||||
func (pp PatchPrompt) Valid() error {
|
||||
return validate.Struct(pp)
|
||||
}
|
||||
@@ -0,0 +1,111 @@
|
||||
package models
|
||||
|
||||
import (
|
||||
"testing"
|
||||
|
||||
"pentagi/pkg/templates"
|
||||
|
||||
"github.com/stretchr/testify/assert"
|
||||
)
|
||||
|
||||
func TestPromptTypeValid(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
validTypes := []struct {
|
||||
name string
|
||||
pt PromptType
|
||||
}{
|
||||
{"primary_agent", PromptType(templates.PromptTypePrimaryAgent)},
|
||||
{"assistant", PromptType(templates.PromptTypeAssistant)},
|
||||
{"pentester", PromptType(templates.PromptTypePentester)},
|
||||
{"coder", PromptType(templates.PromptTypeCoder)},
|
||||
{"searcher", PromptType(templates.PromptTypeSearcher)},
|
||||
{"summarizer", PromptType(templates.PromptTypeSummarizer)},
|
||||
{"reporter", PromptType(templates.PromptTypeReporter)},
|
||||
{"tool_call_id_detector", PromptType(templates.PromptTypeToolCallIDDetector)},
|
||||
}
|
||||
|
||||
for _, tt := range validTypes {
|
||||
tt := tt
|
||||
t.Run("valid_"+tt.name, func(t *testing.T) {
|
||||
t.Parallel()
|
||||
assert.NoError(t, tt.pt.Valid())
|
||||
})
|
||||
}
|
||||
|
||||
invalidTypes := []struct {
|
||||
name string
|
||||
pt PromptType
|
||||
}{
|
||||
{"empty", PromptType("")},
|
||||
{"unknown", PromptType("unknown")},
|
||||
{"typo", PromptType("primarry_agent")},
|
||||
}
|
||||
|
||||
for _, tt := range invalidTypes {
|
||||
tt := tt
|
||||
t.Run("invalid_"+tt.name, func(t *testing.T) {
|
||||
t.Parallel()
|
||||
err := tt.pt.Valid()
|
||||
assert.Error(t, err)
|
||||
assert.Contains(t, err.Error(), "invalid PromptType")
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestPromptTypeString(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
assert.Equal(t, "primary_agent", PromptType(templates.PromptTypePrimaryAgent).String())
|
||||
assert.Equal(t, "assistant", PromptType(templates.PromptTypeAssistant).String())
|
||||
}
|
||||
|
||||
func TestPromptValid(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
validPrompt := Prompt{
|
||||
Type: PromptType(templates.PromptTypePrimaryAgent),
|
||||
Prompt: "You are a security assistant.",
|
||||
}
|
||||
|
||||
t.Run("valid prompt", func(t *testing.T) {
|
||||
t.Parallel()
|
||||
assert.NoError(t, validPrompt.Valid())
|
||||
})
|
||||
|
||||
t.Run("invalid type", func(t *testing.T) {
|
||||
t.Parallel()
|
||||
p := validPrompt
|
||||
p.Type = PromptType("invalid")
|
||||
assert.Error(t, p.Valid())
|
||||
})
|
||||
|
||||
t.Run("missing prompt text", func(t *testing.T) {
|
||||
t.Parallel()
|
||||
p := validPrompt
|
||||
p.Prompt = ""
|
||||
assert.Error(t, p.Valid())
|
||||
})
|
||||
}
|
||||
|
||||
func TestPromptTableName(t *testing.T) {
|
||||
t.Parallel()
|
||||
p := &Prompt{}
|
||||
assert.Equal(t, "prompts", p.TableName())
|
||||
}
|
||||
|
||||
func TestPatchPromptValid(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
t.Run("valid patch", func(t *testing.T) {
|
||||
t.Parallel()
|
||||
pp := PatchPrompt{Prompt: "updated prompt text"}
|
||||
assert.NoError(t, pp.Valid())
|
||||
})
|
||||
|
||||
t.Run("missing prompt", func(t *testing.T) {
|
||||
t.Parallel()
|
||||
pp := PatchPrompt{Prompt: ""}
|
||||
assert.Error(t, pp.Valid())
|
||||
})
|
||||
}
|
||||
@@ -0,0 +1,133 @@
|
||||
package models
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"time"
|
||||
|
||||
"pentagi/pkg/providers/provider"
|
||||
|
||||
"github.com/jinzhu/gorm"
|
||||
)
|
||||
|
||||
type ProviderType provider.ProviderType
|
||||
|
||||
func (s ProviderType) String() string {
|
||||
return string(s)
|
||||
}
|
||||
|
||||
// Valid is function to control input/output data
|
||||
func (s ProviderType) Valid() error {
|
||||
providerType := provider.ProviderType(s)
|
||||
switch providerType {
|
||||
case provider.ProviderOpenAI,
|
||||
provider.ProviderAnthropic,
|
||||
provider.ProviderGemini,
|
||||
provider.ProviderBedrock,
|
||||
provider.ProviderOllama,
|
||||
provider.ProviderCustom,
|
||||
provider.ProviderDeepSeek,
|
||||
provider.ProviderGLM,
|
||||
provider.ProviderKimi,
|
||||
provider.ProviderQwen:
|
||||
return nil
|
||||
default:
|
||||
return fmt.Errorf("invalid ProviderType: %s", s)
|
||||
}
|
||||
}
|
||||
|
||||
// Validate is function to use callback to control input/output data
|
||||
func (s ProviderType) Validate(db *gorm.DB) {
|
||||
if err := s.Valid(); err != nil {
|
||||
db.AddError(err)
|
||||
}
|
||||
}
|
||||
|
||||
// Provider is model to contain provider configuration information
|
||||
// nolint:lll
|
||||
type Provider struct {
|
||||
ID uint64 `form:"id" json:"id" validate:"min=0,numeric" gorm:"type:BIGINT;NOT NULL;PRIMARY_KEY;AUTO_INCREMENT"`
|
||||
UserID uint64 `form:"user_id" json:"user_id" validate:"min=0,numeric" gorm:"type:BIGINT;NOT NULL"`
|
||||
Type ProviderType `form:"type" json:"type" validate:"valid,required" gorm:"type:PROVIDER_TYPE;NOT NULL"`
|
||||
Name string `form:"name" json:"name" validate:"required" gorm:"type:TEXT;NOT NULL"`
|
||||
Config json.RawMessage `form:"config" json:"config" validate:"required" gorm:"type:JSON;NOT NULL"`
|
||||
CreatedAt time.Time `form:"created_at,omitempty" json:"created_at,omitempty" validate:"omitempty" gorm:"type:TIMESTAMPTZ;default:CURRENT_TIMESTAMP"`
|
||||
UpdatedAt time.Time `form:"updated_at,omitempty" json:"updated_at,omitempty" validate:"omitempty" gorm:"type:TIMESTAMPTZ;default:CURRENT_TIMESTAMP"`
|
||||
DeletedAt *time.Time `form:"deleted_at,omitempty" json:"deleted_at,omitempty" validate:"omitempty" sql:"index" gorm:"type:TIMESTAMPTZ"`
|
||||
}
|
||||
|
||||
// TableName returns the table name string to guaranty use correct table
|
||||
func (p *Provider) TableName() string {
|
||||
return "providers"
|
||||
}
|
||||
|
||||
// Valid is function to control input/output data
|
||||
func (p Provider) Valid() error {
|
||||
return validate.Struct(p)
|
||||
}
|
||||
|
||||
// Validate is function to use callback to control input/output data
|
||||
func (p Provider) Validate(db *gorm.DB) {
|
||||
if err := p.Valid(); err != nil {
|
||||
db.AddError(err)
|
||||
}
|
||||
}
|
||||
|
||||
// CreateProvider is model to contain provider creation payload
|
||||
// nolint:lll
|
||||
type CreateProvider struct {
|
||||
Config json.RawMessage `form:"config" json:"config" validate:"required" example:"{}"`
|
||||
}
|
||||
|
||||
// Valid is function to control input/output data
|
||||
func (cp CreateProvider) Valid() error {
|
||||
return validate.Struct(cp)
|
||||
}
|
||||
|
||||
// PatchProvider is model to contain provider patching payload
|
||||
// nolint:lll
|
||||
type PatchProvider struct {
|
||||
Name *string `form:"name,omitempty" json:"name,omitempty" validate:"omitempty" example:"updated provider name"`
|
||||
Config *json.RawMessage `form:"config,omitempty" json:"config,omitempty" validate:"omitempty" example:"{}"`
|
||||
}
|
||||
|
||||
// Valid is function to control input/output data
|
||||
func (pp PatchProvider) Valid() error {
|
||||
return validate.Struct(pp)
|
||||
}
|
||||
|
||||
// ModelPriceInfo is model to contain price information for a model
|
||||
// nolint:lll
|
||||
type ModelPriceInfo struct {
|
||||
Input float64 `form:"input" json:"input" validate:"omitempty,numeric,min=0" example:"1.1"`
|
||||
Output float64 `form:"output" json:"output" validate:"omitempty,numeric,min=0" example:"3.0"`
|
||||
CacheRead float64 `form:"cache_read" json:"cache_read" validate:"omitempty,numeric,min=0" example:"0.1"`
|
||||
CacheWrite float64 `form:"cache_write" json:"cache_write" validate:"omitempty,numeric,min=0" example:"0.3"`
|
||||
}
|
||||
|
||||
// Valid is function to control input/output data
|
||||
func (mpi ModelPriceInfo) Valid() error {
|
||||
return validate.Struct(mpi)
|
||||
}
|
||||
|
||||
// ModelInfo is model to contain model short information for display
|
||||
// nolint:lll
|
||||
type ModelInfo struct {
|
||||
Name string `form:"name" json:"name" validate:"required" example:"gpt-4o"`
|
||||
AgentTypes []string `form:"agent_types" json:"agent_types" validate:"required"`
|
||||
PriceInfo *ModelPriceInfo `form:"price_info" json:"price_info" validate:"omitempty,valid"`
|
||||
}
|
||||
|
||||
// ProviderInfo is model to contain provider short information for display
|
||||
// nolint:lll
|
||||
type ProviderInfo struct {
|
||||
Name string `form:"name" json:"name" validate:"required" example:"my openai provider"`
|
||||
Type ProviderType `form:"type" json:"type" validate:"valid,required" example:"openai"`
|
||||
DefaultModel string `form:"default_model" json:"default_model" validate:"required" example:"gpt-4o"`
|
||||
Models []ModelInfo `form:"models" json:"models" validate:"required"`
|
||||
}
|
||||
|
||||
// Valid is function to control input/output data
|
||||
func (p ProviderInfo) Valid() error {
|
||||
return validate.Struct(p)
|
||||
}
|
||||
@@ -0,0 +1,183 @@
|
||||
package models
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"testing"
|
||||
|
||||
"github.com/stretchr/testify/assert"
|
||||
)
|
||||
|
||||
func TestProviderTypeValid(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
validTypes := []struct {
|
||||
name string
|
||||
pt ProviderType
|
||||
}{
|
||||
{"openai", ProviderType("openai")},
|
||||
{"anthropic", ProviderType("anthropic")},
|
||||
{"gemini", ProviderType("gemini")},
|
||||
{"bedrock", ProviderType("bedrock")},
|
||||
{"ollama", ProviderType("ollama")},
|
||||
{"custom", ProviderType("custom")},
|
||||
{"deepseek", ProviderType("deepseek")},
|
||||
{"glm", ProviderType("glm")},
|
||||
{"kimi", ProviderType("kimi")},
|
||||
{"qwen", ProviderType("qwen")},
|
||||
}
|
||||
|
||||
for _, tt := range validTypes {
|
||||
tt := tt
|
||||
t.Run("valid_"+tt.name, func(t *testing.T) {
|
||||
t.Parallel()
|
||||
assert.NoError(t, tt.pt.Valid())
|
||||
})
|
||||
}
|
||||
|
||||
invalidTypes := []struct {
|
||||
name string
|
||||
pt ProviderType
|
||||
}{
|
||||
{"empty", ProviderType("")},
|
||||
{"unknown", ProviderType("unknown")},
|
||||
{"gpt4", ProviderType("gpt4")},
|
||||
{"azure", ProviderType("azure")},
|
||||
}
|
||||
|
||||
for _, tt := range invalidTypes {
|
||||
tt := tt
|
||||
t.Run("invalid_"+tt.name, func(t *testing.T) {
|
||||
t.Parallel()
|
||||
err := tt.pt.Valid()
|
||||
assert.Error(t, err)
|
||||
assert.Contains(t, err.Error(), "invalid ProviderType")
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestProviderTypeString(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
assert.Equal(t, "openai", ProviderType("openai").String())
|
||||
assert.Equal(t, "anthropic", ProviderType("anthropic").String())
|
||||
}
|
||||
|
||||
func TestProviderValid(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
validProvider := Provider{
|
||||
UserID: 1,
|
||||
Type: ProviderType("openai"),
|
||||
Name: "my-openai",
|
||||
Config: json.RawMessage(`{"api_key":"test"}`),
|
||||
}
|
||||
|
||||
t.Run("valid provider", func(t *testing.T) {
|
||||
t.Parallel()
|
||||
assert.NoError(t, validProvider.Valid())
|
||||
})
|
||||
|
||||
t.Run("missing name", func(t *testing.T) {
|
||||
t.Parallel()
|
||||
p := validProvider
|
||||
p.Name = ""
|
||||
assert.Error(t, p.Valid())
|
||||
})
|
||||
|
||||
t.Run("missing config", func(t *testing.T) {
|
||||
t.Parallel()
|
||||
p := validProvider
|
||||
p.Config = nil
|
||||
assert.Error(t, p.Valid())
|
||||
})
|
||||
|
||||
t.Run("invalid type", func(t *testing.T) {
|
||||
t.Parallel()
|
||||
p := validProvider
|
||||
p.Type = ProviderType("invalid")
|
||||
assert.Error(t, p.Valid())
|
||||
})
|
||||
}
|
||||
|
||||
func TestProviderTableName(t *testing.T) {
|
||||
t.Parallel()
|
||||
p := &Provider{}
|
||||
assert.Equal(t, "providers", p.TableName())
|
||||
}
|
||||
|
||||
func TestCreateProviderValid(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
t.Run("valid create provider", func(t *testing.T) {
|
||||
t.Parallel()
|
||||
cp := CreateProvider{Config: json.RawMessage(`{"key":"val"}`)}
|
||||
assert.NoError(t, cp.Valid())
|
||||
})
|
||||
|
||||
t.Run("missing config", func(t *testing.T) {
|
||||
t.Parallel()
|
||||
cp := CreateProvider{Config: nil}
|
||||
assert.Error(t, cp.Valid())
|
||||
})
|
||||
}
|
||||
|
||||
func TestPatchProviderValid(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
t.Run("valid patch with name", func(t *testing.T) {
|
||||
t.Parallel()
|
||||
name := "updated"
|
||||
pp := PatchProvider{Name: &name}
|
||||
assert.NoError(t, pp.Valid())
|
||||
})
|
||||
|
||||
t.Run("valid empty patch", func(t *testing.T) {
|
||||
t.Parallel()
|
||||
pp := PatchProvider{}
|
||||
assert.NoError(t, pp.Valid())
|
||||
})
|
||||
}
|
||||
|
||||
func TestProviderInfoValid(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
validInfo := ProviderInfo{
|
||||
Name: "my-provider",
|
||||
Type: ProviderType("openai"),
|
||||
DefaultModel: "gpt-4o",
|
||||
Models: []ModelInfo{{Name: "gpt-4o", AgentTypes: []string{"primary_agent"}}},
|
||||
}
|
||||
|
||||
t.Run("valid provider info", func(t *testing.T) {
|
||||
t.Parallel()
|
||||
assert.NoError(t, validInfo.Valid())
|
||||
})
|
||||
|
||||
t.Run("missing name", func(t *testing.T) {
|
||||
t.Parallel()
|
||||
pi := validInfo
|
||||
pi.Name = ""
|
||||
assert.Error(t, pi.Valid())
|
||||
})
|
||||
|
||||
t.Run("invalid type", func(t *testing.T) {
|
||||
t.Parallel()
|
||||
pi := validInfo
|
||||
pi.Type = ProviderType("invalid")
|
||||
assert.Error(t, pi.Valid())
|
||||
})
|
||||
|
||||
t.Run("missing default model", func(t *testing.T) {
|
||||
t.Parallel()
|
||||
pi := validInfo
|
||||
pi.DefaultModel = ""
|
||||
assert.Error(t, pi.Valid())
|
||||
})
|
||||
|
||||
t.Run("missing models", func(t *testing.T) {
|
||||
t.Parallel()
|
||||
pi := validInfo
|
||||
pi.Models = nil
|
||||
assert.Error(t, pi.Valid())
|
||||
})
|
||||
}
|
||||
@@ -0,0 +1,91 @@
|
||||
package models
|
||||
|
||||
import "time"
|
||||
|
||||
// UserResource is the GORM model for a user-owned resource (file or directory)
|
||||
// whose metadata lives in PostgreSQL and whose content (for files) is stored as
|
||||
// a .blob file on disk keyed by MD5 hash.
|
||||
//
|
||||
// nolint:lll
|
||||
type UserResource struct {
|
||||
ID uint64 `form:"id" json:"id" gorm:"type:BIGINT;NOT NULL;PRIMARY_KEY;AUTO_INCREMENT"`
|
||||
UserID uint64 `form:"user_id" json:"user_id" gorm:"type:BIGINT;NOT NULL;index"`
|
||||
Hash string `form:"hash" json:"hash" gorm:"type:TEXT;NOT NULL;default:''"`
|
||||
Name string `form:"name" json:"name" gorm:"type:TEXT;NOT NULL"`
|
||||
Path string `form:"path" json:"path" gorm:"type:TEXT;NOT NULL"`
|
||||
Size int64 `form:"size" json:"size" gorm:"type:BIGINT;NOT NULL;default:0"`
|
||||
IsDir bool `form:"is_dir" json:"is_dir" gorm:"type:BOOLEAN;NOT NULL;default:false"`
|
||||
CreatedAt time.Time `form:"created_at" json:"created_at"`
|
||||
UpdatedAt time.Time `form:"updated_at" json:"updated_at"`
|
||||
}
|
||||
|
||||
// ResourceEntry is the REST response representation of a single resource.
|
||||
type ResourceEntry struct {
|
||||
ID uint64 `json:"id"`
|
||||
UserID uint64 `json:"user_id"`
|
||||
Name string `json:"name"`
|
||||
Path string `json:"path"`
|
||||
Size int64 `json:"size"`
|
||||
IsDir bool `json:"is_dir"`
|
||||
CreatedAt time.Time `json:"created_at"`
|
||||
UpdatedAt time.Time `json:"updated_at"`
|
||||
}
|
||||
|
||||
// ResourceList is the REST list response for resource operations.
|
||||
type ResourceList struct {
|
||||
Items []ResourceEntry `json:"items"`
|
||||
Total uint64 `json:"total"`
|
||||
}
|
||||
|
||||
// MkdirResourceRequest is the request body for creating a virtual directory.
|
||||
type MkdirResourceRequest struct {
|
||||
Path string `json:"path" binding:"required"`
|
||||
}
|
||||
|
||||
// MoveResourceRequest is the request body for moving or renaming a resource.
|
||||
//
|
||||
// At least one of Source or Sources must be non-empty after deduplication.
|
||||
//
|
||||
// - Single source: Destination is the exact target path (file) or root
|
||||
// directory (dir tree), identical to the original single-source behaviour.
|
||||
// - Multiple sources: Destination is treated as a base directory; each
|
||||
// source's base name is appended automatically, e.g. moving "a.txt" and
|
||||
// "b.txt" to "archive" produces "archive/a.txt" and "archive/b.txt".
|
||||
// All moves happen atomically inside a single DB transaction.
|
||||
type MoveResourceRequest struct {
|
||||
// Source is kept for backward compatibility (single source).
|
||||
// Combined with Sources when both are provided; duplicates are removed.
|
||||
Source string `json:"source"`
|
||||
// Sources is a list of virtual resource paths to move.
|
||||
// Combined with Source when both are provided; duplicates are removed.
|
||||
Sources []string `json:"sources"`
|
||||
// Destination is the exact target path (single source) or base directory
|
||||
// (multiple sources). Empty string means the root directory.
|
||||
Destination string `json:"destination"`
|
||||
// Force overwrites existing resources at the target paths when true.
|
||||
Force bool `json:"force"`
|
||||
}
|
||||
|
||||
// CopyResourceRequest is the request body for copying a resource.
|
||||
//
|
||||
// At least one of Source or Sources must be non-empty after deduplication.
|
||||
//
|
||||
// - Single source: Destination is the exact target path (file) or root
|
||||
// directory (dir tree), identical to the original single-source behaviour.
|
||||
// - Multiple sources: Destination is treated as a base directory; each
|
||||
// source's base name is appended automatically, e.g. copying "a.txt" and
|
||||
// "b.txt" to "backup" produces "backup/a.txt" and "backup/b.txt".
|
||||
// All copies happen atomically inside a single DB transaction.
|
||||
type CopyResourceRequest struct {
|
||||
// Source is kept for backward compatibility (single source).
|
||||
// Combined with Sources when both are provided; duplicates are removed.
|
||||
Source string `json:"source"`
|
||||
// Sources is a list of virtual resource paths to copy.
|
||||
// Combined with Source when both are provided; duplicates are removed.
|
||||
Sources []string `json:"sources"`
|
||||
// Destination is the exact target path (single source) or base directory
|
||||
// (multiple sources). Required.
|
||||
Destination string `json:"destination" binding:"required"`
|
||||
// Force overwrites existing resources at the target paths when true.
|
||||
Force bool `json:"force"`
|
||||
}
|
||||
@@ -0,0 +1,74 @@
|
||||
package models
|
||||
|
||||
import "github.com/jinzhu/gorm"
|
||||
|
||||
// Role is model to contain user role information
|
||||
// nolint:lll
|
||||
type Role struct {
|
||||
ID uint64 `form:"id" json:"id" validate:"min=0,numeric" gorm:"type:BIGINT;NOT NULL;PRIMARY_KEY;AUTO_INCREMENT"`
|
||||
Name string `form:"name" json:"name" validate:"max=50,required" gorm:"type:TEXT;NOT NULL;UNIQUE_INDEX"`
|
||||
}
|
||||
|
||||
// TableName returns the table name string to guaranty use correct table
|
||||
func (r *Role) TableName() string {
|
||||
return "roles"
|
||||
}
|
||||
|
||||
// Valid is function to control input/output data
|
||||
func (r Role) Valid() error {
|
||||
return validate.Struct(r)
|
||||
}
|
||||
|
||||
// Validate is function to use callback to control input/output data
|
||||
func (r Role) Validate(db *gorm.DB) {
|
||||
if err := r.Valid(); err != nil {
|
||||
db.AddError(err)
|
||||
}
|
||||
}
|
||||
|
||||
// Privilege is model to contain user privileges
|
||||
// nolint:lll
|
||||
type Privilege struct {
|
||||
ID uint64 `form:"id" json:"id" validate:"min=0,numeric" gorm:"type:BIGINT;NOT NULL;PRIMARY_KEY;AUTO_INCREMENT"`
|
||||
RoleID uint64 `form:"role_id" json:"role_id" validate:"min=0,numeric" gorm:"type:BIGINT;NOT NULL"`
|
||||
Name string `form:"name" json:"name" validate:"max=70,required" gorm:"type:TEXT;NOT NULL"`
|
||||
}
|
||||
|
||||
// TableName returns the table name string to guaranty use correct table
|
||||
func (p *Privilege) TableName() string {
|
||||
return "privileges"
|
||||
}
|
||||
|
||||
// Valid is function to control input/output data
|
||||
func (p Privilege) Valid() error {
|
||||
return validate.Struct(p)
|
||||
}
|
||||
|
||||
// RolePrivileges is model to contain user role privileges
|
||||
// nolint:lll
|
||||
type RolePrivileges struct {
|
||||
Privileges []Privilege `form:"privileges" json:"privileges" validate:"required" gorm:"foreignkey:RoleID;association_autoupdate:false;association_autocreate:false"`
|
||||
Role `form:"" json:""`
|
||||
}
|
||||
|
||||
// TableName returns the table name string to guaranty use correct table
|
||||
func (rp *RolePrivileges) TableName() string {
|
||||
return "roles"
|
||||
}
|
||||
|
||||
// Valid is function to control input/output data
|
||||
func (rp RolePrivileges) Valid() error {
|
||||
for i := range rp.Privileges {
|
||||
if err := rp.Privileges[i].Valid(); err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
return rp.Role.Valid()
|
||||
}
|
||||
|
||||
// Validate is function to use callback to control input/output data
|
||||
func (rp RolePrivileges) Validate(db *gorm.DB) {
|
||||
if err := rp.Valid(); err != nil {
|
||||
db.AddError(err)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,36 @@
|
||||
package models
|
||||
|
||||
import (
|
||||
"time"
|
||||
|
||||
"github.com/jinzhu/gorm"
|
||||
)
|
||||
|
||||
// Screenshot is model to contain screenshot information
|
||||
// nolint:lll
|
||||
type Screenshot struct {
|
||||
ID uint64 `form:"id" json:"id" validate:"min=0,numeric" gorm:"type:BIGINT;NOT NULL;PRIMARY_KEY;AUTO_INCREMENT"`
|
||||
Name string `form:"name" json:"name" validate:"required" gorm:"type:TEXT;NOT NULL"`
|
||||
URL string `form:"url" json:"url" validate:"required" gorm:"type:TEXT;NOT NULL"`
|
||||
FlowID uint64 `form:"flow_id" json:"flow_id" validate:"min=0,numeric,required" gorm:"type:BIGINT;NOT NULL"`
|
||||
TaskID *uint64 `form:"task_id,omitempty" json:"task_id,omitempty" validate:"omitnil,min=0" gorm:"type:BIGINT"`
|
||||
SubtaskID *uint64 `form:"subtask_id,omitempty" json:"subtask_id,omitempty" validate:"omitnil,min=0" gorm:"type:BIGINT"`
|
||||
CreatedAt time.Time `form:"created_at,omitempty" json:"created_at,omitempty" validate:"omitempty" gorm:"type:TIMESTAMPTZ;default:CURRENT_TIMESTAMP"`
|
||||
}
|
||||
|
||||
// TableName returns the table name string to guaranty use correct table
|
||||
func (s *Screenshot) TableName() string {
|
||||
return "screenshots"
|
||||
}
|
||||
|
||||
// Valid is function to control input/output data
|
||||
func (s Screenshot) Valid() error {
|
||||
return validate.Struct(s)
|
||||
}
|
||||
|
||||
// Validate is function to use callback to control input/output data
|
||||
func (s Screenshot) Validate(db *gorm.DB) {
|
||||
if err := s.Valid(); err != nil {
|
||||
db.AddError(err)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,79 @@
|
||||
package models
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"time"
|
||||
|
||||
"github.com/jinzhu/gorm"
|
||||
)
|
||||
|
||||
type SearchEngineType string
|
||||
|
||||
const (
|
||||
SearchEngineTypeGoogle SearchEngineType = "google"
|
||||
SearchEngineTypeDuckduckgo SearchEngineType = "duckduckgo"
|
||||
SearchEngineTypeTavily SearchEngineType = "tavily"
|
||||
SearchEngineTypeTraversaal SearchEngineType = "traversaal"
|
||||
SearchEngineTypePerplexity SearchEngineType = "perplexity"
|
||||
SearchEngineTypeBrowser SearchEngineType = "browser"
|
||||
SearchEngineTypeSploitus SearchEngineType = "sploitus"
|
||||
)
|
||||
|
||||
func (s SearchEngineType) String() string {
|
||||
return string(s)
|
||||
}
|
||||
|
||||
// Valid is function to control input/output data
|
||||
func (s SearchEngineType) Valid() error {
|
||||
switch s {
|
||||
case SearchEngineTypeGoogle,
|
||||
SearchEngineTypeDuckduckgo,
|
||||
SearchEngineTypeTavily,
|
||||
SearchEngineTypeTraversaal,
|
||||
SearchEngineTypePerplexity,
|
||||
SearchEngineTypeBrowser,
|
||||
SearchEngineTypeSploitus:
|
||||
return nil
|
||||
default:
|
||||
return fmt.Errorf("invalid SearchEngineType: %s", s)
|
||||
}
|
||||
}
|
||||
|
||||
// Validate is function to use callback to control input/output data
|
||||
func (s SearchEngineType) Validate(db *gorm.DB) {
|
||||
if err := s.Valid(); err != nil {
|
||||
db.AddError(err)
|
||||
}
|
||||
}
|
||||
|
||||
// Searchlog is model to contain search action information in the internet or local network
|
||||
// nolint:lll
|
||||
type Searchlog struct {
|
||||
ID uint64 `form:"id" json:"id" validate:"min=0,numeric" gorm:"type:BIGINT;NOT NULL;PRIMARY_KEY;AUTO_INCREMENT"`
|
||||
Initiator MsgchainType `json:"initiator" validate:"valid,required" gorm:"type:MSGCHAIN_TYPE;NOT NULL"`
|
||||
Executor MsgchainType `json:"executor" validate:"valid,required" gorm:"type:MSGCHAIN_TYPE;NOT NULL"`
|
||||
Engine SearchEngineType `json:"engine" validate:"valid,required" gorm:"type:SEARCHENGINE_TYPE;NOT NULL"`
|
||||
Query string `json:"query" validate:"required" gorm:"type:TEXT;NOT NULL"`
|
||||
Result string `json:"result" validate:"omitempty" gorm:"type:TEXT;NOT NULL;default:''"`
|
||||
FlowID uint64 `form:"flow_id" json:"flow_id" validate:"min=0,numeric,required" gorm:"type:BIGINT;NOT NULL"`
|
||||
TaskID *uint64 `form:"task_id,omitempty" json:"task_id,omitempty" validate:"omitnil,min=0" gorm:"type:BIGINT;NOT NULL"`
|
||||
SubtaskID *uint64 `form:"subtask_id,omitempty" json:"subtask_id,omitempty" validate:"omitnil,min=0" gorm:"type:BIGINT;NOT NULL"`
|
||||
CreatedAt time.Time `form:"created_at,omitempty" json:"created_at,omitempty" validate:"omitempty" gorm:"type:TIMESTAMPTZ;default:CURRENT_TIMESTAMP"`
|
||||
}
|
||||
|
||||
// TableName returns the table name string to guaranty use correct table
|
||||
func (ml *Searchlog) TableName() string {
|
||||
return "searchlogs"
|
||||
}
|
||||
|
||||
// Valid is function to control input/output data
|
||||
func (ml Searchlog) Valid() error {
|
||||
return validate.Struct(ml)
|
||||
}
|
||||
|
||||
// Validate is function to use callback to control input/output data
|
||||
func (ml Searchlog) Validate(db *gorm.DB) {
|
||||
if err := ml.Valid(); err != nil {
|
||||
db.AddError(err)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,17 @@
|
||||
package models
|
||||
|
||||
// Settings is model to contain application settings information
|
||||
// nolint:lll
|
||||
type Settings struct {
|
||||
Debug bool `json:"debug" example:"false"`
|
||||
AskUser bool `json:"ask_user" example:"false"`
|
||||
Version string `json:"version" example:"v1.0.0"`
|
||||
DockerInside bool `json:"docker_inside" example:"false"`
|
||||
IsDevelopMode bool `json:"is_develop_mode" example:"false"`
|
||||
AssistantUseAgents bool `json:"assistant_use_agents" example:"false"`
|
||||
}
|
||||
|
||||
// Valid is function to control input/output data
|
||||
func (s Settings) Valid() error {
|
||||
return validate.Struct(s)
|
||||
}
|
||||
@@ -0,0 +1,74 @@
|
||||
package models
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"time"
|
||||
|
||||
"github.com/jinzhu/gorm"
|
||||
)
|
||||
|
||||
type SubtaskStatus string
|
||||
|
||||
const (
|
||||
SubtaskStatusCreated SubtaskStatus = "created"
|
||||
SubtaskStatusRunning SubtaskStatus = "running"
|
||||
SubtaskStatusWaiting SubtaskStatus = "waiting"
|
||||
SubtaskStatusFinished SubtaskStatus = "finished"
|
||||
SubtaskStatusFailed SubtaskStatus = "failed"
|
||||
)
|
||||
|
||||
func (s SubtaskStatus) String() string {
|
||||
return string(s)
|
||||
}
|
||||
|
||||
// Valid is function to control input/output data
|
||||
func (s SubtaskStatus) Valid() error {
|
||||
switch s {
|
||||
case SubtaskStatusCreated,
|
||||
SubtaskStatusRunning,
|
||||
SubtaskStatusWaiting,
|
||||
SubtaskStatusFinished,
|
||||
SubtaskStatusFailed:
|
||||
return nil
|
||||
default:
|
||||
return fmt.Errorf("invalid SubtaskStatus: %s", s)
|
||||
}
|
||||
}
|
||||
|
||||
// Validate is function to use callback to control input/output data
|
||||
func (s SubtaskStatus) Validate(db *gorm.DB) {
|
||||
if err := s.Valid(); err != nil {
|
||||
db.AddError(err)
|
||||
}
|
||||
}
|
||||
|
||||
// Subtask is model to contain subtask information
|
||||
// nolint:lll
|
||||
type Subtask struct {
|
||||
ID uint64 `form:"id" json:"id" validate:"min=0,numeric" gorm:"type:BIGINT;NOT NULL;PRIMARY_KEY;AUTO_INCREMENT"`
|
||||
Status SubtaskStatus `form:"status" json:"status" validate:"valid,required" gorm:"type:SUBTASK_STATUS;NOT NULL;default:'created'"`
|
||||
Title string `form:"title" json:"title" validate:"required" gorm:"type:TEXT;NOT NULL"`
|
||||
Description string `form:"description" json:"description" validate:"required" gorm:"type:TEXT;NOT NULL"`
|
||||
Context string `form:"context" json:"context" validate:"omitempty" gorm:"type:TEXT;NOT NULL;default:''"`
|
||||
Result string `form:"result" json:"result" validate:"omitempty" gorm:"type:TEXT;NOT NULL;default:''"`
|
||||
TaskID uint64 `form:"task_id" json:"task_id" validate:"min=0,numeric,required" gorm:"type:BIGINT;NOT NULL"`
|
||||
CreatedAt time.Time `form:"created_at,omitempty" json:"created_at,omitempty" validate:"omitempty" gorm:"type:TIMESTAMPTZ;default:CURRENT_TIMESTAMP"`
|
||||
UpdatedAt time.Time `form:"updated_at,omitempty" json:"updated_at,omitempty" validate:"omitempty" gorm:"type:TIMESTAMPTZ;default:CURRENT_TIMESTAMP"`
|
||||
}
|
||||
|
||||
// TableName returns the table name string to guaranty use correct table
|
||||
func (s *Subtask) TableName() string {
|
||||
return "subtasks"
|
||||
}
|
||||
|
||||
// Valid is function to control input/output data
|
||||
func (s Subtask) Valid() error {
|
||||
return validate.Struct(s)
|
||||
}
|
||||
|
||||
// Validate is function to use callback to control input/output data
|
||||
func (s Subtask) Validate(db *gorm.DB) {
|
||||
if err := s.Valid(); err != nil {
|
||||
db.AddError(err)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,102 @@
|
||||
package models
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"time"
|
||||
|
||||
"github.com/jinzhu/gorm"
|
||||
)
|
||||
|
||||
type TaskStatus string
|
||||
|
||||
const (
|
||||
TaskStatusCreated TaskStatus = "created"
|
||||
TaskStatusRunning TaskStatus = "running"
|
||||
TaskStatusWaiting TaskStatus = "waiting"
|
||||
TaskStatusFinished TaskStatus = "finished"
|
||||
TaskStatusFailed TaskStatus = "failed"
|
||||
)
|
||||
|
||||
func (s TaskStatus) String() string {
|
||||
return string(s)
|
||||
}
|
||||
|
||||
// Valid is function to control input/output data
|
||||
func (s TaskStatus) Valid() error {
|
||||
switch s {
|
||||
case TaskStatusCreated,
|
||||
TaskStatusRunning,
|
||||
TaskStatusWaiting,
|
||||
TaskStatusFinished,
|
||||
TaskStatusFailed:
|
||||
return nil
|
||||
default:
|
||||
return fmt.Errorf("invalid TaskStatus: %s", s)
|
||||
}
|
||||
}
|
||||
|
||||
// Validate is function to use callback to control input/output data
|
||||
func (s TaskStatus) Validate(db *gorm.DB) {
|
||||
if err := s.Valid(); err != nil {
|
||||
db.AddError(err)
|
||||
}
|
||||
}
|
||||
|
||||
// Task is model to contain task information
|
||||
// nolint:lll
|
||||
type Task struct {
|
||||
ID uint64 `form:"id" json:"id" validate:"min=0,numeric" gorm:"type:BIGINT;NOT NULL;PRIMARY_KEY;AUTO_INCREMENT"`
|
||||
Status TaskStatus `form:"status" json:"status" validate:"valid,required" gorm:"type:TASK_STATUS;NOT NULL;default:'created'"`
|
||||
Title string `form:"title" json:"title" validate:"required" gorm:"type:TEXT;NOT NULL;default:'untitled'"`
|
||||
Input string `form:"input" json:"input" validate:"required" gorm:"type:TEXT;NOT NULL"`
|
||||
Result string `form:"result" json:"result" validate:"omitempty" gorm:"type:TEXT;NOT NULL;default:''"`
|
||||
FlowID uint64 `form:"flow_id" json:"flow_id" validate:"min=0,numeric,required" gorm:"type:BIGINT;NOT NULL"`
|
||||
CreatedAt time.Time `form:"created_at,omitempty" json:"created_at,omitempty" validate:"omitempty" gorm:"type:TIMESTAMPTZ;default:CURRENT_TIMESTAMP"`
|
||||
UpdatedAt time.Time `form:"updated_at,omitempty" json:"updated_at,omitempty" validate:"omitempty" gorm:"type:TIMESTAMPTZ;default:CURRENT_TIMESTAMP"`
|
||||
}
|
||||
|
||||
// TableName returns the table name string to guaranty use correct table
|
||||
func (t *Task) TableName() string {
|
||||
return "tasks"
|
||||
}
|
||||
|
||||
// Valid is function to control input/output data
|
||||
func (t Task) Valid() error {
|
||||
return validate.Struct(t)
|
||||
}
|
||||
|
||||
// Validate is function to use callback to control input/output data
|
||||
func (t Task) Validate(db *gorm.DB) {
|
||||
if err := t.Valid(); err != nil {
|
||||
db.AddError(err)
|
||||
}
|
||||
}
|
||||
|
||||
// TaskSubtasks is model to contain task and linked subtasks information
|
||||
// nolint:lll
|
||||
type TaskSubtasks struct {
|
||||
Subtasks []Subtask `form:"subtasks" json:"subtasks" validate:"required" gorm:"foreignkey:TaskID;association_autoupdate:false;association_autocreate:false"`
|
||||
Task `form:"" json:""`
|
||||
}
|
||||
|
||||
// TableName returns the table name string to guaranty use correct table
|
||||
func (ts *TaskSubtasks) TableName() string {
|
||||
return "tasks"
|
||||
}
|
||||
|
||||
// Valid is function to control input/output data
|
||||
func (ts TaskSubtasks) Valid() error {
|
||||
for i := range ts.Subtasks {
|
||||
if err := ts.Subtasks[i].Valid(); err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
return ts.Task.Valid()
|
||||
}
|
||||
|
||||
// Validate is function to use callback to control input/output data
|
||||
func (ts TaskSubtasks) Validate(db *gorm.DB) {
|
||||
if err := ts.Valid(); err != nil {
|
||||
db.AddError(err)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,67 @@
|
||||
package models
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"time"
|
||||
|
||||
"github.com/jinzhu/gorm"
|
||||
)
|
||||
|
||||
type TermlogType string
|
||||
|
||||
const (
|
||||
TermlogTypeStdin TermlogType = "stdin"
|
||||
TermlogTypeStdout TermlogType = "stdout"
|
||||
TermlogTypeStderr TermlogType = "stderr"
|
||||
)
|
||||
|
||||
func (s TermlogType) String() string {
|
||||
return string(s)
|
||||
}
|
||||
|
||||
// Valid is function to control input/output data
|
||||
func (s TermlogType) Valid() error {
|
||||
switch s {
|
||||
case TermlogTypeStdin, TermlogTypeStdout, TermlogTypeStderr:
|
||||
return nil
|
||||
default:
|
||||
return fmt.Errorf("invalid TermlogType: %s", s)
|
||||
}
|
||||
}
|
||||
|
||||
// Validate is function to use callback to control input/output data
|
||||
func (s TermlogType) Validate(db *gorm.DB) {
|
||||
if err := s.Valid(); err != nil {
|
||||
db.AddError(err)
|
||||
}
|
||||
}
|
||||
|
||||
// Termlog is model to contain termlog information
|
||||
// nolint:lll
|
||||
type Termlog struct {
|
||||
ID uint64 `form:"id" json:"id" validate:"min=0,numeric" gorm:"type:BIGINT;NOT NULL;PRIMARY_KEY;AUTO_INCREMENT"`
|
||||
Type TermlogType `form:"type" json:"type" validate:"valid,required" gorm:"type:TERMLOG_TYPE;NOT NULL"`
|
||||
Text string `form:"text" json:"text" validate:"required" gorm:"type:TEXT;NOT NULL"`
|
||||
ContainerID uint64 `form:"container_id" json:"container_id" validate:"min=0,numeric,required" gorm:"type:BIGINT;NOT NULL"`
|
||||
FlowID uint64 `form:"flow_id" json:"flow_id" validate:"min=0,numeric,required" gorm:"type:BIGINT;NOT NULL"`
|
||||
TaskID *uint64 `form:"task_id,omitempty" json:"task_id,omitempty" validate:"omitnil,min=0" gorm:"type:BIGINT"`
|
||||
SubtaskID *uint64 `form:"subtask_id,omitempty" json:"subtask_id,omitempty" validate:"omitnil,min=0" gorm:"type:BIGINT"`
|
||||
CreatedAt time.Time `form:"created_at,omitempty" json:"created_at,omitempty" validate:"omitempty" gorm:"type:TIMESTAMPTZ;default:CURRENT_TIMESTAMP"`
|
||||
}
|
||||
|
||||
// TableName returns the table name string to guaranty use correct table
|
||||
func (tl *Termlog) TableName() string {
|
||||
return "termlogs"
|
||||
}
|
||||
|
||||
// Valid is function to control input/output data
|
||||
func (tl Termlog) Valid() error {
|
||||
return validate.Struct(tl)
|
||||
}
|
||||
|
||||
// Validate is function to use callback to control input/output data
|
||||
func (tl Termlog) Validate(db *gorm.DB) {
|
||||
if err := tl.Valid(); err != nil {
|
||||
db.AddError(err)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,75 @@
|
||||
package models
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"time"
|
||||
|
||||
"github.com/jinzhu/gorm"
|
||||
)
|
||||
|
||||
type ToolcallStatus string
|
||||
|
||||
const (
|
||||
ToolcallStatusReceived ToolcallStatus = "received"
|
||||
ToolcallStatusRunning ToolcallStatus = "running"
|
||||
ToolcallStatusFinished ToolcallStatus = "finished"
|
||||
ToolcallStatusFailed ToolcallStatus = "failed"
|
||||
)
|
||||
|
||||
func (s ToolcallStatus) String() string {
|
||||
return string(s)
|
||||
}
|
||||
|
||||
// Valid is function to control input/output data
|
||||
func (s ToolcallStatus) Valid() error {
|
||||
switch s {
|
||||
case ToolcallStatusReceived,
|
||||
ToolcallStatusRunning,
|
||||
ToolcallStatusFinished,
|
||||
ToolcallStatusFailed:
|
||||
return nil
|
||||
default:
|
||||
return fmt.Errorf("invalid ToolcallStatus: %s", s)
|
||||
}
|
||||
}
|
||||
|
||||
// Validate is function to use callback to control input/output data
|
||||
func (s ToolcallStatus) Validate(db *gorm.DB) {
|
||||
if err := s.Valid(); err != nil {
|
||||
db.AddError(err)
|
||||
}
|
||||
}
|
||||
|
||||
// Toolcall is model to contain tool call information
|
||||
// nolint:lll
|
||||
type Toolcall struct {
|
||||
ID uint64 `form:"id" json:"id" validate:"min=0,numeric" gorm:"type:BIGINT;NOT NULL;PRIMARY_KEY;AUTO_INCREMENT"`
|
||||
CallID string `form:"call_id" json:"call_id" validate:"required" gorm:"type:TEXT;NOT NULL"`
|
||||
Status ToolcallStatus `form:"status" json:"status" validate:"valid,required" gorm:"type:TOOLCALL_STATUS;NOT NULL;default:'received'"`
|
||||
Name string `form:"name" json:"name" validate:"required" gorm:"type:TEXT;NOT NULL"`
|
||||
Args string `form:"args" json:"args" validate:"required" gorm:"type:JSON;NOT NULL"`
|
||||
Result string `form:"result" json:"result" validate:"omitempty" gorm:"type:TEXT;NOT NULL;default:''"`
|
||||
DurationSeconds float64 `form:"duration_seconds" json:"duration_seconds" validate:"min=0" gorm:"type:DOUBLE PRECISION;NOT NULL;default:0"`
|
||||
FlowID uint64 `form:"flow_id" json:"flow_id" validate:"min=0,numeric,required" gorm:"type:BIGINT;NOT NULL"`
|
||||
TaskID *uint64 `form:"task_id,omitempty" json:"task_id,omitempty" validate:"omitnil,min=0" gorm:"type:BIGINT"`
|
||||
SubtaskID *uint64 `form:"subtask_id,omitempty" json:"subtask_id,omitempty" validate:"omitnil,min=0" gorm:"type:BIGINT"`
|
||||
CreatedAt time.Time `form:"created_at,omitempty" json:"created_at,omitempty" validate:"omitempty" gorm:"type:TIMESTAMPTZ;default:CURRENT_TIMESTAMP"`
|
||||
UpdatedAt time.Time `form:"updated_at,omitempty" json:"updated_at,omitempty" validate:"omitempty" gorm:"type:TIMESTAMPTZ;default:CURRENT_TIMESTAMP"`
|
||||
}
|
||||
|
||||
// TableName returns the table name string to guaranty use correct table
|
||||
func (tc *Toolcall) TableName() string {
|
||||
return "toolcalls"
|
||||
}
|
||||
|
||||
// Valid is function to control input/output data
|
||||
func (tc Toolcall) Valid() error {
|
||||
return validate.Struct(tc)
|
||||
}
|
||||
|
||||
// Validate is function to use callback to control input/output data
|
||||
func (tc Toolcall) Validate(db *gorm.DB) {
|
||||
if err := tc.Valid(); err != nil {
|
||||
db.AddError(err)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,304 @@
|
||||
package models
|
||||
|
||||
import (
|
||||
"database/sql/driver"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"time"
|
||||
|
||||
"github.com/jinzhu/gorm"
|
||||
)
|
||||
|
||||
const RoleUser = 2
|
||||
|
||||
type UserStatus string
|
||||
|
||||
const (
|
||||
UserStatusCreated UserStatus = "created"
|
||||
UserStatusActive UserStatus = "active"
|
||||
UserStatusBlocked UserStatus = "blocked"
|
||||
)
|
||||
|
||||
func (s UserStatus) String() string {
|
||||
return string(s)
|
||||
}
|
||||
|
||||
// Valid is function to control input/output data
|
||||
func (s UserStatus) Valid() error {
|
||||
switch s {
|
||||
case UserStatusCreated, UserStatusActive, UserStatusBlocked:
|
||||
return nil
|
||||
default:
|
||||
return fmt.Errorf("invalid UserStatus: %s", s)
|
||||
}
|
||||
}
|
||||
|
||||
// Validate is function to use callback to control input/output data
|
||||
func (s UserStatus) Validate(db *gorm.DB) {
|
||||
if err := s.Valid(); err != nil {
|
||||
db.AddError(err)
|
||||
}
|
||||
}
|
||||
|
||||
type UserType string
|
||||
|
||||
const (
|
||||
UserTypeLocal UserType = "local"
|
||||
UserTypeOAuth UserType = "oauth"
|
||||
UserTypeAPI UserType = "api"
|
||||
)
|
||||
|
||||
func (s UserType) String() string {
|
||||
return string(s)
|
||||
}
|
||||
|
||||
// Valid is function to control input/output data
|
||||
func (s UserType) Valid() error {
|
||||
switch s {
|
||||
case UserTypeLocal, UserTypeOAuth, UserTypeAPI:
|
||||
return nil
|
||||
default:
|
||||
return fmt.Errorf("invalid UserType: %s", s)
|
||||
}
|
||||
}
|
||||
|
||||
// Validate is function to use callback to control input/output data
|
||||
func (s UserType) Validate(db *gorm.DB) {
|
||||
if err := s.Valid(); err != nil {
|
||||
db.AddError(err)
|
||||
}
|
||||
}
|
||||
|
||||
// User is model to contain user information
|
||||
// nolint:lll
|
||||
type User struct {
|
||||
ID uint64 `form:"id" json:"id" validate:"min=0,numeric" gorm:"type:BIGINT;NOT NULL;PRIMARY_KEY;AUTO_INCREMENT"`
|
||||
Hash string `form:"hash" json:"hash" validate:"len=32,hexadecimal,lowercase,omitempty" gorm:"type:TEXT;NOT NULL;UNIQUE_INDEX;default:MD5(RANDOM()::text)"`
|
||||
Type UserType `form:"type" json:"type" validate:"valid,required" gorm:"type:USER_TYPE;NOT NULL;default:'local'"`
|
||||
Mail string `form:"mail" json:"mail" validate:"max=50,vmail,required" gorm:"type:TEXT;NOT NULL;UNIQUE_INDEX"`
|
||||
Name string `form:"name" json:"name" validate:"max=70,omitempty" gorm:"type:TEXT;NOT NULL;default:''"`
|
||||
Status UserStatus `form:"status" json:"status" validate:"valid,required" gorm:"type:USER_STATUS;NOT NULL;default:'created'"`
|
||||
RoleID uint64 `form:"role_id" json:"role_id" validate:"min=0,numeric,required" gorm:"type:BIGINT;NOT NULL;default:2"`
|
||||
PasswordChangeRequired bool `form:"password_change_required" json:"password_change_required" gorm:"type:BOOL;NOT NULL;default:false"`
|
||||
Provider *string `form:"provider,omitempty" json:"provider,omitempty" validate:"omitempty" gorm:"type:TEXT"`
|
||||
CreatedAt time.Time `form:"created_at" json:"created_at" validate:"omitempty" gorm:"type:TIMESTAMPTZ;NOT NULL;default:CURRENT_TIMESTAMP"`
|
||||
}
|
||||
|
||||
// TableName returns the table name string to guaranty use correct table
|
||||
func (u *User) TableName() string {
|
||||
return "users"
|
||||
}
|
||||
|
||||
// Valid is function to control input/output data
|
||||
func (u User) Valid() error {
|
||||
return validate.Struct(u)
|
||||
}
|
||||
|
||||
// Validate is function to use callback to control input/output data
|
||||
func (u User) Validate(db *gorm.DB) {
|
||||
if err := u.Valid(); err != nil {
|
||||
db.AddError(err)
|
||||
}
|
||||
}
|
||||
|
||||
// UserPassword is model to contain user information
|
||||
type UserPassword struct {
|
||||
Password string `form:"password" json:"password" validate:"max=100,required" gorm:"column:password;type:TEXT"`
|
||||
User `form:"" json:""`
|
||||
}
|
||||
|
||||
// TableName returns the table name string to guaranty use correct table
|
||||
func (up *UserPassword) TableName() string {
|
||||
return "users"
|
||||
}
|
||||
|
||||
// Valid is function to control input/output data
|
||||
func (up UserPassword) Valid() error {
|
||||
if err := up.User.Valid(); err != nil {
|
||||
return err
|
||||
}
|
||||
return validate.Struct(up)
|
||||
}
|
||||
|
||||
// Validate is function to use callback to control input/output data
|
||||
func (up UserPassword) Validate(db *gorm.DB) {
|
||||
if err := up.Valid(); err != nil {
|
||||
db.AddError(err)
|
||||
}
|
||||
}
|
||||
|
||||
// Login is model to contain user information on Login procedure
|
||||
// nolint:lll
|
||||
type Login struct {
|
||||
Mail string `form:"mail" json:"mail" validate:"max=50,required" gorm:"type:TEXT;NOT NULL;UNIQUE_INDEX"`
|
||||
Password string `form:"password" json:"password" validate:"min=4,max=100,required" gorm:"type:TEXT"`
|
||||
}
|
||||
|
||||
// TableName returns the table name string to guaranty use correct table
|
||||
func (sin *Login) TableName() string {
|
||||
return "users"
|
||||
}
|
||||
|
||||
// Valid is function to control input/output data
|
||||
func (sin Login) Valid() error {
|
||||
return validate.Struct(sin)
|
||||
}
|
||||
|
||||
// Validate is function to use callback to control input/output data
|
||||
func (sin Login) Validate(db *gorm.DB) {
|
||||
if err := sin.Valid(); err != nil {
|
||||
db.AddError(err)
|
||||
}
|
||||
}
|
||||
|
||||
// AuthCallback is model to contain auth data information from external OAuth application
|
||||
type AuthCallback struct {
|
||||
Code string `form:"code" json:"code" validate:"required"`
|
||||
IdToken string `form:"id_token" json:"id_token" validate:"required,jwt"`
|
||||
Scope string `form:"scope" json:"scope" validate:"required,oauth_min_scope"`
|
||||
State string `form:"state" json:"state" validate:"required"`
|
||||
}
|
||||
|
||||
// Valid is function to control input/output data
|
||||
func (au AuthCallback) Valid() error {
|
||||
return validate.Struct(au)
|
||||
}
|
||||
|
||||
// Password is model to contain user password to change it
|
||||
// nolint:lll
|
||||
type Password struct {
|
||||
CurrentPassword string `form:"current_password" json:"current_password" validate:"nefield=Password,min=5,max=100,required" gorm:"-"`
|
||||
Password string `form:"password" json:"password" validate:"stpass,max=100,required" gorm:"type:TEXT"`
|
||||
ConfirmPassword string `form:"confirm_password" json:"confirm_password" validate:"eqfield=Password" gorm:"-"`
|
||||
}
|
||||
|
||||
// TableName returns the table name string to guaranty use correct table
|
||||
func (p *Password) TableName() string {
|
||||
return "users"
|
||||
}
|
||||
|
||||
// Valid is function to control input/output data
|
||||
func (p Password) Valid() error {
|
||||
return validate.Struct(p)
|
||||
}
|
||||
|
||||
// Validate is function to use callback to control input/output data
|
||||
func (p Password) Validate(db *gorm.DB) {
|
||||
if err := p.Valid(); err != nil {
|
||||
db.AddError(err)
|
||||
}
|
||||
}
|
||||
|
||||
// UserRole is model to contain user information linked with user role
|
||||
// nolint:lll
|
||||
type UserRole struct {
|
||||
Role Role `form:"role,omitempty" json:"role,omitempty" gorm:"association_autoupdate:false;association_autocreate:false"`
|
||||
User `form:"" json:""`
|
||||
}
|
||||
|
||||
// Valid is function to control input/output data
|
||||
func (ur UserRole) Valid() error {
|
||||
if err := ur.Role.Valid(); err != nil {
|
||||
return err
|
||||
}
|
||||
return ur.User.Valid()
|
||||
}
|
||||
|
||||
// Validate is function to use callback to control input/output data
|
||||
func (ur UserRole) Validate(db *gorm.DB) {
|
||||
if err := ur.Valid(); err != nil {
|
||||
db.AddError(err)
|
||||
}
|
||||
}
|
||||
|
||||
// UserRole is model to contain user information linked with user role
|
||||
// nolint:lll
|
||||
type UserRolePrivileges struct {
|
||||
Role RolePrivileges `form:"role,omitempty" json:"role,omitempty" gorm:"association_autoupdate:false;association_autocreate:false"`
|
||||
User `form:"" json:""`
|
||||
}
|
||||
|
||||
// Valid is function to control input/output data
|
||||
func (urp UserRolePrivileges) Valid() error {
|
||||
if err := urp.Role.Valid(); err != nil {
|
||||
return err
|
||||
}
|
||||
return urp.User.Valid()
|
||||
}
|
||||
|
||||
// Validate is function to use callback to control input/output data
|
||||
func (urp UserRolePrivileges) Validate(db *gorm.DB) {
|
||||
if err := urp.Valid(); err != nil {
|
||||
db.AddError(err)
|
||||
}
|
||||
}
|
||||
|
||||
// UserPreferencesOptions is model to contain user preferences as JSON
|
||||
type UserPreferencesOptions struct {
|
||||
FavoriteFlows []int64 `json:"favoriteFlows"`
|
||||
}
|
||||
|
||||
// Value implements driver.Valuer interface for database write
|
||||
func (upo UserPreferencesOptions) Value() (driver.Value, error) {
|
||||
return json.Marshal(upo)
|
||||
}
|
||||
|
||||
// Scan implements sql.Scanner interface for database read
|
||||
func (upo *UserPreferencesOptions) Scan(value any) error {
|
||||
if value == nil {
|
||||
*upo = UserPreferencesOptions{FavoriteFlows: []int64{}}
|
||||
return nil
|
||||
}
|
||||
|
||||
bytes, ok := value.([]byte)
|
||||
if !ok {
|
||||
return fmt.Errorf("failed to scan UserPreferencesOptions: expected []byte, got %T", value)
|
||||
}
|
||||
|
||||
return json.Unmarshal(bytes, upo)
|
||||
}
|
||||
|
||||
// UserPreferences is model to contain user preferences information
|
||||
type UserPreferences struct {
|
||||
ID uint64 `json:"id" gorm:"type:BIGINT;NOT NULL;PRIMARY_KEY;AUTO_INCREMENT"`
|
||||
UserID uint64 `json:"user_id" gorm:"type:BIGINT;NOT NULL;UNIQUE_INDEX"`
|
||||
Preferences UserPreferencesOptions `json:"preferences" gorm:"type:JSONB;NOT NULL"`
|
||||
CreatedAt time.Time `json:"created_at" gorm:"type:TIMESTAMPTZ;NOT NULL;default:CURRENT_TIMESTAMP"`
|
||||
UpdatedAt time.Time `json:"updated_at" gorm:"type:TIMESTAMPTZ;NOT NULL;default:CURRENT_TIMESTAMP"`
|
||||
}
|
||||
|
||||
// TableName returns the table name string to guaranty use correct table
|
||||
func (up *UserPreferences) TableName() string {
|
||||
return "user_preferences"
|
||||
}
|
||||
|
||||
// Valid is function to control input/output data
|
||||
func (up UserPreferences) Valid() error {
|
||||
if up.UserID == 0 {
|
||||
return fmt.Errorf("user_id is required")
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// Validate is function to use callback to control input/output data
|
||||
func (up UserPreferences) Validate(db *gorm.DB) {
|
||||
if err := up.Valid(); err != nil {
|
||||
db.AddError(err)
|
||||
}
|
||||
}
|
||||
|
||||
// NewUserPreferences creates a new UserPreferences with default values
|
||||
func NewUserPreferences(userID uint64) *UserPreferences {
|
||||
return &UserPreferences{
|
||||
UserID: userID,
|
||||
Preferences: UserPreferencesOptions{
|
||||
FavoriteFlows: []int64{},
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
// UserWithPreferences is model to combine User and UserPreferences for transactional creation
|
||||
type UserWithPreferences struct {
|
||||
User User
|
||||
Preferences UserPreferences
|
||||
}
|
||||
@@ -0,0 +1,562 @@
|
||||
package models
|
||||
|
||||
import (
|
||||
"testing"
|
||||
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
func TestUserStatusValid(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
tests := []struct {
|
||||
name string
|
||||
status UserStatus
|
||||
wantErr bool
|
||||
}{
|
||||
{"valid created", UserStatusCreated, false},
|
||||
{"valid active", UserStatusActive, false},
|
||||
{"valid blocked", UserStatusBlocked, false},
|
||||
{"invalid empty", UserStatus(""), true},
|
||||
{"invalid unknown", UserStatus("unknown"), true},
|
||||
{"invalid suspended", UserStatus("suspended"), true},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
tt := tt
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
t.Parallel()
|
||||
err := tt.status.Valid()
|
||||
if tt.wantErr {
|
||||
assert.Error(t, err)
|
||||
assert.Contains(t, err.Error(), "invalid UserStatus")
|
||||
} else {
|
||||
assert.NoError(t, err)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestUserStatusString(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
assert.Equal(t, "created", UserStatusCreated.String())
|
||||
assert.Equal(t, "active", UserStatusActive.String())
|
||||
assert.Equal(t, "blocked", UserStatusBlocked.String())
|
||||
}
|
||||
|
||||
func TestUserTypeValid(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
tests := []struct {
|
||||
name string
|
||||
userType UserType
|
||||
wantErr bool
|
||||
}{
|
||||
{"valid local", UserTypeLocal, false},
|
||||
{"valid oauth", UserTypeOAuth, false},
|
||||
{"valid api", UserTypeAPI, false},
|
||||
{"invalid empty", UserType(""), true},
|
||||
{"invalid unknown", UserType("unknown"), true},
|
||||
{"invalid saml", UserType("saml"), true},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
tt := tt
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
t.Parallel()
|
||||
err := tt.userType.Valid()
|
||||
if tt.wantErr {
|
||||
assert.Error(t, err)
|
||||
assert.Contains(t, err.Error(), "invalid UserType")
|
||||
} else {
|
||||
assert.NoError(t, err)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestUserTypeString(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
assert.Equal(t, "local", UserTypeLocal.String())
|
||||
assert.Equal(t, "oauth", UserTypeOAuth.String())
|
||||
assert.Equal(t, "api", UserTypeAPI.String())
|
||||
}
|
||||
|
||||
func TestLoginValid(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
tests := []struct {
|
||||
name string
|
||||
login Login
|
||||
wantErr bool
|
||||
}{
|
||||
{
|
||||
name: "valid login",
|
||||
login: Login{Mail: "test@example.com", Password: "password123"},
|
||||
wantErr: false,
|
||||
},
|
||||
{
|
||||
name: "valid admin mail",
|
||||
login: Login{Mail: "admin", Password: "password123"},
|
||||
wantErr: false,
|
||||
},
|
||||
{
|
||||
name: "empty mail",
|
||||
login: Login{Mail: "", Password: "password123"},
|
||||
wantErr: true,
|
||||
},
|
||||
{
|
||||
name: "empty password",
|
||||
login: Login{Mail: "test@example.com", Password: ""},
|
||||
wantErr: true,
|
||||
},
|
||||
{
|
||||
name: "password too short",
|
||||
login: Login{Mail: "test@example.com", Password: "ab"},
|
||||
wantErr: true,
|
||||
},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
tt := tt
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
t.Parallel()
|
||||
err := tt.login.Valid()
|
||||
if tt.wantErr {
|
||||
assert.Error(t, err)
|
||||
} else {
|
||||
assert.NoError(t, err)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestPasswordValid(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
tests := []struct {
|
||||
name string
|
||||
pw Password
|
||||
wantErr bool
|
||||
}{
|
||||
{
|
||||
name: "valid strong password with special chars",
|
||||
pw: Password{
|
||||
CurrentPassword: "OldPass1!abc",
|
||||
Password: "NewPass1!abc",
|
||||
ConfirmPassword: "NewPass1!abc",
|
||||
},
|
||||
wantErr: false,
|
||||
},
|
||||
{
|
||||
name: "valid long password over 15 chars",
|
||||
pw: Password{
|
||||
CurrentPassword: "oldpasswordvalue",
|
||||
Password: "newpasswordvalue1",
|
||||
ConfirmPassword: "newpasswordvalue1",
|
||||
},
|
||||
wantErr: false,
|
||||
},
|
||||
{
|
||||
name: "confirm password mismatch",
|
||||
pw: Password{
|
||||
CurrentPassword: "OldPass1!abc",
|
||||
Password: "NewPass1!abc",
|
||||
ConfirmPassword: "DifferentPass1!",
|
||||
},
|
||||
wantErr: true,
|
||||
},
|
||||
{
|
||||
name: "current equals new password",
|
||||
pw: Password{
|
||||
CurrentPassword: "SamePass1!abc",
|
||||
Password: "SamePass1!abc",
|
||||
ConfirmPassword: "SamePass1!abc",
|
||||
},
|
||||
wantErr: true,
|
||||
},
|
||||
{
|
||||
name: "weak password no special chars",
|
||||
pw: Password{
|
||||
CurrentPassword: "OldPass1!abc",
|
||||
Password: "newpass1",
|
||||
ConfirmPassword: "newpass1",
|
||||
},
|
||||
wantErr: true,
|
||||
},
|
||||
{
|
||||
name: "empty current password",
|
||||
pw: Password{
|
||||
CurrentPassword: "",
|
||||
Password: "NewPass1!abc",
|
||||
ConfirmPassword: "NewPass1!abc",
|
||||
},
|
||||
wantErr: true,
|
||||
},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
tt := tt
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
t.Parallel()
|
||||
err := tt.pw.Valid()
|
||||
if tt.wantErr {
|
||||
assert.Error(t, err)
|
||||
} else {
|
||||
assert.NoError(t, err)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestUserValid(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
validUser := User{
|
||||
ID: 1,
|
||||
Hash: "abcdef1234567890abcdef1234567890",
|
||||
Type: UserTypeLocal,
|
||||
Mail: "test@example.com",
|
||||
Status: UserStatusActive,
|
||||
RoleID: RoleUser,
|
||||
}
|
||||
|
||||
t.Run("valid user", func(t *testing.T) {
|
||||
t.Parallel()
|
||||
assert.NoError(t, validUser.Valid())
|
||||
})
|
||||
|
||||
t.Run("missing mail", func(t *testing.T) {
|
||||
t.Parallel()
|
||||
u := validUser
|
||||
u.Mail = ""
|
||||
assert.Error(t, u.Valid())
|
||||
})
|
||||
|
||||
t.Run("invalid user type", func(t *testing.T) {
|
||||
t.Parallel()
|
||||
u := validUser
|
||||
u.Type = UserType("invalid")
|
||||
assert.Error(t, u.Valid())
|
||||
})
|
||||
|
||||
t.Run("invalid user status", func(t *testing.T) {
|
||||
t.Parallel()
|
||||
u := validUser
|
||||
u.Status = UserStatus("invalid")
|
||||
assert.Error(t, u.Valid())
|
||||
})
|
||||
|
||||
t.Run("invalid hash length", func(t *testing.T) {
|
||||
t.Parallel()
|
||||
u := validUser
|
||||
u.Hash = "tooshort"
|
||||
assert.Error(t, u.Valid())
|
||||
})
|
||||
}
|
||||
|
||||
func TestUserTableName(t *testing.T) {
|
||||
t.Parallel()
|
||||
u := &User{}
|
||||
assert.Equal(t, "users", u.TableName())
|
||||
}
|
||||
|
||||
func TestUserPasswordValid(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
t.Run("valid user password", func(t *testing.T) {
|
||||
t.Parallel()
|
||||
up := UserPassword{
|
||||
Password: "somepassword",
|
||||
User: User{
|
||||
ID: 1,
|
||||
Hash: "abcdef1234567890abcdef1234567890",
|
||||
Type: UserTypeLocal,
|
||||
Mail: "test@example.com",
|
||||
Status: UserStatusActive,
|
||||
RoleID: RoleUser,
|
||||
},
|
||||
}
|
||||
assert.NoError(t, up.Valid())
|
||||
})
|
||||
|
||||
t.Run("empty password", func(t *testing.T) {
|
||||
t.Parallel()
|
||||
up := UserPassword{
|
||||
Password: "",
|
||||
User: User{
|
||||
ID: 1,
|
||||
Hash: "abcdef1234567890abcdef1234567890",
|
||||
Type: UserTypeLocal,
|
||||
Mail: "test@example.com",
|
||||
Status: UserStatusActive,
|
||||
RoleID: RoleUser,
|
||||
},
|
||||
}
|
||||
assert.Error(t, up.Valid())
|
||||
})
|
||||
|
||||
t.Run("invalid user in user password", func(t *testing.T) {
|
||||
t.Parallel()
|
||||
up := UserPassword{
|
||||
Password: "somepassword",
|
||||
User: User{
|
||||
Mail: "",
|
||||
},
|
||||
}
|
||||
assert.Error(t, up.Valid())
|
||||
})
|
||||
}
|
||||
|
||||
func TestUserPasswordTableName(t *testing.T) {
|
||||
t.Parallel()
|
||||
up := &UserPassword{}
|
||||
assert.Equal(t, "users", up.TableName())
|
||||
}
|
||||
|
||||
func TestLoginTableName(t *testing.T) {
|
||||
t.Parallel()
|
||||
l := &Login{}
|
||||
assert.Equal(t, "users", l.TableName())
|
||||
}
|
||||
|
||||
func TestPasswordTableName(t *testing.T) {
|
||||
t.Parallel()
|
||||
p := &Password{}
|
||||
assert.Equal(t, "users", p.TableName())
|
||||
}
|
||||
|
||||
func TestUserPreferencesOptionsValueScan(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
t.Run("value and scan round trip", func(t *testing.T) {
|
||||
t.Parallel()
|
||||
original := UserPreferencesOptions{FavoriteFlows: []int64{1, 2, 3}}
|
||||
val, err := original.Value()
|
||||
require.NoError(t, err)
|
||||
|
||||
var scanned UserPreferencesOptions
|
||||
switch v := val.(type) {
|
||||
case string:
|
||||
err = scanned.Scan([]byte(v))
|
||||
case []byte:
|
||||
err = scanned.Scan(v)
|
||||
default:
|
||||
t.Fatalf("unexpected Value() type: %T", val)
|
||||
}
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, original.FavoriteFlows, scanned.FavoriteFlows)
|
||||
})
|
||||
|
||||
t.Run("scan nil value", func(t *testing.T) {
|
||||
t.Parallel()
|
||||
var upo UserPreferencesOptions
|
||||
err := upo.Scan(nil)
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, []int64{}, upo.FavoriteFlows)
|
||||
})
|
||||
|
||||
t.Run("scan unsupported type", func(t *testing.T) {
|
||||
t.Parallel()
|
||||
var upo UserPreferencesOptions
|
||||
err := upo.Scan(12345)
|
||||
assert.Error(t, err)
|
||||
assert.Contains(t, err.Error(), "expected []byte")
|
||||
})
|
||||
|
||||
t.Run("scan invalid json", func(t *testing.T) {
|
||||
t.Parallel()
|
||||
var upo UserPreferencesOptions
|
||||
err := upo.Scan([]byte("not json"))
|
||||
assert.Error(t, err)
|
||||
})
|
||||
|
||||
t.Run("value with empty flows", func(t *testing.T) {
|
||||
t.Parallel()
|
||||
upo := UserPreferencesOptions{FavoriteFlows: []int64{}}
|
||||
val, err := upo.Value()
|
||||
require.NoError(t, err)
|
||||
var valStr string
|
||||
switch v := val.(type) {
|
||||
case string:
|
||||
valStr = v
|
||||
case []byte:
|
||||
valStr = string(v)
|
||||
}
|
||||
assert.Contains(t, valStr, "favoriteFlows")
|
||||
})
|
||||
}
|
||||
|
||||
func TestUserPreferencesValid(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
t.Run("valid preferences", func(t *testing.T) {
|
||||
t.Parallel()
|
||||
up := UserPreferences{UserID: 1}
|
||||
assert.NoError(t, up.Valid())
|
||||
})
|
||||
|
||||
t.Run("zero user id", func(t *testing.T) {
|
||||
t.Parallel()
|
||||
up := UserPreferences{UserID: 0}
|
||||
err := up.Valid()
|
||||
assert.Error(t, err)
|
||||
assert.Contains(t, err.Error(), "user_id")
|
||||
})
|
||||
}
|
||||
|
||||
func TestUserPreferencesTableName(t *testing.T) {
|
||||
t.Parallel()
|
||||
up := &UserPreferences{}
|
||||
assert.Equal(t, "user_preferences", up.TableName())
|
||||
}
|
||||
|
||||
func TestNewUserPreferences(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
up := NewUserPreferences(42)
|
||||
assert.Equal(t, uint64(42), up.UserID)
|
||||
assert.NotNil(t, up.Preferences.FavoriteFlows)
|
||||
assert.Empty(t, up.Preferences.FavoriteFlows)
|
||||
}
|
||||
|
||||
func TestAuthCallbackValid(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
// JWT token with 3 dot-separated base64 segments (header.payload.signature)
|
||||
validJWT := "eyJhbGciOiJSUzI1NiIsInR5cCI6IkpXVCJ9.eyJzdWIiOiIxMjM0NTY3ODkwIn0.POstGetfAytaZS82wHcjoTyoqhMyxXiWdR7Nn7A29DNSl0EiXLdwJ6xC6AfgZWF1bOsS_TuYI3OG85AmiExREkrS6tDfTQ2B3WXlrr-wp5AokiRbz3_oB4OxG-W9KcEEbDRcZc0nH3L7LzYptiy1PtAylQGxHTWZXtGz4ht0bAecBgmpdgXMguEIcoqPJ1n3pIWk_dUZegpqx0Lka21H6XxUTxiy8OcaarA8zdnPUnV6AmNP3ecFawIFYdvJB_cm-GvpCSbr8G8y_Mllj8f4x9nBH8pQux89_6gUY618iYv7tuPWBFfEbLxtF2pZS6YC1aSfLQxaOoaBSTNRg"
|
||||
|
||||
t.Run("valid callback", func(t *testing.T) {
|
||||
t.Parallel()
|
||||
ac := AuthCallback{
|
||||
Code: "auth-code-123",
|
||||
IdToken: validJWT,
|
||||
Scope: "openid email profile",
|
||||
State: "random-state-value",
|
||||
}
|
||||
assert.NoError(t, ac.Valid())
|
||||
})
|
||||
|
||||
t.Run("missing code", func(t *testing.T) {
|
||||
t.Parallel()
|
||||
ac := AuthCallback{
|
||||
Code: "",
|
||||
IdToken: validJWT,
|
||||
Scope: "openid email",
|
||||
State: "state123",
|
||||
}
|
||||
assert.Error(t, ac.Valid())
|
||||
})
|
||||
|
||||
t.Run("scope missing openid", func(t *testing.T) {
|
||||
t.Parallel()
|
||||
ac := AuthCallback{
|
||||
Code: "code123",
|
||||
IdToken: validJWT,
|
||||
Scope: "email profile",
|
||||
State: "state123",
|
||||
}
|
||||
assert.Error(t, ac.Valid())
|
||||
})
|
||||
|
||||
t.Run("invalid id token not jwt", func(t *testing.T) {
|
||||
t.Parallel()
|
||||
ac := AuthCallback{
|
||||
Code: "code123",
|
||||
IdToken: "not-a-jwt",
|
||||
Scope: "openid email",
|
||||
State: "state123",
|
||||
}
|
||||
assert.Error(t, ac.Valid())
|
||||
})
|
||||
}
|
||||
|
||||
func TestUserRoleValid(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
validUserForRole := User{
|
||||
Hash: "abcdef1234567890abcdef1234567890",
|
||||
Type: UserTypeLocal,
|
||||
Mail: "test@example.com",
|
||||
Status: UserStatusActive,
|
||||
RoleID: RoleUser,
|
||||
}
|
||||
|
||||
t.Run("valid user role", func(t *testing.T) {
|
||||
t.Parallel()
|
||||
ur := UserRole{
|
||||
Role: Role{ID: 1, Name: "admin"},
|
||||
User: validUserForRole,
|
||||
}
|
||||
assert.NoError(t, ur.Valid())
|
||||
})
|
||||
|
||||
t.Run("invalid role", func(t *testing.T) {
|
||||
t.Parallel()
|
||||
ur := UserRole{
|
||||
Role: Role{Name: ""},
|
||||
User: validUserForRole,
|
||||
}
|
||||
assert.Error(t, ur.Valid())
|
||||
})
|
||||
|
||||
t.Run("invalid user", func(t *testing.T) {
|
||||
t.Parallel()
|
||||
ur := UserRole{
|
||||
Role: Role{ID: 1, Name: "admin"},
|
||||
User: User{Mail: ""},
|
||||
}
|
||||
assert.Error(t, ur.Valid())
|
||||
})
|
||||
}
|
||||
|
||||
func TestUserRolePrivilegesValid(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
validUserForRole := User{
|
||||
Hash: "abcdef1234567890abcdef1234567890",
|
||||
Type: UserTypeLocal,
|
||||
Mail: "test@example.com",
|
||||
Status: UserStatusActive,
|
||||
RoleID: RoleUser,
|
||||
}
|
||||
|
||||
t.Run("valid user role privileges", func(t *testing.T) {
|
||||
t.Parallel()
|
||||
urp := UserRolePrivileges{
|
||||
Role: RolePrivileges{
|
||||
Privileges: []Privilege{{Name: "read"}},
|
||||
Role: Role{ID: 1, Name: "admin"},
|
||||
},
|
||||
User: validUserForRole,
|
||||
}
|
||||
assert.NoError(t, urp.Valid())
|
||||
})
|
||||
|
||||
t.Run("invalid role privileges", func(t *testing.T) {
|
||||
t.Parallel()
|
||||
urp := UserRolePrivileges{
|
||||
Role: RolePrivileges{
|
||||
Privileges: []Privilege{{Name: ""}},
|
||||
Role: Role{ID: 1, Name: "admin"},
|
||||
},
|
||||
User: validUserForRole,
|
||||
}
|
||||
assert.Error(t, urp.Valid())
|
||||
})
|
||||
|
||||
t.Run("invalid user", func(t *testing.T) {
|
||||
t.Parallel()
|
||||
urp := UserRolePrivileges{
|
||||
Role: RolePrivileges{
|
||||
Privileges: []Privilege{{Name: "read"}},
|
||||
Role: Role{ID: 1, Name: "admin"},
|
||||
},
|
||||
User: User{Mail: ""},
|
||||
}
|
||||
assert.Error(t, urp.Valid())
|
||||
})
|
||||
}
|
||||
@@ -0,0 +1,69 @@
|
||||
package models
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"time"
|
||||
|
||||
"github.com/jinzhu/gorm"
|
||||
)
|
||||
|
||||
type VecstoreActionType string
|
||||
|
||||
const (
|
||||
VecstoreActionTypeRetrieve VecstoreActionType = "retrieve"
|
||||
VecstoreActionTypeStore VecstoreActionType = "store"
|
||||
)
|
||||
|
||||
func (s VecstoreActionType) String() string {
|
||||
return string(s)
|
||||
}
|
||||
|
||||
// Valid is function to control input/output data
|
||||
func (s VecstoreActionType) Valid() error {
|
||||
switch s {
|
||||
case VecstoreActionTypeRetrieve, VecstoreActionTypeStore:
|
||||
return nil
|
||||
default:
|
||||
return fmt.Errorf("invalid VecstoreActionType: %s", s)
|
||||
}
|
||||
}
|
||||
|
||||
// Validate is function to use callback to control input/output data
|
||||
func (s VecstoreActionType) Validate(db *gorm.DB) {
|
||||
if err := s.Valid(); err != nil {
|
||||
db.AddError(err)
|
||||
}
|
||||
}
|
||||
|
||||
// Vecstorelog is model to contain vecstore action information
|
||||
// nolint:lll
|
||||
type Vecstorelog struct {
|
||||
ID uint64 `form:"id" json:"id" validate:"min=0,numeric" gorm:"type:BIGINT;NOT NULL;PRIMARY_KEY;AUTO_INCREMENT"`
|
||||
Initiator MsgchainType `json:"initiator" validate:"valid,required" gorm:"type:MSGCHAIN_TYPE;NOT NULL"`
|
||||
Executor MsgchainType `json:"executor" validate:"valid,required" gorm:"type:MSGCHAIN_TYPE;NOT NULL"`
|
||||
Filter string `json:"filter" validate:"required" gorm:"type:JSON;NOT NULL"`
|
||||
Query string `json:"query" validate:"required" gorm:"type:TEXT;NOT NULL"`
|
||||
Action VecstoreActionType `json:"action" validate:"valid,required" gorm:"type:VECSTORE_ACTION_TYPE;NOT NULL"`
|
||||
Result string `json:"result" validate:"omitempty" gorm:"type:TEXT;NOT NULL"`
|
||||
FlowID uint64 `form:"flow_id" json:"flow_id" validate:"min=0,numeric,required" gorm:"type:BIGINT;NOT NULL"`
|
||||
TaskID *uint64 `form:"task_id,omitempty" json:"task_id,omitempty" validate:"omitnil,min=0" gorm:"type:BIGINT;NOT NULL"`
|
||||
SubtaskID *uint64 `form:"subtask_id,omitempty" json:"subtask_id,omitempty" validate:"omitnil,min=0" gorm:"type:BIGINT;NOT NULL"`
|
||||
CreatedAt time.Time `form:"created_at,omitempty" json:"created_at,omitempty" validate:"omitempty" gorm:"type:TIMESTAMPTZ;default:CURRENT_TIMESTAMP"`
|
||||
}
|
||||
|
||||
// TableName returns the table name string to guaranty use correct table
|
||||
func (ml *Vecstorelog) TableName() string {
|
||||
return "vecstorelogs"
|
||||
}
|
||||
|
||||
// Valid is function to control input/output data
|
||||
func (ml Vecstorelog) Valid() error {
|
||||
return validate.Struct(ml)
|
||||
}
|
||||
|
||||
// Validate is function to use callback to control input/output data
|
||||
func (ml Vecstorelog) Validate(db *gorm.DB) {
|
||||
if err := ml.Valid(); err != nil {
|
||||
db.AddError(err)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,64 @@
|
||||
package oauth
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
|
||||
"golang.org/x/oauth2"
|
||||
)
|
||||
|
||||
type OAuthEmailResolver func(ctx context.Context, nonce string, token *oauth2.Token) (string, error)
|
||||
|
||||
type OAuthClient interface {
|
||||
ProviderName() string
|
||||
ResolveEmail(ctx context.Context, nonce string, token *oauth2.Token) (string, error)
|
||||
TokenSource(ctx context.Context, token *oauth2.Token) oauth2.TokenSource
|
||||
Exchange(ctx context.Context, code string, opts ...oauth2.AuthCodeOption) (*oauth2.Token, error)
|
||||
RefreshToken(ctx context.Context, token string) (*oauth2.Token, error)
|
||||
AuthCodeURL(state string, opts ...oauth2.AuthCodeOption) string
|
||||
}
|
||||
|
||||
type oauthClient struct {
|
||||
name string
|
||||
verifier string
|
||||
conf *oauth2.Config
|
||||
emailResolver OAuthEmailResolver
|
||||
}
|
||||
|
||||
func NewOAuthClient(name string, conf *oauth2.Config, emailResolver OAuthEmailResolver) OAuthClient {
|
||||
return &oauthClient{
|
||||
name: name,
|
||||
verifier: oauth2.GenerateVerifier(),
|
||||
conf: conf,
|
||||
emailResolver: emailResolver,
|
||||
}
|
||||
}
|
||||
|
||||
func (o *oauthClient) ProviderName() string {
|
||||
return o.name
|
||||
}
|
||||
|
||||
func (o *oauthClient) ResolveEmail(ctx context.Context, nonce string, token *oauth2.Token) (string, error) {
|
||||
if o.emailResolver == nil {
|
||||
return "", fmt.Errorf("email resolver is not set")
|
||||
}
|
||||
return o.emailResolver(ctx, nonce, token)
|
||||
}
|
||||
|
||||
func (o *oauthClient) TokenSource(ctx context.Context, token *oauth2.Token) oauth2.TokenSource {
|
||||
return o.conf.TokenSource(ctx, token)
|
||||
}
|
||||
|
||||
func (o *oauthClient) Exchange(ctx context.Context, code string, opts ...oauth2.AuthCodeOption) (*oauth2.Token, error) {
|
||||
opts = append(opts, oauth2.VerifierOption(o.verifier))
|
||||
return o.conf.Exchange(ctx, code, opts...)
|
||||
}
|
||||
|
||||
func (o *oauthClient) RefreshToken(ctx context.Context, token string) (*oauth2.Token, error) {
|
||||
return o.conf.TokenSource(ctx, &oauth2.Token{RefreshToken: token}).Token()
|
||||
}
|
||||
|
||||
func (o *oauthClient) AuthCodeURL(state string, opts ...oauth2.AuthCodeOption) string {
|
||||
opts = append(opts, oauth2.S256ChallengeOption(o.verifier))
|
||||
return o.conf.AuthCodeURL(state, opts...)
|
||||
}
|
||||
@@ -0,0 +1,70 @@
|
||||
package oauth
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"io"
|
||||
"net/http"
|
||||
|
||||
"golang.org/x/oauth2"
|
||||
"golang.org/x/oauth2/github"
|
||||
)
|
||||
|
||||
type githubEmail struct {
|
||||
Email string `json:"email"`
|
||||
Primary bool `json:"primary"`
|
||||
Verified bool `json:"verified"`
|
||||
Visibility string `json:"visibility"`
|
||||
}
|
||||
|
||||
func githubEmailResolver(ctx context.Context, nonce string, token *oauth2.Token) (string, error) {
|
||||
req, err := http.NewRequestWithContext(ctx, http.MethodGet, "https://api.github.com/user/emails", nil)
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
req.Header.Set("Authorization", fmt.Sprintf("token %s", token.AccessToken))
|
||||
|
||||
resp, err := http.DefaultClient.Do(req)
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
defer resp.Body.Close()
|
||||
|
||||
body, err := io.ReadAll(resp.Body)
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
|
||||
emails := []githubEmail{}
|
||||
if err := json.Unmarshal(body, &emails); err != nil {
|
||||
return "", err
|
||||
}
|
||||
|
||||
for _, email := range emails {
|
||||
if email.Verified && email.Primary {
|
||||
return email.Email, nil
|
||||
}
|
||||
}
|
||||
|
||||
for _, email := range emails {
|
||||
if email.Verified {
|
||||
return email.Email, nil
|
||||
}
|
||||
}
|
||||
|
||||
return "", fmt.Errorf("no verified primary email found")
|
||||
}
|
||||
|
||||
func NewGithubOAuthClient(clientID, clientSecret, redirectURL string) OAuthClient {
|
||||
return NewOAuthClient("github", &oauth2.Config{
|
||||
ClientID: clientID,
|
||||
ClientSecret: clientSecret,
|
||||
RedirectURL: redirectURL,
|
||||
Scopes: []string{
|
||||
"user:email",
|
||||
"openid",
|
||||
},
|
||||
Endpoint: github.Endpoint,
|
||||
}, githubEmailResolver)
|
||||
}
|
||||
@@ -0,0 +1,76 @@
|
||||
package oauth
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
|
||||
"github.com/coreos/go-oidc/v3/oidc"
|
||||
"golang.org/x/oauth2"
|
||||
"golang.org/x/oauth2/google"
|
||||
)
|
||||
|
||||
type googleTokenClaims struct {
|
||||
Nonce string `json:"nonce"`
|
||||
Email string `json:"email"`
|
||||
EmailVerified bool `json:"email_verified"`
|
||||
}
|
||||
|
||||
func newGoogleEmailResolver(clientID string) OAuthEmailResolver {
|
||||
return func(ctx context.Context, nonce string, token *oauth2.Token) (string, error) {
|
||||
provider, err := oidc.NewProvider(ctx, "https://accounts.google.com")
|
||||
if err != nil {
|
||||
return "", fmt.Errorf("could not create Google OpenID client: %w", err)
|
||||
}
|
||||
|
||||
oidToken, ok := token.Extra("id_token").(string)
|
||||
if !ok {
|
||||
return "", fmt.Errorf("id_token is not present in the token")
|
||||
}
|
||||
|
||||
verifier := provider.Verifier(&oidc.Config{ClientID: clientID})
|
||||
idToken, err := verifier.Verify(ctx, oidToken)
|
||||
if err != nil {
|
||||
return "", fmt.Errorf("could not verify Google ID Token: %w", err)
|
||||
}
|
||||
|
||||
if idToken.Nonce != nonce {
|
||||
return "", fmt.Errorf("nonce mismatch in Google ID Token")
|
||||
}
|
||||
|
||||
if err = idToken.VerifyAccessToken(token.AccessToken); err != nil {
|
||||
return "", fmt.Errorf("failed to verify Google Access Token: %w", err)
|
||||
}
|
||||
|
||||
claims := googleTokenClaims{}
|
||||
if err := idToken.Claims(&claims); err != nil {
|
||||
return "", fmt.Errorf("failed to parse Google ID Token claims: %w", err)
|
||||
}
|
||||
|
||||
if claims.Nonce != nonce {
|
||||
return "", fmt.Errorf("nonce mismatch in Google ID Token claims")
|
||||
}
|
||||
|
||||
if !claims.EmailVerified {
|
||||
return "", fmt.Errorf("email not verified in Google ID Token claims")
|
||||
}
|
||||
|
||||
if claims.Email == "" {
|
||||
return "", fmt.Errorf("email is empty in Google ID Token claims")
|
||||
}
|
||||
|
||||
return claims.Email, nil
|
||||
}
|
||||
}
|
||||
|
||||
func NewGoogleOAuthClient(clientID, clientSecret, redirectURL string) OAuthClient {
|
||||
return NewOAuthClient("google", &oauth2.Config{
|
||||
ClientID: clientID,
|
||||
ClientSecret: clientSecret,
|
||||
RedirectURL: redirectURL,
|
||||
Scopes: []string{
|
||||
"https://www.googleapis.com/auth/userinfo.email",
|
||||
"openid",
|
||||
},
|
||||
Endpoint: google.Endpoint,
|
||||
}, newGoogleEmailResolver(clientID))
|
||||
}
|
||||
@@ -0,0 +1,370 @@
|
||||
package rdb
|
||||
|
||||
import (
|
||||
"crypto/md5" //nolint:gosec
|
||||
"encoding/hex"
|
||||
"errors"
|
||||
"slices"
|
||||
"strconv"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"github.com/google/uuid"
|
||||
"github.com/jinzhu/gorm"
|
||||
"golang.org/x/crypto/bcrypt"
|
||||
)
|
||||
|
||||
// TableFilter is auxiliary struct to contain method of filtering
|
||||
//
|
||||
//nolint:lll
|
||||
type TableFilter struct {
|
||||
Value any `form:"value" json:"value,omitempty" binding:"required" swaggertype:"object"`
|
||||
Field string `form:"field" json:"field,omitempty" binding:"required"`
|
||||
Operator string `form:"operator" json:"operator,omitempty" binding:"oneof='<' '<=' '>=' '>' '=' '!=' 'like' 'not like' 'in',omitempty" default:"like" enums:"<,<=,>=,>,=,!=,like,not like,in"`
|
||||
}
|
||||
|
||||
// TableSort is auxiliary struct to contain method of sorting
|
||||
type TableSort struct {
|
||||
Prop string `form:"prop" json:"prop,omitempty" binding:"omitempty"`
|
||||
Order string `form:"order" json:"order,omitempty" binding:"oneof=ascending descending,required_with=Prop,omitempty" enums:"ascending,descending"`
|
||||
}
|
||||
|
||||
// TableQuery is main struct to contain input params
|
||||
//
|
||||
//nolint:lll
|
||||
type TableQuery struct {
|
||||
// Number of page (since 1)
|
||||
Page int `form:"page" json:"page" binding:"min=1,required" default:"1" minimum:"1"`
|
||||
// Amount items per page (min -1, max 1000, -1 means unlimited)
|
||||
Size int `form:"pageSize" json:"pageSize" binding:"min=-1,max=1000" default:"5" minimum:"-1" maximum:"1000"`
|
||||
// Type of request
|
||||
Type string `form:"type" json:"type" binding:"oneof=sort filter init page size,required" default:"init" enums:"sort,filter,init,page,size"`
|
||||
// Sorting result on server e.g. {"prop":"...","order":"..."}
|
||||
// field order is "ascending" or "descending" value
|
||||
// order is required if prop is not empty
|
||||
Sort []TableSort `form:"sort[]" json:"sort[],omitempty" binding:"omitempty,dive" swaggertype:"array,string"`
|
||||
// Filtering result on server e.g. {"value":[...],"field":"...","operator":"..."}
|
||||
// field is the unique identifier of the table column, different for each endpoint
|
||||
// value should be integer or string or array type, "value":123 or "value":"string" or "value":[123,456]
|
||||
// operator value should be one of <,<=,>=,>,=,!=,like,not like,in
|
||||
// default operator value is 'like' or '=' if field is 'id' or '*_id' or '*_at'
|
||||
Filters []TableFilter `form:"filters[]" json:"filters[],omitempty" binding:"omitempty,dive" swaggertype:"array,string"`
|
||||
// Field to group results by
|
||||
Group string `form:"group" json:"group,omitempty" binding:"omitempty" swaggertype:"string"`
|
||||
// non input arguments
|
||||
table string `form:"-" json:"-"`
|
||||
groupField string `form:"-" json:"-"`
|
||||
sqlMappers map[string]any `form:"-" json:"-"`
|
||||
sqlFind func(out any) func(*gorm.DB) *gorm.DB `form:"-" json:"-"`
|
||||
sqlFilters []func(*gorm.DB) *gorm.DB `form:"-" json:"-"`
|
||||
sqlOrders []func(*gorm.DB) *gorm.DB `form:"-" json:"-"`
|
||||
}
|
||||
|
||||
// Init is function to set table name and sql mapping to data columns
|
||||
func (q *TableQuery) Init(table string, sqlMappers map[string]any) error {
|
||||
q.table = table
|
||||
q.sqlFind = func(out any) func(db *gorm.DB) *gorm.DB {
|
||||
return func(db *gorm.DB) *gorm.DB {
|
||||
return db.Find(out)
|
||||
}
|
||||
}
|
||||
q.sqlMappers = make(map[string]any)
|
||||
q.sqlOrders = append(q.sqlOrders, func(db *gorm.DB) *gorm.DB {
|
||||
return db.Order("id DESC")
|
||||
})
|
||||
for k, v := range sqlMappers {
|
||||
switch t := v.(type) {
|
||||
case string:
|
||||
t = q.DoConditionFormat(t)
|
||||
if isNumbericField(k) {
|
||||
q.sqlMappers[k] = t
|
||||
} else {
|
||||
q.sqlMappers[k] = "LOWER(" + t + "::text)"
|
||||
}
|
||||
case func(q *TableQuery, db *gorm.DB, value any) *gorm.DB:
|
||||
q.sqlMappers[k] = t
|
||||
default:
|
||||
continue
|
||||
}
|
||||
}
|
||||
if q.Group != "" {
|
||||
var ok bool
|
||||
q.groupField, ok = q.sqlMappers[q.Group].(string)
|
||||
if !ok {
|
||||
return errors.New("wrong field for grouping")
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// DoConditionFormat is auxiliary function to prepare condition to the table
|
||||
func (q *TableQuery) DoConditionFormat(cond string) string {
|
||||
cond = strings.ReplaceAll(cond, "{{type}}", q.Type)
|
||||
cond = strings.ReplaceAll(cond, "{{table}}", q.table)
|
||||
cond = strings.ReplaceAll(cond, "{{page}}", strconv.Itoa(q.Page))
|
||||
cond = strings.ReplaceAll(cond, "{{size}}", strconv.Itoa(q.Size))
|
||||
return cond
|
||||
}
|
||||
|
||||
// SetFilters is function to set custom filters to build target SQL query
|
||||
func (q *TableQuery) SetFilters(sqlFilters []func(*gorm.DB) *gorm.DB) {
|
||||
q.sqlFilters = sqlFilters
|
||||
}
|
||||
|
||||
// SetFind is function to set custom find function to build target SQL query
|
||||
func (q *TableQuery) SetFind(find func(out any) func(*gorm.DB) *gorm.DB) {
|
||||
q.sqlFind = find
|
||||
}
|
||||
|
||||
// SetOrders is function to set custom ordering to build target SQL query
|
||||
func (q *TableQuery) SetOrders(sqlOrders []func(*gorm.DB) *gorm.DB) {
|
||||
q.sqlOrders = sqlOrders
|
||||
}
|
||||
|
||||
// Mappers is getter for private field (SQL find funcction to use it in custom query)
|
||||
func (q *TableQuery) Find(out any) func(*gorm.DB) *gorm.DB {
|
||||
return q.sqlFind(out)
|
||||
}
|
||||
|
||||
// Mappers is getter for private field (SQL mappers fields to table ones)
|
||||
func (q *TableQuery) Mappers() map[string]any {
|
||||
return q.sqlMappers
|
||||
}
|
||||
|
||||
// Table is getter for private field (table name)
|
||||
func (q *TableQuery) Table() string {
|
||||
return q.table
|
||||
}
|
||||
|
||||
// Ordering is function to get order of data rows according with input params
|
||||
func (q *TableQuery) Ordering() func(db *gorm.DB) *gorm.DB {
|
||||
var sortItems []TableSort
|
||||
|
||||
for _, sort := range q.Sort {
|
||||
var t TableSort
|
||||
|
||||
switch sort.Order {
|
||||
case "ascending":
|
||||
t.Order = "ASC"
|
||||
case "descending":
|
||||
t.Order = "DESC"
|
||||
}
|
||||
|
||||
if v, ok := q.sqlMappers[sort.Prop]; ok {
|
||||
if s, ok := v.(string); ok {
|
||||
t.Prop = s
|
||||
}
|
||||
}
|
||||
|
||||
if t.Prop != "" && t.Order != "" {
|
||||
sortItems = append(sortItems, t)
|
||||
}
|
||||
}
|
||||
|
||||
return func(db *gorm.DB) *gorm.DB {
|
||||
for _, sort := range sortItems {
|
||||
// sort.Prop comes from server-side whitelist (q.sqlMappers)
|
||||
// sort.Order is validated to be only "ASC" or "DESC"
|
||||
db = db.Order(sort.Prop + " " + sort.Order)
|
||||
}
|
||||
for _, order := range q.sqlOrders {
|
||||
db = order(db)
|
||||
}
|
||||
return db
|
||||
}
|
||||
}
|
||||
|
||||
// Paginate is function to navigate between pages according with input params
|
||||
func (q *TableQuery) Paginate() func(db *gorm.DB) *gorm.DB {
|
||||
return func(db *gorm.DB) *gorm.DB {
|
||||
if q.Page <= 0 && q.Size >= 0 {
|
||||
return db.Limit(q.Size)
|
||||
} else if q.Page > 0 && q.Size >= 0 {
|
||||
offset := (q.Page - 1) * q.Size
|
||||
return db.Offset(offset).Limit(q.Size)
|
||||
}
|
||||
return db
|
||||
}
|
||||
}
|
||||
|
||||
// GroupBy is function to group results by some field
|
||||
func (q *TableQuery) GroupBy(total *uint64, result any) func(db *gorm.DB) *gorm.DB {
|
||||
return func(db *gorm.DB) *gorm.DB {
|
||||
return db.Group(q.groupField).Where(q.groupField+" IS NOT NULL").Count(total).Pluck(q.groupField, result)
|
||||
}
|
||||
}
|
||||
|
||||
// DataFilter is function to build main data filter from filters input params
|
||||
func (q *TableQuery) DataFilter() func(db *gorm.DB) *gorm.DB {
|
||||
type item struct {
|
||||
op string
|
||||
v any
|
||||
}
|
||||
|
||||
fl := make(map[string][]item)
|
||||
setFilter := func(field, operator string, value any) {
|
||||
if operator == "" {
|
||||
operator = "like" // nolint:goconst
|
||||
}
|
||||
fvalue := []item{}
|
||||
if fv, ok := fl[field]; ok {
|
||||
fvalue = fv
|
||||
}
|
||||
switch tvalue := value.(type) {
|
||||
case string, float64, bool:
|
||||
fl[field] = append(fvalue, item{operator, tvalue})
|
||||
case []any:
|
||||
fl[field] = append(fvalue, item{operator, tvalue})
|
||||
}
|
||||
}
|
||||
patchOperator := func(f *TableFilter) {
|
||||
switch f.Operator {
|
||||
case "<", "<=", ">=", ">", "=", "!=", "in":
|
||||
case "not like":
|
||||
if isNumbericField(f.Field) {
|
||||
f.Operator = "!="
|
||||
}
|
||||
default:
|
||||
f.Operator = "like"
|
||||
if isNumbericField(f.Field) {
|
||||
f.Operator = "="
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
for _, f := range q.Filters {
|
||||
f := f
|
||||
|
||||
patchOperator(&f)
|
||||
if _, ok := q.sqlMappers[f.Field]; ok {
|
||||
if v, ok := f.Value.(string); ok && v != "" {
|
||||
vs := v
|
||||
if slices.Contains([]string{"like", "not like"}, f.Operator) {
|
||||
vs = "%" + strings.ToLower(vs) + "%"
|
||||
}
|
||||
setFilter(f.Field, f.Operator, vs)
|
||||
}
|
||||
if v, ok := f.Value.(float64); ok {
|
||||
setFilter(f.Field, f.Operator, v)
|
||||
}
|
||||
if v, ok := f.Value.(bool); ok {
|
||||
setFilter(f.Field, f.Operator, v)
|
||||
}
|
||||
if v, ok := f.Value.([]any); ok && len(v) != 0 {
|
||||
var vi []any
|
||||
for _, ti := range v {
|
||||
if ts, ok := ti.(string); ok {
|
||||
vi = append(vi, strings.ToLower(ts))
|
||||
}
|
||||
if ts, ok := ti.(float64); ok {
|
||||
vi = append(vi, ts)
|
||||
}
|
||||
if ts, ok := ti.(bool); ok {
|
||||
vi = append(vi, ts)
|
||||
}
|
||||
}
|
||||
if len(vi) != 0 {
|
||||
setFilter(f.Field, "in", vi)
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
return func(db *gorm.DB) *gorm.DB {
|
||||
doFilter := func(db *gorm.DB, k, s string, v any) *gorm.DB {
|
||||
switch t := q.sqlMappers[k].(type) {
|
||||
case string:
|
||||
return db.Where(t+s, v)
|
||||
case func(q *TableQuery, db *gorm.DB, value any) *gorm.DB:
|
||||
return t(q, db, v)
|
||||
default:
|
||||
return db
|
||||
}
|
||||
}
|
||||
for k, f := range fl {
|
||||
for _, it := range f {
|
||||
if _, ok := it.v.([]any); ok {
|
||||
db = doFilter(db, k, " "+it.op+" (?)", it.v)
|
||||
} else {
|
||||
db = doFilter(db, k, " "+it.op+" ?", it.v)
|
||||
}
|
||||
}
|
||||
}
|
||||
for _, filter := range q.sqlFilters {
|
||||
db = filter(db)
|
||||
}
|
||||
return db
|
||||
}
|
||||
}
|
||||
|
||||
// Query is function to retrieve table data according with input params
|
||||
func (q *TableQuery) Query(db *gorm.DB, result any,
|
||||
funcs ...func(*gorm.DB) *gorm.DB) (uint64, error) {
|
||||
var total uint64
|
||||
err := ApplyToChainDB(
|
||||
ApplyToChainDB(db.Table(q.Table()), funcs...).Scopes(q.DataFilter()).Count(&total),
|
||||
q.Ordering(),
|
||||
q.Paginate(),
|
||||
q.Find(result),
|
||||
).Error
|
||||
return uint64(total), err
|
||||
}
|
||||
|
||||
// QueryGrouped is function to retrieve grouped data according with input params
|
||||
func (q *TableQuery) QueryGrouped(db *gorm.DB, result any,
|
||||
funcs ...func(*gorm.DB) *gorm.DB) (uint64, error) {
|
||||
if _, ok := q.sqlMappers[q.Group]; !ok {
|
||||
return 0, errors.New("group field not found")
|
||||
}
|
||||
|
||||
var total uint64
|
||||
err := ApplyToChainDB(
|
||||
ApplyToChainDB(db.Table(q.Table()), funcs...).Scopes(q.DataFilter()),
|
||||
q.GroupBy(&total, result),
|
||||
).Error
|
||||
return uint64(total), err
|
||||
}
|
||||
|
||||
// ApplyToChainDB is function to extend gorm method chaining by custom functions
|
||||
func ApplyToChainDB(db *gorm.DB, funcs ...func(*gorm.DB) *gorm.DB) (tx *gorm.DB) {
|
||||
for _, f := range funcs {
|
||||
db = f(db)
|
||||
}
|
||||
return db
|
||||
}
|
||||
|
||||
// EncryptPassword is function to prepare user data as a password
|
||||
func EncryptPassword(password string) (hpass []byte, err error) {
|
||||
hpass, err = bcrypt.GenerateFromPassword([]byte(password), bcrypt.DefaultCost)
|
||||
return
|
||||
}
|
||||
|
||||
// MakeMD5Hash is function to generate common hash by value
|
||||
func MakeMD5Hash(value, salt string) string {
|
||||
currentTime := time.Now().Format("2006-01-02 15:04:05.000000000")
|
||||
hash := md5.Sum([]byte(currentTime + value + salt)) // nolint:gosec
|
||||
return hex.EncodeToString(hash[:])
|
||||
}
|
||||
|
||||
// MakeUserHash is function to generate user hash from name
|
||||
func MakeUserHash(name string) string {
|
||||
currentTime := time.Now().Format("2006-01-02 15:04:05.000000000")
|
||||
return MakeMD5Hash(name+currentTime, "248a8bd896595be1319e65c308a903c568afdb9b")
|
||||
}
|
||||
|
||||
// MakeUuidStrFromHash is function to convert format view from hash to UUID
|
||||
func MakeUuidStrFromHash(hash string) (string, error) {
|
||||
hashBytes, err := hex.DecodeString(hash)
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
userIdUuid, err := uuid.FromBytes(hashBytes)
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
return userIdUuid.String(), nil
|
||||
}
|
||||
|
||||
func isNumbericField(field string) bool {
|
||||
return strings.HasSuffix(field, "_id") || strings.HasSuffix(field, "_at") || field == "id"
|
||||
}
|
||||
@@ -0,0 +1,171 @@
|
||||
package response
|
||||
|
||||
// general
|
||||
|
||||
var ErrInternal = NewHttpError(500, "Internal", "internal server error")
|
||||
var ErrInternalDBNotFound = NewHttpError(500, "Internal.DBNotFound", "db not found")
|
||||
var ErrInternalServiceNotFound = NewHttpError(500, "Internal.ServiceNotFound", "service not found")
|
||||
var ErrInternalDBEncryptorNotFound = NewHttpError(500, "Internal.DBEncryptorNotFound", "DBEncryptor not found")
|
||||
var ErrNotPermitted = NewHttpError(403, "NotPermitted", "action not permitted")
|
||||
var ErrAuthRequired = NewHttpError(403, "AuthRequired", "auth required")
|
||||
var ErrLocalUserRequired = NewHttpError(403, "LocalUserRequired", "local user required")
|
||||
var ErrPrivilegesRequired = NewHttpError(403, "PrivilegesRequired", "some privileges required")
|
||||
var ErrAdminRequired = NewHttpError(403, "AdminRequired", "admin required")
|
||||
var ErrSuperRequired = NewHttpError(403, "SuperRequired", "super admin required")
|
||||
|
||||
// auth
|
||||
|
||||
var ErrAuthInvalidLoginRequest = NewHttpError(400, "Auth.InvalidLoginRequest", "invalid login data")
|
||||
var ErrAuthInvalidAuthorizeQuery = NewHttpError(400, "Auth.InvalidAuthorizeQuery", "invalid authorize query")
|
||||
var ErrAuthInvalidLoginCallbackRequest = NewHttpError(400, "Auth.InvalidLoginCallbackRequest", "invalid login callback data")
|
||||
var ErrAuthInvalidAuthorizationState = NewHttpError(400, "Auth.InvalidAuthorizationState", "invalid authorization state data")
|
||||
var ErrAuthInvalidSwitchServiceHash = NewHttpError(400, "Auth.InvalidSwitchServiceHash", "invalid switch service hash input data")
|
||||
var ErrAuthInvalidAuthorizationNonce = NewHttpError(400, "Auth.InvalidAuthorizationNonce", "invalid authorization nonce data")
|
||||
var ErrAuthInvalidCredentials = NewHttpError(401, "Auth.InvalidCredentials", "invalid login or password")
|
||||
var ErrAuthInvalidUserData = NewHttpError(500, "Auth.InvalidUserData", "invalid user data")
|
||||
var ErrAuthInactiveUser = NewHttpError(403, "Auth.InactiveUser", "user is inactive")
|
||||
var ErrAuthExchangeTokenFail = NewHttpError(403, "Auth.ExchangeTokenFail", "error on exchanging token")
|
||||
var ErrAuthTokenExpired = NewHttpError(403, "Auth.TokenExpired", "token is expired")
|
||||
var ErrAuthVerificationTokenFail = NewHttpError(403, "Auth.VerificationTokenFail", "error on verifying token")
|
||||
var ErrAuthInvalidServiceData = NewHttpError(500, "Auth.InvalidServiceData", "invalid service data")
|
||||
var ErrAuthInvalidTenantData = NewHttpError(500, "Auth.InvalidTenantData", "invalid tenant data")
|
||||
|
||||
// info
|
||||
|
||||
var ErrInfoUserNotFound = NewHttpError(404, "Info.UserNotFound", "user not found")
|
||||
var ErrInfoInvalidUserData = NewHttpError(500, "Info.InvalidUserData", "invalid user data")
|
||||
var ErrInfoInvalidServiceData = NewHttpError(500, "Info.InvalidServiceData", "invalid service data")
|
||||
|
||||
// users
|
||||
|
||||
var ErrUsersNotFound = NewHttpError(404, "Users.NotFound", "user not found")
|
||||
var ErrUsersInvalidData = NewHttpError(500, "Users.InvalidData", "invalid user data")
|
||||
var ErrUsersInvalidRequest = NewHttpError(400, "Users.InvalidRequest", "invalid user request data")
|
||||
var ErrChangePasswordCurrentUserInvalidPassword = NewHttpError(400, "Users.ChangePasswordCurrentUser.InvalidPassword", "failed to validate user password")
|
||||
var ErrChangePasswordCurrentUserInvalidCurrentPassword = NewHttpError(403, "Users.ChangePasswordCurrentUser.InvalidCurrentPassword", "invalid current password")
|
||||
var ErrChangePasswordCurrentUserInvalidNewPassword = NewHttpError(400, "Users.ChangePasswordCurrentUser.InvalidNewPassword", "invalid new password form data")
|
||||
var ErrGetUserModelsNotFound = NewHttpError(404, "Users.GetUser.ModelsNotFound", "user linked models not found")
|
||||
var ErrCreateUserInvalidUser = NewHttpError(400, "Users.CreateUser.InvalidUser", "failed to validate user")
|
||||
var ErrPatchUserModelsNotFound = NewHttpError(404, "Users.PatchUser.ModelsNotFound", "user linked models not found")
|
||||
var ErrDeleteUserModelsNotFound = NewHttpError(404, "Users.DeleteUser.ModelsNotFound", "user linked models not found")
|
||||
|
||||
// roles
|
||||
|
||||
var ErrRolesInvalidRequest = NewHttpError(400, "Roles.InvalidRequest", "invalid role request data")
|
||||
var ErrRolesInvalidData = NewHttpError(500, "Roles.InvalidData", "invalid role data")
|
||||
var ErrRolesNotFound = NewHttpError(404, "Roles.NotFound", "role not found")
|
||||
|
||||
// prompts
|
||||
|
||||
var ErrPromptsInvalidRequest = NewHttpError(400, "Prompts.InvalidRequest", "invalid prompt request data")
|
||||
var ErrPromptsInvalidData = NewHttpError(500, "Prompts.InvalidData", "invalid prompt data")
|
||||
var ErrPromptsNotFound = NewHttpError(404, "Prompts.NotFound", "prompt not found")
|
||||
|
||||
// screenshots
|
||||
|
||||
var ErrScreenshotsInvalidRequest = NewHttpError(400, "Screenshots.InvalidRequest", "invalid screenshot request data")
|
||||
var ErrScreenshotsNotFound = NewHttpError(404, "Screenshots.NotFound", "screenshot not found")
|
||||
var ErrScreenshotsInvalidData = NewHttpError(500, "Screenshots.InvalidData", "invalid screenshot data")
|
||||
|
||||
// containers
|
||||
|
||||
var ErrContainersInvalidRequest = NewHttpError(400, "Containers.InvalidRequest", "invalid container request data")
|
||||
var ErrContainersNotFound = NewHttpError(404, "Containers.NotFound", "container not found")
|
||||
var ErrContainersInvalidData = NewHttpError(500, "Containers.InvalidData", "invalid container data")
|
||||
|
||||
// agentlogs
|
||||
|
||||
var ErrAgentlogsInvalidRequest = NewHttpError(400, "Agentlogs.InvalidRequest", "invalid agentlog request data")
|
||||
var ErrAgentlogsInvalidData = NewHttpError(500, "Agentlogs.InvalidData", "invalid agentlog data")
|
||||
|
||||
// assistantlogs
|
||||
|
||||
var ErrAssistantlogsInvalidRequest = NewHttpError(400, "Assistantlogs.InvalidRequest", "invalid assistantlog request data")
|
||||
var ErrAssistantlogsInvalidData = NewHttpError(500, "Assistantlogs.InvalidData", "invalid assistantlog data")
|
||||
|
||||
// msglogs
|
||||
|
||||
var ErrMsglogsInvalidRequest = NewHttpError(400, "Msglogs.InvalidRequest", "invalid msglog request data")
|
||||
var ErrMsglogsInvalidData = NewHttpError(500, "Msglogs.InvalidData", "invalid msglog data")
|
||||
|
||||
// searchlogs
|
||||
|
||||
var ErrSearchlogsInvalidRequest = NewHttpError(400, "Searchlogs.InvalidRequest", "invalid searchlog request data")
|
||||
var ErrSearchlogsInvalidData = NewHttpError(500, "Searchlogs.InvalidData", "invalid searchlog data")
|
||||
|
||||
// termlogs
|
||||
|
||||
var ErrTermlogsInvalidRequest = NewHttpError(400, "Termlogs.InvalidRequest", "invalid termlog request data")
|
||||
var ErrTermlogsInvalidData = NewHttpError(500, "Termlogs.InvalidData", "invalid termlog data")
|
||||
|
||||
// vecstorelogs
|
||||
|
||||
var ErrVecstorelogsInvalidRequest = NewHttpError(400, "Vecstorelogs.InvalidRequest", "invalid vecstorelog request data")
|
||||
var ErrVecstorelogsInvalidData = NewHttpError(500, "Vecstorelogs.InvalidData", "invalid vecstorelog data")
|
||||
|
||||
// flows
|
||||
|
||||
var ErrFlowsInvalidRequest = NewHttpError(400, "Flows.InvalidRequest", "invalid flow request data")
|
||||
var ErrFlowsNotFound = NewHttpError(404, "Flows.NotFound", "flow not found")
|
||||
var ErrFlowsInvalidData = NewHttpError(500, "Flows.InvalidData", "invalid flow data")
|
||||
|
||||
// flow files
|
||||
|
||||
var ErrFlowFilesInvalidRequest = NewHttpError(400, "FlowFiles.InvalidRequest", "invalid flow file request data")
|
||||
var ErrFlowFilesNotFound = NewHttpError(404, "FlowFiles.NotFound", "flow file not found")
|
||||
var ErrFlowFilesInvalidData = NewHttpError(400, "FlowFiles.InvalidData", "invalid flow file data")
|
||||
var ErrFlowFilesAlreadyExists = NewHttpError(409, "FlowFiles.AlreadyExists", "flow file already exists")
|
||||
var ErrFlowFilesContainerNotRunning = NewHttpError(400, "FlowFiles.ContainerNotRunning", "container is not running")
|
||||
|
||||
// tasks
|
||||
|
||||
var ErrTasksInvalidRequest = NewHttpError(400, "Tasks.InvalidRequest", "invalid task request data")
|
||||
var ErrTasksNotFound = NewHttpError(404, "Tasks.NotFound", "task not found")
|
||||
var ErrTasksInvalidData = NewHttpError(500, "Tasks.InvalidData", "invalid task data")
|
||||
|
||||
// subtasks
|
||||
|
||||
var ErrSubtasksInvalidRequest = NewHttpError(400, "Subtasks.InvalidRequest", "invalid subtask request data")
|
||||
var ErrSubtasksNotFound = NewHttpError(404, "Subtasks.NotFound", "subtask not found")
|
||||
var ErrSubtasksInvalidData = NewHttpError(500, "Subtasks.InvalidData", "invalid subtask data")
|
||||
|
||||
// assistants
|
||||
|
||||
var ErrAssistantsInvalidRequest = NewHttpError(400, "Assistants.InvalidRequest", "invalid assistant request data")
|
||||
var ErrAssistantsNotFound = NewHttpError(404, "Assistants.NotFound", "assistant not found")
|
||||
var ErrAssistantsInvalidData = NewHttpError(500, "Assistants.InvalidData", "invalid assistant data")
|
||||
|
||||
// resources
|
||||
|
||||
var ErrResourcesInvalidRequest = NewHttpError(400, "Resources.InvalidRequest", "invalid resource request data")
|
||||
var ErrResourcesNotFound = NewHttpError(404, "Resources.NotFound", "resource not found")
|
||||
var ErrResourcesAlreadyExists = NewHttpError(409, "Resources.AlreadyExists", "resource already exists")
|
||||
var ErrResourcesInvalidData = NewHttpError(400, "Resources.InvalidData", "invalid resource data")
|
||||
var ErrResourcesConflict = NewHttpError(409, "Resources.Conflict", "resource conflict: use force=true to merge")
|
||||
|
||||
// knowledge
|
||||
|
||||
var ErrKnowledgeInvalidRequest = NewHttpError(400, "Knowledge.InvalidRequest", "invalid knowledge request data")
|
||||
var ErrKnowledgeNotFound = NewHttpError(404, "Knowledge.NotFound", "knowledge document not found")
|
||||
var ErrKnowledgeUnauthorized = NewHttpError(403, "Knowledge.Unauthorized", "not authorized to manage this knowledge document")
|
||||
var ErrKnowledgeStoreUnavail = NewHttpError(503, "Knowledge.StoreUnavailable", "knowledge store (embedding provider) is not configured")
|
||||
var ErrKnowledgeInvalidData = NewHttpError(500, "Knowledge.InvalidData", "invalid knowledge document data")
|
||||
|
||||
// toolcalls
|
||||
|
||||
var ErrToolcallsInvalidRequest = NewHttpError(400, "Toolcalls.InvalidRequest", "invalid toolcall request data")
|
||||
var ErrToolcallsNotFound = NewHttpError(404, "Toolcalls.NotFound", "toolcall not found")
|
||||
var ErrToolcallsInvalidData = NewHttpError(500, "Toolcalls.InvalidData", "invalid toolcall data")
|
||||
|
||||
// anonymize
|
||||
|
||||
var ErrAnonymizeInvalidRequest = NewHttpError(400, "Anonymize.InvalidRequest", "invalid anonymize request data")
|
||||
var ErrAnonymizeUnavailable = NewHttpError(503, "Anonymize.Unavailable", "anonymizer is not configured")
|
||||
|
||||
// tokens
|
||||
|
||||
var ErrTokenCreationDisabled = NewHttpError(400, "Token.CreationDisabled", "token creation is disabled with default configuration")
|
||||
var ErrTokenNotFound = NewHttpError(404, "Token.NotFound", "token not found")
|
||||
var ErrTokenUnauthorized = NewHttpError(403, "Token.Unauthorized", "not authorized to manage this token")
|
||||
var ErrTokenInvalidRequest = NewHttpError(400, "Token.InvalidRequest", "invalid token request data")
|
||||
var ErrTokenInvalidData = NewHttpError(500, "Token.InvalidData", "invalid token data")
|
||||
@@ -0,0 +1,75 @@
|
||||
package response
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
|
||||
"pentagi/pkg/server/logger"
|
||||
"pentagi/pkg/version"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/sirupsen/logrus"
|
||||
)
|
||||
|
||||
type HttpError struct {
|
||||
message string
|
||||
code string
|
||||
httpCode int
|
||||
}
|
||||
|
||||
func (h *HttpError) Code() string {
|
||||
return h.code
|
||||
}
|
||||
|
||||
func (h *HttpError) HttpCode() int {
|
||||
return h.httpCode
|
||||
}
|
||||
|
||||
func (h *HttpError) Msg() string {
|
||||
return h.message
|
||||
}
|
||||
|
||||
func NewHttpError(httpCode int, code, message string) *HttpError {
|
||||
return &HttpError{httpCode: httpCode, message: message, code: code}
|
||||
}
|
||||
|
||||
func (h *HttpError) Error() string {
|
||||
return fmt.Sprintf("%s: %s", h.code, h.message)
|
||||
}
|
||||
|
||||
func Error(c *gin.Context, err *HttpError, original error) {
|
||||
body := gin.H{
|
||||
"status": "error",
|
||||
"code": err.Code(),
|
||||
"msg": err.Msg(),
|
||||
}
|
||||
|
||||
if version.IsDevelopMode() && original != nil {
|
||||
body["error"] = original.Error()
|
||||
}
|
||||
|
||||
fields := logrus.Fields{
|
||||
"code": err.HttpCode(),
|
||||
"message": err.Msg(),
|
||||
}
|
||||
logger.FromContext(c).WithFields(fields).WithError(original).Error("api error")
|
||||
|
||||
c.AbortWithStatusJSON(err.HttpCode(), body)
|
||||
}
|
||||
|
||||
func Success(c *gin.Context, code int, data any) {
|
||||
c.JSON(code, gin.H{"status": "success", "data": data})
|
||||
}
|
||||
|
||||
//lint:ignore U1000 successResp
|
||||
type successResp struct {
|
||||
Status string `json:"status" example:"success"`
|
||||
Data any `json:"data" swaggertype:"object"`
|
||||
} // @name SuccessResponse
|
||||
|
||||
//lint:ignore U1000 errorResp
|
||||
type errorResp struct {
|
||||
Status string `json:"status" example:"error"`
|
||||
Code string `json:"code" example:"Internal"`
|
||||
Msg string `json:"msg,omitempty" example:"internal server error"`
|
||||
Error string `json:"error,omitempty" example:"original server error message"`
|
||||
} // @name ErrorResponse
|
||||
@@ -0,0 +1,434 @@
|
||||
package response
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"testing"
|
||||
|
||||
"pentagi/pkg/version"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
func init() {
|
||||
gin.SetMode(gin.TestMode)
|
||||
}
|
||||
|
||||
func TestNewHttpError(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
err := NewHttpError(404, "NotFound", "resource not found")
|
||||
assert.Equal(t, 404, err.HttpCode())
|
||||
assert.Equal(t, "NotFound", err.Code())
|
||||
assert.Equal(t, "resource not found", err.Msg())
|
||||
}
|
||||
|
||||
func TestHttpError_Error(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
err := NewHttpError(500, "Internal", "something broke")
|
||||
assert.Equal(t, "Internal: something broke", err.Error())
|
||||
}
|
||||
|
||||
func TestHttpError_ImplementsError(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
var err error = NewHttpError(400, "Bad", "bad request")
|
||||
assert.Error(t, err)
|
||||
assert.Contains(t, err.Error(), "Bad")
|
||||
}
|
||||
|
||||
func TestPredefinedErrors(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
tests := []struct {
|
||||
name string
|
||||
err *HttpError
|
||||
httpCode int
|
||||
code string
|
||||
}{
|
||||
// General errors
|
||||
{"ErrInternal", ErrInternal, 500, "Internal"},
|
||||
{"ErrInternalDBNotFound", ErrInternalDBNotFound, 500, "Internal.DBNotFound"},
|
||||
{"ErrInternalServiceNotFound", ErrInternalServiceNotFound, 500, "Internal.ServiceNotFound"},
|
||||
{"ErrInternalDBEncryptorNotFound", ErrInternalDBEncryptorNotFound, 500, "Internal.DBEncryptorNotFound"},
|
||||
{"ErrNotPermitted", ErrNotPermitted, 403, "NotPermitted"},
|
||||
{"ErrAuthRequired", ErrAuthRequired, 403, "AuthRequired"},
|
||||
{"ErrLocalUserRequired", ErrLocalUserRequired, 403, "LocalUserRequired"},
|
||||
{"ErrPrivilegesRequired", ErrPrivilegesRequired, 403, "PrivilegesRequired"},
|
||||
{"ErrAdminRequired", ErrAdminRequired, 403, "AdminRequired"},
|
||||
{"ErrSuperRequired", ErrSuperRequired, 403, "SuperRequired"},
|
||||
|
||||
// Auth errors
|
||||
{"ErrAuthInvalidLoginRequest", ErrAuthInvalidLoginRequest, 400, "Auth.InvalidLoginRequest"},
|
||||
{"ErrAuthInvalidAuthorizeQuery", ErrAuthInvalidAuthorizeQuery, 400, "Auth.InvalidAuthorizeQuery"},
|
||||
{"ErrAuthInvalidLoginCallbackRequest", ErrAuthInvalidLoginCallbackRequest, 400, "Auth.InvalidLoginCallbackRequest"},
|
||||
{"ErrAuthInvalidAuthorizationState", ErrAuthInvalidAuthorizationState, 400, "Auth.InvalidAuthorizationState"},
|
||||
{"ErrAuthInvalidSwitchServiceHash", ErrAuthInvalidSwitchServiceHash, 400, "Auth.InvalidSwitchServiceHash"},
|
||||
{"ErrAuthInvalidAuthorizationNonce", ErrAuthInvalidAuthorizationNonce, 400, "Auth.InvalidAuthorizationNonce"},
|
||||
{"ErrAuthInvalidCredentials", ErrAuthInvalidCredentials, 401, "Auth.InvalidCredentials"},
|
||||
{"ErrAuthInvalidUserData", ErrAuthInvalidUserData, 500, "Auth.InvalidUserData"},
|
||||
{"ErrAuthInactiveUser", ErrAuthInactiveUser, 403, "Auth.InactiveUser"},
|
||||
{"ErrAuthExchangeTokenFail", ErrAuthExchangeTokenFail, 403, "Auth.ExchangeTokenFail"},
|
||||
{"ErrAuthTokenExpired", ErrAuthTokenExpired, 403, "Auth.TokenExpired"},
|
||||
{"ErrAuthVerificationTokenFail", ErrAuthVerificationTokenFail, 403, "Auth.VerificationTokenFail"},
|
||||
{"ErrAuthInvalidServiceData", ErrAuthInvalidServiceData, 500, "Auth.InvalidServiceData"},
|
||||
{"ErrAuthInvalidTenantData", ErrAuthInvalidTenantData, 500, "Auth.InvalidTenantData"},
|
||||
|
||||
// Info errors
|
||||
{"ErrInfoUserNotFound", ErrInfoUserNotFound, 404, "Info.UserNotFound"},
|
||||
{"ErrInfoInvalidUserData", ErrInfoInvalidUserData, 500, "Info.InvalidUserData"},
|
||||
{"ErrInfoInvalidServiceData", ErrInfoInvalidServiceData, 500, "Info.InvalidServiceData"},
|
||||
|
||||
// Users errors
|
||||
{"ErrUsersNotFound", ErrUsersNotFound, 404, "Users.NotFound"},
|
||||
{"ErrUsersInvalidData", ErrUsersInvalidData, 500, "Users.InvalidData"},
|
||||
{"ErrUsersInvalidRequest", ErrUsersInvalidRequest, 400, "Users.InvalidRequest"},
|
||||
{"ErrChangePasswordCurrentUserInvalidPassword", ErrChangePasswordCurrentUserInvalidPassword, 400, "Users.ChangePasswordCurrentUser.InvalidPassword"},
|
||||
{"ErrChangePasswordCurrentUserInvalidCurrentPassword", ErrChangePasswordCurrentUserInvalidCurrentPassword, 403, "Users.ChangePasswordCurrentUser.InvalidCurrentPassword"},
|
||||
{"ErrChangePasswordCurrentUserInvalidNewPassword", ErrChangePasswordCurrentUserInvalidNewPassword, 400, "Users.ChangePasswordCurrentUser.InvalidNewPassword"},
|
||||
{"ErrGetUserModelsNotFound", ErrGetUserModelsNotFound, 404, "Users.GetUser.ModelsNotFound"},
|
||||
{"ErrCreateUserInvalidUser", ErrCreateUserInvalidUser, 400, "Users.CreateUser.InvalidUser"},
|
||||
{"ErrPatchUserModelsNotFound", ErrPatchUserModelsNotFound, 404, "Users.PatchUser.ModelsNotFound"},
|
||||
{"ErrDeleteUserModelsNotFound", ErrDeleteUserModelsNotFound, 404, "Users.DeleteUser.ModelsNotFound"},
|
||||
|
||||
// Roles errors
|
||||
{"ErrRolesInvalidRequest", ErrRolesInvalidRequest, 400, "Roles.InvalidRequest"},
|
||||
{"ErrRolesInvalidData", ErrRolesInvalidData, 500, "Roles.InvalidData"},
|
||||
{"ErrRolesNotFound", ErrRolesNotFound, 404, "Roles.NotFound"},
|
||||
|
||||
// Prompts errors
|
||||
{"ErrPromptsInvalidRequest", ErrPromptsInvalidRequest, 400, "Prompts.InvalidRequest"},
|
||||
{"ErrPromptsInvalidData", ErrPromptsInvalidData, 500, "Prompts.InvalidData"},
|
||||
{"ErrPromptsNotFound", ErrPromptsNotFound, 404, "Prompts.NotFound"},
|
||||
|
||||
// Screenshots errors
|
||||
{"ErrScreenshotsInvalidRequest", ErrScreenshotsInvalidRequest, 400, "Screenshots.InvalidRequest"},
|
||||
{"ErrScreenshotsNotFound", ErrScreenshotsNotFound, 404, "Screenshots.NotFound"},
|
||||
{"ErrScreenshotsInvalidData", ErrScreenshotsInvalidData, 500, "Screenshots.InvalidData"},
|
||||
|
||||
// Containers errors
|
||||
{"ErrContainersInvalidRequest", ErrContainersInvalidRequest, 400, "Containers.InvalidRequest"},
|
||||
{"ErrContainersNotFound", ErrContainersNotFound, 404, "Containers.NotFound"},
|
||||
{"ErrContainersInvalidData", ErrContainersInvalidData, 500, "Containers.InvalidData"},
|
||||
|
||||
// Agentlogs errors
|
||||
{"ErrAgentlogsInvalidRequest", ErrAgentlogsInvalidRequest, 400, "Agentlogs.InvalidRequest"},
|
||||
{"ErrAgentlogsInvalidData", ErrAgentlogsInvalidData, 500, "Agentlogs.InvalidData"},
|
||||
|
||||
// Assistantlogs errors
|
||||
{"ErrAssistantlogsInvalidRequest", ErrAssistantlogsInvalidRequest, 400, "Assistantlogs.InvalidRequest"},
|
||||
{"ErrAssistantlogsInvalidData", ErrAssistantlogsInvalidData, 500, "Assistantlogs.InvalidData"},
|
||||
|
||||
// Msglogs errors
|
||||
{"ErrMsglogsInvalidRequest", ErrMsglogsInvalidRequest, 400, "Msglogs.InvalidRequest"},
|
||||
{"ErrMsglogsInvalidData", ErrMsglogsInvalidData, 500, "Msglogs.InvalidData"},
|
||||
|
||||
// Searchlogs errors
|
||||
{"ErrSearchlogsInvalidRequest", ErrSearchlogsInvalidRequest, 400, "Searchlogs.InvalidRequest"},
|
||||
{"ErrSearchlogsInvalidData", ErrSearchlogsInvalidData, 500, "Searchlogs.InvalidData"},
|
||||
|
||||
// Termlogs errors
|
||||
{"ErrTermlogsInvalidRequest", ErrTermlogsInvalidRequest, 400, "Termlogs.InvalidRequest"},
|
||||
{"ErrTermlogsInvalidData", ErrTermlogsInvalidData, 500, "Termlogs.InvalidData"},
|
||||
|
||||
// Vecstorelogs errors
|
||||
{"ErrVecstorelogsInvalidRequest", ErrVecstorelogsInvalidRequest, 400, "Vecstorelogs.InvalidRequest"},
|
||||
{"ErrVecstorelogsInvalidData", ErrVecstorelogsInvalidData, 500, "Vecstorelogs.InvalidData"},
|
||||
|
||||
// Flows errors
|
||||
{"ErrFlowsInvalidRequest", ErrFlowsInvalidRequest, 400, "Flows.InvalidRequest"},
|
||||
{"ErrFlowsNotFound", ErrFlowsNotFound, 404, "Flows.NotFound"},
|
||||
{"ErrFlowsInvalidData", ErrFlowsInvalidData, 500, "Flows.InvalidData"},
|
||||
|
||||
// Tasks errors
|
||||
{"ErrTasksInvalidRequest", ErrTasksInvalidRequest, 400, "Tasks.InvalidRequest"},
|
||||
{"ErrTasksNotFound", ErrTasksNotFound, 404, "Tasks.NotFound"},
|
||||
{"ErrTasksInvalidData", ErrTasksInvalidData, 500, "Tasks.InvalidData"},
|
||||
|
||||
// Subtasks errors
|
||||
{"ErrSubtasksInvalidRequest", ErrSubtasksInvalidRequest, 400, "Subtasks.InvalidRequest"},
|
||||
{"ErrSubtasksNotFound", ErrSubtasksNotFound, 404, "Subtasks.NotFound"},
|
||||
{"ErrSubtasksInvalidData", ErrSubtasksInvalidData, 500, "Subtasks.InvalidData"},
|
||||
|
||||
// Assistants errors
|
||||
{"ErrAssistantsInvalidRequest", ErrAssistantsInvalidRequest, 400, "Assistants.InvalidRequest"},
|
||||
{"ErrAssistantsNotFound", ErrAssistantsNotFound, 404, "Assistants.NotFound"},
|
||||
{"ErrAssistantsInvalidData", ErrAssistantsInvalidData, 500, "Assistants.InvalidData"},
|
||||
|
||||
// Tokens errors
|
||||
{"ErrTokenCreationDisabled", ErrTokenCreationDisabled, 400, "Token.CreationDisabled"},
|
||||
{"ErrTokenNotFound", ErrTokenNotFound, 404, "Token.NotFound"},
|
||||
{"ErrTokenUnauthorized", ErrTokenUnauthorized, 403, "Token.Unauthorized"},
|
||||
{"ErrTokenInvalidRequest", ErrTokenInvalidRequest, 400, "Token.InvalidRequest"},
|
||||
{"ErrTokenInvalidData", ErrTokenInvalidData, 500, "Token.InvalidData"},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
assert.Equal(t, tt.httpCode, tt.err.HttpCode())
|
||||
assert.Equal(t, tt.code, tt.err.Code())
|
||||
assert.NotEmpty(t, tt.err.Msg())
|
||||
assert.NotEmpty(t, tt.err.Error())
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestSuccessResponse(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
w := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(w)
|
||||
|
||||
data := map[string]string{"id": "123"}
|
||||
Success(c, http.StatusOK, data)
|
||||
|
||||
assert.Equal(t, http.StatusOK, w.Code)
|
||||
|
||||
var body map[string]any
|
||||
err := json.Unmarshal(w.Body.Bytes(), &body)
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, "success", body["status"])
|
||||
assert.NotNil(t, body["data"])
|
||||
}
|
||||
|
||||
func TestSuccessResponse_Created(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
w := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(w)
|
||||
|
||||
Success(c, http.StatusCreated, gin.H{"name": "test"})
|
||||
|
||||
assert.Equal(t, http.StatusCreated, w.Code)
|
||||
}
|
||||
|
||||
func TestErrorResponse(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
w := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(w)
|
||||
c.Request = httptest.NewRequest(http.MethodGet, "/test", nil)
|
||||
|
||||
Error(c, ErrInternal, errors.New("db connection failed"))
|
||||
|
||||
assert.Equal(t, http.StatusInternalServerError, w.Code)
|
||||
|
||||
var body map[string]any
|
||||
err := json.Unmarshal(w.Body.Bytes(), &body)
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, "error", body["status"])
|
||||
assert.Equal(t, "Internal", body["code"])
|
||||
assert.Equal(t, "internal server error", body["msg"])
|
||||
}
|
||||
|
||||
func TestErrorResponse_DevMode(t *testing.T) {
|
||||
// Save original version and restore after test
|
||||
oldVer := version.PackageVer
|
||||
defer func() { version.PackageVer = oldVer }()
|
||||
|
||||
// Enable dev mode
|
||||
version.PackageVer = ""
|
||||
|
||||
w := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(w)
|
||||
c.Request = httptest.NewRequest(http.MethodGet, "/test", nil)
|
||||
|
||||
originalErr := errors.New("detailed error info")
|
||||
Error(c, ErrInternal, originalErr)
|
||||
|
||||
var body map[string]any
|
||||
err := json.Unmarshal(w.Body.Bytes(), &body)
|
||||
require.NoError(t, err)
|
||||
|
||||
// In dev mode, original error should be included
|
||||
assert.Equal(t, "detailed error info", body["error"])
|
||||
}
|
||||
|
||||
func TestErrorResponse_ProductionMode(t *testing.T) {
|
||||
// Save original version and restore after test
|
||||
oldVer := version.PackageVer
|
||||
defer func() { version.PackageVer = oldVer }()
|
||||
|
||||
// Set production mode
|
||||
version.PackageVer = "1.0.0"
|
||||
|
||||
w := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(w)
|
||||
c.Request = httptest.NewRequest(http.MethodGet, "/test", nil)
|
||||
|
||||
Error(c, ErrInternal, errors.New("should not appear"))
|
||||
|
||||
var body map[string]any
|
||||
err := json.Unmarshal(w.Body.Bytes(), &body)
|
||||
require.NoError(t, err)
|
||||
|
||||
// In production mode, original error should NOT be included
|
||||
_, hasError := body["error"]
|
||||
assert.False(t, hasError)
|
||||
}
|
||||
|
||||
func TestErrorResponse_NilOriginalError(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
w := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(w)
|
||||
c.Request = httptest.NewRequest(http.MethodGet, "/test", nil)
|
||||
|
||||
Error(c, ErrNotPermitted, nil)
|
||||
|
||||
assert.Equal(t, http.StatusForbidden, w.Code)
|
||||
|
||||
var body map[string]any
|
||||
err := json.Unmarshal(w.Body.Bytes(), &body)
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, "NotPermitted", body["code"])
|
||||
}
|
||||
|
||||
func TestHttpError_MultipleInstancesIndependent(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
err1 := NewHttpError(404, "NotFound", "resource 1 not found")
|
||||
err2 := NewHttpError(404, "NotFound", "resource 2 not found")
|
||||
|
||||
// Verify they are independent instances
|
||||
assert.NotEqual(t, err1.Msg(), err2.Msg())
|
||||
assert.Equal(t, err1.Code(), err2.Code())
|
||||
assert.Equal(t, err1.HttpCode(), err2.HttpCode())
|
||||
}
|
||||
|
||||
func TestSuccessResponse_EmptyData(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
w := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(w)
|
||||
|
||||
Success(c, http.StatusOK, nil)
|
||||
|
||||
assert.Equal(t, http.StatusOK, w.Code)
|
||||
|
||||
var body map[string]any
|
||||
err := json.Unmarshal(w.Body.Bytes(), &body)
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, "success", body["status"])
|
||||
assert.Nil(t, body["data"])
|
||||
}
|
||||
|
||||
func TestSuccessResponse_ComplexData(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
w := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(w)
|
||||
|
||||
data := gin.H{
|
||||
"users": []gin.H{
|
||||
{"id": 1, "name": "Alice"},
|
||||
{"id": 2, "name": "Bob"},
|
||||
},
|
||||
"count": 2,
|
||||
"meta": gin.H{
|
||||
"page": 1,
|
||||
"total": 100,
|
||||
},
|
||||
}
|
||||
|
||||
Success(c, http.StatusOK, data)
|
||||
|
||||
assert.Equal(t, http.StatusOK, w.Code)
|
||||
|
||||
var body map[string]any
|
||||
err := json.Unmarshal(w.Body.Bytes(), &body)
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, "success", body["status"])
|
||||
|
||||
responseData, ok := body["data"].(map[string]any)
|
||||
require.True(t, ok)
|
||||
assert.Equal(t, float64(2), responseData["count"])
|
||||
}
|
||||
|
||||
func TestErrorResponse_DifferentHttpCodes(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
tests := []struct {
|
||||
name string
|
||||
err *HttpError
|
||||
expected int
|
||||
}{
|
||||
{"400 Bad Request", ErrPromptsInvalidRequest, http.StatusBadRequest},
|
||||
{"401 Unauthorized", ErrAuthInvalidCredentials, http.StatusUnauthorized},
|
||||
{"403 Forbidden", ErrNotPermitted, http.StatusForbidden},
|
||||
{"404 Not Found", ErrUsersNotFound, http.StatusNotFound},
|
||||
{"500 Internal", ErrInternal, http.StatusInternalServerError},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
w := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(w)
|
||||
c.Request = httptest.NewRequest(http.MethodGet, "/test", nil)
|
||||
|
||||
Error(c, tt.err, nil)
|
||||
|
||||
assert.Equal(t, tt.expected, w.Code)
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestErrorResponse_ResponseStructure(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
w := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(w)
|
||||
c.Request = httptest.NewRequest(http.MethodGet, "/test", nil)
|
||||
|
||||
Error(c, ErrUsersNotFound, nil)
|
||||
|
||||
var body map[string]any
|
||||
err := json.Unmarshal(w.Body.Bytes(), &body)
|
||||
require.NoError(t, err)
|
||||
|
||||
// Verify required fields
|
||||
assert.Equal(t, "error", body["status"])
|
||||
assert.Equal(t, "Users.NotFound", body["code"])
|
||||
assert.Equal(t, "user not found", body["msg"])
|
||||
|
||||
// Verify error field is not present in non-dev mode
|
||||
_, hasError := body["error"]
|
||||
assert.False(t, hasError)
|
||||
}
|
||||
|
||||
func TestSuccessResponse_StatusCodes(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
tests := []struct {
|
||||
name string
|
||||
statusCode int
|
||||
expectedCode int
|
||||
}{
|
||||
{"200 OK", http.StatusOK, 200},
|
||||
{"201 Created", http.StatusCreated, 201},
|
||||
{"202 Accepted", http.StatusAccepted, 202},
|
||||
{"204 No Content", http.StatusNoContent, 204},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
w := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(w)
|
||||
|
||||
Success(c, tt.statusCode, gin.H{"test": "data"})
|
||||
|
||||
assert.Equal(t, tt.expectedCode, w.Code)
|
||||
})
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,695 @@
|
||||
package router
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/gob"
|
||||
"net"
|
||||
"net/http"
|
||||
"net/http/httputil"
|
||||
"net/url"
|
||||
"os"
|
||||
"path"
|
||||
"path/filepath"
|
||||
"slices"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"pentagi/pkg/config"
|
||||
"pentagi/pkg/controller"
|
||||
"pentagi/pkg/database"
|
||||
"pentagi/pkg/database/knowledge"
|
||||
"pentagi/pkg/docker"
|
||||
"pentagi/pkg/graph/subscriptions"
|
||||
"pentagi/pkg/providers"
|
||||
"pentagi/pkg/server/auth"
|
||||
"pentagi/pkg/server/logger"
|
||||
"pentagi/pkg/server/oauth"
|
||||
"pentagi/pkg/server/services"
|
||||
|
||||
_ "pentagi/pkg/server/docs" // swagger docs
|
||||
|
||||
"github.com/gin-contrib/cors"
|
||||
"github.com/gin-contrib/sessions"
|
||||
"github.com/gin-contrib/sessions/cookie"
|
||||
"github.com/gin-contrib/static"
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/jinzhu/gorm"
|
||||
"github.com/sirupsen/logrus"
|
||||
ginSwagger "github.com/swaggo/gin-swagger"
|
||||
"github.com/swaggo/gin-swagger/swaggerFiles"
|
||||
"github.com/vxcontrol/cloud/anonymizer"
|
||||
"github.com/vxcontrol/cloud/anonymizer/patterns"
|
||||
"github.com/vxcontrol/langchaingo/vectorstores/pgvector"
|
||||
)
|
||||
|
||||
const baseURL = "/api/v1"
|
||||
|
||||
const corsAllowGoogleOAuth = "https://accounts.google.com"
|
||||
|
||||
// frontendRoutes defines the list of URI prefixes that should be handled by the frontend SPA.
|
||||
// Add new frontend base routes here if they are added in the frontend router (e.g., in App.tsx).
|
||||
var frontendRoutes = []string{
|
||||
"/chat",
|
||||
"/oauth",
|
||||
"/login",
|
||||
"/flows",
|
||||
"/settings",
|
||||
"/templates",
|
||||
"/resources",
|
||||
"/knowledges",
|
||||
"/dashboard",
|
||||
}
|
||||
|
||||
// @title PentAGI Swagger API
|
||||
// @version 1.0
|
||||
// @description Swagger API for Penetration Testing Advanced General Intelligence PentAGI.
|
||||
// @termsOfService http://swagger.io/terms/
|
||||
|
||||
// @contact.url https://pentagi.com
|
||||
// @contact.name PentAGI Development Team
|
||||
// @contact.email team@pentagi.com
|
||||
|
||||
// @license.name MIT
|
||||
// @license.url https://opensource.org/license/mit
|
||||
|
||||
// @query.collection.format multi
|
||||
|
||||
// @securityDefinitions.apikey BearerAuth
|
||||
// @in header
|
||||
// @name Authorization
|
||||
// @description Type "Bearer" followed by a space and JWT token.
|
||||
|
||||
// @BasePath /api/v1
|
||||
func NewRouter(
|
||||
db *database.Queries,
|
||||
orm *gorm.DB,
|
||||
cfg *config.Config,
|
||||
providers providers.ProviderController,
|
||||
controller controller.FlowController,
|
||||
subscriptions subscriptions.SubscriptionsController,
|
||||
dockerClient docker.DockerClient,
|
||||
) *gin.Engine {
|
||||
gin.SetMode(gin.ReleaseMode)
|
||||
if cfg.Debug {
|
||||
gin.SetMode(gin.DebugMode)
|
||||
}
|
||||
|
||||
gob.Register([]string{})
|
||||
|
||||
tokenCache := auth.NewTokenCache(orm)
|
||||
userCache := auth.NewUserCache(orm)
|
||||
authMiddleware := auth.NewAuthMiddleware(baseURL, cfg.CookieSigningSalt, tokenCache, userCache)
|
||||
oauthClients := make(map[string]oauth.OAuthClient)
|
||||
oauthLoginCallbackURL := "/auth/login-callback"
|
||||
|
||||
publicURL, err := url.Parse(cfg.PublicURL)
|
||||
if err == nil {
|
||||
publicURL.Path = path.Join(baseURL, oauthLoginCallbackURL)
|
||||
}
|
||||
|
||||
if publicURL != nil && cfg.OAuthGoogleClientID != "" && cfg.OAuthGoogleClientSecret != "" {
|
||||
googleClient := oauth.NewGoogleOAuthClient(
|
||||
cfg.OAuthGoogleClientID,
|
||||
cfg.OAuthGoogleClientSecret,
|
||||
publicURL.String(),
|
||||
)
|
||||
oauthClients[googleClient.ProviderName()] = googleClient
|
||||
}
|
||||
|
||||
if publicURL != nil && cfg.OAuthGithubClientID != "" && cfg.OAuthGithubClientSecret != "" {
|
||||
githubClient := oauth.NewGithubOAuthClient(
|
||||
cfg.OAuthGithubClientID,
|
||||
cfg.OAuthGithubClientSecret,
|
||||
publicURL.String(),
|
||||
)
|
||||
oauthClients[githubClient.ProviderName()] = githubClient
|
||||
}
|
||||
|
||||
// ---- Knowledge (pgvector) store -----------------------------------------
|
||||
// Shared by both the GraphQL and REST layers.
|
||||
// Store and embedder are nil when no embedding provider is configured;
|
||||
// the knowledge store handles that gracefully (embedding-dependent ops error).
|
||||
embedder := providers.Embedder()
|
||||
var pgStore *pgvector.Store
|
||||
if embedder.IsAvailable() {
|
||||
opts := []pgvector.Option{
|
||||
pgvector.WithEmbedder(embedder),
|
||||
pgvector.WithCollectionName("langchain"),
|
||||
}
|
||||
if cfg.PgxPool != nil {
|
||||
opts = append(opts, pgvector.WithConn(cfg.PgxPool))
|
||||
} else {
|
||||
opts = append(opts, pgvector.WithConnectionURL(cfg.DatabaseURL))
|
||||
}
|
||||
if s, err := pgvector.New(context.Background(), opts...); err == nil {
|
||||
pgStore = &s
|
||||
} else {
|
||||
logrus.WithError(err).Warn("failed to initialise pgvector store for knowledge API; embedding operations will be unavailable")
|
||||
}
|
||||
}
|
||||
var knowledgeStore knowledge.KnowledgeStore
|
||||
knowledgeStore = knowledge.NewKnowledgeStore(db, pgStore, embedder, subscriptions.NewKnowledgePublisher, cfg.EmbeddingMaxTextBytes)
|
||||
|
||||
// ---- Anonymizer replacer ------------------------------------------------
|
||||
// Shared singleton used by the GraphQL anonymizeText mutation.
|
||||
// Falls back to a no-op nil replacer on failure so the rest of the server still starts correctly.
|
||||
var textReplacer anonymizer.Replacer
|
||||
if allPatterns, err := patterns.LoadPatterns(patterns.PatternListTypeAll); err != nil {
|
||||
logrus.WithError(err).Warn("failed to load anonymizer patterns; anonymizeText mutation will be unavailable")
|
||||
} else {
|
||||
allPatterns.Patterns = append(allPatterns.Patterns, cfg.GetSecretPatterns()...)
|
||||
|
||||
if r, err := anonymizer.NewReplacer(allPatterns.Regexes(), allPatterns.Names()); err != nil {
|
||||
logrus.WithError(err).Warn("failed to create anonymizer replacer; anonymizeText mutation will be unavailable")
|
||||
} else {
|
||||
textReplacer = r
|
||||
}
|
||||
}
|
||||
|
||||
// services
|
||||
authService := services.NewAuthService(
|
||||
services.AuthServiceConfig{
|
||||
BaseURL: baseURL,
|
||||
LoginCallbackURL: oauthLoginCallbackURL,
|
||||
SessionTimeout: 4 * 60 * 60, // 4 hours
|
||||
},
|
||||
orm,
|
||||
oauthClients,
|
||||
)
|
||||
userService := services.NewUserService(orm, userCache)
|
||||
roleService := services.NewRoleService(orm)
|
||||
providerService := services.NewProviderService(providers)
|
||||
settingsService := services.NewSettingsService(cfg)
|
||||
flowService := services.NewFlowService(orm, providers, controller, subscriptions)
|
||||
flowFileService := services.NewFlowFileService(orm, cfg.DataDir, dockerClient, subscriptions)
|
||||
resourceService := services.NewResourceService(orm, cfg.DataDir, subscriptions)
|
||||
taskService := services.NewTaskService(orm)
|
||||
subtaskService := services.NewSubtaskService(orm)
|
||||
containerService := services.NewContainerService(orm)
|
||||
toolcallService := services.NewToolcallService(orm)
|
||||
assistantService := services.NewAssistantService(orm, providers, controller, subscriptions)
|
||||
agentlogService := services.NewAgentlogService(orm)
|
||||
assistantlogService := services.NewAssistantlogService(orm)
|
||||
msglogService := services.NewMsglogService(orm)
|
||||
searchlogService := services.NewSearchlogService(orm)
|
||||
vecstorelogService := services.NewVecstorelogService(orm)
|
||||
termlogService := services.NewTermlogService(orm)
|
||||
screenshotService := services.NewScreenshotService(orm, cfg.DataDir)
|
||||
promptService := services.NewPromptService(orm)
|
||||
analyticsService := services.NewAnalyticsService(orm)
|
||||
tokenService := services.NewTokenService(orm, cfg.CookieSigningSalt, tokenCache, subscriptions)
|
||||
knowledgeService := services.NewKnowledgeService(orm, knowledgeStore)
|
||||
anonymizerService := services.NewAnonymizerService(textReplacer)
|
||||
graphqlService := services.NewGraphqlService(
|
||||
db, cfg, baseURL, cfg.CorsOrigins, tokenCache, providers, controller, subscriptions, knowledgeStore, textReplacer,
|
||||
)
|
||||
|
||||
router := gin.Default()
|
||||
|
||||
// Setup Cross-Origin Resource Sharing policy
|
||||
config := cors.DefaultConfig()
|
||||
if !slices.Contains(cfg.CorsOrigins, "*") {
|
||||
config.AllowCredentials = true
|
||||
}
|
||||
config.AllowWildcard = true
|
||||
config.AllowWebSockets = true
|
||||
config.AllowPrivateNetwork = true
|
||||
|
||||
// Add OAuth provider origins to CORS allowed origins
|
||||
allowedOrigins := make([]string, len(cfg.CorsOrigins))
|
||||
copy(allowedOrigins, cfg.CorsOrigins)
|
||||
|
||||
// Google OAuth uses POST callback from accounts.google.com
|
||||
if cfg.OAuthGoogleClientID != "" && cfg.OAuthGoogleClientSecret != "" {
|
||||
if !slices.Contains(allowedOrigins, corsAllowGoogleOAuth) && !slices.Contains(cfg.CorsOrigins, "*") {
|
||||
allowedOrigins = append(allowedOrigins, corsAllowGoogleOAuth)
|
||||
logrus.Infof("Added %s to CORS allowed origins for Google OAuth", corsAllowGoogleOAuth)
|
||||
}
|
||||
}
|
||||
|
||||
config.AllowOrigins = allowedOrigins
|
||||
config.AllowMethods = []string{"GET", "POST", "PUT", "DELETE", "OPTIONS"}
|
||||
if err := config.Validate(); err != nil {
|
||||
logrus.WithError(err).Error("failed to validate cors config")
|
||||
} else {
|
||||
router.Use(cors.New(config))
|
||||
}
|
||||
|
||||
router.Use(gin.Recovery())
|
||||
router.Use(logger.WithGinLogger("pentagi-api"))
|
||||
|
||||
cookieStore := cookie.NewStore(auth.MakeCookieStoreKey(cfg.CookieSigningSalt)...)
|
||||
router.Use(sessions.Sessions("auth", cookieStore))
|
||||
|
||||
api := router.Group(baseURL)
|
||||
api.Use(noCacheMiddleware())
|
||||
|
||||
// Special case for local user own password change
|
||||
changePasswordGroup := api.Group("/user")
|
||||
changePasswordGroup.Use(authMiddleware.AuthUserRequired)
|
||||
changePasswordGroup.Use(localUserRequired())
|
||||
changePasswordGroup.PUT("/password", userService.ChangePasswordCurrentUser)
|
||||
|
||||
publicGroup := api.Group("/")
|
||||
publicGroup.Use(authMiddleware.TryAuth)
|
||||
{
|
||||
publicGroup.GET("/info", authService.Info)
|
||||
|
||||
developerGroup := publicGroup.Group("/")
|
||||
{
|
||||
developerGroup.GET("/graphql/playground", graphqlService.ServeGraphqlPlayground)
|
||||
developerGroup.GET("/swagger/*any", ginSwagger.WrapHandler(swaggerFiles.Handler))
|
||||
}
|
||||
|
||||
authGroup := publicGroup.Group("/auth")
|
||||
{
|
||||
authGroup.POST("/login", authService.AuthLogin)
|
||||
authGroup.GET("/logout", authService.AuthLogout)
|
||||
authGroup.GET("/authorize", authService.AuthAuthorize)
|
||||
authGroup.GET("/login-callback", authService.AuthLoginGetCallback)
|
||||
authGroup.POST("/login-callback", authService.AuthLoginPostCallback)
|
||||
authGroup.POST("/logout-callback", authService.AuthLogoutCallback)
|
||||
}
|
||||
}
|
||||
|
||||
privateGroup := api.Group("/")
|
||||
privateGroup.Use(authMiddleware.AuthTokenRequired)
|
||||
{
|
||||
setGraphqlGroup(privateGroup, graphqlService)
|
||||
|
||||
setKnowledgeGroup(privateGroup, knowledgeService)
|
||||
setProvidersGroup(privateGroup, providerService)
|
||||
setSettingsGroup(privateGroup, settingsService)
|
||||
setFlowsGroup(privateGroup, flowService)
|
||||
setFlowFilesGroup(privateGroup, flowFileService)
|
||||
setResourcesGroup(privateGroup, resourceService)
|
||||
setTasksGroup(privateGroup, taskService)
|
||||
setSubtasksGroup(privateGroup, subtaskService)
|
||||
setContainersGroup(privateGroup, containerService)
|
||||
setToolcallsGroup(privateGroup, toolcallService)
|
||||
setAssistantsGroup(privateGroup, assistantService)
|
||||
setAgentlogsGroup(privateGroup, agentlogService)
|
||||
setAssistantlogsGroup(privateGroup, assistantlogService)
|
||||
setMsglogsGroup(privateGroup, msglogService)
|
||||
setTermlogsGroup(privateGroup, termlogService)
|
||||
setSearchlogsGroup(privateGroup, searchlogService)
|
||||
setVecstorelogsGroup(privateGroup, vecstorelogService)
|
||||
setScreenshotsGroup(privateGroup, screenshotService)
|
||||
setPromptsGroup(privateGroup, promptService)
|
||||
setAnonymizeGroup(privateGroup, anonymizerService)
|
||||
setAnalyticsGroup(privateGroup, analyticsService)
|
||||
}
|
||||
|
||||
privateUserGroup := api.Group("/")
|
||||
privateUserGroup.Use(authMiddleware.AuthUserRequired)
|
||||
{
|
||||
setRolesGroup(privateGroup, roleService)
|
||||
setUsersGroup(privateGroup, userService)
|
||||
setTokensGroup(privateGroup, tokenService)
|
||||
}
|
||||
|
||||
if cfg.StaticURL != nil && cfg.StaticURL.Scheme != "" && cfg.StaticURL.Host != "" {
|
||||
router.NoRoute(func() gin.HandlerFunc {
|
||||
return func(c *gin.Context) {
|
||||
director := func(req *http.Request) {
|
||||
*req = *c.Request
|
||||
req.URL.Scheme = cfg.StaticURL.Scheme
|
||||
req.URL.Host = cfg.StaticURL.Host
|
||||
}
|
||||
dialer := &net.Dialer{
|
||||
Timeout: 30 * time.Second,
|
||||
KeepAlive: 30 * time.Second,
|
||||
}
|
||||
httpTransport := &http.Transport{
|
||||
DialContext: dialer.DialContext,
|
||||
ForceAttemptHTTP2: true,
|
||||
MaxIdleConns: 20,
|
||||
IdleConnTimeout: 60 * time.Second,
|
||||
TLSHandshakeTimeout: 10 * time.Second,
|
||||
ExpectContinueTimeout: 1 * time.Second,
|
||||
}
|
||||
|
||||
proxy := &httputil.ReverseProxy{
|
||||
Director: director,
|
||||
Transport: httpTransport,
|
||||
}
|
||||
proxy.ServeHTTP(c.Writer, c.Request)
|
||||
}
|
||||
}())
|
||||
} else {
|
||||
router.Use(static.Serve("/", static.LocalFile(cfg.StaticDir, true)))
|
||||
|
||||
indexExists := true
|
||||
indexPath := filepath.Join(cfg.StaticDir, "index.html")
|
||||
if _, err := os.Stat(indexPath); err != nil {
|
||||
indexExists = false
|
||||
}
|
||||
|
||||
router.NoRoute(func(c *gin.Context) {
|
||||
if c.Request.Method == "GET" && !strings.HasPrefix(c.Request.URL.Path, baseURL) {
|
||||
isFrontendRoute := false
|
||||
path := c.Request.URL.Path
|
||||
for _, prefix := range frontendRoutes {
|
||||
if path == prefix || strings.HasPrefix(path, prefix+"/") {
|
||||
isFrontendRoute = true
|
||||
break
|
||||
}
|
||||
}
|
||||
|
||||
if isFrontendRoute && indexExists {
|
||||
c.File(indexPath)
|
||||
return
|
||||
}
|
||||
}
|
||||
|
||||
c.Redirect(http.StatusMovedPermanently, "/")
|
||||
})
|
||||
}
|
||||
|
||||
return router
|
||||
}
|
||||
|
||||
func setKnowledgeGroup(parent *gin.RouterGroup, svc *services.KnowledgeService) {
|
||||
kg := parent.Group("/knowledge")
|
||||
{
|
||||
kg.GET("/", svc.ListDocuments)
|
||||
kg.GET("/:id", svc.GetDocument)
|
||||
kg.POST("/", svc.CreateDocument)
|
||||
kg.POST("/search", svc.SearchDocuments)
|
||||
kg.PUT("/:id", svc.UpdateDocument)
|
||||
kg.DELETE("/:id", svc.DeleteDocument)
|
||||
}
|
||||
}
|
||||
|
||||
func setProvidersGroup(parent *gin.RouterGroup, svc *services.ProviderService) {
|
||||
providersGroup := parent.Group("/providers")
|
||||
{
|
||||
providersGroup.GET("/", svc.GetProviders)
|
||||
}
|
||||
}
|
||||
|
||||
func setSettingsGroup(parent *gin.RouterGroup, svc *services.SettingsService) {
|
||||
settingsGroup := parent.Group("/settings")
|
||||
{
|
||||
settingsGroup.GET("/", svc.GetSettings)
|
||||
}
|
||||
}
|
||||
|
||||
func setGraphqlGroup(parent *gin.RouterGroup, svc *services.GraphqlService) {
|
||||
graphqlGroup := parent.Group("/")
|
||||
{
|
||||
graphqlGroup.Any("/graphql", svc.ServeGraphql)
|
||||
}
|
||||
}
|
||||
|
||||
func setSubtasksGroup(parent *gin.RouterGroup, svc *services.SubtaskService) {
|
||||
flowSubtasksViewGroup := parent.Group("/flows/:flowID/subtasks")
|
||||
{
|
||||
flowSubtasksViewGroup.GET("/", svc.GetFlowSubtasks)
|
||||
}
|
||||
|
||||
flowTaskSubtasksViewGroup := parent.Group("/flows/:flowID/tasks/:taskID/subtasks")
|
||||
{
|
||||
flowTaskSubtasksViewGroup.GET("/", svc.GetFlowTaskSubtasks)
|
||||
flowTaskSubtasksViewGroup.GET("/:subtaskID", svc.GetFlowTaskSubtask)
|
||||
}
|
||||
}
|
||||
|
||||
func setTasksGroup(parent *gin.RouterGroup, svc *services.TaskService) {
|
||||
flowTaskViewGroup := parent.Group("/flows/:flowID/tasks")
|
||||
{
|
||||
flowTaskViewGroup.GET("/", svc.GetFlowTasks)
|
||||
flowTaskViewGroup.GET("/:taskID", svc.GetFlowTask)
|
||||
flowTaskViewGroup.GET("/:taskID/graph", svc.GetFlowTaskGraph)
|
||||
}
|
||||
}
|
||||
|
||||
func setFlowsGroup(parent *gin.RouterGroup, svc *services.FlowService) {
|
||||
flowCreateGroup := parent.Group("/flows")
|
||||
{
|
||||
flowCreateGroup.POST("/", svc.CreateFlow)
|
||||
}
|
||||
|
||||
flowDeleteGroup := parent.Group("/flows")
|
||||
{
|
||||
flowDeleteGroup.DELETE("/:flowID", svc.DeleteFlow)
|
||||
}
|
||||
|
||||
flowEditGroup := parent.Group("/flows")
|
||||
{
|
||||
flowEditGroup.PUT("/:flowID", svc.PatchFlow)
|
||||
}
|
||||
|
||||
flowsViewGroup := parent.Group("/flows")
|
||||
{
|
||||
flowsViewGroup.GET("/", svc.GetFlows)
|
||||
flowsViewGroup.GET("/:flowID", svc.GetFlow)
|
||||
flowsViewGroup.GET("/:flowID/graph", svc.GetFlowGraph)
|
||||
}
|
||||
}
|
||||
|
||||
func setFlowFilesGroup(parent *gin.RouterGroup, svc *services.FlowFileService) {
|
||||
flowFilesGroup := parent.Group("/flows/:flowID/files")
|
||||
{
|
||||
flowFilesGroup.GET("/", svc.GetFlowFiles)
|
||||
flowFilesGroup.GET("/container", svc.GetFlowContainerFiles)
|
||||
flowFilesGroup.POST("/", svc.UploadFlowFiles)
|
||||
flowFilesGroup.DELETE("/", svc.DeleteFlowFile)
|
||||
flowFilesGroup.GET("/download", svc.DownloadFlowFile)
|
||||
flowFilesGroup.POST("/pull", svc.PullFlowFiles)
|
||||
flowFilesGroup.POST("/resources", svc.AddResourcesToFlow)
|
||||
flowFilesGroup.POST("/to-resources", svc.AddResourceFromFlow)
|
||||
}
|
||||
}
|
||||
|
||||
func setResourcesGroup(parent *gin.RouterGroup, svc *services.ResourceService) {
|
||||
rg := parent.Group("/resources")
|
||||
{
|
||||
rg.GET("/", svc.ListResources)
|
||||
rg.POST("/", svc.UploadResources)
|
||||
rg.POST("/mkdir", svc.MkdirResource)
|
||||
rg.PUT("/move", svc.MoveResource)
|
||||
rg.POST("/copy", svc.CopyResource)
|
||||
rg.DELETE("/", svc.DeleteResource)
|
||||
rg.GET("/download", svc.DownloadResource)
|
||||
}
|
||||
}
|
||||
|
||||
func setContainersGroup(parent *gin.RouterGroup, svc *services.ContainerService) {
|
||||
containersViewGroup := parent.Group("/containers")
|
||||
{
|
||||
containersViewGroup.GET("/", svc.GetContainers)
|
||||
}
|
||||
|
||||
flowContainersViewGroup := parent.Group("/flows/:flowID/containers")
|
||||
{
|
||||
flowContainersViewGroup.GET("/", svc.GetFlowContainers)
|
||||
flowContainersViewGroup.GET("/:containerID", svc.GetFlowContainer)
|
||||
}
|
||||
}
|
||||
|
||||
func setToolcallsGroup(parent *gin.RouterGroup, svc *services.ToolcallService) {
|
||||
toolcallsViewGroup := parent.Group("/toolcalls")
|
||||
{
|
||||
toolcallsViewGroup.GET("/", svc.GetToolcalls)
|
||||
}
|
||||
|
||||
flowToolcallsViewGroup := parent.Group("/flows/:flowID/toolcalls")
|
||||
{
|
||||
flowToolcallsViewGroup.GET("/", svc.GetFlowToolcalls)
|
||||
flowToolcallsViewGroup.GET("/:toolcallID", svc.GetFlowToolcall)
|
||||
}
|
||||
}
|
||||
|
||||
func setAssistantsGroup(parent *gin.RouterGroup, svc *services.AssistantService) {
|
||||
flowCreateGroup := parent.Group("/flows/:flowID/assistants")
|
||||
{
|
||||
flowCreateGroup.POST("/", svc.CreateFlowAssistant)
|
||||
}
|
||||
|
||||
flowDeleteGroup := parent.Group("/flows/:flowID/assistants")
|
||||
{
|
||||
flowDeleteGroup.DELETE("/:assistantID", svc.DeleteAssistant)
|
||||
}
|
||||
|
||||
flowEditGroup := parent.Group("/flows/:flowID/assistants")
|
||||
{
|
||||
flowEditGroup.PUT("/:assistantID", svc.PatchAssistant)
|
||||
}
|
||||
|
||||
flowsViewGroup := parent.Group("/flows/:flowID/assistants")
|
||||
{
|
||||
flowsViewGroup.GET("/", svc.GetFlowAssistants)
|
||||
flowsViewGroup.GET("/:assistantID", svc.GetFlowAssistant)
|
||||
}
|
||||
}
|
||||
|
||||
func setAgentlogsGroup(parent *gin.RouterGroup, svc *services.AgentlogService) {
|
||||
agentlogsViewGroup := parent.Group("/agentlogs")
|
||||
{
|
||||
agentlogsViewGroup.GET("/", svc.GetAgentlogs)
|
||||
}
|
||||
|
||||
flowAgentlogsViewGroup := parent.Group("/flows/:flowID/agentlogs")
|
||||
{
|
||||
flowAgentlogsViewGroup.GET("/", svc.GetFlowAgentlogs)
|
||||
}
|
||||
}
|
||||
|
||||
func setAssistantlogsGroup(parent *gin.RouterGroup, svc *services.AssistantlogService) {
|
||||
assistantlogsViewGroup := parent.Group("/assistantlogs")
|
||||
{
|
||||
assistantlogsViewGroup.GET("/", svc.GetAssistantlogs)
|
||||
}
|
||||
|
||||
flowAssistantlogsViewGroup := parent.Group("/flows/:flowID/assistantlogs")
|
||||
{
|
||||
flowAssistantlogsViewGroup.GET("/", svc.GetFlowAssistantlogs)
|
||||
}
|
||||
}
|
||||
|
||||
func setMsglogsGroup(parent *gin.RouterGroup, svc *services.MsglogService) {
|
||||
msglogsViewGroup := parent.Group("/msglogs")
|
||||
{
|
||||
msglogsViewGroup.GET("/", svc.GetMsglogs)
|
||||
}
|
||||
|
||||
flowMsglogsViewGroup := parent.Group("/flows/:flowID/msglogs")
|
||||
{
|
||||
flowMsglogsViewGroup.GET("/", svc.GetFlowMsglogs)
|
||||
}
|
||||
}
|
||||
|
||||
func setSearchlogsGroup(parent *gin.RouterGroup, svc *services.SearchlogService) {
|
||||
searchlogsViewGroup := parent.Group("/searchlogs")
|
||||
{
|
||||
searchlogsViewGroup.GET("/", svc.GetSearchlogs)
|
||||
}
|
||||
|
||||
flowSearchlogsViewGroup := parent.Group("/flows/:flowID/searchlogs")
|
||||
{
|
||||
flowSearchlogsViewGroup.GET("/", svc.GetFlowSearchlogs)
|
||||
}
|
||||
}
|
||||
|
||||
func setTermlogsGroup(parent *gin.RouterGroup, svc *services.TermlogService) {
|
||||
termlogsViewGroup := parent.Group("/termlogs")
|
||||
{
|
||||
termlogsViewGroup.GET("/", svc.GetTermlogs)
|
||||
}
|
||||
|
||||
flowTermlogsViewGroup := parent.Group("/flows/:flowID/termlogs")
|
||||
{
|
||||
flowTermlogsViewGroup.GET("/", svc.GetFlowTermlogs)
|
||||
}
|
||||
}
|
||||
|
||||
func setVecstorelogsGroup(parent *gin.RouterGroup, svc *services.VecstorelogService) {
|
||||
vecstorelogsViewGroup := parent.Group("/vecstorelogs")
|
||||
{
|
||||
vecstorelogsViewGroup.GET("/", svc.GetVecstorelogs)
|
||||
}
|
||||
|
||||
flowVecstorelogsViewGroup := parent.Group("/flows/:flowID/vecstorelogs")
|
||||
{
|
||||
flowVecstorelogsViewGroup.GET("/", svc.GetFlowVecstorelogs)
|
||||
}
|
||||
}
|
||||
|
||||
func setScreenshotsGroup(parent *gin.RouterGroup, svc *services.ScreenshotService) {
|
||||
screenshotsViewGroup := parent.Group("/screenshots")
|
||||
{
|
||||
screenshotsViewGroup.GET("/", svc.GetScreenshots)
|
||||
}
|
||||
|
||||
flowScreenshotsViewGroup := parent.Group("/flows/:flowID/screenshots")
|
||||
{
|
||||
flowScreenshotsViewGroup.GET("/", svc.GetFlowScreenshots)
|
||||
flowScreenshotsViewGroup.GET("/:screenshotID", svc.GetFlowScreenshot)
|
||||
flowScreenshotsViewGroup.GET("/:screenshotID/file", svc.GetFlowScreenshotFile)
|
||||
}
|
||||
}
|
||||
|
||||
func setAnonymizeGroup(parent *gin.RouterGroup, svc *services.AnonymizerService) {
|
||||
group := parent.Group("/anonymize")
|
||||
{
|
||||
group.POST("/text", svc.AnonymizeText)
|
||||
}
|
||||
}
|
||||
|
||||
func setPromptsGroup(parent *gin.RouterGroup, svc *services.PromptService) {
|
||||
promptsViewGroup := parent.Group("/prompts")
|
||||
{
|
||||
promptsViewGroup.GET("/", svc.GetPrompts)
|
||||
promptsViewGroup.GET("/:promptType", svc.GetPrompt)
|
||||
}
|
||||
|
||||
promptsEditGroup := parent.Group("/prompts")
|
||||
{
|
||||
promptsEditGroup.PUT("/:promptType", svc.PatchPrompt)
|
||||
promptsEditGroup.POST("/:promptType/default", svc.ResetPrompt)
|
||||
promptsEditGroup.DELETE("/:promptType", svc.DeletePrompt)
|
||||
}
|
||||
}
|
||||
|
||||
func setRolesGroup(parent *gin.RouterGroup, svc *services.RoleService) {
|
||||
rolesViewGroup := parent.Group("/roles")
|
||||
{
|
||||
rolesViewGroup.GET("/", svc.GetRoles)
|
||||
rolesViewGroup.GET("/:roleID", svc.GetRole)
|
||||
}
|
||||
}
|
||||
|
||||
func setUsersGroup(parent *gin.RouterGroup, svc *services.UserService) {
|
||||
usersCreateGroup := parent.Group("/users")
|
||||
{
|
||||
usersCreateGroup.POST("/", svc.CreateUser)
|
||||
}
|
||||
|
||||
usersDeleteGroup := parent.Group("/users")
|
||||
{
|
||||
usersDeleteGroup.DELETE("/:hash", svc.DeleteUser)
|
||||
}
|
||||
|
||||
usersEditGroup := parent.Group("/users")
|
||||
{
|
||||
usersEditGroup.PUT("/:hash", svc.PatchUser)
|
||||
}
|
||||
|
||||
usersViewGroup := parent.Group("/users")
|
||||
{
|
||||
usersViewGroup.GET("/", svc.GetUsers)
|
||||
usersViewGroup.GET("/:hash", svc.GetUser)
|
||||
}
|
||||
|
||||
userViewGroup := parent.Group("/user")
|
||||
{
|
||||
userViewGroup.GET("/", svc.GetCurrentUser)
|
||||
}
|
||||
}
|
||||
|
||||
func setAnalyticsGroup(parent *gin.RouterGroup, svc *services.AnalyticsService) {
|
||||
// System-wide analytics
|
||||
usageViewGroup := parent.Group("/usage")
|
||||
{
|
||||
usageViewGroup.GET("/", svc.GetSystemUsage)
|
||||
usageViewGroup.GET("/:period", svc.GetPeriodUsage)
|
||||
}
|
||||
|
||||
// Flow-specific analytics
|
||||
flowUsageViewGroup := parent.Group("/flows/:flowID/usage")
|
||||
{
|
||||
flowUsageViewGroup.GET("/", svc.GetFlowUsage)
|
||||
}
|
||||
}
|
||||
|
||||
func setTokensGroup(parent *gin.RouterGroup, svc *services.TokenService) {
|
||||
tokensGroup := parent.Group("/tokens")
|
||||
{
|
||||
tokensGroup.POST("/", svc.CreateToken)
|
||||
tokensGroup.GET("/", svc.ListTokens)
|
||||
tokensGroup.GET("/:tokenID", svc.GetToken)
|
||||
tokensGroup.PUT("/:tokenID", svc.UpdateToken)
|
||||
tokensGroup.DELETE("/:tokenID", svc.DeleteToken)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,220 @@
|
||||
package services
|
||||
|
||||
import (
|
||||
"errors"
|
||||
"net/http"
|
||||
"slices"
|
||||
"strconv"
|
||||
|
||||
"pentagi/pkg/server/logger"
|
||||
"pentagi/pkg/server/models"
|
||||
"pentagi/pkg/server/rdb"
|
||||
"pentagi/pkg/server/response"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/jinzhu/gorm"
|
||||
)
|
||||
|
||||
type agentlogs struct {
|
||||
AgentLogs []models.Agentlog `json:"agentlogs"`
|
||||
Total uint64 `json:"total"`
|
||||
}
|
||||
|
||||
type agentlogsGrouped struct {
|
||||
Grouped []string `json:"grouped"`
|
||||
Total uint64 `json:"total"`
|
||||
}
|
||||
|
||||
var agentlogsSQLMappers = map[string]any{
|
||||
"id": "{{table}}.id",
|
||||
"initiator": "{{table}}.initiator",
|
||||
"executor": "{{table}}.executor",
|
||||
"task": "{{table}}.task",
|
||||
"result": "{{table}}.result",
|
||||
"flow_id": "{{table}}.flow_id",
|
||||
"task_id": "{{table}}.task_id",
|
||||
"subtask_id": "{{table}}.subtask_id",
|
||||
"created_at": "{{table}}.created_at",
|
||||
"data": "({{table}}.task || ' ' || {{table}}.result)",
|
||||
}
|
||||
|
||||
type AgentlogService struct {
|
||||
db *gorm.DB
|
||||
}
|
||||
|
||||
func NewAgentlogService(db *gorm.DB) *AgentlogService {
|
||||
return &AgentlogService{
|
||||
db: db,
|
||||
}
|
||||
}
|
||||
|
||||
// GetAgentlogs is a function to return agentlogs list
|
||||
// @Summary Retrieve agentlogs list
|
||||
// @Tags Agentlogs
|
||||
// @Produce json
|
||||
// @Security BearerAuth
|
||||
// @Param request query rdb.TableQuery true "query table params"
|
||||
// @Success 200 {object} response.successResp{data=agentlogs} "agentlogs list received successful"
|
||||
// @Failure 400 {object} response.errorResp "invalid query request data"
|
||||
// @Failure 403 {object} response.errorResp "getting agentlogs not permitted"
|
||||
// @Failure 500 {object} response.errorResp "internal error on getting agentlogs"
|
||||
// @Router /agentlogs/ [get]
|
||||
func (s *AgentlogService) GetAgentlogs(c *gin.Context) {
|
||||
var (
|
||||
err error
|
||||
query rdb.TableQuery
|
||||
resp agentlogs
|
||||
)
|
||||
|
||||
if err = c.ShouldBindQuery(&query); err != nil {
|
||||
logger.FromContext(c).WithError(err).Errorf("error binding query")
|
||||
response.Error(c, response.ErrAgentlogsInvalidRequest, err)
|
||||
return
|
||||
}
|
||||
|
||||
uid := c.GetUint64("uid")
|
||||
privs := c.GetStringSlice("prm")
|
||||
var scope func(db *gorm.DB) *gorm.DB
|
||||
if slices.Contains(privs, "agentlogs.admin") {
|
||||
scope = func(db *gorm.DB) *gorm.DB {
|
||||
return db.
|
||||
Joins("INNER JOIN flows f ON f.id = flow_id")
|
||||
}
|
||||
} else if slices.Contains(privs, "agentlogs.view") {
|
||||
scope = func(db *gorm.DB) *gorm.DB {
|
||||
return db.
|
||||
Joins("INNER JOIN flows f ON f.id = flow_id").
|
||||
Where("f.user_id = ?", uid)
|
||||
}
|
||||
} else {
|
||||
logger.FromContext(c).Errorf("error filtering user role permissions: permission not found")
|
||||
response.Error(c, response.ErrNotPermitted, nil)
|
||||
return
|
||||
}
|
||||
|
||||
query.Init("agentlogs", agentlogsSQLMappers)
|
||||
|
||||
if query.Group != "" {
|
||||
if _, ok := agentlogsSQLMappers[query.Group]; !ok {
|
||||
logger.FromContext(c).Errorf("error finding agentlogs grouped: group field not found")
|
||||
response.Error(c, response.ErrAgentlogsInvalidRequest, errors.New("group field not found"))
|
||||
return
|
||||
}
|
||||
|
||||
var respGrouped agentlogsGrouped
|
||||
if respGrouped.Total, err = query.QueryGrouped(s.db, &respGrouped.Grouped, scope); err != nil {
|
||||
logger.FromContext(c).WithError(err).Errorf("error finding agentlogs grouped")
|
||||
response.Error(c, response.ErrInternal, err)
|
||||
return
|
||||
}
|
||||
|
||||
response.Success(c, http.StatusOK, respGrouped)
|
||||
return
|
||||
}
|
||||
|
||||
if resp.Total, err = query.Query(s.db, &resp.AgentLogs, scope); err != nil {
|
||||
logger.FromContext(c).WithError(err).Errorf("error finding agentlogs")
|
||||
response.Error(c, response.ErrInternal, err)
|
||||
return
|
||||
}
|
||||
|
||||
for i := 0; i < len(resp.AgentLogs); i++ {
|
||||
if err = resp.AgentLogs[i].Valid(); err != nil {
|
||||
logger.FromContext(c).WithError(err).Errorf("error validating agentlog data '%d'", resp.AgentLogs[i].ID)
|
||||
response.Error(c, response.ErrAgentlogsInvalidData, err)
|
||||
return
|
||||
}
|
||||
}
|
||||
|
||||
response.Success(c, http.StatusOK, resp)
|
||||
}
|
||||
|
||||
// GetFlowAgentlogs is a function to return agentlogs list by flow id
|
||||
// @Summary Retrieve agentlogs list by flow id
|
||||
// @Tags Agentlogs
|
||||
// @Produce json
|
||||
// @Security BearerAuth
|
||||
// @Param flowID path int true "flow id" minimum(0)
|
||||
// @Param request query rdb.TableQuery true "query table params"
|
||||
// @Success 200 {object} response.successResp{data=agentlogs} "agentlogs list received successful"
|
||||
// @Failure 400 {object} response.errorResp "invalid query request data"
|
||||
// @Failure 403 {object} response.errorResp "getting agentlogs not permitted"
|
||||
// @Failure 500 {object} response.errorResp "internal error on getting agentlogs"
|
||||
// @Router /flows/{flowID}/agentlogs/ [get]
|
||||
func (s *AgentlogService) GetFlowAgentlogs(c *gin.Context) {
|
||||
var (
|
||||
err error
|
||||
flowID uint64
|
||||
query rdb.TableQuery
|
||||
resp agentlogs
|
||||
)
|
||||
|
||||
if flowID, err = strconv.ParseUint(c.Param("flowID"), 10, 64); err != nil {
|
||||
logger.FromContext(c).WithError(err).Errorf("error parsing flow id")
|
||||
response.Error(c, response.ErrAgentlogsInvalidRequest, err)
|
||||
return
|
||||
}
|
||||
|
||||
if err = c.ShouldBindQuery(&query); err != nil {
|
||||
logger.FromContext(c).WithError(err).Errorf("error binding query")
|
||||
response.Error(c, response.ErrAgentlogsInvalidRequest, err)
|
||||
return
|
||||
}
|
||||
|
||||
uid := c.GetUint64("uid")
|
||||
privs := c.GetStringSlice("prm")
|
||||
var scope func(db *gorm.DB) *gorm.DB
|
||||
if slices.Contains(privs, "agentlogs.admin") {
|
||||
scope = func(db *gorm.DB) *gorm.DB {
|
||||
return db.
|
||||
Joins("INNER JOIN flows f ON f.id = flow_id").
|
||||
Where("f.id = ?", flowID)
|
||||
}
|
||||
} else if slices.Contains(privs, "agentlogs.view") {
|
||||
scope = func(db *gorm.DB) *gorm.DB {
|
||||
return db.
|
||||
Joins("INNER JOIN flows f ON f.id = flow_id").
|
||||
Where("f.id = ? AND f.user_id = ?", flowID, uid)
|
||||
}
|
||||
} else {
|
||||
logger.FromContext(c).Errorf("error filtering user role permissions: permission not found")
|
||||
response.Error(c, response.ErrNotPermitted, nil)
|
||||
return
|
||||
}
|
||||
|
||||
query.Init("agentlogs", agentlogsSQLMappers)
|
||||
|
||||
if query.Group != "" {
|
||||
if _, ok := agentlogsSQLMappers[query.Group]; !ok {
|
||||
logger.FromContext(c).Errorf("error finding agentlogs grouped: group field not found")
|
||||
response.Error(c, response.ErrAgentlogsInvalidRequest, errors.New("group field not found"))
|
||||
return
|
||||
}
|
||||
|
||||
var respGrouped agentlogsGrouped
|
||||
if respGrouped.Total, err = query.QueryGrouped(s.db, &respGrouped.Grouped, scope); err != nil {
|
||||
logger.FromContext(c).WithError(err).Errorf("error finding agentlogs grouped")
|
||||
response.Error(c, response.ErrInternal, err)
|
||||
return
|
||||
}
|
||||
|
||||
response.Success(c, http.StatusOK, respGrouped)
|
||||
return
|
||||
}
|
||||
|
||||
if resp.Total, err = query.Query(s.db, &resp.AgentLogs, scope); err != nil {
|
||||
logger.FromContext(c).WithError(err).Errorf("error finding agentlogs")
|
||||
response.Error(c, response.ErrInternal, err)
|
||||
return
|
||||
}
|
||||
|
||||
for i := 0; i < len(resp.AgentLogs); i++ {
|
||||
if err = resp.AgentLogs[i].Valid(); err != nil {
|
||||
logger.FromContext(c).WithError(err).Errorf("error validating agentlog data '%d'", resp.AgentLogs[i].ID)
|
||||
response.Error(c, response.ErrAgentlogsInvalidData, err)
|
||||
return
|
||||
}
|
||||
}
|
||||
|
||||
response.Success(c, http.StatusOK, resp)
|
||||
}
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,67 @@
|
||||
package services
|
||||
|
||||
import (
|
||||
"net/http"
|
||||
"slices"
|
||||
|
||||
"pentagi/pkg/server/logger"
|
||||
"pentagi/pkg/server/response"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/vxcontrol/cloud/anonymizer"
|
||||
)
|
||||
|
||||
type anonymizeTextRequest struct {
|
||||
Text string `json:"text" binding:"required"`
|
||||
}
|
||||
|
||||
type anonymizeTextResponse struct {
|
||||
Text string `json:"text"`
|
||||
}
|
||||
|
||||
// AnonymizerService handles text anonymization via the REST API.
|
||||
type AnonymizerService struct {
|
||||
replacer anonymizer.Replacer
|
||||
}
|
||||
|
||||
func NewAnonymizerService(replacer anonymizer.Replacer) *AnonymizerService {
|
||||
return &AnonymizerService{replacer: replacer}
|
||||
}
|
||||
|
||||
// AnonymizeText replaces sensitive data patterns in the provided text.
|
||||
// @Summary Anonymize text
|
||||
// @Description Replace sensitive data patterns (credentials, keys, PII, etc.) in the provided text with safe placeholders.
|
||||
// @Tags Anonymize
|
||||
// @Accept json
|
||||
// @Produce json
|
||||
// @Security BearerAuth
|
||||
// @Param json body anonymizeTextRequest true "Text to anonymize"
|
||||
// @Success 200 {object} response.successResp{data=anonymizeTextResponse} "anonymized text returned successfully"
|
||||
// @Failure 400 {object} response.errorResp "invalid request body"
|
||||
// @Failure 403 {object} response.errorResp "anonymize.call permission required"
|
||||
// @Failure 503 {object} response.errorResp "anonymizer is not configured"
|
||||
// @Router /anonymize/text [post]
|
||||
func (s *AnonymizerService) AnonymizeText(c *gin.Context) {
|
||||
privs := c.GetStringSlice("prm")
|
||||
if !slices.Contains(privs, "anonymize.call") {
|
||||
logger.FromContext(c).Errorf("error filtering user role permissions: permission not found")
|
||||
response.Error(c, response.ErrNotPermitted, nil)
|
||||
return
|
||||
}
|
||||
|
||||
if s.replacer == nil {
|
||||
response.Error(c, response.ErrAnonymizeUnavailable, nil)
|
||||
return
|
||||
}
|
||||
|
||||
var req anonymizeTextRequest
|
||||
if err := c.ShouldBindJSON(&req); err != nil {
|
||||
logger.FromContext(c).WithError(err).Error("error binding anonymize request")
|
||||
response.Error(c, response.ErrAnonymizeInvalidRequest, err)
|
||||
return
|
||||
}
|
||||
|
||||
response.Success(c, http.StatusOK, anonymizeTextResponse{
|
||||
Text: s.replacer.ReplaceString(req.Text),
|
||||
})
|
||||
}
|
||||
@@ -0,0 +1,424 @@
|
||||
package services
|
||||
|
||||
import (
|
||||
"errors"
|
||||
"net/http"
|
||||
"time"
|
||||
|
||||
"pentagi/pkg/database"
|
||||
"pentagi/pkg/graph/subscriptions"
|
||||
"pentagi/pkg/server/auth"
|
||||
"pentagi/pkg/server/logger"
|
||||
"pentagi/pkg/server/models"
|
||||
"pentagi/pkg/server/response"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/jinzhu/gorm"
|
||||
)
|
||||
|
||||
type tokens struct {
|
||||
Tokens []models.APIToken `json:"tokens"`
|
||||
Total uint64 `json:"total"`
|
||||
}
|
||||
|
||||
// TokenService handles API token management
|
||||
type TokenService struct {
|
||||
db *gorm.DB
|
||||
globalSalt string
|
||||
tokenCache *auth.TokenCache
|
||||
ss subscriptions.SubscriptionsController
|
||||
}
|
||||
|
||||
// NewTokenService creates a new TokenService instance
|
||||
func NewTokenService(
|
||||
db *gorm.DB,
|
||||
globalSalt string,
|
||||
tokenCache *auth.TokenCache,
|
||||
ss subscriptions.SubscriptionsController,
|
||||
) *TokenService {
|
||||
return &TokenService{
|
||||
db: db,
|
||||
globalSalt: globalSalt,
|
||||
tokenCache: tokenCache,
|
||||
ss: ss,
|
||||
}
|
||||
}
|
||||
|
||||
// CreateToken creates a new API token
|
||||
// @Summary Create new API token for automation
|
||||
// @Tags Tokens
|
||||
// @Accept json
|
||||
// @Produce json
|
||||
// @Param json body models.CreateAPITokenRequest true "Token creation request"
|
||||
// @Success 201 {object} response.successResp{data=models.APITokenWithSecret} "token created successful"
|
||||
// @Failure 400 {object} response.errorResp "invalid token request or default salt"
|
||||
// @Failure 403 {object} response.errorResp "creating token not permitted"
|
||||
// @Failure 500 {object} response.errorResp "internal error on creating token"
|
||||
// @Router /tokens [post]
|
||||
func (s *TokenService) CreateToken(c *gin.Context) {
|
||||
// check for default salt
|
||||
if s.globalSalt == "" || s.globalSalt == "salt" {
|
||||
logger.FromContext(c).Errorf("token creation attempted with default salt")
|
||||
response.Error(c, response.ErrTokenCreationDisabled, errors.New("token creation is disabled with default salt"))
|
||||
return
|
||||
}
|
||||
|
||||
uid := c.GetUint64("uid")
|
||||
rid := c.GetUint64("rid")
|
||||
uhash := c.GetString("uhash")
|
||||
|
||||
var req models.CreateAPITokenRequest
|
||||
if err := c.ShouldBindJSON(&req); err != nil {
|
||||
logger.FromContext(c).WithError(err).Errorf("error binding JSON")
|
||||
response.Error(c, response.ErrTokenInvalidRequest, err)
|
||||
return
|
||||
}
|
||||
if err := req.Valid(); err != nil {
|
||||
logger.FromContext(c).WithError(err).Errorf("error validating JSON")
|
||||
response.Error(c, response.ErrTokenInvalidRequest, err)
|
||||
return
|
||||
}
|
||||
|
||||
// check if name is unique for this user (if provided)
|
||||
if req.Name != nil && *req.Name != "" {
|
||||
var existing models.APIToken
|
||||
err := s.db.
|
||||
Where("user_id = ? AND name = ? AND deleted_at IS NULL", uid, *req.Name).
|
||||
First(&existing).
|
||||
Error
|
||||
if err == nil {
|
||||
logger.FromContext(c).Errorf("token with name '%s' already exists for user %d", *req.Name, uid)
|
||||
response.Error(c, response.ErrTokenInvalidRequest, errors.New("token with this name already exists"))
|
||||
return
|
||||
}
|
||||
}
|
||||
|
||||
// generate token_id
|
||||
tokenID, err := auth.GenerateTokenID()
|
||||
if err != nil {
|
||||
logger.FromContext(c).WithError(err).Errorf("error generating token ID")
|
||||
response.Error(c, response.ErrInternal, err)
|
||||
return
|
||||
}
|
||||
|
||||
// create JWT claims
|
||||
claims := auth.MakeAPITokenClaims(tokenID, uhash, uid, rid, req.TTL)
|
||||
|
||||
// sign token
|
||||
token, err := auth.MakeAPIToken(s.globalSalt, claims)
|
||||
if err != nil {
|
||||
logger.FromContext(c).WithError(err).Errorf("error signing token")
|
||||
response.Error(c, response.ErrInternal, err)
|
||||
return
|
||||
}
|
||||
|
||||
// save to database
|
||||
apiToken := models.APIToken{
|
||||
TokenID: tokenID,
|
||||
UserID: uid,
|
||||
RoleID: rid,
|
||||
Name: req.Name,
|
||||
TTL: req.TTL,
|
||||
Status: models.TokenStatusActive,
|
||||
}
|
||||
|
||||
if err := s.db.Create(&apiToken).Error; err != nil {
|
||||
logger.FromContext(c).WithError(err).Errorf("error creating token in database")
|
||||
response.Error(c, response.ErrInternal, err)
|
||||
return
|
||||
}
|
||||
|
||||
result := models.APITokenWithSecret{
|
||||
APIToken: apiToken,
|
||||
Token: token,
|
||||
}
|
||||
|
||||
// invalidate cache for negative caching results
|
||||
s.tokenCache.Invalidate(apiToken.TokenID)
|
||||
s.tokenCache.InvalidateUser(apiToken.UserID)
|
||||
|
||||
if s.ss != nil {
|
||||
publisher := s.ss.NewAPITokenPublisher(int64(apiToken.UserID))
|
||||
publisher.APITokenCreated(c, database.APITokenWithSecret{
|
||||
ApiToken: convertAPITokenToDatabase(apiToken),
|
||||
Token: token,
|
||||
})
|
||||
}
|
||||
|
||||
response.Success(c, http.StatusCreated, result)
|
||||
}
|
||||
|
||||
// ListTokens returns a list of tokens (user sees only their own, admin sees all)
|
||||
// @Summary List API tokens
|
||||
// @Tags Tokens
|
||||
// @Produce json
|
||||
// @Success 200 {object} response.successResp{data=tokens} "tokens retrieved successful"
|
||||
// @Failure 403 {object} response.errorResp "listing tokens not permitted"
|
||||
// @Failure 500 {object} response.errorResp "internal error on listing tokens"
|
||||
// @Router /tokens [get]
|
||||
func (s *TokenService) ListTokens(c *gin.Context) {
|
||||
uid := c.GetUint64("uid")
|
||||
prms := c.GetStringSlice("prm")
|
||||
|
||||
query := s.db.Where("deleted_at IS NULL")
|
||||
|
||||
// check if user has admin privilege
|
||||
hasAdmin := auth.LookupPerm(prms, "settings.tokens.admin")
|
||||
if !hasAdmin {
|
||||
// regular user sees only their own tokens
|
||||
query = query.Where("user_id = ?", uid)
|
||||
}
|
||||
|
||||
var tokenList []models.APIToken
|
||||
var total uint64
|
||||
|
||||
if err := query.Order("created_at DESC").Find(&tokenList).Count(&total).Error; err != nil {
|
||||
logger.FromContext(c).WithError(err).Errorf("error finding tokens")
|
||||
response.Error(c, response.ErrInternal, err)
|
||||
return
|
||||
}
|
||||
|
||||
for i := range tokenList {
|
||||
token := &tokenList[i]
|
||||
isExpired := token.CreatedAt.Add(time.Duration(token.TTL) * time.Second).Before(time.Now())
|
||||
if token.Status == models.TokenStatusActive && isExpired {
|
||||
token.Status = models.TokenStatusExpired
|
||||
}
|
||||
}
|
||||
|
||||
result := tokens{
|
||||
Tokens: tokenList,
|
||||
Total: total,
|
||||
}
|
||||
|
||||
response.Success(c, http.StatusOK, result)
|
||||
}
|
||||
|
||||
// GetToken returns information about a specific token
|
||||
// @Summary Get API token details
|
||||
// @Tags Tokens
|
||||
// @Produce json
|
||||
// @Param tokenID path string true "Token ID"
|
||||
// @Success 200 {object} response.successResp{data=models.APIToken} "token retrieved successful"
|
||||
// @Failure 403 {object} response.errorResp "accessing token not permitted"
|
||||
// @Failure 404 {object} response.errorResp "token not found"
|
||||
// @Failure 500 {object} response.errorResp "internal error on getting token"
|
||||
// @Router /tokens/{tokenID} [get]
|
||||
func (s *TokenService) GetToken(c *gin.Context) {
|
||||
uid := c.GetUint64("uid")
|
||||
prms := c.GetStringSlice("prm")
|
||||
tokenID := c.Param("tokenID")
|
||||
|
||||
var token models.APIToken
|
||||
if err := s.db.Where("token_id = ? AND deleted_at IS NULL", tokenID).First(&token).Error; err != nil {
|
||||
logger.FromContext(c).WithError(err).Errorf("error finding token")
|
||||
if errors.Is(err, gorm.ErrRecordNotFound) {
|
||||
response.Error(c, response.ErrTokenNotFound, err)
|
||||
} else {
|
||||
response.Error(c, response.ErrInternal, err)
|
||||
}
|
||||
return
|
||||
}
|
||||
|
||||
// check authorization
|
||||
hasAdmin := auth.LookupPerm(prms, "settings.tokens.admin")
|
||||
if !hasAdmin && token.UserID != uid {
|
||||
logger.FromContext(c).Errorf("user %d attempted to access token of user %d", uid, token.UserID)
|
||||
response.Error(c, response.ErrTokenUnauthorized, errors.New("not authorized to access this token"))
|
||||
return
|
||||
}
|
||||
|
||||
isExpired := token.CreatedAt.Add(time.Duration(token.TTL) * time.Second).Before(time.Now())
|
||||
if token.Status == models.TokenStatusActive && isExpired {
|
||||
token.Status = models.TokenStatusExpired
|
||||
}
|
||||
|
||||
if err := token.Valid(); err != nil {
|
||||
logger.FromContext(c).WithError(err).Errorf("error validating token data")
|
||||
response.Error(c, response.ErrTokenInvalidData, err)
|
||||
return
|
||||
}
|
||||
|
||||
response.Success(c, http.StatusOK, token)
|
||||
}
|
||||
|
||||
// UpdateToken updates name and/or status of a token
|
||||
// @Summary Update API token
|
||||
// @Tags Tokens
|
||||
// @Accept json
|
||||
// @Produce json
|
||||
// @Param tokenID path string true "Token ID"
|
||||
// @Param json body models.UpdateAPITokenRequest true "Token update request"
|
||||
// @Success 200 {object} response.successResp{data=models.APIToken} "token updated successful"
|
||||
// @Failure 400 {object} response.errorResp "invalid update request"
|
||||
// @Failure 403 {object} response.errorResp "updating token not permitted"
|
||||
// @Failure 404 {object} response.errorResp "token not found"
|
||||
// @Failure 500 {object} response.errorResp "internal error on updating token"
|
||||
// @Router /tokens/{tokenID} [put]
|
||||
func (s *TokenService) UpdateToken(c *gin.Context) {
|
||||
uid := c.GetUint64("uid")
|
||||
prms := c.GetStringSlice("prm")
|
||||
tokenID := c.Param("tokenID")
|
||||
|
||||
var req models.UpdateAPITokenRequest
|
||||
if err := c.ShouldBindJSON(&req); err != nil {
|
||||
logger.FromContext(c).WithError(err).Errorf("error binding JSON")
|
||||
response.Error(c, response.ErrTokenInvalidRequest, err)
|
||||
return
|
||||
}
|
||||
if err := req.Valid(); err != nil {
|
||||
logger.FromContext(c).WithError(err).Errorf("error validating JSON")
|
||||
response.Error(c, response.ErrTokenInvalidRequest, err)
|
||||
return
|
||||
}
|
||||
|
||||
var token models.APIToken
|
||||
if err := s.db.Where("token_id = ? AND deleted_at IS NULL", tokenID).First(&token).Error; err != nil {
|
||||
logger.FromContext(c).WithError(err).Errorf("error finding token")
|
||||
if errors.Is(err, gorm.ErrRecordNotFound) {
|
||||
response.Error(c, response.ErrTokenNotFound, err)
|
||||
} else {
|
||||
response.Error(c, response.ErrInternal, err)
|
||||
}
|
||||
return
|
||||
}
|
||||
|
||||
// check authorization
|
||||
hasAdmin := auth.LookupPerm(prms, "settings.tokens.admin")
|
||||
if !hasAdmin && token.UserID != uid {
|
||||
logger.FromContext(c).Errorf("user %d attempted to update token of user %d", uid, token.UserID)
|
||||
response.Error(c, response.ErrTokenUnauthorized, errors.New("not authorized to update this token"))
|
||||
return
|
||||
}
|
||||
|
||||
// update fields
|
||||
updates := make(map[string]any)
|
||||
if req.Name != nil {
|
||||
// check uniqueness if name is changing
|
||||
if token.Name == nil || *token.Name != *req.Name {
|
||||
if *req.Name != "" {
|
||||
var existing models.APIToken
|
||||
err := s.db.
|
||||
Where("user_id = ? AND name = ? AND token_id != ? AND deleted_at IS NULL", token.UserID, *req.Name, tokenID).
|
||||
First(&existing).
|
||||
Error
|
||||
if err == nil {
|
||||
logger.FromContext(c).Errorf("token with name '%s' already exists for user %d", *req.Name, token.UserID)
|
||||
response.Error(c, response.ErrTokenInvalidRequest, errors.New("token with this name already exists"))
|
||||
return
|
||||
}
|
||||
}
|
||||
}
|
||||
updates["name"] = req.Name
|
||||
}
|
||||
switch req.Status {
|
||||
case models.TokenStatusActive:
|
||||
updates["status"] = models.TokenStatusActive
|
||||
case models.TokenStatusRevoked:
|
||||
updates["status"] = models.TokenStatusRevoked
|
||||
case models.TokenStatusExpired:
|
||||
updates["status"] = models.TokenStatusRevoked
|
||||
}
|
||||
|
||||
if len(updates) > 0 {
|
||||
if err := s.db.Model(&token).Updates(updates).Error; err != nil {
|
||||
logger.FromContext(c).WithError(err).Errorf("error updating token")
|
||||
response.Error(c, response.ErrInternal, err)
|
||||
return
|
||||
}
|
||||
|
||||
// invalidate cache if status changed
|
||||
if req.Status != "" {
|
||||
s.tokenCache.Invalidate(tokenID)
|
||||
// also invalidate all tokens for this user (in case of role change or security event)
|
||||
s.tokenCache.InvalidateUser(token.UserID)
|
||||
}
|
||||
|
||||
// reload token
|
||||
if err := s.db.Where("token_id = ?", tokenID).First(&token).Error; err != nil {
|
||||
logger.FromContext(c).WithError(err).Errorf("error reloading token")
|
||||
response.Error(c, response.ErrInternal, err)
|
||||
return
|
||||
}
|
||||
}
|
||||
|
||||
isExpired := token.CreatedAt.Add(time.Duration(token.TTL) * time.Second).Before(time.Now())
|
||||
if token.Status == models.TokenStatusActive && isExpired {
|
||||
token.Status = models.TokenStatusExpired
|
||||
}
|
||||
|
||||
if s.ss != nil {
|
||||
publisher := s.ss.NewAPITokenPublisher(int64(token.UserID))
|
||||
publisher.APITokenUpdated(c, convertAPITokenToDatabase(token))
|
||||
}
|
||||
|
||||
response.Success(c, http.StatusOK, token)
|
||||
}
|
||||
|
||||
// DeleteToken performs soft delete of a token
|
||||
// @Summary Delete API token
|
||||
// @Tags Tokens
|
||||
// @Produce json
|
||||
// @Param tokenID path string true "Token ID"
|
||||
// @Success 200 {object} response.successResp "token deleted successful"
|
||||
// @Failure 403 {object} response.errorResp "deleting token not permitted"
|
||||
// @Failure 404 {object} response.errorResp "token not found"
|
||||
// @Failure 500 {object} response.errorResp "internal error on deleting token"
|
||||
// @Router /tokens/{tokenID} [delete]
|
||||
func (s *TokenService) DeleteToken(c *gin.Context) {
|
||||
uid := c.GetUint64("uid")
|
||||
prms := c.GetStringSlice("prm")
|
||||
tokenID := c.Param("tokenID")
|
||||
|
||||
var token models.APIToken
|
||||
if err := s.db.Where("token_id = ? AND deleted_at IS NULL", tokenID).First(&token).Error; err != nil {
|
||||
logger.FromContext(c).WithError(err).Errorf("error finding token")
|
||||
if errors.Is(err, gorm.ErrRecordNotFound) {
|
||||
response.Error(c, response.ErrTokenNotFound, err)
|
||||
} else {
|
||||
response.Error(c, response.ErrInternal, err)
|
||||
}
|
||||
return
|
||||
}
|
||||
|
||||
// check authorization
|
||||
hasAdmin := auth.LookupPerm(prms, "settings.tokens.admin")
|
||||
if !hasAdmin && token.UserID != uid {
|
||||
logger.FromContext(c).Errorf("user %d attempted to delete token of user %d", uid, token.UserID)
|
||||
response.Error(c, response.ErrTokenUnauthorized, errors.New("not authorized to delete this token"))
|
||||
return
|
||||
}
|
||||
|
||||
// soft delete
|
||||
if err := s.db.Delete(&token).Error; err != nil {
|
||||
logger.FromContext(c).WithError(err).Errorf("error deleting token")
|
||||
response.Error(c, response.ErrInternal, err)
|
||||
return
|
||||
}
|
||||
|
||||
// invalidate cache for this token and all user's tokens
|
||||
s.tokenCache.Invalidate(tokenID)
|
||||
s.tokenCache.InvalidateUser(token.UserID)
|
||||
|
||||
if s.ss != nil {
|
||||
publisher := s.ss.NewAPITokenPublisher(int64(token.UserID))
|
||||
publisher.APITokenDeleted(c, convertAPITokenToDatabase(token))
|
||||
}
|
||||
|
||||
response.Success(c, http.StatusOK, gin.H{"message": "token deleted successfully"})
|
||||
}
|
||||
|
||||
func convertAPITokenToDatabase(apiToken models.APIToken) database.ApiToken {
|
||||
return database.ApiToken{
|
||||
ID: int64(apiToken.ID),
|
||||
TokenID: apiToken.TokenID,
|
||||
UserID: int64(apiToken.UserID),
|
||||
RoleID: int64(apiToken.RoleID),
|
||||
Name: database.StringToNullString(*apiToken.Name),
|
||||
Ttl: int64(apiToken.TTL),
|
||||
Status: database.TokenStatus(apiToken.Status),
|
||||
CreatedAt: database.TimeToNullTime(apiToken.CreatedAt),
|
||||
UpdatedAt: database.TimeToNullTime(apiToken.UpdatedAt),
|
||||
DeletedAt: database.PtrTimeToNullTime(apiToken.DeletedAt),
|
||||
}
|
||||
}
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,219 @@
|
||||
package services
|
||||
|
||||
import (
|
||||
"errors"
|
||||
"net/http"
|
||||
"slices"
|
||||
"strconv"
|
||||
|
||||
"pentagi/pkg/server/logger"
|
||||
"pentagi/pkg/server/models"
|
||||
"pentagi/pkg/server/rdb"
|
||||
"pentagi/pkg/server/response"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/jinzhu/gorm"
|
||||
)
|
||||
|
||||
type assistantlogs struct {
|
||||
AssistantLogs []models.Assistantlog `json:"assistantlogs"`
|
||||
Total uint64 `json:"total"`
|
||||
}
|
||||
|
||||
type assistantlogsGrouped struct {
|
||||
Grouped []string `json:"grouped"`
|
||||
Total uint64 `json:"total"`
|
||||
}
|
||||
|
||||
var assistantlogsSQLMappers = map[string]any{
|
||||
"id": "{{table}}.id",
|
||||
"type": "{{table}}.type",
|
||||
"message": "{{table}}.message",
|
||||
"result": "{{table}}.result",
|
||||
"result_format": "{{table}}.result_format",
|
||||
"flow_id": "{{table}}.flow_id",
|
||||
"assistant_id": "{{table}}.assistant_id",
|
||||
"created_at": "{{table}}.created_at",
|
||||
"data": "({{table}}.type || ' ' || {{table}}.message || ' ' || {{table}}.result)",
|
||||
}
|
||||
|
||||
type AssistantlogService struct {
|
||||
db *gorm.DB
|
||||
}
|
||||
|
||||
func NewAssistantlogService(db *gorm.DB) *AssistantlogService {
|
||||
return &AssistantlogService{
|
||||
db: db,
|
||||
}
|
||||
}
|
||||
|
||||
// GetAssistantlogs is a function to return assistantlogs list
|
||||
// @Summary Retrieve assistantlogs list
|
||||
// @Tags Assistantlogs
|
||||
// @Produce json
|
||||
// @Security BearerAuth
|
||||
// @Param request query rdb.TableQuery true "query table params"
|
||||
// @Success 200 {object} response.successResp{data=assistantlogs} "assistantlogs list received successful"
|
||||
// @Failure 400 {object} response.errorResp "invalid query request data"
|
||||
// @Failure 403 {object} response.errorResp "getting assistantlogs not permitted"
|
||||
// @Failure 500 {object} response.errorResp "internal error on getting assistantlogs"
|
||||
// @Router /assistantlogs/ [get]
|
||||
func (s *AssistantlogService) GetAssistantlogs(c *gin.Context) {
|
||||
var (
|
||||
err error
|
||||
query rdb.TableQuery
|
||||
resp assistantlogs
|
||||
)
|
||||
|
||||
if err = c.ShouldBindQuery(&query); err != nil {
|
||||
logger.FromContext(c).WithError(err).Errorf("error binding query")
|
||||
response.Error(c, response.ErrAssistantlogsInvalidRequest, err)
|
||||
return
|
||||
}
|
||||
|
||||
uid := c.GetUint64("uid")
|
||||
privs := c.GetStringSlice("prm")
|
||||
var scope func(db *gorm.DB) *gorm.DB
|
||||
if slices.Contains(privs, "assistantlogs.admin") {
|
||||
scope = func(db *gorm.DB) *gorm.DB {
|
||||
return db.
|
||||
Joins("INNER JOIN flows f ON f.id = flow_id")
|
||||
}
|
||||
} else if slices.Contains(privs, "assistantlogs.view") {
|
||||
scope = func(db *gorm.DB) *gorm.DB {
|
||||
return db.
|
||||
Joins("INNER JOIN flows f ON f.id = flow_id").
|
||||
Where("f.user_id = ?", uid)
|
||||
}
|
||||
} else {
|
||||
logger.FromContext(c).Errorf("error filtering user role permissions: permission not found")
|
||||
response.Error(c, response.ErrNotPermitted, nil)
|
||||
return
|
||||
}
|
||||
|
||||
query.Init("assistantlogs", assistantlogsSQLMappers)
|
||||
|
||||
if query.Group != "" {
|
||||
if _, ok := assistantlogsSQLMappers[query.Group]; !ok {
|
||||
logger.FromContext(c).Errorf("error finding assistantlogs grouped: group field not found")
|
||||
response.Error(c, response.ErrAssistantlogsInvalidRequest, errors.New("group field not found"))
|
||||
return
|
||||
}
|
||||
|
||||
var respGrouped assistantlogsGrouped
|
||||
if respGrouped.Total, err = query.QueryGrouped(s.db, &respGrouped.Grouped, scope); err != nil {
|
||||
logger.FromContext(c).WithError(err).Errorf("error finding assistantlogs grouped")
|
||||
response.Error(c, response.ErrInternal, err)
|
||||
return
|
||||
}
|
||||
|
||||
response.Success(c, http.StatusOK, respGrouped)
|
||||
return
|
||||
}
|
||||
|
||||
if resp.Total, err = query.Query(s.db, &resp.AssistantLogs, scope); err != nil {
|
||||
logger.FromContext(c).WithError(err).Errorf("error finding assistantlogs")
|
||||
response.Error(c, response.ErrInternal, err)
|
||||
return
|
||||
}
|
||||
|
||||
for i := 0; i < len(resp.AssistantLogs); i++ {
|
||||
if err = resp.AssistantLogs[i].Valid(); err != nil {
|
||||
logger.FromContext(c).WithError(err).Errorf("error validating assistantlog data '%d'", resp.AssistantLogs[i].ID)
|
||||
response.Error(c, response.ErrAssistantlogsInvalidData, err)
|
||||
return
|
||||
}
|
||||
}
|
||||
|
||||
response.Success(c, http.StatusOK, resp)
|
||||
}
|
||||
|
||||
// GetFlowAssistantlogs is a function to return assistantlogs list by flow id
|
||||
// @Summary Retrieve assistantlogs list by flow id
|
||||
// @Tags Assistantlogs
|
||||
// @Produce json
|
||||
// @Security BearerAuth
|
||||
// @Param flowID path int true "flow id" minimum(0)
|
||||
// @Param request query rdb.TableQuery true "query table params"
|
||||
// @Success 200 {object} response.successResp{data=assistantlogs} "assistantlogs list received successful"
|
||||
// @Failure 400 {object} response.errorResp "invalid query request data"
|
||||
// @Failure 403 {object} response.errorResp "getting assistantlogs not permitted"
|
||||
// @Failure 500 {object} response.errorResp "internal error on getting assistantlogs"
|
||||
// @Router /flows/{flowID}/assistantlogs/ [get]
|
||||
func (s *AssistantlogService) GetFlowAssistantlogs(c *gin.Context) {
|
||||
var (
|
||||
err error
|
||||
flowID uint64
|
||||
query rdb.TableQuery
|
||||
resp assistantlogs
|
||||
)
|
||||
|
||||
if flowID, err = strconv.ParseUint(c.Param("flowID"), 10, 64); err != nil {
|
||||
logger.FromContext(c).WithError(err).Errorf("error parsing flow id")
|
||||
response.Error(c, response.ErrAssistantlogsInvalidRequest, err)
|
||||
return
|
||||
}
|
||||
|
||||
if err = c.ShouldBindQuery(&query); err != nil {
|
||||
logger.FromContext(c).WithError(err).Errorf("error binding query")
|
||||
response.Error(c, response.ErrAssistantlogsInvalidRequest, err)
|
||||
return
|
||||
}
|
||||
|
||||
uid := c.GetUint64("uid")
|
||||
privs := c.GetStringSlice("prm")
|
||||
var scope func(db *gorm.DB) *gorm.DB
|
||||
if slices.Contains(privs, "assistantlogs.admin") {
|
||||
scope = func(db *gorm.DB) *gorm.DB {
|
||||
return db.
|
||||
Joins("INNER JOIN flows f ON f.id = flow_id").
|
||||
Where("f.id = ?", flowID)
|
||||
}
|
||||
} else if slices.Contains(privs, "assistantlogs.view") {
|
||||
scope = func(db *gorm.DB) *gorm.DB {
|
||||
return db.
|
||||
Joins("INNER JOIN flows f ON f.id = flow_id").
|
||||
Where("f.id = ? AND f.user_id = ?", flowID, uid)
|
||||
}
|
||||
} else {
|
||||
logger.FromContext(c).Errorf("error filtering user role permissions: permission not found")
|
||||
response.Error(c, response.ErrNotPermitted, nil)
|
||||
return
|
||||
}
|
||||
|
||||
query.Init("assistantlogs", assistantlogsSQLMappers)
|
||||
|
||||
if query.Group != "" {
|
||||
if _, ok := assistantlogsSQLMappers[query.Group]; !ok {
|
||||
logger.FromContext(c).Errorf("error finding assistantlogs grouped: group field not found")
|
||||
response.Error(c, response.ErrAssistantlogsInvalidRequest, errors.New("group field not found"))
|
||||
return
|
||||
}
|
||||
|
||||
var respGrouped assistantlogsGrouped
|
||||
if respGrouped.Total, err = query.QueryGrouped(s.db, &respGrouped.Grouped, scope); err != nil {
|
||||
logger.FromContext(c).WithError(err).Errorf("error finding assistantlogs grouped")
|
||||
response.Error(c, response.ErrInternal, err)
|
||||
return
|
||||
}
|
||||
|
||||
response.Success(c, http.StatusOK, respGrouped)
|
||||
return
|
||||
}
|
||||
|
||||
if resp.Total, err = query.Query(s.db, &resp.AssistantLogs, scope); err != nil {
|
||||
logger.FromContext(c).WithError(err).Errorf("error finding assistantlogs")
|
||||
response.Error(c, response.ErrInternal, err)
|
||||
return
|
||||
}
|
||||
|
||||
for i := 0; i < len(resp.AssistantLogs); i++ {
|
||||
if err = resp.AssistantLogs[i].Valid(); err != nil {
|
||||
logger.FromContext(c).WithError(err).Errorf("error validating assistantlog data '%d'", resp.AssistantLogs[i].ID)
|
||||
response.Error(c, response.ErrAssistantlogsInvalidData, err)
|
||||
return
|
||||
}
|
||||
}
|
||||
|
||||
response.Success(c, http.StatusOK, resp)
|
||||
}
|
||||
@@ -0,0 +1,684 @@
|
||||
package services
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"fmt"
|
||||
"net/http"
|
||||
"slices"
|
||||
"strconv"
|
||||
|
||||
"pentagi/pkg/controller"
|
||||
"pentagi/pkg/database"
|
||||
"pentagi/pkg/graph/subscriptions"
|
||||
"pentagi/pkg/providers"
|
||||
"pentagi/pkg/providers/provider"
|
||||
"pentagi/pkg/server/logger"
|
||||
"pentagi/pkg/server/models"
|
||||
"pentagi/pkg/server/rdb"
|
||||
"pentagi/pkg/server/response"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/jinzhu/gorm"
|
||||
)
|
||||
|
||||
type assistants struct {
|
||||
Assistants []models.Assistant `json:"assistants"`
|
||||
Total uint64 `json:"total"`
|
||||
}
|
||||
|
||||
type assistantsGrouped struct {
|
||||
Grouped []string `json:"grouped"`
|
||||
Total uint64 `json:"total"`
|
||||
}
|
||||
|
||||
var assistantsSQLMappers = map[string]any{
|
||||
"id": "{{table}}.id",
|
||||
"status": "{{table}}.status",
|
||||
"title": "{{table}}.title",
|
||||
"model": "{{table}}.model",
|
||||
"model_provider_name": "{{table}}.model_provider_name",
|
||||
"model_provider_type": "{{table}}.model_provider_type",
|
||||
"language": "{{table}}.language",
|
||||
"flow_id": "{{table}}.flow_id",
|
||||
"msgchain_id": "{{table}}.msgchain_id",
|
||||
"created_at": "{{table}}.created_at",
|
||||
"updated_at": "{{table}}.updated_at",
|
||||
"data": "({{table}}.status || ' ' || {{table}}.title || ' ' || {{table}}.flow_id)",
|
||||
}
|
||||
|
||||
type AssistantService struct {
|
||||
db *gorm.DB
|
||||
pc providers.ProviderController
|
||||
fc controller.FlowController
|
||||
ss subscriptions.SubscriptionsController
|
||||
}
|
||||
|
||||
func NewAssistantService(
|
||||
db *gorm.DB,
|
||||
pc providers.ProviderController,
|
||||
fc controller.FlowController,
|
||||
ss subscriptions.SubscriptionsController,
|
||||
) *AssistantService {
|
||||
return &AssistantService{
|
||||
db: db,
|
||||
pc: pc,
|
||||
fc: fc,
|
||||
ss: ss,
|
||||
}
|
||||
}
|
||||
|
||||
// GetAssistants is a function to return assistants list
|
||||
// @Summary Retrieve assistants list
|
||||
// @Tags Assistants
|
||||
// @Produce json
|
||||
// @Security BearerAuth
|
||||
// @Param flowID path int true "flow id" minimum(0)
|
||||
// @Param request query rdb.TableQuery true "query table params"
|
||||
// @Success 200 {object} response.successResp{data=assistants} "assistants list received successful"
|
||||
// @Failure 400 {object} response.errorResp "invalid query request data"
|
||||
// @Failure 403 {object} response.errorResp "getting assistants not permitted"
|
||||
// @Failure 500 {object} response.errorResp "internal error on getting assistants"
|
||||
// @Router /flows/{flowID}/assistants/ [get]
|
||||
func (s *AssistantService) GetFlowAssistants(c *gin.Context) {
|
||||
var (
|
||||
err error
|
||||
flowID uint64
|
||||
query rdb.TableQuery
|
||||
resp assistants
|
||||
)
|
||||
|
||||
if err = c.ShouldBindQuery(&query); err != nil {
|
||||
logger.FromContext(c).WithError(err).Errorf("error binding query")
|
||||
response.Error(c, response.ErrAssistantsInvalidRequest, err)
|
||||
return
|
||||
}
|
||||
|
||||
if flowID, err = strconv.ParseUint(c.Param("flowID"), 10, 64); err != nil {
|
||||
logger.FromContext(c).WithError(err).Errorf("error parsing flow id")
|
||||
response.Error(c, response.ErrAssistantsInvalidRequest, err)
|
||||
return
|
||||
}
|
||||
|
||||
uid := c.GetUint64("uid")
|
||||
privs := c.GetStringSlice("prm")
|
||||
var scope func(db *gorm.DB) *gorm.DB
|
||||
if slices.Contains(privs, "assistants.admin") {
|
||||
scope = func(db *gorm.DB) *gorm.DB {
|
||||
return db.
|
||||
Joins("INNER JOIN flows f ON f.id = assistants.flow_id").
|
||||
Where("f.id = ?", flowID)
|
||||
}
|
||||
} else if slices.Contains(privs, "assistants.view") {
|
||||
scope = func(db *gorm.DB) *gorm.DB {
|
||||
return db.
|
||||
Joins("INNER JOIN flows f ON f.id = assistants.flow_id").
|
||||
Where("f.id = ? AND f.user_id = ?", flowID, uid)
|
||||
}
|
||||
} else {
|
||||
logger.FromContext(c).Errorf("error filtering user role permissions: permission not found")
|
||||
response.Error(c, response.ErrNotPermitted, nil)
|
||||
return
|
||||
}
|
||||
|
||||
query.Init("assistants", assistantsSQLMappers)
|
||||
|
||||
if query.Group != "" {
|
||||
if _, ok := assistantsSQLMappers[query.Group]; !ok {
|
||||
logger.FromContext(c).Errorf("error finding assistants grouped: group field not found")
|
||||
response.Error(c, response.ErrAssistantsInvalidRequest, errors.New("group field not found"))
|
||||
return
|
||||
}
|
||||
|
||||
var respGrouped assistantsGrouped
|
||||
if respGrouped.Total, err = query.QueryGrouped(s.db, &respGrouped.Grouped, scope); err != nil {
|
||||
logger.FromContext(c).WithError(err).Errorf("error finding assistants grouped")
|
||||
response.Error(c, response.ErrInternal, err)
|
||||
return
|
||||
}
|
||||
|
||||
response.Success(c, http.StatusOK, respGrouped)
|
||||
return
|
||||
}
|
||||
|
||||
if resp.Total, err = query.Query(s.db, &resp.Assistants, scope); err != nil {
|
||||
logger.FromContext(c).WithError(err).Errorf("error finding assistants")
|
||||
response.Error(c, response.ErrInternal, err)
|
||||
return
|
||||
}
|
||||
|
||||
for i := 0; i < len(resp.Assistants); i++ {
|
||||
if err = resp.Assistants[i].Valid(); err != nil {
|
||||
logger.FromContext(c).WithError(err).Errorf("error validating assistant data '%d'", resp.Assistants[i].ID)
|
||||
response.Error(c, response.ErrAssistantsInvalidData, err)
|
||||
return
|
||||
}
|
||||
}
|
||||
|
||||
response.Success(c, http.StatusOK, resp)
|
||||
}
|
||||
|
||||
// GetFlowAssistant is a function to return flow assistant by id
|
||||
// @Summary Retrieve flow assistant by id
|
||||
// @Tags Assistants
|
||||
// @Produce json
|
||||
// @Security BearerAuth
|
||||
// @Param flowID path int true "flow id" minimum(0)
|
||||
// @Param assistantID path int true "assistant id" minimum(0)
|
||||
// @Success 200 {object} response.successResp{data=models.Assistant} "flow assistant received successful"
|
||||
// @Failure 403 {object} response.errorResp "getting flow assistant not permitted"
|
||||
// @Failure 404 {object} response.errorResp "flow assistant not found"
|
||||
// @Failure 500 {object} response.errorResp "internal error on getting flow assistant"
|
||||
// @Router /flows/{flowID}/assistants/{assistantID} [get]
|
||||
func (s *AssistantService) GetFlowAssistant(c *gin.Context) {
|
||||
var (
|
||||
err error
|
||||
flowID uint64
|
||||
assistantID uint64
|
||||
resp models.Assistant
|
||||
)
|
||||
|
||||
if flowID, err = strconv.ParseUint(c.Param("flowID"), 10, 64); err != nil {
|
||||
logger.FromContext(c).WithError(err).Errorf("error parsing flow id")
|
||||
response.Error(c, response.ErrAssistantsInvalidRequest, err)
|
||||
return
|
||||
}
|
||||
|
||||
if assistantID, err = strconv.ParseUint(c.Param("assistantID"), 10, 64); err != nil {
|
||||
logger.FromContext(c).WithError(err).Errorf("error parsing assistant id")
|
||||
response.Error(c, response.ErrAssistantsInvalidRequest, err)
|
||||
return
|
||||
}
|
||||
|
||||
uid := c.GetUint64("uid")
|
||||
privs := c.GetStringSlice("prm")
|
||||
var scope func(db *gorm.DB) *gorm.DB
|
||||
if slices.Contains(privs, "assistants.admin") {
|
||||
scope = func(db *gorm.DB) *gorm.DB {
|
||||
return db.
|
||||
Joins("INNER JOIN flows f ON f.id = assistants.flow_id").
|
||||
Where("f.id = ?", flowID)
|
||||
}
|
||||
} else if slices.Contains(privs, "assistants.view") {
|
||||
scope = func(db *gorm.DB) *gorm.DB {
|
||||
return db.
|
||||
Joins("INNER JOIN flows f ON f.id = assistants.flow_id").
|
||||
Where("f.id = ? AND f.user_id = ?", flowID, uid)
|
||||
}
|
||||
} else {
|
||||
logger.FromContext(c).Errorf("error filtering user role permissions: permission not found")
|
||||
response.Error(c, response.ErrNotPermitted, nil)
|
||||
return
|
||||
}
|
||||
|
||||
err = s.db.Model(&resp).
|
||||
Scopes(scope).
|
||||
Where("assistants.id = ?", assistantID).
|
||||
Take(&resp).Error
|
||||
if err != nil {
|
||||
logger.FromContext(c).WithError(err).Errorf("error on getting flow assistant by id")
|
||||
if gorm.IsRecordNotFoundError(err) {
|
||||
response.Error(c, response.ErrAssistantsNotFound, err)
|
||||
} else {
|
||||
response.Error(c, response.ErrInternal, err)
|
||||
}
|
||||
return
|
||||
}
|
||||
|
||||
response.Success(c, http.StatusOK, resp)
|
||||
}
|
||||
|
||||
// CreateFlowAssistant is a function to create new assistant with custom functions
|
||||
// @Summary Create new assistant with custom functions
|
||||
// @Tags Assistants
|
||||
// @Accept json
|
||||
// @Produce json
|
||||
// @Security BearerAuth
|
||||
// @Param flowID path int true "flow id (use 0 to create a new flow together with the assistant)" minimum(0)
|
||||
// @Param json body models.CreateAssistant true "assistant model to create"
|
||||
// @Success 201 {object} response.successResp{data=models.AssistantFlow} "assistant created successful"
|
||||
// @Failure 400 {object} response.errorResp "invalid assistant request data"
|
||||
// @Failure 403 {object} response.errorResp "creating assistant not permitted"
|
||||
// @Failure 500 {object} response.errorResp "internal error on creating assistant"
|
||||
// @Router /flows/{flowID}/assistants/ [post]
|
||||
func (s *AssistantService) CreateFlowAssistant(c *gin.Context) {
|
||||
var (
|
||||
err error
|
||||
flowID uint64
|
||||
resp models.AssistantFlow
|
||||
createAssistant models.CreateAssistant
|
||||
)
|
||||
|
||||
if err := c.ShouldBindJSON(&createAssistant); err != nil {
|
||||
logger.FromContext(c).WithError(err).Errorf("error binding JSON")
|
||||
response.Error(c, response.ErrAssistantsInvalidRequest, err)
|
||||
return
|
||||
}
|
||||
|
||||
if err := createAssistant.Valid(); err != nil {
|
||||
logger.FromContext(c).WithError(err).Errorf("error validating assistant data")
|
||||
response.Error(c, response.ErrAssistantsInvalidData, err)
|
||||
return
|
||||
}
|
||||
|
||||
if flowID, err = strconv.ParseUint(c.Param("flowID"), 10, 64); err != nil {
|
||||
logger.FromContext(c).WithError(err).Errorf("error parsing flow id")
|
||||
response.Error(c, response.ErrAssistantsInvalidRequest, err)
|
||||
return
|
||||
}
|
||||
|
||||
uid := c.GetUint64("uid")
|
||||
privs := c.GetStringSlice("prm")
|
||||
|
||||
if flowID == 0 {
|
||||
// Creating a new flow together with the assistant requires both permissions.
|
||||
if !slices.Contains(privs, "assistants.create") || !slices.Contains(privs, "flows.create") {
|
||||
logger.FromContext(c).Errorf("error filtering user role permissions: permission not found")
|
||||
response.Error(c, response.ErrNotPermitted, nil)
|
||||
return
|
||||
}
|
||||
} else {
|
||||
if !slices.Contains(privs, "assistants.create") {
|
||||
logger.FromContext(c).Errorf("error filtering user role permissions: permission not found")
|
||||
response.Error(c, response.ErrNotPermitted, nil)
|
||||
return
|
||||
}
|
||||
|
||||
// Verify the flow exists and belongs to the user (unless admin).
|
||||
var flow models.Flow
|
||||
var flowScope func(db *gorm.DB) *gorm.DB
|
||||
if slices.Contains(privs, "flows.admin") {
|
||||
flowScope = func(db *gorm.DB) *gorm.DB { return db.Where("id = ?", flowID) }
|
||||
} else {
|
||||
flowScope = func(db *gorm.DB) *gorm.DB { return db.Where("id = ? AND user_id = ?", flowID, uid) }
|
||||
}
|
||||
if err = s.db.Model(&flow).Scopes(flowScope).Take(&flow).Error; err != nil {
|
||||
logger.FromContext(c).WithError(err).Errorf("error getting flow by id")
|
||||
if gorm.IsRecordNotFoundError(err) {
|
||||
response.Error(c, response.ErrFlowsNotFound, err)
|
||||
} else {
|
||||
response.Error(c, response.ErrInternal, err)
|
||||
}
|
||||
return
|
||||
}
|
||||
}
|
||||
|
||||
prvname := provider.ProviderName(createAssistant.Provider)
|
||||
|
||||
prv, err := s.pc.GetProvider(c, prvname, int64(uid))
|
||||
if err != nil {
|
||||
logger.FromContext(c).WithError(err).Errorf("error getting provider: not found")
|
||||
response.Error(c, response.ErrInternal, err)
|
||||
return
|
||||
}
|
||||
prvtype := prv.Type()
|
||||
|
||||
dbResources, err := validateServiceResources(s.db, uid, privs, createAssistant.ResourceIDs)
|
||||
if err != nil {
|
||||
logger.FromContext(c).WithError(err).Errorf("error validating resource ids")
|
||||
response.Error(c, response.ErrAssistantsInvalidRequest, err)
|
||||
return
|
||||
}
|
||||
|
||||
aw, err := s.fc.CreateAssistant(
|
||||
c,
|
||||
int64(uid),
|
||||
int64(flowID),
|
||||
createAssistant.Input,
|
||||
createAssistant.UseAgents,
|
||||
prvname,
|
||||
prvtype,
|
||||
createAssistant.Functions,
|
||||
dbResources,
|
||||
)
|
||||
if err != nil {
|
||||
logger.FromContext(c).WithError(err).Errorf("error creating assistant")
|
||||
response.Error(c, response.ErrInternal, err)
|
||||
return
|
||||
}
|
||||
|
||||
if err = s.db.Model(&resp.Assistant).Where("id = ?", aw.GetAssistantID()).Take(&resp.Assistant).Error; err != nil {
|
||||
logger.FromContext(c).WithError(err).Errorf("error getting assistant by id")
|
||||
response.Error(c, response.ErrInternal, err)
|
||||
return
|
||||
}
|
||||
|
||||
if err = s.db.Model(&resp.Flow).Where("id = ?", resp.Assistant.FlowID).Take(&resp.Flow).Error; err != nil {
|
||||
logger.FromContext(c).WithError(err).Errorf("error getting flow by id")
|
||||
response.Error(c, response.ErrInternal, err)
|
||||
return
|
||||
}
|
||||
|
||||
response.Success(c, http.StatusCreated, resp)
|
||||
}
|
||||
|
||||
// PatchAssistant is a function to patch assistant
|
||||
// @Summary Patch assistant
|
||||
// @Tags Assistants
|
||||
// @Accept json
|
||||
// @Produce json
|
||||
// @Security BearerAuth
|
||||
// @Param flowID path int true "flow id" minimum(0)
|
||||
// @Param assistantID path int true "assistant id" minimum(0)
|
||||
// @Param json body models.PatchAssistant true "assistant model to patch"
|
||||
// @Success 200 {object} response.successResp{data=models.AssistantFlow} "assistant patched successful"
|
||||
// @Failure 400 {object} response.errorResp "invalid assistant request data"
|
||||
// @Failure 403 {object} response.errorResp "patching assistant not permitted"
|
||||
// @Failure 500 {object} response.errorResp "internal error on patching assistant"
|
||||
// @Router /flows/{flowID}/assistants/{assistantID} [put]
|
||||
func (s *AssistantService) PatchAssistant(c *gin.Context) {
|
||||
var (
|
||||
err error
|
||||
flowID uint64
|
||||
assistant models.AssistantFlow
|
||||
assistantID uint64
|
||||
patchAssistant models.PatchAssistant
|
||||
)
|
||||
|
||||
if err := c.ShouldBindJSON(&patchAssistant); err != nil {
|
||||
logger.FromContext(c).WithError(err).Errorf("error binding JSON")
|
||||
response.Error(c, response.ErrAssistantsInvalidRequest, err)
|
||||
return
|
||||
}
|
||||
|
||||
if err := patchAssistant.Valid(); err != nil {
|
||||
logger.FromContext(c).WithError(err).Errorf("error validating assistant data")
|
||||
response.Error(c, response.ErrAssistantsInvalidData, err)
|
||||
return
|
||||
}
|
||||
|
||||
flowID, err = strconv.ParseUint(c.Param("flowID"), 10, 64)
|
||||
if err != nil {
|
||||
logger.FromContext(c).WithError(err).Errorf("error parsing flow id")
|
||||
response.Error(c, response.ErrAssistantsInvalidRequest, err)
|
||||
return
|
||||
}
|
||||
|
||||
assistantID, err = strconv.ParseUint(c.Param("assistantID"), 10, 64)
|
||||
if err != nil {
|
||||
logger.FromContext(c).WithError(err).Errorf("error parsing assistant id")
|
||||
response.Error(c, response.ErrAssistantsInvalidRequest, err)
|
||||
return
|
||||
}
|
||||
|
||||
uid := c.GetUint64("uid")
|
||||
privs := c.GetStringSlice("prm")
|
||||
var scope func(db *gorm.DB) *gorm.DB
|
||||
if slices.Contains(privs, "assistants.admin") {
|
||||
scope = func(db *gorm.DB) *gorm.DB {
|
||||
return db.Where("id = ?", assistantID)
|
||||
}
|
||||
} else if slices.Contains(privs, "assistants.edit") {
|
||||
scope = func(db *gorm.DB) *gorm.DB {
|
||||
return db.Where("assistants.id = ? AND assistants.flow_id = ?", assistantID, flowID).
|
||||
Joins("INNER JOIN flows f ON f.id = assistants.flow_id").
|
||||
Where("f.user_id = ?", uid)
|
||||
}
|
||||
} else {
|
||||
logger.FromContext(c).Errorf("error filtering user role permissions: permission not found")
|
||||
response.Error(c, response.ErrNotPermitted, nil)
|
||||
return
|
||||
}
|
||||
|
||||
if err = s.db.Model(&assistant).Scopes(scope).Take(&assistant).Error; err != nil {
|
||||
logger.FromContext(c).WithError(err).Errorf("error getting assistant by id")
|
||||
if gorm.IsRecordNotFoundError(err) {
|
||||
response.Error(c, response.ErrAssistantsNotFound, err)
|
||||
} else {
|
||||
response.Error(c, response.ErrInternal, err)
|
||||
}
|
||||
return
|
||||
}
|
||||
|
||||
fw, err := s.fc.GetFlow(c, int64(flowID))
|
||||
if err != nil {
|
||||
logger.FromContext(c).WithError(err).Errorf("error getting flow by id in flow controller")
|
||||
response.Error(c, response.ErrInternal, err)
|
||||
return
|
||||
}
|
||||
|
||||
aw, err := fw.GetAssistant(c, int64(assistantID))
|
||||
if err != nil {
|
||||
logger.FromContext(c).WithError(err).Errorf("error getting assistant by id in flow controller")
|
||||
response.Error(c, response.ErrInternal, err)
|
||||
return
|
||||
}
|
||||
|
||||
switch patchAssistant.Action {
|
||||
case "stop":
|
||||
if err := aw.Stop(c); err != nil {
|
||||
logger.FromContext(c).WithError(err).Errorf("error stopping assistant")
|
||||
response.Error(c, response.ErrInternal, err)
|
||||
return
|
||||
}
|
||||
case "input":
|
||||
if patchAssistant.Input == nil || *patchAssistant.Input == "" {
|
||||
logger.FromContext(c).Errorf("error sending input to assistant: input is empty")
|
||||
response.Error(c, response.ErrAssistantsInvalidRequest, nil)
|
||||
return
|
||||
}
|
||||
|
||||
dbResources, err := validateServiceResources(s.db, uid, privs, patchAssistant.ResourceIDs)
|
||||
if err != nil {
|
||||
logger.FromContext(c).WithError(err).Errorf("error validating resource ids")
|
||||
response.Error(c, response.ErrAssistantsInvalidRequest, err)
|
||||
return
|
||||
}
|
||||
|
||||
if err := aw.PutInput(c, *patchAssistant.Input, patchAssistant.UseAgents, dbResources); err != nil {
|
||||
logger.FromContext(c).WithError(err).Errorf("error sending input to assistant")
|
||||
response.Error(c, response.ErrInternal, err)
|
||||
return
|
||||
}
|
||||
default:
|
||||
logger.FromContext(c).Errorf("error filtering assistant action")
|
||||
response.Error(c, response.ErrAssistantsInvalidRequest, nil)
|
||||
return
|
||||
}
|
||||
|
||||
if err = s.db.Model(&assistant).Scopes(scope).Take(&assistant).Error; err != nil {
|
||||
logger.FromContext(c).WithError(err).Errorf("error getting assistant by id")
|
||||
if gorm.IsRecordNotFoundError(err) {
|
||||
response.Error(c, response.ErrAssistantsNotFound, err)
|
||||
} else {
|
||||
response.Error(c, response.ErrInternal, err)
|
||||
}
|
||||
return
|
||||
}
|
||||
|
||||
assistantDB, err := convertAssistantToDatabase(assistant.Assistant)
|
||||
if err != nil {
|
||||
logger.FromContext(c).WithError(err).Errorf("error converting assistant to database")
|
||||
response.Error(c, response.ErrInternal, err)
|
||||
return
|
||||
}
|
||||
|
||||
if s.ss != nil {
|
||||
publisher := s.ss.NewFlowPublisher(int64(assistant.Flow.UserID), int64(assistant.FlowID))
|
||||
publisher.AssistantUpdated(c, assistantDB)
|
||||
}
|
||||
|
||||
response.Success(c, http.StatusOK, assistant)
|
||||
}
|
||||
|
||||
// DeleteAssistant is a function to delete assistant by id
|
||||
// @Summary Delete assistant by id
|
||||
// @Tags Assistants
|
||||
// @Security BearerAuth
|
||||
// @Param flowID path int true "flow id" minimum(0)
|
||||
// @Param assistantID path int true "assistant id" minimum(0)
|
||||
// @Success 200 {object} response.successResp{data=models.AssistantFlow} "assistant deleted successful"
|
||||
// @Failure 403 {object} response.errorResp "deleting assistant not permitted"
|
||||
// @Failure 404 {object} response.errorResp "assistant not found"
|
||||
// @Failure 500 {object} response.errorResp "internal error on deleting assistant"
|
||||
// @Router /flows/{flowID}/assistants/{assistantID} [delete]
|
||||
func (s *AssistantService) DeleteAssistant(c *gin.Context) {
|
||||
var (
|
||||
err error
|
||||
flowID uint64
|
||||
assistant models.AssistantFlow
|
||||
assistantID uint64
|
||||
)
|
||||
|
||||
flowID, err = strconv.ParseUint(c.Param("flowID"), 10, 64)
|
||||
if err != nil {
|
||||
logger.FromContext(c).WithError(err).Errorf("error parsing flow id")
|
||||
response.Error(c, response.ErrAssistantsInvalidRequest, err)
|
||||
return
|
||||
}
|
||||
|
||||
assistantID, err = strconv.ParseUint(c.Param("assistantID"), 10, 64)
|
||||
if err != nil {
|
||||
logger.FromContext(c).WithError(err).Errorf("error parsing assistant id")
|
||||
response.Error(c, response.ErrAssistantsInvalidRequest, err)
|
||||
return
|
||||
}
|
||||
|
||||
uid := c.GetUint64("uid")
|
||||
privs := c.GetStringSlice("prm")
|
||||
var scope func(db *gorm.DB) *gorm.DB
|
||||
if slices.Contains(privs, "assistants.admin") {
|
||||
scope = func(db *gorm.DB) *gorm.DB {
|
||||
return db.Where("id = ?", assistantID)
|
||||
}
|
||||
} else if slices.Contains(privs, "assistants.delete") {
|
||||
scope = func(db *gorm.DB) *gorm.DB {
|
||||
return db.Where("assistants.id = ? AND assistants.flow_id = ?", assistantID, flowID).
|
||||
Joins("INNER JOIN flows f ON f.id = assistants.flow_id").
|
||||
Where("f.user_id = ?", uid)
|
||||
}
|
||||
} else {
|
||||
logger.FromContext(c).Errorf("error filtering user role permissions: permission not found")
|
||||
response.Error(c, response.ErrNotPermitted, nil)
|
||||
return
|
||||
}
|
||||
|
||||
if err = s.db.Model(&assistant).Scopes(scope).Take(&assistant).Error; err != nil {
|
||||
logger.FromContext(c).WithError(err).Errorf("error getting assistant by id")
|
||||
if gorm.IsRecordNotFoundError(err) {
|
||||
response.Error(c, response.ErrAssistantsNotFound, err)
|
||||
} else {
|
||||
response.Error(c, response.ErrInternal, err)
|
||||
}
|
||||
return
|
||||
}
|
||||
|
||||
fw, err := s.fc.GetFlow(c, int64(flowID))
|
||||
if err != nil {
|
||||
logger.FromContext(c).WithError(err).Errorf("error getting flow by id in flow controller")
|
||||
response.Error(c, response.ErrInternal, err)
|
||||
return
|
||||
}
|
||||
|
||||
aw, err := fw.GetAssistant(c, int64(assistantID))
|
||||
if err != nil {
|
||||
logger.FromContext(c).WithError(err).Errorf("error getting assistant by id in flow controller")
|
||||
response.Error(c, response.ErrInternal, err)
|
||||
return
|
||||
}
|
||||
|
||||
if err := aw.Finish(c); err != nil {
|
||||
logger.FromContext(c).WithError(err).Errorf("error stopping assistant")
|
||||
response.Error(c, response.ErrInternal, err)
|
||||
return
|
||||
}
|
||||
|
||||
if err = s.db.Scopes(scope).Delete(&assistant).Error; err != nil {
|
||||
logger.FromContext(c).WithError(err).Errorf("error deleting assistant by id")
|
||||
if gorm.IsRecordNotFoundError(err) {
|
||||
response.Error(c, response.ErrAssistantsNotFound, err)
|
||||
} else {
|
||||
response.Error(c, response.ErrInternal, err)
|
||||
}
|
||||
return
|
||||
}
|
||||
|
||||
assistantDB, err := convertAssistantToDatabase(assistant.Assistant)
|
||||
if err != nil {
|
||||
logger.FromContext(c).WithError(err).Errorf("error converting assistant to database")
|
||||
response.Error(c, response.ErrInternal, err)
|
||||
return
|
||||
}
|
||||
|
||||
if s.ss != nil {
|
||||
publisher := s.ss.NewFlowPublisher(int64(assistant.Flow.UserID), int64(assistant.FlowID))
|
||||
publisher.AssistantDeleted(c, assistantDB)
|
||||
}
|
||||
|
||||
response.Success(c, http.StatusOK, assistant)
|
||||
}
|
||||
|
||||
func convertAssistantToDatabase(assistant models.Assistant) (database.Assistant, error) {
|
||||
functions, err := json.Marshal(assistant.Functions)
|
||||
if err != nil {
|
||||
return database.Assistant{}, err
|
||||
}
|
||||
|
||||
return database.Assistant{
|
||||
ID: int64(assistant.ID),
|
||||
Status: database.AssistantStatus(assistant.Status),
|
||||
Title: assistant.Title,
|
||||
Model: assistant.Model,
|
||||
ModelProviderName: assistant.ModelProviderName,
|
||||
Language: assistant.Language,
|
||||
Functions: functions,
|
||||
TraceID: database.PtrStringToNullString(assistant.TraceID),
|
||||
FlowID: int64(assistant.FlowID),
|
||||
UseAgents: assistant.UseAgents,
|
||||
MsgchainID: database.Uint64ToNullInt64(assistant.MsgchainID),
|
||||
CreatedAt: database.TimeToNullTime(assistant.CreatedAt),
|
||||
UpdatedAt: database.TimeToNullTime(assistant.UpdatedAt),
|
||||
DeletedAt: database.PtrTimeToNullTime(assistant.DeletedAt),
|
||||
ModelProviderType: database.ProviderType(assistant.ModelProviderType),
|
||||
ToolCallIDTemplate: assistant.ToolCallIDTemplate,
|
||||
}, nil
|
||||
}
|
||||
|
||||
// validateServiceResources fetches and validates user resource ownership for REST handlers.
|
||||
// It mirrors the logic in validateUserResources from the graph package but works with *gorm.DB.
|
||||
// privs must contain at least "resources.view" or "resources.admin"; otherwise permission is denied.
|
||||
// "resources.admin" bypasses the user_id ownership check.
|
||||
func validateServiceResources(db *gorm.DB, uid uint64, privs []string, ids []uint64) ([]database.UserResource, error) {
|
||||
if len(ids) == 0 {
|
||||
return nil, nil
|
||||
}
|
||||
|
||||
isAdmin := slices.Contains(privs, "resources.admin")
|
||||
if !isAdmin && !slices.Contains(privs, "resources.view") && len(ids) > 0 {
|
||||
return nil, fmt.Errorf("permission 'resources.view' required to use resource IDs")
|
||||
}
|
||||
|
||||
var recs []models.UserResource
|
||||
if err := db.Model(&models.UserResource{}).Where("id IN (?)", ids).Find(&recs).Error; err != nil {
|
||||
return nil, fmt.Errorf("failed to fetch resources: %w", err)
|
||||
}
|
||||
|
||||
found := make(map[uint64]models.UserResource, len(recs))
|
||||
for _, r := range recs {
|
||||
found[r.ID] = r
|
||||
}
|
||||
|
||||
result := make([]database.UserResource, 0, len(ids))
|
||||
for _, id := range ids {
|
||||
r, ok := found[id]
|
||||
if !ok {
|
||||
return nil, fmt.Errorf("resource %d not found", id)
|
||||
}
|
||||
if !isAdmin && r.UserID != uid {
|
||||
return nil, fmt.Errorf("resource %d not accessible", id)
|
||||
}
|
||||
result = append(result, database.UserResource{
|
||||
ID: int64(r.ID),
|
||||
UserID: int64(r.UserID),
|
||||
Hash: r.Hash,
|
||||
Name: r.Name,
|
||||
Path: r.Path,
|
||||
Size: r.Size,
|
||||
IsDir: r.IsDir,
|
||||
CreatedAt: database.TimeToNullTime(r.CreatedAt),
|
||||
UpdatedAt: database.TimeToNullTime(r.UpdatedAt),
|
||||
})
|
||||
}
|
||||
|
||||
return result, nil
|
||||
}
|
||||
@@ -0,0 +1,889 @@
|
||||
package services
|
||||
|
||||
import (
|
||||
"crypto/hmac"
|
||||
"crypto/rand"
|
||||
"crypto/sha256"
|
||||
"encoding/base64"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"fmt"
|
||||
"io"
|
||||
"net/http"
|
||||
"net/url"
|
||||
"path"
|
||||
"slices"
|
||||
"strconv"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"pentagi/pkg/server/logger"
|
||||
"pentagi/pkg/server/models"
|
||||
"pentagi/pkg/server/oauth"
|
||||
"pentagi/pkg/server/rdb"
|
||||
"pentagi/pkg/server/response"
|
||||
"pentagi/pkg/version"
|
||||
|
||||
"github.com/gin-contrib/sessions"
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/jinzhu/gorm"
|
||||
"github.com/sirupsen/logrus"
|
||||
"golang.org/x/crypto/bcrypt"
|
||||
"golang.org/x/oauth2"
|
||||
)
|
||||
|
||||
const (
|
||||
authStateCookieName = "state"
|
||||
authNonceCookieName = "nonce"
|
||||
authStateRequestTTL = 5 * time.Minute
|
||||
)
|
||||
|
||||
type AuthServiceConfig struct {
|
||||
BaseURL string
|
||||
LoginCallbackURL string
|
||||
SessionTimeout int // in seconds
|
||||
}
|
||||
|
||||
type AuthService struct {
|
||||
cfg AuthServiceConfig
|
||||
db *gorm.DB
|
||||
key []byte
|
||||
oauth map[string]oauth.OAuthClient
|
||||
}
|
||||
|
||||
func NewAuthService(
|
||||
cfg AuthServiceConfig,
|
||||
db *gorm.DB,
|
||||
oauth map[string]oauth.OAuthClient,
|
||||
) *AuthService {
|
||||
var count int
|
||||
err := db.Model(&models.User{}).Where("type = 'local'").Count(&count).Error
|
||||
if err != nil {
|
||||
logrus.WithError(err).Errorf("error getting local users count")
|
||||
}
|
||||
|
||||
key, err := randBytes(32)
|
||||
if err != nil {
|
||||
logrus.WithError(err).Errorf("error generating key")
|
||||
}
|
||||
|
||||
return &AuthService{
|
||||
cfg: cfg,
|
||||
db: db,
|
||||
key: key,
|
||||
oauth: oauth,
|
||||
}
|
||||
}
|
||||
|
||||
// AuthLogin is function to login user in the system
|
||||
// @Summary Login user into system
|
||||
// @Tags Public
|
||||
// @Accept json
|
||||
// @Produce json
|
||||
// @Param json body models.Login true "Login form JSON data"
|
||||
// @Success 200 {object} response.successResp "login successful"
|
||||
// @Failure 400 {object} response.errorResp "invalid login data"
|
||||
// @Failure 401 {object} response.errorResp "invalid login or password"
|
||||
// @Failure 403 {object} response.errorResp "login not permitted"
|
||||
// @Failure 500 {object} response.errorResp "internal error on login"
|
||||
// @Router /auth/login [post]
|
||||
func (s *AuthService) AuthLogin(c *gin.Context) {
|
||||
var data models.Login
|
||||
if err := c.ShouldBindJSON(&data); err != nil || data.Valid() != nil {
|
||||
if err == nil {
|
||||
err = data.Valid()
|
||||
}
|
||||
logger.FromContext(c).WithError(err).Errorf("error validating request data")
|
||||
response.Error(c, response.ErrAuthInvalidLoginRequest, err)
|
||||
return
|
||||
}
|
||||
|
||||
var user models.UserPassword
|
||||
if err := s.db.Take(&user, "mail = ? AND password IS NOT NULL", data.Mail).Error; err != nil {
|
||||
logrus.WithError(err).Errorf("error getting user by mail '%s'", data.Mail)
|
||||
response.Error(c, response.ErrAuthInvalidCredentials, err)
|
||||
return
|
||||
} else if err = user.Valid(); err != nil {
|
||||
logger.FromContext(c).WithError(err).Errorf("error validating user data '%s'", user.Hash)
|
||||
response.Error(c, response.ErrAuthInvalidUserData, err)
|
||||
return
|
||||
} else if user.RoleID == 100 {
|
||||
logger.FromContext(c).WithError(err).Errorf("can't authorize external user '%s'", user.Hash)
|
||||
response.Error(c, response.ErrAuthInvalidUserData, fmt.Errorf("user is external"))
|
||||
return
|
||||
}
|
||||
|
||||
if err := bcrypt.CompareHashAndPassword([]byte(user.Password), []byte(data.Password)); err != nil {
|
||||
logger.FromContext(c).Errorf("error matching user input password")
|
||||
response.Error(c, response.ErrAuthInvalidCredentials, err)
|
||||
return
|
||||
}
|
||||
|
||||
if user.Status != "active" {
|
||||
logger.FromContext(c).Errorf("error checking active state for user '%s'", user.Status)
|
||||
response.Error(c, response.ErrAuthInactiveUser, fmt.Errorf("user is inactive"))
|
||||
return
|
||||
}
|
||||
|
||||
var privs []string
|
||||
err := s.db.Table("privileges").
|
||||
Where("role_id = ?", user.RoleID).
|
||||
Pluck("name", &privs).Error
|
||||
if err != nil {
|
||||
logger.FromContext(c).WithError(err).Errorf("error getting user privileges list '%s'", user.Hash)
|
||||
response.Error(c, response.ErrAuthInvalidServiceData, err)
|
||||
return
|
||||
}
|
||||
|
||||
uuid, err := rdb.MakeUuidStrFromHash(user.Hash)
|
||||
if err != nil {
|
||||
logger.FromContext(c).WithError(err).Errorf("error validating user data '%s'", user.Hash)
|
||||
response.Error(c, response.ErrAuthInvalidUserData, err)
|
||||
return
|
||||
}
|
||||
|
||||
expires := s.cfg.SessionTimeout
|
||||
session := sessions.Default(c)
|
||||
session.Set("uid", user.ID)
|
||||
session.Set("uhash", user.Hash)
|
||||
session.Set("rid", user.RoleID)
|
||||
session.Set("tid", models.UserTypeLocal.String())
|
||||
session.Set("prm", privs)
|
||||
session.Set("gtm", time.Now().Unix())
|
||||
session.Set("exp", time.Now().Add(time.Duration(expires)*time.Second).Unix())
|
||||
session.Set("uuid", uuid)
|
||||
session.Set("uname", user.Name)
|
||||
session.Options(sessions.Options{
|
||||
HttpOnly: true,
|
||||
Secure: c.Request.TLS != nil,
|
||||
SameSite: http.SameSiteLaxMode,
|
||||
Path: s.cfg.BaseURL,
|
||||
MaxAge: int(expires),
|
||||
})
|
||||
if err := session.Save(); err != nil {
|
||||
logger.FromContext(c).WithError(err).Errorf("error saving session")
|
||||
response.Error(c, response.ErrInternal, err)
|
||||
return
|
||||
}
|
||||
|
||||
logger.FromContext(c).
|
||||
WithFields(logrus.Fields{
|
||||
"age": expires,
|
||||
"uid": user.ID,
|
||||
"uhash": user.Hash,
|
||||
"rid": user.RoleID,
|
||||
"tid": session.Get("tid"),
|
||||
"gtm": session.Get("gtm"),
|
||||
"exp": session.Get("exp"),
|
||||
"prm": session.Get("prm"),
|
||||
}).
|
||||
Infof("user made successful local login for '%s'", data.Mail)
|
||||
|
||||
response.Success(c, http.StatusOK, struct{}{})
|
||||
}
|
||||
|
||||
func (s *AuthService) refreshCookie(c *gin.Context, resp *info, privs []string) error {
|
||||
session := sessions.Default(c)
|
||||
expires := int(s.cfg.SessionTimeout)
|
||||
session.Set("prm", privs)
|
||||
session.Set("gtm", time.Now().Unix())
|
||||
session.Set("exp", time.Now().Add(time.Duration(expires)*time.Second).Unix())
|
||||
resp.Privs = privs
|
||||
|
||||
session.Set("uid", resp.User.ID)
|
||||
session.Set("uhash", resp.User.Hash)
|
||||
session.Set("rid", resp.User.RoleID)
|
||||
session.Set("tid", resp.User.Type.String())
|
||||
session.Options(sessions.Options{
|
||||
HttpOnly: true,
|
||||
Secure: c.Request.TLS != nil,
|
||||
SameSite: http.SameSiteLaxMode,
|
||||
Path: s.cfg.BaseURL,
|
||||
MaxAge: expires,
|
||||
})
|
||||
if err := session.Save(); err != nil {
|
||||
logger.FromContext(c).WithError(err).Errorf("error saving session")
|
||||
return err
|
||||
}
|
||||
|
||||
logger.FromContext(c).
|
||||
WithFields(logrus.Fields{
|
||||
"age": expires,
|
||||
"uid": resp.User.ID,
|
||||
"uhash": resp.User.Hash,
|
||||
"rid": resp.User.RoleID,
|
||||
"tid": session.Get("tid"),
|
||||
"gtm": session.Get("gtm"),
|
||||
"exp": session.Get("exp"),
|
||||
"prm": session.Get("prm"),
|
||||
}).
|
||||
Infof("session was refreshed for '%s' '%s'", resp.User.Mail, resp.User.Name)
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
// AuthAuthorize is function to login user in OAuth2 external system
|
||||
// @Summary Login user into OAuth2 external system via HTTP redirect
|
||||
// @Tags Public
|
||||
// @Produce json
|
||||
// @Param return_uri query string false "URI to redirect user there after login" default(/)
|
||||
// @Param provider query string false "OAuth provider name (google, github, etc.)" default(google) enums:"google,github"
|
||||
// @Success 307 "redirect to SSO login page"
|
||||
// @Failure 400 {object} response.errorResp "invalid autorizarion query"
|
||||
// @Failure 403 {object} response.errorResp "authorize not permitted"
|
||||
// @Failure 500 {object} response.errorResp "internal error on autorizarion"
|
||||
// @Router /auth/authorize [get]
|
||||
func (s *AuthService) AuthAuthorize(c *gin.Context) {
|
||||
stateData := map[string]string{
|
||||
"exp": strconv.FormatInt(time.Now().Add(authStateRequestTTL).Unix(), 10),
|
||||
}
|
||||
|
||||
queryReturnURI := c.Query("return_uri")
|
||||
if queryReturnURI != "" {
|
||||
returnURL, err := url.Parse(queryReturnURI)
|
||||
if err != nil {
|
||||
logger.FromContext(c).WithError(err).Errorf("failed to parse return url argument '%s'", queryReturnURI)
|
||||
response.Error(c, response.ErrAuthInvalidAuthorizeQuery, err)
|
||||
return
|
||||
}
|
||||
returnURL.Path = path.Clean(path.Join("/", returnURL.Path))
|
||||
stateData["return_uri"] = returnURL.RequestURI()
|
||||
}
|
||||
|
||||
provider := c.Query("provider")
|
||||
oauthClient, ok := s.oauth[provider]
|
||||
if !ok {
|
||||
logger.FromContext(c).Errorf("external OAuth2 provider '%s' is not initialized", provider)
|
||||
err := fmt.Errorf("provider not initialized")
|
||||
response.Error(c, response.ErrNotPermitted, err)
|
||||
return
|
||||
}
|
||||
stateData["provider"] = provider
|
||||
|
||||
stateUniq, err := randBase64String(16)
|
||||
if err != nil {
|
||||
logger.FromContext(c).WithError(err).Errorf("failed to generate state random data")
|
||||
response.Error(c, response.ErrInternal, err)
|
||||
return
|
||||
}
|
||||
stateData["uniq"] = stateUniq
|
||||
|
||||
nonce, err := randBase64String(16)
|
||||
if err != nil {
|
||||
logger.FromContext(c).WithError(err).Errorf("failed to generate nonce random data")
|
||||
response.Error(c, response.ErrInternal, err)
|
||||
return
|
||||
}
|
||||
|
||||
stateJSON, err := json.Marshal(stateData)
|
||||
if err != nil {
|
||||
logger.FromContext(c).WithError(err).Errorf("failed to marshal state json data")
|
||||
response.Error(c, response.ErrInternal, err)
|
||||
return
|
||||
}
|
||||
mac := hmac.New(sha256.New, s.key)
|
||||
mac.Write(stateJSON)
|
||||
signature := mac.Sum(nil)
|
||||
|
||||
signedStateJSON := append(signature, stateJSON...)
|
||||
state := base64.RawURLEncoding.EncodeToString(signedStateJSON)
|
||||
|
||||
// Google OAuth uses POST callback which requires SameSite=None for cross-site requests
|
||||
// GitHub and other providers use GET callback which works with SameSite=Lax
|
||||
sameSiteMode := http.SameSiteLaxMode
|
||||
if provider == "google" {
|
||||
sameSiteMode = http.SameSiteNoneMode
|
||||
}
|
||||
|
||||
maxAge := int(authStateRequestTTL / time.Second)
|
||||
s.setCallbackCookie(c.Writer, c.Request, authStateCookieName, state, maxAge, sameSiteMode)
|
||||
s.setCallbackCookie(c.Writer, c.Request, authNonceCookieName, nonce, maxAge, sameSiteMode)
|
||||
|
||||
authOpts := []oauth2.AuthCodeOption{
|
||||
oauth2.SetAuthURLParam("nonce", nonce),
|
||||
oauth2.SetAuthURLParam("response_mode", "form_post"),
|
||||
oauth2.SetAuthURLParam("response_type", "code id_token"),
|
||||
}
|
||||
http.Redirect(c.Writer, c.Request,
|
||||
oauthClient.AuthCodeURL(state, authOpts...),
|
||||
http.StatusTemporaryRedirect)
|
||||
}
|
||||
|
||||
// AuthLoginGetCallback is function to catch login callback from OAuth application with code only
|
||||
// @Summary Login user from external OAuth application
|
||||
// @Tags Public
|
||||
// @Accept json
|
||||
// @Produce json
|
||||
// @Param code query string false "Auth code from OAuth provider to exchange token"
|
||||
// @Success 303 "redirect to registered return_uri path in the state"
|
||||
// @Failure 400 {object} response.errorResp "invalid login data"
|
||||
// @Failure 401 {object} response.errorResp "invalid login or password"
|
||||
// @Failure 403 {object} response.errorResp "login not permitted"
|
||||
// @Failure 500 {object} response.errorResp "internal error on login"
|
||||
// @Router /auth/login-callback [get]
|
||||
func (s *AuthService) AuthLoginGetCallback(c *gin.Context) {
|
||||
code := c.Query("code")
|
||||
if code == "" {
|
||||
response.Error(c, response.ErrAuthInvalidLoginCallbackRequest, fmt.Errorf("code is required"))
|
||||
return
|
||||
}
|
||||
|
||||
state, err := c.Request.Cookie(authStateCookieName)
|
||||
if err != nil {
|
||||
logger.FromContext(c).WithError(err).Errorf("error getting state from cookie")
|
||||
response.Error(c, response.ErrAuthInvalidAuthorizationState, err)
|
||||
return
|
||||
}
|
||||
|
||||
queryState := c.Query("state")
|
||||
if queryState == "" {
|
||||
logger.FromContext(c).Errorf("error missing state parameter in OAuth callback")
|
||||
response.Error(c, response.ErrAuthInvalidAuthorizationState, fmt.Errorf("state parameter is required"))
|
||||
return
|
||||
}
|
||||
|
||||
if queryState != state.Value {
|
||||
logger.FromContext(c).Errorf("error matching received state to stored one")
|
||||
response.Error(c, response.ErrAuthInvalidAuthorizationState, nil)
|
||||
return
|
||||
}
|
||||
|
||||
stateData, err := s.parseState(c, state.Value)
|
||||
if err != nil {
|
||||
return
|
||||
}
|
||||
|
||||
s.authLoginCallback(c, stateData, code)
|
||||
}
|
||||
|
||||
// AuthLoginPostCallback is function to catch login callback from OAuth application
|
||||
// @Summary Login user from external OAuth application
|
||||
// @Tags Public
|
||||
// @Accept json
|
||||
// @Produce json
|
||||
// @Param json body models.AuthCallback true "Auth form JSON data"
|
||||
// @Success 303 "redirect to registered return_uri path in the state"
|
||||
// @Failure 400 {object} response.errorResp "invalid login data"
|
||||
// @Failure 401 {object} response.errorResp "invalid login or password"
|
||||
// @Failure 403 {object} response.errorResp "login not permitted"
|
||||
// @Failure 500 {object} response.errorResp "internal error on login"
|
||||
// @Router /auth/login-callback [post]
|
||||
func (s *AuthService) AuthLoginPostCallback(c *gin.Context) {
|
||||
var (
|
||||
data models.AuthCallback
|
||||
err error
|
||||
)
|
||||
|
||||
if err = c.ShouldBind(&data); err != nil || data.Valid() != nil {
|
||||
if err == nil {
|
||||
err = data.Valid()
|
||||
}
|
||||
logger.FromContext(c).WithError(err).Errorf("error validating request data")
|
||||
response.Error(c, response.ErrAuthInvalidLoginCallbackRequest, err)
|
||||
return
|
||||
}
|
||||
|
||||
state, err := c.Request.Cookie(authStateCookieName)
|
||||
if err != nil {
|
||||
logger.FromContext(c).WithError(err).Errorf("error getting state from cookie")
|
||||
response.Error(c, response.ErrAuthInvalidAuthorizationState, err)
|
||||
return
|
||||
}
|
||||
|
||||
if data.State != state.Value {
|
||||
logger.FromContext(c).Errorf("error matching received state to stored one")
|
||||
response.Error(c, response.ErrAuthInvalidAuthorizationState, nil)
|
||||
return
|
||||
}
|
||||
|
||||
stateData, err := s.parseState(c, state.Value)
|
||||
if err != nil {
|
||||
return
|
||||
}
|
||||
|
||||
s.authLoginCallback(c, stateData, data.Code)
|
||||
}
|
||||
|
||||
// AuthLogoutCallback is function to catch logout callback from OAuth application
|
||||
// @Summary Logout current user from external OAuth application
|
||||
// @Tags Public
|
||||
// @Accept json
|
||||
// @Success 303 {object} response.successResp "logout successful"
|
||||
// @Router /auth/logout-callback [post]
|
||||
func (s *AuthService) AuthLogoutCallback(c *gin.Context) {
|
||||
s.resetSession(c)
|
||||
http.Redirect(c.Writer, c.Request, "/", http.StatusSeeOther)
|
||||
}
|
||||
|
||||
// AuthLogout is function to logout current user
|
||||
// @Summary Logout current user via HTTP redirect
|
||||
// @Tags Public
|
||||
// @Produce json
|
||||
// @Param return_uri query string false "URI to redirect user there after logout" default(/)
|
||||
// @Success 307 "redirect to input return_uri path"
|
||||
// @Router /auth/logout [get]
|
||||
func (s *AuthService) AuthLogout(c *gin.Context) {
|
||||
returnURI := "/"
|
||||
if returnURL, err := url.Parse(c.Query("return_uri")); err == nil {
|
||||
if uri := returnURL.RequestURI(); uri != "" {
|
||||
returnURI = path.Clean(path.Join("/", uri))
|
||||
}
|
||||
}
|
||||
|
||||
session := sessions.Default(c)
|
||||
logger.FromContext(c).
|
||||
WithFields(logrus.Fields{
|
||||
"uid": session.Get("uid"),
|
||||
"uhash": session.Get("uhash"),
|
||||
"rid": session.Get("rid"),
|
||||
"tid": session.Get("tid"),
|
||||
"gtm": session.Get("gtm"),
|
||||
"exp": session.Get("exp"),
|
||||
"prm": session.Get("prm"),
|
||||
}).
|
||||
Info("user made successful logout")
|
||||
|
||||
s.resetSession(c)
|
||||
http.Redirect(c.Writer, c.Request, returnURI, http.StatusTemporaryRedirect)
|
||||
}
|
||||
|
||||
func (s *AuthService) authLoginCallback(c *gin.Context, stateData map[string]string, code string) {
|
||||
var (
|
||||
privs []string
|
||||
role models.Role
|
||||
user models.User
|
||||
)
|
||||
|
||||
provider := stateData["provider"]
|
||||
oauthClient, ok := s.oauth[provider]
|
||||
if !ok {
|
||||
logger.FromContext(c).Errorf("external OAuth2 provider '%s' is not initialized", provider)
|
||||
response.Error(c, response.ErrNotPermitted, fmt.Errorf("provider not initialized"))
|
||||
return
|
||||
}
|
||||
|
||||
ctx := c.Request.Context()
|
||||
|
||||
oauth2Token, err := oauthClient.Exchange(ctx, code)
|
||||
if err != nil {
|
||||
logger.FromContext(c).WithError(err).Errorf("failed to exchange token")
|
||||
response.Error(c, response.ErrAuthExchangeTokenFail, err)
|
||||
return
|
||||
}
|
||||
|
||||
if !oauth2Token.Valid() {
|
||||
logger.FromContext(c).Errorf("failed to validate OAuth2 token")
|
||||
response.Error(c, response.ErrAuthVerificationTokenFail, nil)
|
||||
return
|
||||
}
|
||||
|
||||
nonce, err := c.Request.Cookie(authNonceCookieName)
|
||||
if err != nil {
|
||||
logger.FromContext(c).WithError(err).Errorf("error getting nonce from cookie")
|
||||
response.Error(c, response.ErrAuthInvalidAuthorizationNonce, err)
|
||||
return
|
||||
}
|
||||
|
||||
email, err := oauthClient.ResolveEmail(ctx, nonce.Value, oauth2Token)
|
||||
if err != nil {
|
||||
logger.FromContext(c).WithError(err).Errorf("failed to resolve email")
|
||||
response.Error(c, response.ErrAuthInvalidUserData, err)
|
||||
return
|
||||
}
|
||||
|
||||
if !strings.Contains(email, "@") {
|
||||
logger.FromContext(c).Errorf("invalid email format '%s'", email)
|
||||
response.Error(c, response.ErrAuthInvalidUserData, fmt.Errorf("invalid email format"))
|
||||
return
|
||||
}
|
||||
|
||||
username := strings.Split(email, "@")[0]
|
||||
if username == "" {
|
||||
logger.FromContext(c).Errorf("empty username from email '%s'", email)
|
||||
response.Error(c, response.ErrAuthInvalidUserData, fmt.Errorf("empty username"))
|
||||
return
|
||||
}
|
||||
|
||||
err = s.db.Take(&role, "id = ?", models.RoleUser).Error
|
||||
if err != nil {
|
||||
logger.FromContext(c).WithError(err).Errorf("error getting user role '%d'", models.RoleUser)
|
||||
response.Error(c, response.ErrAuthInvalidServiceData, err)
|
||||
return
|
||||
}
|
||||
|
||||
err = s.db.Table("privileges").
|
||||
Where("role_id = ?", models.RoleUser).
|
||||
Pluck("name", &privs).Error
|
||||
if err != nil {
|
||||
logger.FromContext(c).WithError(err).Errorf("error getting user privileges list '%s'", user.Hash)
|
||||
response.Error(c, response.ErrAuthInvalidServiceData, err)
|
||||
return
|
||||
}
|
||||
|
||||
filterQuery := "mail = ? AND type = ?"
|
||||
if err = s.db.Take(&user, filterQuery, email, models.UserTypeOAuth).Error; err != nil {
|
||||
if errors.Is(err, gorm.ErrRecordNotFound) {
|
||||
user = models.User{
|
||||
Hash: rdb.MakeUserHash(email),
|
||||
Mail: email,
|
||||
Name: username,
|
||||
RoleID: models.RoleUser,
|
||||
Status: "active",
|
||||
Type: models.UserTypeOAuth,
|
||||
}
|
||||
|
||||
tx := s.db.Begin()
|
||||
if tx.Error != nil {
|
||||
logger.FromContext(c).WithError(tx.Error).Errorf("error starting transaction")
|
||||
response.Error(c, response.ErrInternal, tx.Error)
|
||||
return
|
||||
}
|
||||
|
||||
if err = tx.Create(&user).Error; err != nil {
|
||||
tx.Rollback()
|
||||
logger.FromContext(c).WithError(err).Errorf("error creating user")
|
||||
response.Error(c, response.ErrInternal, err)
|
||||
return
|
||||
}
|
||||
|
||||
preferences := models.NewUserPreferences(user.ID)
|
||||
if err = tx.Create(preferences).Error; err != nil {
|
||||
tx.Rollback()
|
||||
logger.FromContext(c).WithError(err).Errorf("error creating user preferences")
|
||||
response.Error(c, response.ErrInternal, err)
|
||||
return
|
||||
}
|
||||
|
||||
if err = tx.Commit().Error; err != nil {
|
||||
logger.FromContext(c).WithError(err).Errorf("error committing transaction")
|
||||
response.Error(c, response.ErrInternal, err)
|
||||
return
|
||||
}
|
||||
} else {
|
||||
logger.FromContext(c).WithError(err).Errorf("error searching user by email '%s'", email)
|
||||
response.Error(c, response.ErrInternal, err)
|
||||
return
|
||||
}
|
||||
} else if err = user.Valid(); err != nil {
|
||||
logger.FromContext(c).WithError(err).Errorf("error validating user data '%s'", user.Hash)
|
||||
response.Error(c, response.ErrAuthInvalidUserData, err)
|
||||
return
|
||||
}
|
||||
|
||||
if user.Status != "active" {
|
||||
logger.FromContext(c).Errorf("error checking active state for user '%s'", user.Status)
|
||||
response.Error(c, response.ErrAuthInactiveUser, fmt.Errorf("user is inactive"))
|
||||
return
|
||||
}
|
||||
|
||||
expires := s.cfg.SessionTimeout
|
||||
gtm := time.Now().Unix()
|
||||
exp := time.Now().Add(time.Duration(expires) * time.Second).Unix()
|
||||
session := sessions.Default(c)
|
||||
session.Set("uid", user.ID)
|
||||
session.Set("uhash", user.Hash)
|
||||
session.Set("rid", user.RoleID)
|
||||
session.Set("tid", user.Type.String())
|
||||
session.Set("prm", privs)
|
||||
session.Set("gtm", gtm)
|
||||
session.Set("exp", exp)
|
||||
session.Set("uuid", user.Mail)
|
||||
session.Set("uname", user.Name)
|
||||
session.Options(sessions.Options{
|
||||
HttpOnly: true,
|
||||
Secure: c.Request.TLS != nil,
|
||||
SameSite: http.SameSiteLaxMode,
|
||||
Path: s.cfg.BaseURL,
|
||||
MaxAge: expires,
|
||||
})
|
||||
if err := session.Save(); err != nil {
|
||||
logger.FromContext(c).WithError(err).Errorf("error saving session")
|
||||
response.Error(c, response.ErrInternal, err)
|
||||
return
|
||||
}
|
||||
|
||||
// delete temporary cookies
|
||||
// Google OAuth uses POST callback which requires SameSite=None for cross-site requests
|
||||
// GitHub and other providers use GET callback which works with SameSite=Lax
|
||||
sameSiteMode := http.SameSiteLaxMode
|
||||
if stateData["provider"] == "google" {
|
||||
sameSiteMode = http.SameSiteNoneMode
|
||||
}
|
||||
s.setCallbackCookie(c.Writer, c.Request, authStateCookieName, "", 0, sameSiteMode)
|
||||
s.setCallbackCookie(c.Writer, c.Request, authNonceCookieName, "", 0, sameSiteMode)
|
||||
|
||||
logger.FromContext(c).
|
||||
WithFields(logrus.Fields{
|
||||
"age": expires,
|
||||
"uid": user.ID,
|
||||
"uhash": user.Hash,
|
||||
"rid": user.RoleID,
|
||||
"tid": user.Type,
|
||||
"gtm": session.Get("gtm"),
|
||||
"exp": session.Get("exp"),
|
||||
"prm": session.Get("prm"),
|
||||
}).
|
||||
Infof("user made successful SSO login for '%s' '%s'", user.Mail, user.Name)
|
||||
|
||||
if returnURI := stateData["return_uri"]; returnURI == "" {
|
||||
response.Success(c, http.StatusOK, nil)
|
||||
} else {
|
||||
u, err := url.Parse(returnURI)
|
||||
if err != nil {
|
||||
response.Success(c, http.StatusOK, nil)
|
||||
return
|
||||
}
|
||||
query := u.Query()
|
||||
query.Add("status", "success")
|
||||
u.RawQuery = query.Encode()
|
||||
http.Redirect(c.Writer, c.Request, u.RequestURI(), http.StatusSeeOther)
|
||||
}
|
||||
}
|
||||
|
||||
func (s *AuthService) parseState(c *gin.Context, state string) (map[string]string, error) {
|
||||
var stateData map[string]string
|
||||
|
||||
stateJSON, err := base64.RawURLEncoding.DecodeString(state)
|
||||
if err != nil {
|
||||
logger.FromContext(c).WithError(err).Errorf("error on getting state as a base64")
|
||||
response.Error(c, response.ErrAuthInvalidAuthorizationState, err)
|
||||
return nil, err
|
||||
}
|
||||
|
||||
signatureLen := 32
|
||||
if len(stateJSON) <= signatureLen {
|
||||
logger.FromContext(c).Errorf("error on parsing state from json data")
|
||||
err := fmt.Errorf("unexpected state length")
|
||||
response.Error(c, response.ErrAuthInvalidAuthorizationState, err)
|
||||
return nil, err
|
||||
}
|
||||
stateSignature := stateJSON[:signatureLen]
|
||||
stateJSON = stateJSON[signatureLen:]
|
||||
|
||||
mac := hmac.New(sha256.New, s.key)
|
||||
mac.Write(stateJSON)
|
||||
signature := mac.Sum(nil)
|
||||
|
||||
if !hmac.Equal(stateSignature, signature) {
|
||||
logger.FromContext(c).Errorf("error on matching signature")
|
||||
err := fmt.Errorf("mismatch state signature")
|
||||
response.Error(c, response.ErrAuthInvalidAuthorizationState, err)
|
||||
return nil, err
|
||||
}
|
||||
|
||||
if err := json.Unmarshal(stateJSON, &stateData); err != nil {
|
||||
logger.FromContext(c).WithError(err).Errorf("error on parsing state from json data")
|
||||
response.Error(c, response.ErrAuthInvalidAuthorizationState, err)
|
||||
return nil, err
|
||||
}
|
||||
|
||||
expStr, ok := stateData["exp"]
|
||||
if !ok || expStr == "" {
|
||||
err := fmt.Errorf("missing required field: exp")
|
||||
logger.FromContext(c).WithError(err).Errorf("error on validating state data")
|
||||
response.Error(c, response.ErrAuthInvalidAuthorizationState, err)
|
||||
return nil, err
|
||||
}
|
||||
|
||||
if _, ok := stateData["provider"]; !ok {
|
||||
err := fmt.Errorf("missing required field: provider")
|
||||
logger.FromContext(c).WithError(err).Errorf("error on validating state data")
|
||||
response.Error(c, response.ErrAuthInvalidAuthorizationState, err)
|
||||
return nil, err
|
||||
}
|
||||
|
||||
exp, err := strconv.ParseInt(expStr, 10, 64)
|
||||
if err != nil {
|
||||
logger.FromContext(c).WithError(err).Errorf("error on parsing expiration time")
|
||||
response.Error(c, response.ErrAuthInvalidAuthorizationState, err)
|
||||
return nil, err
|
||||
}
|
||||
|
||||
if time.Now().Unix() > exp {
|
||||
logger.FromContext(c).Errorf("error on checking expiration time")
|
||||
err := fmt.Errorf("state signature expired")
|
||||
response.Error(c, response.ErrAuthTokenExpired, err)
|
||||
return nil, err
|
||||
}
|
||||
|
||||
return stateData, nil
|
||||
}
|
||||
|
||||
func (s *AuthService) setCallbackCookie(
|
||||
w http.ResponseWriter, r *http.Request,
|
||||
name, value string, maxAge int,
|
||||
sameSite http.SameSite,
|
||||
) {
|
||||
// Check both direct TLS and X-Forwarded-Proto header (for reverse proxy setups)
|
||||
useTLS := r.TLS != nil || r.Header.Get("X-Forwarded-Proto") == "https"
|
||||
|
||||
c := &http.Cookie{
|
||||
Name: name,
|
||||
Value: value,
|
||||
HttpOnly: true,
|
||||
Secure: useTLS,
|
||||
SameSite: sameSite,
|
||||
Path: path.Join(s.cfg.BaseURL, s.cfg.LoginCallbackURL),
|
||||
MaxAge: maxAge,
|
||||
}
|
||||
http.SetCookie(w, c)
|
||||
}
|
||||
|
||||
func (s *AuthService) resetSession(c *gin.Context) {
|
||||
now := time.Now().Add(-1 * time.Second)
|
||||
session := sessions.Default(c)
|
||||
session.Set("gtm", now.Unix())
|
||||
session.Set("exp", now.Unix())
|
||||
session.Options(sessions.Options{
|
||||
HttpOnly: true,
|
||||
Secure: c.Request.TLS != nil,
|
||||
SameSite: http.SameSiteLaxMode,
|
||||
Path: s.cfg.BaseURL,
|
||||
MaxAge: -1,
|
||||
})
|
||||
if err := session.Save(); err != nil {
|
||||
logger.FromContext(c).WithError(err).Errorf("error resetting session")
|
||||
}
|
||||
}
|
||||
|
||||
// randBase64String is function to generate random base64 with set length (bytes)
|
||||
func randBase64String(nByte int) (string, error) {
|
||||
b := make([]byte, nByte)
|
||||
if _, err := io.ReadFull(rand.Reader, b); err != nil {
|
||||
return "", err
|
||||
}
|
||||
return base64.RawURLEncoding.EncodeToString(b), nil
|
||||
}
|
||||
|
||||
// randBytes is function to generate random bytes with set length (bytes)
|
||||
func randBytes(nByte int) ([]byte, error) {
|
||||
b := make([]byte, nByte)
|
||||
if _, err := io.ReadFull(rand.Reader, b); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return b, nil
|
||||
}
|
||||
|
||||
type info struct {
|
||||
Type string `json:"type"`
|
||||
Develop bool `json:"develop"`
|
||||
User models.User `json:"user"`
|
||||
Role models.Role `json:"role"`
|
||||
Providers []string `json:"providers"`
|
||||
Privs []string `json:"privileges"`
|
||||
OAuth bool `json:"oauth"`
|
||||
IssuedAt time.Time `json:"issued_at"`
|
||||
ExpiresAt time.Time `json:"expires_at"`
|
||||
}
|
||||
|
||||
// Info is function to return settings and current information about system and config
|
||||
// @Summary Retrieve current user and system settings
|
||||
// @Tags Public
|
||||
// @Produce json
|
||||
// @Security BearerAuth
|
||||
// @Param refresh_cookie query boolean false "boolean arg to refresh current cookie, use explicit false"
|
||||
// @Success 200 {object} response.successResp{data=info} "info received successful"
|
||||
// @Failure 403 {object} response.errorResp "getting info not permitted"
|
||||
// @Failure 404 {object} response.errorResp "user not found"
|
||||
// @Failure 500 {object} response.errorResp "internal error on getting information about system and config"
|
||||
// @Router /info [get]
|
||||
func (s *AuthService) Info(c *gin.Context) {
|
||||
var resp info
|
||||
|
||||
logger.FromContext(c).WithFields(logrus.Fields(c.Keys)).Trace("AuthService.Info")
|
||||
now := time.Now().Unix()
|
||||
uhash := c.GetString("uhash")
|
||||
uid := c.GetUint64("uid")
|
||||
tid := c.GetString("tid")
|
||||
exp := c.GetInt64("exp")
|
||||
gtm := c.GetInt64("gtm")
|
||||
cpt := c.GetString("cpt")
|
||||
privs := c.GetStringSlice("prm")
|
||||
|
||||
resp.Privs = privs
|
||||
resp.IssuedAt = time.Unix(gtm, 0).UTC()
|
||||
resp.ExpiresAt = time.Unix(exp, 0).UTC()
|
||||
resp.Develop = version.IsDevelopMode()
|
||||
resp.OAuth = tid == models.UserTypeOAuth.String()
|
||||
for name := range s.oauth {
|
||||
resp.Providers = append(resp.Providers, name)
|
||||
}
|
||||
|
||||
logger.FromContext(c).WithFields(logrus.Fields(
|
||||
map[string]any{
|
||||
"exp": exp,
|
||||
"gtm": gtm,
|
||||
"uhash": uhash,
|
||||
"now": now,
|
||||
"cpt": cpt,
|
||||
"uid": uid,
|
||||
"tid": tid,
|
||||
},
|
||||
)).Trace("AuthService.Info")
|
||||
|
||||
if uhash == "" || exp == 0 || gtm == 0 || now > exp {
|
||||
resp.Type = "guest"
|
||||
resp.Privs = []string{}
|
||||
response.Success(c, http.StatusOK, resp)
|
||||
return
|
||||
}
|
||||
|
||||
err := s.db.Take(&resp.User, "id = ?", uid).Related(&resp.Role).Error
|
||||
if err != nil {
|
||||
response.Error(c, response.ErrInfoUserNotFound, err)
|
||||
return
|
||||
} else if err = resp.User.Valid(); err != nil {
|
||||
logger.FromContext(c).WithError(err).Errorf("error validating user data '%s'", resp.User.Hash)
|
||||
response.Error(c, response.ErrInfoInvalidUserData, err)
|
||||
return
|
||||
}
|
||||
if err = s.db.Table("privileges").Where("role_id = ?", resp.User.RoleID).Pluck("name", &privs).Error; err != nil {
|
||||
logger.FromContext(c).WithError(err).Errorf("error getting user privileges list '%s'", resp.User.Hash)
|
||||
response.Error(c, response.ErrInfoInvalidUserData, err)
|
||||
return
|
||||
}
|
||||
|
||||
if cpt == "automation" {
|
||||
resp.Type = models.UserTypeAPI.String()
|
||||
// filter out privileges that are not supported for API tokens
|
||||
privs = slices.DeleteFunc(privs, func(priv string) bool {
|
||||
return strings.HasPrefix(priv, "users.") ||
|
||||
strings.HasPrefix(priv, "roles.") ||
|
||||
strings.HasPrefix(priv, "settings.user.") ||
|
||||
strings.HasPrefix(priv, "settings.tokens.")
|
||||
})
|
||||
resp.Privs = privs
|
||||
response.Success(c, http.StatusOK, resp)
|
||||
return
|
||||
}
|
||||
|
||||
resp.Type = "user"
|
||||
|
||||
// check 5 minutes timeout to refresh current token
|
||||
var fiveMins int64 = 5 * 60
|
||||
if now >= gtm+fiveMins && c.Query("refresh_cookie") != "false" {
|
||||
if err = s.refreshCookie(c, &resp, privs); err != nil {
|
||||
logger.FromContext(c).WithError(err).Errorf("failed to refresh token")
|
||||
// raise error when there is elapsing last five minutes
|
||||
if now >= gtm+int64(s.cfg.SessionTimeout)-fiveMins {
|
||||
response.Error(c, response.ErrAuthRequired, err)
|
||||
return
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// raise error when user has no permissions in the session auth cookie
|
||||
if resp.Type != "guest" && resp.Privs == nil {
|
||||
logger.FromContext(c).
|
||||
WithFields(logrus.Fields{
|
||||
"uid": resp.User.ID,
|
||||
"rid": resp.User.RoleID,
|
||||
"tid": resp.User.Type,
|
||||
}).
|
||||
Errorf("failed to get user privileges for '%s' '%s'", resp.User.Mail, resp.User.Name)
|
||||
response.Error(c, response.ErrInternal, err)
|
||||
return
|
||||
}
|
||||
|
||||
response.Success(c, http.StatusOK, resp)
|
||||
}
|
||||
@@ -0,0 +1,287 @@
|
||||
package services
|
||||
|
||||
import (
|
||||
"errors"
|
||||
"net/http"
|
||||
"slices"
|
||||
"strconv"
|
||||
|
||||
"pentagi/pkg/server/logger"
|
||||
"pentagi/pkg/server/models"
|
||||
"pentagi/pkg/server/rdb"
|
||||
"pentagi/pkg/server/response"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/jinzhu/gorm"
|
||||
)
|
||||
|
||||
type containers struct {
|
||||
Containers []models.Container `json:"containers"`
|
||||
Total uint64 `json:"total"`
|
||||
}
|
||||
|
||||
type containersGrouped struct {
|
||||
Grouped []string `json:"grouped"`
|
||||
Total uint64 `json:"total"`
|
||||
}
|
||||
|
||||
var containersSQLMappers = map[string]any{
|
||||
"id": "{{table}}.id",
|
||||
"type": "{{table}}.type",
|
||||
"name": "{{table}}.name",
|
||||
"image": "{{table}}.image",
|
||||
"status": "{{table}}.status",
|
||||
"local_id": "{{table}}.local_id",
|
||||
"local_dir": "{{table}}.local_dir",
|
||||
"flow_id": "{{table}}.flow_id",
|
||||
"created_at": "{{table}}.created_at",
|
||||
"updated_at": "{{table}}.updated_at",
|
||||
"data": "({{table}}.type || ' ' || {{table}}.name || ' ' || {{table}}.status || ' ' || {{table}}.local_id || ' ' || {{table}}.local_dir)",
|
||||
}
|
||||
|
||||
type ContainerService struct {
|
||||
db *gorm.DB
|
||||
}
|
||||
|
||||
func NewContainerService(db *gorm.DB) *ContainerService {
|
||||
return &ContainerService{
|
||||
db: db,
|
||||
}
|
||||
}
|
||||
|
||||
// GetContainers is a function to return containers list
|
||||
// @Summary Retrieve containers list
|
||||
// @Tags Containers
|
||||
// @Produce json
|
||||
// @Security BearerAuth
|
||||
// @Param request query rdb.TableQuery true "query table params"
|
||||
// @Success 200 {object} response.successResp{data=containers} "containers list received successful"
|
||||
// @Failure 400 {object} response.errorResp "invalid query request data"
|
||||
// @Failure 403 {object} response.errorResp "getting containers not permitted"
|
||||
// @Failure 500 {object} response.errorResp "internal error on getting containers"
|
||||
// @Router /containers/ [get]
|
||||
func (s *ContainerService) GetContainers(c *gin.Context) {
|
||||
var (
|
||||
err error
|
||||
query rdb.TableQuery
|
||||
resp containers
|
||||
)
|
||||
|
||||
if err = c.ShouldBindQuery(&query); err != nil {
|
||||
logger.FromContext(c).WithError(err).Errorf("error binding query")
|
||||
response.Error(c, response.ErrContainersInvalidRequest, err)
|
||||
return
|
||||
}
|
||||
|
||||
uid := c.GetUint64("uid")
|
||||
privs := c.GetStringSlice("prm")
|
||||
var scope func(db *gorm.DB) *gorm.DB
|
||||
if slices.Contains(privs, "containers.admin") {
|
||||
scope = func(db *gorm.DB) *gorm.DB {
|
||||
return db.
|
||||
Joins("INNER JOIN flows f ON f.id = containers.flow_id")
|
||||
}
|
||||
} else if slices.Contains(privs, "containers.view") {
|
||||
scope = func(db *gorm.DB) *gorm.DB {
|
||||
return db.
|
||||
Joins("INNER JOIN flows f ON f.id = containers.flow_id").
|
||||
Where("f.user_id = ?", uid)
|
||||
}
|
||||
} else {
|
||||
logger.FromContext(c).Errorf("error filtering user role permissions: permission not found")
|
||||
response.Error(c, response.ErrNotPermitted, nil)
|
||||
return
|
||||
}
|
||||
|
||||
query.Init("containers", containersSQLMappers)
|
||||
|
||||
if query.Group != "" {
|
||||
if _, ok := containersSQLMappers[query.Group]; !ok {
|
||||
logger.FromContext(c).Errorf("error finding containers grouped: group field not found")
|
||||
response.Error(c, response.ErrContainersInvalidRequest, errors.New("group field not found"))
|
||||
return
|
||||
}
|
||||
|
||||
var respGrouped containersGrouped
|
||||
if respGrouped.Total, err = query.QueryGrouped(s.db, &respGrouped.Grouped, scope); err != nil {
|
||||
logger.FromContext(c).WithError(err).Errorf("error finding containers grouped")
|
||||
response.Error(c, response.ErrInternal, err)
|
||||
return
|
||||
}
|
||||
|
||||
response.Success(c, http.StatusOK, respGrouped)
|
||||
return
|
||||
}
|
||||
|
||||
if resp.Total, err = query.Query(s.db, &resp.Containers, scope); err != nil {
|
||||
logger.FromContext(c).WithError(err).Errorf("error finding containers")
|
||||
response.Error(c, response.ErrInternal, err)
|
||||
return
|
||||
}
|
||||
|
||||
for i := 0; i < len(resp.Containers); i++ {
|
||||
if err = resp.Containers[i].Valid(); err != nil {
|
||||
logger.FromContext(c).WithError(err).Errorf("error validating container data '%d'", resp.Containers[i].ID)
|
||||
response.Error(c, response.ErrContainersInvalidData, err)
|
||||
return
|
||||
}
|
||||
}
|
||||
|
||||
response.Success(c, http.StatusOK, resp)
|
||||
}
|
||||
|
||||
// GetFlowContainers is a function to return containers list by flow id
|
||||
// @Summary Retrieve containers list by flow id
|
||||
// @Tags Containers
|
||||
// @Produce json
|
||||
// @Security BearerAuth
|
||||
// @Param flowID path int true "flow id" minimum(0)
|
||||
// @Param request query rdb.TableQuery true "query table params"
|
||||
// @Success 200 {object} response.successResp{data=containers} "containers list received successful"
|
||||
// @Failure 400 {object} response.errorResp "invalid query request data"
|
||||
// @Failure 403 {object} response.errorResp "getting containers not permitted"
|
||||
// @Failure 500 {object} response.errorResp "internal error on getting containers"
|
||||
// @Router /flows/{flowID}/containers/ [get]
|
||||
func (s *ContainerService) GetFlowContainers(c *gin.Context) {
|
||||
var (
|
||||
err error
|
||||
flowID uint64
|
||||
query rdb.TableQuery
|
||||
resp containers
|
||||
)
|
||||
|
||||
if flowID, err = strconv.ParseUint(c.Param("flowID"), 10, 64); err != nil {
|
||||
logger.FromContext(c).WithError(err).Errorf("error parsing flow id")
|
||||
response.Error(c, response.ErrContainersInvalidRequest, err)
|
||||
return
|
||||
}
|
||||
|
||||
if err = c.ShouldBindQuery(&query); err != nil {
|
||||
logger.FromContext(c).WithError(err).Errorf("error binding query")
|
||||
response.Error(c, response.ErrContainersInvalidRequest, err)
|
||||
return
|
||||
}
|
||||
|
||||
uid := c.GetUint64("uid")
|
||||
privs := c.GetStringSlice("prm")
|
||||
var scope func(db *gorm.DB) *gorm.DB
|
||||
if slices.Contains(privs, "containers.admin") {
|
||||
scope = func(db *gorm.DB) *gorm.DB {
|
||||
return db.
|
||||
Joins("INNER JOIN flows f ON f.id = containers.flow_id").
|
||||
Where("f.id = ?", flowID)
|
||||
}
|
||||
} else if slices.Contains(privs, "containers.view") {
|
||||
scope = func(db *gorm.DB) *gorm.DB {
|
||||
return db.
|
||||
Joins("INNER JOIN flows f ON f.id = containers.flow_id").
|
||||
Where("f.id = ? AND f.user_id = ?", flowID, uid)
|
||||
}
|
||||
} else {
|
||||
logger.FromContext(c).Errorf("error filtering user role permissions: permission not found")
|
||||
response.Error(c, response.ErrNotPermitted, nil)
|
||||
return
|
||||
}
|
||||
|
||||
query.Init("containers", containersSQLMappers)
|
||||
|
||||
if query.Group != "" {
|
||||
if _, ok := containersSQLMappers[query.Group]; !ok {
|
||||
logger.FromContext(c).Errorf("error finding containers grouped: group field not found")
|
||||
response.Error(c, response.ErrContainersInvalidRequest, errors.New("group field not found"))
|
||||
return
|
||||
}
|
||||
|
||||
var respGrouped containersGrouped
|
||||
if respGrouped.Total, err = query.QueryGrouped(s.db, &respGrouped.Grouped, scope); err != nil {
|
||||
logger.FromContext(c).WithError(err).Errorf("error finding containers grouped")
|
||||
response.Error(c, response.ErrInternal, err)
|
||||
return
|
||||
}
|
||||
|
||||
response.Success(c, http.StatusOK, respGrouped)
|
||||
return
|
||||
}
|
||||
|
||||
if resp.Total, err = query.Query(s.db, &resp.Containers, scope); err != nil {
|
||||
logger.FromContext(c).WithError(err).Errorf("error finding containers")
|
||||
response.Error(c, response.ErrInternal, err)
|
||||
return
|
||||
}
|
||||
|
||||
for i := 0; i < len(resp.Containers); i++ {
|
||||
if err = resp.Containers[i].Valid(); err != nil {
|
||||
logger.FromContext(c).WithError(err).Errorf("error validating container data '%d'", resp.Containers[i].ID)
|
||||
response.Error(c, response.ErrContainersInvalidData, err)
|
||||
return
|
||||
}
|
||||
}
|
||||
|
||||
response.Success(c, http.StatusOK, resp)
|
||||
}
|
||||
|
||||
// GetFlowContainer is a function to return container info by id and flow id
|
||||
// @Summary Retrieve container info by id and flow id
|
||||
// @Tags Containers
|
||||
// @Produce json
|
||||
// @Security BearerAuth
|
||||
// @Param flowID path int true "flow id" minimum(0)
|
||||
// @Param containerID path int true "container id" minimum(0)
|
||||
// @Success 200 {object} response.successResp{data=models.Container} "container info received successful"
|
||||
// @Failure 403 {object} response.errorResp "getting container not permitted"
|
||||
// @Failure 404 {object} response.errorResp "container not found"
|
||||
// @Failure 500 {object} response.errorResp "internal error on getting container"
|
||||
// @Router /flows/{flowID}/containers/{containerID} [get]
|
||||
func (s *ContainerService) GetFlowContainer(c *gin.Context) {
|
||||
var (
|
||||
err error
|
||||
containerID uint64
|
||||
flowID uint64
|
||||
resp models.Container
|
||||
)
|
||||
|
||||
if flowID, err = strconv.ParseUint(c.Param("flowID"), 10, 64); err != nil {
|
||||
logger.FromContext(c).WithError(err).Errorf("error parsing flow id")
|
||||
response.Error(c, response.ErrContainersInvalidRequest, err)
|
||||
return
|
||||
}
|
||||
if containerID, err = strconv.ParseUint(c.Param("containerID"), 10, 64); err != nil {
|
||||
logger.FromContext(c).WithError(err).Errorf("error parsing container id")
|
||||
response.Error(c, response.ErrContainersInvalidRequest, err)
|
||||
return
|
||||
}
|
||||
|
||||
uid := c.GetUint64("uid")
|
||||
privs := c.GetStringSlice("prm")
|
||||
var scope func(db *gorm.DB) *gorm.DB
|
||||
if slices.Contains(privs, "containers.admin") {
|
||||
scope = func(db *gorm.DB) *gorm.DB {
|
||||
return db.Where("f.id = ?", flowID)
|
||||
}
|
||||
} else if slices.Contains(privs, "containers.view") {
|
||||
scope = func(db *gorm.DB) *gorm.DB {
|
||||
return db.Where("f.id = ? AND f.user_id = ?", flowID, uid)
|
||||
}
|
||||
} else {
|
||||
logger.FromContext(c).Errorf("error filtering user role permissions: permission not found")
|
||||
response.Error(c, response.ErrNotPermitted, nil)
|
||||
return
|
||||
}
|
||||
|
||||
err = s.db.Model(&resp).
|
||||
Joins("INNER JOIN flows f ON f.id = flow_id").
|
||||
Scopes(scope).
|
||||
Where("containers.id = ?", containerID).
|
||||
Take(&resp).Error
|
||||
if err != nil {
|
||||
logger.FromContext(c).WithError(err).Errorf("error on getting container by id")
|
||||
if gorm.IsRecordNotFoundError(err) {
|
||||
response.Error(c, response.ErrContainersNotFound, err)
|
||||
} else {
|
||||
response.Error(c, response.ErrInternal, err)
|
||||
}
|
||||
return
|
||||
}
|
||||
|
||||
response.Success(c, http.StatusOK, resp)
|
||||
}
|
||||
File diff suppressed because it is too large
Load Diff
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,675 @@
|
||||
package services
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"net/http"
|
||||
"slices"
|
||||
"strconv"
|
||||
|
||||
"pentagi/pkg/controller"
|
||||
"pentagi/pkg/database"
|
||||
"pentagi/pkg/graph/subscriptions"
|
||||
"pentagi/pkg/providers"
|
||||
"pentagi/pkg/providers/provider"
|
||||
"pentagi/pkg/server/logger"
|
||||
"pentagi/pkg/server/models"
|
||||
"pentagi/pkg/server/rdb"
|
||||
"pentagi/pkg/server/response"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/jinzhu/gorm"
|
||||
)
|
||||
|
||||
type flows struct {
|
||||
Flows []models.Flow `json:"flows"`
|
||||
Total uint64 `json:"total"`
|
||||
}
|
||||
|
||||
type flowsGrouped struct {
|
||||
Grouped []string `json:"grouped"`
|
||||
Total uint64 `json:"total"`
|
||||
}
|
||||
|
||||
var flowsSQLMappers = map[string]any{
|
||||
"id": "{{table}}.id",
|
||||
"status": "{{table}}.status",
|
||||
"title": "{{table}}.title",
|
||||
"model": "{{table}}.model",
|
||||
"model_provider_name": "{{table}}.model_provider_name",
|
||||
"model_provider_type": "{{table}}.model_provider_type",
|
||||
"language": "{{table}}.language",
|
||||
"created_at": "{{table}}.created_at",
|
||||
"updated_at": "{{table}}.updated_at",
|
||||
"data": "({{table}}.status || ' ' || {{table}}.title || ' ' || {{table}}.model || ' ' || {{table}}.model_provider || ' ' || {{table}}.language)",
|
||||
}
|
||||
|
||||
type FlowService struct {
|
||||
db *gorm.DB
|
||||
pc providers.ProviderController
|
||||
fc controller.FlowController
|
||||
ss subscriptions.SubscriptionsController
|
||||
}
|
||||
|
||||
func NewFlowService(
|
||||
db *gorm.DB,
|
||||
pc providers.ProviderController,
|
||||
fc controller.FlowController,
|
||||
ss subscriptions.SubscriptionsController,
|
||||
) *FlowService {
|
||||
return &FlowService{
|
||||
db: db,
|
||||
pc: pc,
|
||||
fc: fc,
|
||||
ss: ss,
|
||||
}
|
||||
}
|
||||
|
||||
// GetFlows is a function to return flows list
|
||||
// @Summary Retrieve flows list
|
||||
// @Tags Flows
|
||||
// @Produce json
|
||||
// @Security BearerAuth
|
||||
// @Param request query rdb.TableQuery true "query table params"
|
||||
// @Success 200 {object} response.successResp{data=flows} "flows list received successful"
|
||||
// @Failure 400 {object} response.errorResp "invalid query request data"
|
||||
// @Failure 403 {object} response.errorResp "getting flows not permitted"
|
||||
// @Failure 500 {object} response.errorResp "internal error on getting flows"
|
||||
// @Router /flows/ [get]
|
||||
func (s *FlowService) GetFlows(c *gin.Context) {
|
||||
var (
|
||||
err error
|
||||
query rdb.TableQuery
|
||||
resp flows
|
||||
)
|
||||
|
||||
if err = c.ShouldBindQuery(&query); err != nil {
|
||||
logger.FromContext(c).WithError(err).Errorf("error binding query")
|
||||
response.Error(c, response.ErrFlowsInvalidRequest, err)
|
||||
return
|
||||
}
|
||||
|
||||
uid := c.GetUint64("uid")
|
||||
privs := c.GetStringSlice("prm")
|
||||
var scope func(db *gorm.DB) *gorm.DB
|
||||
if slices.Contains(privs, "flows.admin") {
|
||||
scope = func(db *gorm.DB) *gorm.DB {
|
||||
return db
|
||||
}
|
||||
} else if slices.Contains(privs, "flows.view") {
|
||||
scope = func(db *gorm.DB) *gorm.DB {
|
||||
return db.Where("user_id = ?", uid)
|
||||
}
|
||||
} else {
|
||||
logger.FromContext(c).Errorf("error filtering user role permissions: permission not found")
|
||||
response.Error(c, response.ErrNotPermitted, nil)
|
||||
return
|
||||
}
|
||||
|
||||
query.Init("flows", flowsSQLMappers)
|
||||
|
||||
if query.Group != "" {
|
||||
if _, ok := flowsSQLMappers[query.Group]; !ok {
|
||||
logger.FromContext(c).Errorf("error finding flows grouped: group field not found")
|
||||
response.Error(c, response.ErrFlowsInvalidRequest, errors.New("group field not found"))
|
||||
return
|
||||
}
|
||||
|
||||
var respGrouped flowsGrouped
|
||||
if respGrouped.Total, err = query.QueryGrouped(s.db, &respGrouped.Grouped, scope); err != nil {
|
||||
logger.FromContext(c).WithError(err).Errorf("error finding flows grouped")
|
||||
response.Error(c, response.ErrInternal, err)
|
||||
return
|
||||
}
|
||||
|
||||
response.Success(c, http.StatusOK, respGrouped)
|
||||
return
|
||||
}
|
||||
|
||||
if resp.Total, err = query.Query(s.db, &resp.Flows, scope); err != nil {
|
||||
logger.FromContext(c).WithError(err).Errorf("error finding flows")
|
||||
response.Error(c, response.ErrInternal, err)
|
||||
return
|
||||
}
|
||||
|
||||
for i := 0; i < len(resp.Flows); i++ {
|
||||
if err = resp.Flows[i].Valid(); err != nil {
|
||||
logger.FromContext(c).WithError(err).Errorf("error validating flow data '%d'", resp.Flows[i].ID)
|
||||
response.Error(c, response.ErrFlowsInvalidData, err)
|
||||
return
|
||||
}
|
||||
}
|
||||
|
||||
response.Success(c, http.StatusOK, resp)
|
||||
}
|
||||
|
||||
// GetFlow is a function to return flow by id
|
||||
// @Summary Retrieve flow by id
|
||||
// @Tags Flows
|
||||
// @Produce json
|
||||
// @Security BearerAuth
|
||||
// @Param flowID path int true "flow id" minimum(0)
|
||||
// @Success 200 {object} response.successResp{data=models.Flow} "flow received successful"
|
||||
// @Failure 403 {object} response.errorResp "getting flow not permitted"
|
||||
// @Failure 404 {object} response.errorResp "flow not found"
|
||||
// @Failure 500 {object} response.errorResp "internal error on getting flow"
|
||||
// @Router /flows/{flowID} [get]
|
||||
func (s *FlowService) GetFlow(c *gin.Context) {
|
||||
var (
|
||||
err error
|
||||
flowID uint64
|
||||
resp models.Flow
|
||||
)
|
||||
|
||||
if flowID, err = strconv.ParseUint(c.Param("flowID"), 10, 64); err != nil {
|
||||
logger.FromContext(c).WithError(err).Errorf("error parsing flow id")
|
||||
response.Error(c, response.ErrFlowsInvalidRequest, err)
|
||||
return
|
||||
}
|
||||
|
||||
uid := c.GetUint64("uid")
|
||||
privs := c.GetStringSlice("prm")
|
||||
var scope func(db *gorm.DB) *gorm.DB
|
||||
if slices.Contains(privs, "flows.admin") {
|
||||
scope = func(db *gorm.DB) *gorm.DB {
|
||||
return db.Where("id = ?", flowID)
|
||||
}
|
||||
} else if slices.Contains(privs, "flows.view") {
|
||||
scope = func(db *gorm.DB) *gorm.DB {
|
||||
return db.Where("id = ? AND user_id = ?", flowID, uid)
|
||||
}
|
||||
} else {
|
||||
logger.FromContext(c).Errorf("error filtering user role permissions: permission not found")
|
||||
response.Error(c, response.ErrNotPermitted, nil)
|
||||
return
|
||||
}
|
||||
|
||||
if err = s.db.Model(&resp).Scopes(scope).Take(&resp).Error; err != nil {
|
||||
logger.FromContext(c).WithError(err).Errorf("error on getting flow by id")
|
||||
if gorm.IsRecordNotFoundError(err) {
|
||||
response.Error(c, response.ErrFlowsNotFound, err)
|
||||
} else {
|
||||
response.Error(c, response.ErrInternal, err)
|
||||
}
|
||||
return
|
||||
}
|
||||
|
||||
response.Success(c, http.StatusOK, resp)
|
||||
}
|
||||
|
||||
// GetFlowGraph is a function to return flow graph by id
|
||||
// @Summary Retrieve flow graph by id
|
||||
// @Tags Flows
|
||||
// @Produce json
|
||||
// @Security BearerAuth
|
||||
// @Param flowID path int true "flow id" minimum(0)
|
||||
// @Success 200 {object} response.successResp{data=models.FlowTasksSubtasks} "flow graph received successful"
|
||||
// @Failure 403 {object} response.errorResp "getting flow graph not permitted"
|
||||
// @Failure 404 {object} response.errorResp "flow graph not found"
|
||||
// @Failure 500 {object} response.errorResp "internal error on getting flow graph"
|
||||
// @Router /flows/{flowID}/graph [get]
|
||||
func (s *FlowService) GetFlowGraph(c *gin.Context) {
|
||||
var (
|
||||
err error
|
||||
flowID uint64
|
||||
resp models.FlowTasksSubtasks
|
||||
tids []uint64
|
||||
)
|
||||
|
||||
if flowID, err = strconv.ParseUint(c.Param("flowID"), 10, 64); err != nil {
|
||||
logger.FromContext(c).WithError(err).Errorf("error parsing flow id")
|
||||
response.Error(c, response.ErrFlowsInvalidRequest, err)
|
||||
return
|
||||
}
|
||||
|
||||
uid := c.GetUint64("uid")
|
||||
privs := c.GetStringSlice("prm")
|
||||
var scope func(db *gorm.DB) *gorm.DB
|
||||
if slices.Contains(privs, "flows.admin") {
|
||||
scope = func(db *gorm.DB) *gorm.DB {
|
||||
return db.Where("id = ?", flowID)
|
||||
}
|
||||
} else if slices.Contains(privs, "flows.view") {
|
||||
scope = func(db *gorm.DB) *gorm.DB {
|
||||
return db.Where("id = ? AND user_id = ?", flowID, uid)
|
||||
}
|
||||
} else {
|
||||
logger.FromContext(c).Errorf("error filtering user role permissions: permission not found")
|
||||
response.Error(c, response.ErrNotPermitted, nil)
|
||||
return
|
||||
}
|
||||
|
||||
err = s.db.Model(&resp).
|
||||
Scopes(scope).
|
||||
Take(&resp).Error
|
||||
if err != nil {
|
||||
logger.FromContext(c).WithError(err).Errorf("error on getting flow by id")
|
||||
if gorm.IsRecordNotFoundError(err) {
|
||||
response.Error(c, response.ErrFlowsNotFound, err)
|
||||
} else {
|
||||
response.Error(c, response.ErrInternal, err)
|
||||
}
|
||||
return
|
||||
}
|
||||
|
||||
isTasksAdmin := slices.Contains(privs, "tasks.admin")
|
||||
isTasksView := slices.Contains(privs, "tasks.view")
|
||||
if !(resp.UserID == uid && isTasksView) && !(resp.UserID != uid && isTasksAdmin) {
|
||||
response.Success(c, http.StatusOK, resp)
|
||||
return
|
||||
}
|
||||
|
||||
if resp.UserID != uid && !slices.Contains(privs, "tasks.admin") {
|
||||
response.Success(c, http.StatusOK, resp)
|
||||
return
|
||||
}
|
||||
|
||||
err = s.db.Model(&resp).Association("tasks").Find(&resp.Tasks).Error
|
||||
if err != nil {
|
||||
logger.FromContext(c).WithError(err).Errorf("error on getting flow tasks")
|
||||
response.Error(c, response.ErrInternal, err)
|
||||
return
|
||||
}
|
||||
|
||||
isSubtasksAdmin := slices.Contains(privs, "subtasks.admin")
|
||||
isSubtasksView := slices.Contains(privs, "subtasks.view")
|
||||
if !(resp.UserID == uid && isSubtasksView) && !(resp.UserID != uid && isSubtasksAdmin) {
|
||||
response.Success(c, http.StatusOK, resp)
|
||||
return
|
||||
}
|
||||
|
||||
for _, task := range resp.Tasks {
|
||||
tids = append(tids, task.ID)
|
||||
}
|
||||
|
||||
var subtasks []models.Subtask
|
||||
err = s.db.Model(&subtasks).Where("task_id IN (?)", tids).Find(&subtasks).Error
|
||||
if err != nil {
|
||||
logger.FromContext(c).WithError(err).Errorf("error on getting flow subtasks")
|
||||
response.Error(c, response.ErrInternal, err)
|
||||
return
|
||||
}
|
||||
|
||||
tasksSubtasks := map[uint64][]models.Subtask{}
|
||||
for _, subtask := range subtasks {
|
||||
tasksSubtasks[subtask.TaskID] = append(tasksSubtasks[subtask.TaskID], subtask)
|
||||
}
|
||||
|
||||
for i := range resp.Tasks {
|
||||
resp.Tasks[i].Subtasks = tasksSubtasks[resp.Tasks[i].ID]
|
||||
}
|
||||
|
||||
if err = resp.Valid(); err != nil {
|
||||
logger.FromContext(c).WithError(err).Errorf("error validating flow data '%d'", flowID)
|
||||
response.Error(c, response.ErrFlowsInvalidData, err)
|
||||
return
|
||||
}
|
||||
|
||||
response.Success(c, http.StatusOK, resp)
|
||||
}
|
||||
|
||||
// CreateFlow is a function to create new flow with custom functions
|
||||
// @Summary Create new flow with custom functions
|
||||
// @Tags Flows
|
||||
// @Accept json
|
||||
// @Produce json
|
||||
// @Security BearerAuth
|
||||
// @Param json body models.CreateFlow true "flow model to create"
|
||||
// @Success 201 {object} response.successResp{data=models.Flow} "flow created successful"
|
||||
// @Failure 400 {object} response.errorResp "invalid flow request data"
|
||||
// @Failure 403 {object} response.errorResp "creating flow not permitted"
|
||||
// @Failure 500 {object} response.errorResp "internal error on creating flow"
|
||||
// @Router /flows/ [post]
|
||||
func (s *FlowService) CreateFlow(c *gin.Context) {
|
||||
var (
|
||||
err error
|
||||
flow models.Flow
|
||||
createFlow models.CreateFlow
|
||||
)
|
||||
|
||||
if err := c.ShouldBindJSON(&createFlow); err != nil {
|
||||
logger.FromContext(c).WithError(err).Errorf("error binding JSON")
|
||||
response.Error(c, response.ErrFlowsInvalidRequest, err)
|
||||
return
|
||||
}
|
||||
|
||||
privs := c.GetStringSlice("prm")
|
||||
if !slices.Contains(privs, "flows.create") {
|
||||
logger.FromContext(c).Errorf("error filtering user role permissions: permission not found")
|
||||
response.Error(c, response.ErrNotPermitted, nil)
|
||||
return
|
||||
}
|
||||
|
||||
if err := createFlow.Valid(); err != nil {
|
||||
logger.FromContext(c).WithError(err).Errorf("error validating flow data")
|
||||
response.Error(c, response.ErrFlowsInvalidData, err)
|
||||
return
|
||||
}
|
||||
|
||||
uid := c.GetUint64("uid")
|
||||
prvname := provider.ProviderName(createFlow.Provider)
|
||||
|
||||
prv, err := s.pc.GetProvider(c, prvname, int64(uid))
|
||||
if err != nil {
|
||||
logger.FromContext(c).WithError(err).Errorf("error getting provider: not found")
|
||||
response.Error(c, response.ErrInternal, err)
|
||||
return
|
||||
}
|
||||
prvtype := prv.Type()
|
||||
|
||||
dbResources, err := validateServiceResources(s.db, uid, privs, createFlow.ResourceIDs)
|
||||
if err != nil {
|
||||
logger.FromContext(c).WithError(err).Errorf("error validating resource ids")
|
||||
response.Error(c, response.ErrFlowsInvalidRequest, err)
|
||||
return
|
||||
}
|
||||
|
||||
fw, err := s.fc.CreateFlow(c, int64(uid), createFlow.Input, prvname, prvtype, createFlow.Functions, dbResources)
|
||||
if err != nil {
|
||||
logger.FromContext(c).WithError(err).Errorf("error creating flow")
|
||||
response.Error(c, response.ErrInternal, err)
|
||||
return
|
||||
}
|
||||
|
||||
err = s.db.Model(&flow).Where("id = ?", fw.GetFlowID()).Take(&flow).Error
|
||||
if err != nil {
|
||||
logger.FromContext(c).WithError(err).Errorf("error getting flow by id")
|
||||
response.Error(c, response.ErrInternal, err)
|
||||
return
|
||||
}
|
||||
|
||||
response.Success(c, http.StatusCreated, flow)
|
||||
}
|
||||
|
||||
// PatchFlow is a function to patch flow
|
||||
// @Summary Patch flow
|
||||
// @Tags Flows
|
||||
// @Accept json
|
||||
// @Produce json
|
||||
// @Security BearerAuth
|
||||
// @Param flowID path int true "flow id" minimum(0)
|
||||
// @Param json body models.PatchFlow true "flow model to patch"
|
||||
// @Success 200 {object} response.successResp{data=models.Flow} "flow patched successful"
|
||||
// @Failure 400 {object} response.errorResp "invalid flow request data"
|
||||
// @Failure 403 {object} response.errorResp "patching flow not permitted"
|
||||
// @Failure 500 {object} response.errorResp "internal error on patching flow"
|
||||
// @Router /flows/{flowID} [put]
|
||||
func (s *FlowService) PatchFlow(c *gin.Context) {
|
||||
var (
|
||||
err error
|
||||
flow models.Flow
|
||||
flowID uint64
|
||||
patchFlow models.PatchFlow
|
||||
)
|
||||
|
||||
if err := c.ShouldBindJSON(&patchFlow); err != nil {
|
||||
logger.FromContext(c).WithError(err).Errorf("error binding JSON")
|
||||
response.Error(c, response.ErrFlowsInvalidRequest, err)
|
||||
return
|
||||
}
|
||||
|
||||
if err := patchFlow.Valid(); err != nil {
|
||||
logger.FromContext(c).WithError(err).Errorf("error validating flow data")
|
||||
response.Error(c, response.ErrFlowsInvalidData, err)
|
||||
return
|
||||
}
|
||||
|
||||
flowID, err = strconv.ParseUint(c.Param("flowID"), 10, 64)
|
||||
if err != nil {
|
||||
logger.FromContext(c).WithError(err).Errorf("error parsing flow id")
|
||||
response.Error(c, response.ErrFlowsInvalidRequest, err)
|
||||
return
|
||||
}
|
||||
|
||||
uid := c.GetUint64("uid")
|
||||
privs := c.GetStringSlice("prm")
|
||||
var scope func(db *gorm.DB) *gorm.DB
|
||||
if slices.Contains(privs, "flows.admin") {
|
||||
scope = func(db *gorm.DB) *gorm.DB {
|
||||
return db.Where("id = ?", flowID)
|
||||
}
|
||||
} else if slices.Contains(privs, "flows.edit") {
|
||||
scope = func(db *gorm.DB) *gorm.DB {
|
||||
return db.Where("id = ? AND user_id = ?", flowID, uid)
|
||||
}
|
||||
} else {
|
||||
logger.FromContext(c).Errorf("error filtering user role permissions: permission not found")
|
||||
response.Error(c, response.ErrNotPermitted, nil)
|
||||
return
|
||||
}
|
||||
|
||||
if err = s.db.Model(&flow).Scopes(scope).Take(&flow).Error; err != nil {
|
||||
logger.FromContext(c).WithError(err).Errorf("error getting flow by id")
|
||||
if gorm.IsRecordNotFoundError(err) {
|
||||
response.Error(c, response.ErrFlowsNotFound, err)
|
||||
} else {
|
||||
response.Error(c, response.ErrInternal, err)
|
||||
}
|
||||
return
|
||||
}
|
||||
|
||||
fw, err := s.fc.GetFlow(c, int64(flow.ID))
|
||||
if err != nil {
|
||||
logger.FromContext(c).WithError(err).Errorf("error getting flow by id in flow controller")
|
||||
response.Error(c, response.ErrInternal, err)
|
||||
return
|
||||
}
|
||||
|
||||
switch patchFlow.Action {
|
||||
case "stop":
|
||||
if err := fw.Stop(c); err != nil {
|
||||
logger.FromContext(c).WithError(err).Errorf("error stopping flow")
|
||||
response.Error(c, response.ErrInternal, err)
|
||||
return
|
||||
}
|
||||
case "finish":
|
||||
if err := fw.Finish(c); err != nil {
|
||||
logger.FromContext(c).WithError(err).Errorf("error finishing flow")
|
||||
response.Error(c, response.ErrInternal, err)
|
||||
return
|
||||
}
|
||||
case "input":
|
||||
if patchFlow.Input == nil || *patchFlow.Input == "" {
|
||||
logger.FromContext(c).Errorf("error sending input to flow: input is empty")
|
||||
response.Error(c, response.ErrFlowsInvalidRequest, nil)
|
||||
return
|
||||
}
|
||||
|
||||
var prv provider.Provider
|
||||
if patchFlow.Provider != nil && *patchFlow.Provider != "" {
|
||||
prv, err = s.pc.GetProvider(c, provider.ProviderName(*patchFlow.Provider), int64(uid))
|
||||
if err != nil {
|
||||
logger.FromContext(c).WithError(err).Errorf("error getting provider by name")
|
||||
response.Error(c, response.ErrInternal, err)
|
||||
return
|
||||
}
|
||||
}
|
||||
|
||||
dbResources, err := validateServiceResources(s.db, uid, privs, patchFlow.ResourceIDs)
|
||||
if err != nil {
|
||||
logger.FromContext(c).WithError(err).Errorf("error validating resource ids")
|
||||
response.Error(c, response.ErrFlowsInvalidRequest, err)
|
||||
return
|
||||
}
|
||||
|
||||
if err := fw.PutInput(c, *patchFlow.Input, prv, dbResources); err != nil {
|
||||
logger.FromContext(c).WithError(err).Errorf("error sending input to flow")
|
||||
response.Error(c, response.ErrInternal, err)
|
||||
return
|
||||
}
|
||||
case "rename":
|
||||
if patchFlow.Name == nil || *patchFlow.Name == "" {
|
||||
logger.FromContext(c).Errorf("error renaming flow: name is empty")
|
||||
response.Error(c, response.ErrFlowsInvalidRequest, nil)
|
||||
return
|
||||
}
|
||||
if err := fw.Rename(c, *patchFlow.Name); err != nil {
|
||||
logger.FromContext(c).WithError(err).Errorf("error renaming flow")
|
||||
response.Error(c, response.ErrInternal, err)
|
||||
return
|
||||
}
|
||||
default:
|
||||
logger.FromContext(c).Errorf("error filtering flow action")
|
||||
response.Error(c, response.ErrFlowsInvalidRequest, nil)
|
||||
return
|
||||
}
|
||||
|
||||
if err = s.db.Model(&flow).Scopes(scope).Take(&flow).Error; err != nil {
|
||||
logger.FromContext(c).WithError(err).Errorf("error getting flow by id")
|
||||
if gorm.IsRecordNotFoundError(err) {
|
||||
response.Error(c, response.ErrFlowsNotFound, err)
|
||||
} else {
|
||||
response.Error(c, response.ErrInternal, err)
|
||||
}
|
||||
return
|
||||
}
|
||||
|
||||
response.Success(c, http.StatusOK, flow)
|
||||
}
|
||||
|
||||
// DeleteFlow is a function to delete flow by id
|
||||
// @Summary Delete flow by id
|
||||
// @Tags Flows
|
||||
// @Security BearerAuth
|
||||
// @Param flowID path int true "flow id" minimum(0)
|
||||
// @Success 200 {object} response.successResp{data=models.Flow} "flow deleted successful"
|
||||
// @Failure 403 {object} response.errorResp "deleting flow not permitted"
|
||||
// @Failure 404 {object} response.errorResp "flow not found"
|
||||
// @Failure 500 {object} response.errorResp "internal error on deleting flow"
|
||||
// @Router /flows/{flowID} [delete]
|
||||
func (s *FlowService) DeleteFlow(c *gin.Context) {
|
||||
var (
|
||||
err error
|
||||
flow models.Flow
|
||||
flowID uint64
|
||||
)
|
||||
|
||||
flowID, err = strconv.ParseUint(c.Param("flowID"), 10, 64)
|
||||
if err != nil {
|
||||
logger.FromContext(c).WithError(err).Errorf("error parsing flow id")
|
||||
response.Error(c, response.ErrFlowsInvalidRequest, err)
|
||||
return
|
||||
}
|
||||
|
||||
uid := c.GetUint64("uid")
|
||||
privs := c.GetStringSlice("prm")
|
||||
var scope func(db *gorm.DB) *gorm.DB
|
||||
if slices.Contains(privs, "flows.admin") {
|
||||
scope = func(db *gorm.DB) *gorm.DB {
|
||||
return db.Where("id = ?", flowID)
|
||||
}
|
||||
} else if slices.Contains(privs, "flows.delete") {
|
||||
scope = func(db *gorm.DB) *gorm.DB {
|
||||
return db.Where("id = ? AND user_id = ?", flowID, uid)
|
||||
}
|
||||
} else {
|
||||
logger.FromContext(c).Errorf("error filtering user role permissions: permission not found")
|
||||
response.Error(c, response.ErrNotPermitted, nil)
|
||||
return
|
||||
}
|
||||
|
||||
if err = s.db.Model(&flow).Scopes(scope).Take(&flow).Error; err != nil {
|
||||
logger.FromContext(c).WithError(err).Errorf("error getting flow by id")
|
||||
if gorm.IsRecordNotFoundError(err) {
|
||||
response.Error(c, response.ErrFlowsNotFound, err)
|
||||
} else {
|
||||
response.Error(c, response.ErrInternal, err)
|
||||
}
|
||||
return
|
||||
}
|
||||
|
||||
if err := s.fc.FinishFlow(c, int64(flow.ID)); err != nil {
|
||||
logger.FromContext(c).WithError(err).Errorf("error stopping flow")
|
||||
response.Error(c, response.ErrInternal, err)
|
||||
return
|
||||
}
|
||||
|
||||
var containers []models.Container
|
||||
err = s.db.Model(&containers).Where("flow_id = ?", flow.ID).Find(&containers).Error
|
||||
if err != nil {
|
||||
logger.FromContext(c).WithError(err).Errorf("error getting flow containers")
|
||||
response.Error(c, response.ErrInternal, err)
|
||||
return
|
||||
}
|
||||
|
||||
if err = s.db.Scopes(scope).Delete(&flow).Error; err != nil {
|
||||
logger.FromContext(c).WithError(err).Errorf("error deleting flow by id")
|
||||
if gorm.IsRecordNotFoundError(err) {
|
||||
response.Error(c, response.ErrFlowsNotFound, err)
|
||||
} else {
|
||||
response.Error(c, response.ErrInternal, err)
|
||||
}
|
||||
return
|
||||
}
|
||||
|
||||
// Best-effort cleanup: remove accumulated long-term memory documents for this
|
||||
// flow from the vector store. They will never be re-used after the flow is gone.
|
||||
if err := s.db.Exec(
|
||||
`DELETE FROM langchain_pg_embedding
|
||||
WHERE collection_id = (SELECT uuid FROM langchain_pg_collection WHERE name = 'langchain')
|
||||
AND COALESCE(cmetadata ->> 'doc_type', '') = 'memory'
|
||||
AND (cmetadata ->> 'flow_id') = ?`,
|
||||
strconv.FormatUint(flow.ID, 10),
|
||||
).Error; err != nil {
|
||||
logger.FromContext(c).WithError(err).Warnf("failed to clean up memory documents for deleted flow %d", flow.ID)
|
||||
}
|
||||
|
||||
flowDB, err := convertFlowToDatabase(flow)
|
||||
if err != nil {
|
||||
logger.FromContext(c).WithError(err).Errorf("error converting flow to database")
|
||||
response.Error(c, response.ErrInternal, err)
|
||||
return
|
||||
}
|
||||
|
||||
containersDB := make([]database.Container, 0, len(containers))
|
||||
for _, container := range containers {
|
||||
containersDB = append(containersDB, convertContainerToDatabase(container))
|
||||
}
|
||||
|
||||
if s.ss != nil {
|
||||
publisher := s.ss.NewFlowPublisher(int64(flow.UserID), int64(flow.ID))
|
||||
publisher.FlowUpdated(c, flowDB, containersDB)
|
||||
publisher.FlowDeleted(c, flowDB, containersDB)
|
||||
}
|
||||
|
||||
response.Success(c, http.StatusOK, flow)
|
||||
}
|
||||
|
||||
func convertFlowToDatabase(flow models.Flow) (database.Flow, error) {
|
||||
functions, err := json.Marshal(flow.Functions)
|
||||
if err != nil {
|
||||
return database.Flow{}, err
|
||||
}
|
||||
|
||||
return database.Flow{
|
||||
ID: int64(flow.ID),
|
||||
Status: database.FlowStatus(flow.Status),
|
||||
Title: flow.Title,
|
||||
Model: flow.Model,
|
||||
ModelProviderName: flow.ModelProviderName,
|
||||
Language: flow.Language,
|
||||
Functions: functions,
|
||||
UserID: int64(flow.UserID),
|
||||
CreatedAt: database.TimeToNullTime(flow.CreatedAt),
|
||||
UpdatedAt: database.TimeToNullTime(flow.UpdatedAt),
|
||||
DeletedAt: database.PtrTimeToNullTime(flow.DeletedAt),
|
||||
TraceID: database.PtrStringToNullString(flow.TraceID),
|
||||
ModelProviderType: database.ProviderType(flow.ModelProviderType),
|
||||
ToolCallIDTemplate: flow.ToolCallIDTemplate,
|
||||
}, nil
|
||||
}
|
||||
|
||||
func convertContainerToDatabase(container models.Container) database.Container {
|
||||
return database.Container{
|
||||
ID: int64(container.ID),
|
||||
Type: database.ContainerType(container.Type),
|
||||
Name: container.Name,
|
||||
Image: container.Image,
|
||||
Status: database.ContainerStatus(container.Status),
|
||||
LocalID: database.StringToNullString(container.LocalID),
|
||||
LocalDir: database.StringToNullString(container.LocalDir),
|
||||
FlowID: int64(container.FlowID),
|
||||
CreatedAt: database.TimeToNullTime(container.CreatedAt),
|
||||
UpdatedAt: database.TimeToNullTime(container.UpdatedAt),
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,227 @@
|
||||
package services
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"net/http"
|
||||
"slices"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"pentagi/pkg/config"
|
||||
"pentagi/pkg/controller"
|
||||
"pentagi/pkg/database"
|
||||
"pentagi/pkg/database/knowledge"
|
||||
"pentagi/pkg/graph"
|
||||
"pentagi/pkg/graph/subscriptions"
|
||||
"pentagi/pkg/providers"
|
||||
"pentagi/pkg/server/auth"
|
||||
"pentagi/pkg/server/logger"
|
||||
"pentagi/pkg/templates"
|
||||
|
||||
"github.com/99designs/gqlgen/graphql"
|
||||
"github.com/99designs/gqlgen/graphql/handler"
|
||||
"github.com/99designs/gqlgen/graphql/handler/extension"
|
||||
"github.com/99designs/gqlgen/graphql/handler/lru"
|
||||
"github.com/99designs/gqlgen/graphql/handler/transport"
|
||||
"github.com/99designs/gqlgen/graphql/playground"
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/gorilla/websocket"
|
||||
"github.com/sirupsen/logrus"
|
||||
"github.com/vektah/gqlparser/v2/ast"
|
||||
"github.com/vxcontrol/cloud/anonymizer"
|
||||
)
|
||||
|
||||
var (
|
||||
_ graphql.RawParams // @lint:ignore U1000
|
||||
_ graphql.Response // @lint:ignore U1000
|
||||
)
|
||||
|
||||
type GraphqlService struct {
|
||||
srv *handler.Server
|
||||
play http.HandlerFunc
|
||||
}
|
||||
|
||||
type originValidator struct {
|
||||
allowAll bool
|
||||
allowed []string
|
||||
wildcards [][]string
|
||||
wrappers []string
|
||||
}
|
||||
|
||||
func NewGraphqlService(
|
||||
db *database.Queries,
|
||||
cfg *config.Config,
|
||||
baseURL string,
|
||||
origins []string,
|
||||
tokenCache *auth.TokenCache,
|
||||
providers providers.ProviderController,
|
||||
controller controller.FlowController,
|
||||
subscriptions subscriptions.SubscriptionsController,
|
||||
knowledgeStore knowledge.KnowledgeStore,
|
||||
replacer anonymizer.Replacer,
|
||||
) *GraphqlService {
|
||||
srv := handler.New(graph.NewExecutableSchema(graph.Config{Resolvers: &graph.Resolver{
|
||||
DB: db,
|
||||
Config: cfg,
|
||||
Logger: logrus.StandardLogger().WithField("component", "pentagi-gql-bl"),
|
||||
TokenCache: tokenCache,
|
||||
DefaultPrompter: templates.NewDefaultPrompter(),
|
||||
ProvidersCtrl: providers,
|
||||
Controller: controller,
|
||||
Subscriptions: subscriptions,
|
||||
Knowledge: knowledgeStore,
|
||||
Replacer: replacer,
|
||||
}}))
|
||||
|
||||
component := "pentagi-gql"
|
||||
srv.AroundResponses(logger.WithGqlLogger(component))
|
||||
logger := logrus.WithField("component", component)
|
||||
|
||||
srv.AddTransport(transport.Options{})
|
||||
srv.AddTransport(transport.GET{})
|
||||
srv.AddTransport(transport.POST{})
|
||||
srv.AddTransport(transport.MultipartForm{
|
||||
MaxMemory: 32 << 20, // 32MB
|
||||
})
|
||||
|
||||
srv.SetQueryCache(lru.New[*ast.QueryDocument](1000))
|
||||
|
||||
srv.Use(extension.Introspection{})
|
||||
srv.Use(extension.AutomaticPersistedQuery{
|
||||
Cache: lru.New[string](100),
|
||||
})
|
||||
srv.Use(extension.FixedComplexityLimit(20000))
|
||||
|
||||
ov := newOriginValidator(origins)
|
||||
|
||||
// Add transport to support GraphQL subscriptions
|
||||
srv.AddTransport(&transport.Websocket{
|
||||
KeepAlivePingInterval: 10 * time.Second,
|
||||
InitFunc: func(ctx context.Context, initPayload transport.InitPayload) (context.Context, *transport.InitPayload, error) {
|
||||
uid, err := graph.GetUserID(ctx)
|
||||
if err != nil {
|
||||
return nil, nil, fmt.Errorf("unauthorized: invalid user: %v", err)
|
||||
}
|
||||
logger.WithField("uid", uid).Info("graphql websocket upgrade")
|
||||
|
||||
return ctx, &initPayload, nil
|
||||
},
|
||||
Upgrader: websocket.Upgrader{
|
||||
CheckOrigin: func(r *http.Request) (allowed bool) {
|
||||
return ov.validateOrigin(r.Header.Get("Origin"), r.Host)
|
||||
},
|
||||
EnableCompression: true,
|
||||
},
|
||||
})
|
||||
|
||||
return &GraphqlService{
|
||||
srv: srv,
|
||||
play: playground.Handler("GraphQL", baseURL+"/graphql"),
|
||||
}
|
||||
}
|
||||
|
||||
// ServeGraphql is a function to perform graphql requests
|
||||
// @Summary Perform graphql requests
|
||||
// @Tags GraphQL
|
||||
// @Accept json
|
||||
// @Produce json
|
||||
// @Security BearerAuth
|
||||
// @Param json body graphql.RawParams true "graphql request"
|
||||
// @Success 200 {object} graphql.Response "graphql response"
|
||||
// @Failure 400 {object} graphql.Response "invalid graphql request data"
|
||||
// @Failure 403 {object} graphql.Response "unauthorized"
|
||||
// @Failure 500 {object} graphql.Response "internal error on graphql request"
|
||||
// @Router /graphql [post]
|
||||
func (s *GraphqlService) ServeGraphql(c *gin.Context) {
|
||||
uid := c.GetUint64("uid")
|
||||
tid := c.GetString("tid")
|
||||
privs := c.GetStringSlice("prm")
|
||||
|
||||
savedCtx := c.Request.Context()
|
||||
defer func() {
|
||||
c.Request = c.Request.WithContext(savedCtx)
|
||||
}()
|
||||
|
||||
ctx := savedCtx
|
||||
ctx = graph.SetUserID(ctx, uid)
|
||||
ctx = graph.SetUserType(ctx, tid)
|
||||
ctx = graph.SetUserPermissions(ctx, privs)
|
||||
c.Request = c.Request.WithContext(ctx)
|
||||
|
||||
s.srv.ServeHTTP(c.Writer, c.Request)
|
||||
}
|
||||
|
||||
func (s *GraphqlService) ServeGraphqlPlayground(c *gin.Context) {
|
||||
s.play.ServeHTTP(c.Writer, c.Request)
|
||||
}
|
||||
|
||||
func newOriginValidator(origins []string) *originValidator {
|
||||
var wRules [][]string
|
||||
|
||||
for _, o := range origins {
|
||||
if !strings.Contains(o, "*") {
|
||||
continue
|
||||
}
|
||||
|
||||
if c := strings.Count(o, "*"); c > 1 {
|
||||
continue
|
||||
}
|
||||
|
||||
i := strings.Index(o, "*")
|
||||
if i == 0 {
|
||||
wRules = append(wRules, []string{"*", o[1:]})
|
||||
continue
|
||||
}
|
||||
if i == (len(o) - 1) {
|
||||
wRules = append(wRules, []string{o[:i], "*"})
|
||||
continue
|
||||
}
|
||||
|
||||
wRules = append(wRules, []string{o[:i], o[i+1:]})
|
||||
}
|
||||
|
||||
return &originValidator{
|
||||
allowAll: slices.Contains(origins, "*"),
|
||||
allowed: origins,
|
||||
wildcards: wRules,
|
||||
wrappers: []string{"http://", "https://", "ws://", "wss://"},
|
||||
}
|
||||
}
|
||||
|
||||
func (ov *originValidator) validateOrigin(origin, host string) bool {
|
||||
if ov.allowAll {
|
||||
// CORS for origin '*' is allowed
|
||||
return true
|
||||
}
|
||||
|
||||
if len(origin) == 0 {
|
||||
// request is not a CORS request
|
||||
return true
|
||||
}
|
||||
|
||||
for _, wrapper := range ov.wrappers {
|
||||
if origin == wrapper+host {
|
||||
// request is not a CORS request but have origin header
|
||||
return true
|
||||
}
|
||||
}
|
||||
|
||||
if slices.Contains(ov.allowed, origin) {
|
||||
return true
|
||||
}
|
||||
|
||||
for _, w := range ov.wildcards {
|
||||
if w[0] == "*" && strings.HasSuffix(origin, w[1]) {
|
||||
return true
|
||||
}
|
||||
if w[1] == "*" && strings.HasPrefix(origin, w[0]) {
|
||||
return true
|
||||
}
|
||||
if strings.HasPrefix(origin, w[0]) && strings.HasSuffix(origin, w[1]) {
|
||||
return true
|
||||
}
|
||||
}
|
||||
|
||||
return false
|
||||
}
|
||||
@@ -0,0 +1,519 @@
|
||||
package services
|
||||
|
||||
import (
|
||||
"errors"
|
||||
"net/http"
|
||||
"strings"
|
||||
|
||||
knowledgepkg "pentagi/pkg/database/knowledge"
|
||||
gqlmodel "pentagi/pkg/graph/model"
|
||||
"pentagi/pkg/server/auth"
|
||||
"pentagi/pkg/server/logger"
|
||||
"pentagi/pkg/server/models"
|
||||
"pentagi/pkg/server/rdb"
|
||||
"pentagi/pkg/server/response"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/jinzhu/gorm"
|
||||
)
|
||||
|
||||
// knowledgeSQLMappers maps rdb filter field names to the corresponding SQL
|
||||
// expressions against langchain_pg_embedding / its cmetadata JSON column.
|
||||
// Fields ending in _id are treated as numeric by rdb (operator = instead of like),
|
||||
// so numeric JSON fields use custom function mappers that cast cmetadata values.
|
||||
var knowledgeSQLMappers = map[string]any{
|
||||
// Primary key exposed as "id" so TableQuery default ordering still works via SetOrders override.
|
||||
"id": "langchain_pg_embedding.uuid::text",
|
||||
|
||||
// Text fields from cmetadata — rdb wraps them in LOWER(…::text) automatically.
|
||||
"doc_type": "(langchain_pg_embedding.cmetadata ->> 'doc_type')",
|
||||
"question": "(langchain_pg_embedding.cmetadata ->> 'question')",
|
||||
"description": "(langchain_pg_embedding.cmetadata ->> 'description')",
|
||||
"guide_type": "(langchain_pg_embedding.cmetadata ->> 'guide_type')",
|
||||
"answer_type": "(langchain_pg_embedding.cmetadata ->> 'answer_type')",
|
||||
"code_lang": "(langchain_pg_embedding.cmetadata ->> 'code_lang')",
|
||||
|
||||
// Searchable concat of common text fields for generic "data" filter.
|
||||
"data": "((langchain_pg_embedding.cmetadata ->> 'doc_type') || ' ' || " +
|
||||
"COALESCE(langchain_pg_embedding.cmetadata ->> 'question', ''))",
|
||||
|
||||
// Numeric ID fields: stored as JSON numbers; cast to bigint before comparison.
|
||||
"user_id": func(q *rdb.TableQuery, db *gorm.DB, value any) *gorm.DB {
|
||||
return applyKnowledgeIntFilter(db, "user_id", value)
|
||||
},
|
||||
"flow_id": func(q *rdb.TableQuery, db *gorm.DB, value any) *gorm.DB {
|
||||
return applyKnowledgeIntFilter(db, "flow_id", value)
|
||||
},
|
||||
"task_id": func(q *rdb.TableQuery, db *gorm.DB, value any) *gorm.DB {
|
||||
return applyKnowledgeIntFilter(db, "task_id", value)
|
||||
},
|
||||
"subtask_id": func(q *rdb.TableQuery, db *gorm.DB, value any) *gorm.DB {
|
||||
return applyKnowledgeIntFilter(db, "subtask_id", value)
|
||||
},
|
||||
|
||||
// Boolean field: 'true' / anything else (JSON booleans serialise as 'true'/'false').
|
||||
"manual": func(q *rdb.TableQuery, db *gorm.DB, value any) *gorm.DB {
|
||||
if b, ok := value.(bool); ok {
|
||||
if b {
|
||||
return db.Where("(langchain_pg_embedding.cmetadata ->> 'manual') = 'true'")
|
||||
}
|
||||
return db.Where("(langchain_pg_embedding.cmetadata ->> 'manual') IS DISTINCT FROM 'true'")
|
||||
}
|
||||
return db
|
||||
},
|
||||
}
|
||||
|
||||
// applyKnowledgeIntFilter adds a bigint equality / IN filter against a cmetadata integer field.
|
||||
func applyKnowledgeIntFilter(db *gorm.DB, field string, value any) *gorm.DB {
|
||||
expr := "(langchain_pg_embedding.cmetadata ->> '" + field + "')::bigint"
|
||||
switch v := value.(type) {
|
||||
case float64:
|
||||
return db.Where(expr+" = ?", int64(v))
|
||||
case []any:
|
||||
var ids []int64
|
||||
for _, vi := range v {
|
||||
if f, ok := vi.(float64); ok {
|
||||
ids = append(ids, int64(f))
|
||||
}
|
||||
}
|
||||
if len(ids) > 0 {
|
||||
return db.Where(expr+" IN (?)", ids)
|
||||
}
|
||||
}
|
||||
return db
|
||||
}
|
||||
|
||||
// knowledgeGrouped is the response shape for grouped list queries.
|
||||
type knowledgeGrouped struct {
|
||||
Grouped []string `json:"grouped"`
|
||||
Total uint64 `json:"total"`
|
||||
}
|
||||
|
||||
// KnowledgeService exposes the knowledge (pgvector) store via a REST API
|
||||
// that mirrors the GraphQL knowledge operations.
|
||||
//
|
||||
// Authorization model:
|
||||
// - knowledge.admin → read/write any document (no user_id filtering)
|
||||
// - knowledge.view → list/get own documents (filtered by user_id)
|
||||
// - knowledge.create → create a new document
|
||||
// - knowledge.edit → update own document (admin: any)
|
||||
// - knowledge.delete → delete own document (admin: any)
|
||||
// - knowledge.search → semantic search (own docs for regular users)
|
||||
type KnowledgeService struct {
|
||||
db *gorm.DB
|
||||
store knowledgepkg.KnowledgeStore
|
||||
}
|
||||
|
||||
// NewKnowledgeService creates a KnowledgeService.
|
||||
// db is the gorm connection used for paginated list queries via rdb.TableQuery.
|
||||
// store may be nil when the embedding provider is not configured; in that
|
||||
// case embedding-dependent endpoints return 503.
|
||||
func NewKnowledgeService(db *gorm.DB, store knowledgepkg.KnowledgeStore) *KnowledgeService {
|
||||
return &KnowledgeService{db: db, store: store}
|
||||
}
|
||||
|
||||
func isKnowledgeAdmin(c *gin.Context) bool {
|
||||
return auth.LookupPerm(c.GetStringSlice("prm"), "knowledge.admin")
|
||||
}
|
||||
|
||||
// ---- ListDocuments ----------------------------------------------------------
|
||||
|
||||
// ListDocuments returns a paginated, sortable, filterable list of knowledge documents.
|
||||
// It uses the standard rdb.TableQuery protocol (page, pageSize, sort[], filters[])
|
||||
// plus an additional `with_content` boolean query parameter.
|
||||
//
|
||||
// Filterable fields:
|
||||
// - id – UUID (exact match)
|
||||
// - doc_type – answer | guide | code (text, like/=)
|
||||
// - question – question text (text, like)
|
||||
// - description – description text (text, like)
|
||||
// - guide_type – guide sub-type (text, like/=)
|
||||
// - answer_type – answer sub-type (text, like/=)
|
||||
// - code_lang – code language (text, like/=)
|
||||
// - manual – boolean flag (bool, =)
|
||||
// - user_id – owner user ID (int64, =/in)
|
||||
// - flow_id – originating flow ID (int64, =/in)
|
||||
// - task_id – originating task ID (int64, =/in)
|
||||
// - subtask_id – originating subtask ID (int64, =/in)
|
||||
// - data – full-text across doc_type + question (text, like)
|
||||
//
|
||||
// @Summary List knowledge documents
|
||||
// @Tags Knowledge
|
||||
// @Produce json
|
||||
// @Security BearerAuth
|
||||
// @Param request query rdb.TableQuery true "pagination / sort / filter params"
|
||||
// @Param with_content query bool false "include document text in the response"
|
||||
// @Success 200 {object} response.successResp{data=models.KnowledgeDocList}
|
||||
// @Failure 400 {object} response.errorResp "invalid query parameters"
|
||||
// @Failure 403 {object} response.errorResp "not permitted"
|
||||
// @Failure 500 {object} response.errorResp "internal error"
|
||||
// @Router /knowledge/ [get]
|
||||
func (s *KnowledgeService) ListDocuments(c *gin.Context) {
|
||||
uid := int64(c.GetUint64("uid"))
|
||||
admin := isKnowledgeAdmin(c)
|
||||
|
||||
var query rdb.TableQuery
|
||||
if err := c.ShouldBindQuery(&query); err != nil {
|
||||
logger.FromContext(c).WithError(err).Error("error binding list query")
|
||||
response.Error(c, response.ErrKnowledgeInvalidRequest, err)
|
||||
return
|
||||
}
|
||||
|
||||
withContent := c.Query("with_content") == "true" || c.Query("with_content") == "1"
|
||||
|
||||
if err := query.Init("langchain_pg_embedding", knowledgeSQLMappers); err != nil {
|
||||
// Init fails when the requested group field is not a sortable string column
|
||||
// (e.g. a custom function mapper like user_id).
|
||||
logger.FromContext(c).WithError(err).Error("error initialising knowledge query")
|
||||
response.Error(c, response.ErrKnowledgeInvalidRequest, err)
|
||||
return
|
||||
}
|
||||
|
||||
// Base scope: JOIN the collection table, restrict to 'langchain', exclude memory.
|
||||
baseScope := func(db *gorm.DB) *gorm.DB {
|
||||
return db.
|
||||
Joins("JOIN langchain_pg_collection ON langchain_pg_embedding.collection_id = langchain_pg_collection.uuid").
|
||||
Where("langchain_pg_collection.name = ?", "langchain").
|
||||
Where("COALESCE(langchain_pg_embedding.cmetadata ->> 'doc_type', '') NOT IN (?)", []string{"memory"})
|
||||
}
|
||||
|
||||
// Authorization scope: admins see all documents; regular users see only their own.
|
||||
authScope := func(db *gorm.DB) *gorm.DB {
|
||||
if admin {
|
||||
return db
|
||||
}
|
||||
return db.Where("(langchain_pg_embedding.cmetadata ->> 'user_id')::bigint = ?", uid)
|
||||
}
|
||||
|
||||
// ---- Grouped query -------------------------------------------------------
|
||||
// When a group field is requested the response is a list of distinct values
|
||||
// rather than full document rows (mirrors the pattern used in other services).
|
||||
if query.Group != "" {
|
||||
if _, ok := knowledgeSQLMappers[query.Group]; !ok {
|
||||
logger.FromContext(c).Errorf("group field %q not found in knowledge mappers", query.Group)
|
||||
response.Error(c, response.ErrKnowledgeInvalidRequest, errors.New("group field not found"))
|
||||
return
|
||||
}
|
||||
|
||||
var resp knowledgeGrouped
|
||||
var err error
|
||||
if resp.Total, err = query.QueryGrouped(s.db, &resp.Grouped, baseScope, authScope); err != nil {
|
||||
logger.FromContext(c).WithError(err).Error("error querying knowledge documents grouped")
|
||||
response.Error(c, response.ErrInternal, err)
|
||||
return
|
||||
}
|
||||
|
||||
response.Success(c, http.StatusOK, resp)
|
||||
return
|
||||
}
|
||||
|
||||
// ---- Paginated / sorted query --------------------------------------------
|
||||
|
||||
// Override the default "ORDER BY id DESC" — the PK column is uuid, not id.
|
||||
query.SetOrders([]func(*gorm.DB) *gorm.DB{
|
||||
func(db *gorm.DB) *gorm.DB {
|
||||
return db.Order("langchain_pg_embedding.uuid DESC")
|
||||
},
|
||||
})
|
||||
|
||||
// SELECT only the columns the API exposes; never return the embedding vector.
|
||||
query.SetFind(func(out any) func(*gorm.DB) *gorm.DB {
|
||||
docCol := "'' AS document"
|
||||
if withContent {
|
||||
docCol = "COALESCE(langchain_pg_embedding.document, '') AS document"
|
||||
}
|
||||
sel := "langchain_pg_embedding.uuid::text AS id, " +
|
||||
docCol + ", " +
|
||||
"COALESCE(langchain_pg_embedding.cmetadata::text, '{}') AS cmetadata"
|
||||
return func(db *gorm.DB) *gorm.DB {
|
||||
return db.Select(sel).Find(out)
|
||||
}
|
||||
})
|
||||
|
||||
var rows []models.KnowledgeEmbeddingRow
|
||||
total, err := query.Query(s.db, &rows, baseScope, authScope)
|
||||
if err != nil {
|
||||
logger.FromContext(c).WithError(err).Error("error querying knowledge documents")
|
||||
response.Error(c, response.ErrInternal, err)
|
||||
return
|
||||
}
|
||||
|
||||
response.Success(c, http.StatusOK, models.KnowledgeDocListFromRows(rows, total, withContent))
|
||||
}
|
||||
|
||||
// ---- GetDocument ------------------------------------------------------------
|
||||
|
||||
// GetDocument returns a single knowledge document by its UUID.
|
||||
//
|
||||
// @Summary Get knowledge document
|
||||
// @Tags Knowledge
|
||||
// @Produce json
|
||||
// @Security BearerAuth
|
||||
// @Param id path string true "Document UUID"
|
||||
// @Success 200 {object} response.successResp{data=models.KnowledgeDocEntry}
|
||||
// @Failure 403 {object} response.errorResp "not permitted or wrong owner"
|
||||
// @Failure 404 {object} response.errorResp "document not found"
|
||||
// @Failure 500 {object} response.errorResp "internal error"
|
||||
// @Router /knowledge/{id} [get]
|
||||
func (s *KnowledgeService) GetDocument(c *gin.Context) {
|
||||
uid := int64(c.GetUint64("uid"))
|
||||
admin := isKnowledgeAdmin(c)
|
||||
id := c.Param("id")
|
||||
|
||||
if id == "" {
|
||||
response.Error(c, response.ErrKnowledgeInvalidRequest, errors.New("document id is required"))
|
||||
return
|
||||
}
|
||||
|
||||
ctx := c.Request.Context()
|
||||
|
||||
var (
|
||||
doc *gqlmodel.KnowledgeDocument
|
||||
err error
|
||||
)
|
||||
if admin {
|
||||
doc, err = s.store.GetDocument(ctx, id)
|
||||
} else {
|
||||
doc, err = s.store.GetUserDocument(ctx, uid, id)
|
||||
}
|
||||
if err != nil {
|
||||
logger.FromContext(c).WithError(err).Errorf("error getting knowledge document %s", id)
|
||||
response.Error(c, response.ErrKnowledgeNotFound, err)
|
||||
return
|
||||
}
|
||||
|
||||
response.Success(c, http.StatusOK, models.KnowledgeDocFromGQL(doc))
|
||||
}
|
||||
|
||||
// ---- SearchDocuments --------------------------------------------------------
|
||||
|
||||
// SearchDocuments performs a semantic similarity search over the knowledge store.
|
||||
//
|
||||
// @Summary Semantic search in knowledge base
|
||||
// @Tags Knowledge
|
||||
// @Accept json
|
||||
// @Produce json
|
||||
// @Security BearerAuth
|
||||
// @Param json body models.KnowledgeSearchRequest true "Search request"
|
||||
// @Success 200 {object} response.successResp{data=models.KnowledgeSearchResult}
|
||||
// @Failure 400 {object} response.errorResp "invalid request"
|
||||
// @Failure 403 {object} response.errorResp "not permitted"
|
||||
// @Failure 503 {object} response.errorResp "embedding provider not configured"
|
||||
// @Failure 500 {object} response.errorResp "internal error"
|
||||
// @Router /knowledge/search [post]
|
||||
func (s *KnowledgeService) SearchDocuments(c *gin.Context) {
|
||||
uid := int64(c.GetUint64("uid"))
|
||||
admin := isKnowledgeAdmin(c)
|
||||
|
||||
var req models.KnowledgeSearchRequest
|
||||
if err := c.ShouldBindJSON(&req); err != nil {
|
||||
logger.FromContext(c).WithError(err).Error("error binding search request")
|
||||
response.Error(c, response.ErrKnowledgeInvalidRequest, err)
|
||||
return
|
||||
}
|
||||
if err := req.Valid(); err != nil {
|
||||
logger.FromContext(c).WithError(err).Error("invalid search request")
|
||||
response.Error(c, response.ErrKnowledgeInvalidRequest, err)
|
||||
return
|
||||
}
|
||||
|
||||
ctx := c.Request.Context()
|
||||
filter := req.ToGQLFilter()
|
||||
|
||||
var (
|
||||
results []*gqlmodel.KnowledgeDocumentWithScore
|
||||
err error
|
||||
)
|
||||
if admin {
|
||||
results, err = s.store.SearchDocuments(ctx, req.Query, filter, req.Limit)
|
||||
} else {
|
||||
results, err = s.store.SearchUserDocuments(ctx, uid, req.Query, filter, req.Limit)
|
||||
}
|
||||
if err != nil {
|
||||
if isKnowledgeStoreUnavailable(err) {
|
||||
response.Error(c, response.ErrKnowledgeStoreUnavail, err)
|
||||
return
|
||||
}
|
||||
logger.FromContext(c).WithError(err).Error("error searching knowledge documents")
|
||||
response.Error(c, response.ErrInternal, err)
|
||||
return
|
||||
}
|
||||
|
||||
response.Success(c, http.StatusOK, models.KnowledgeSearchResultFromGQL(results))
|
||||
}
|
||||
|
||||
// ---- CreateDocument ---------------------------------------------------------
|
||||
|
||||
// CreateDocument creates a new knowledge document and computes its embedding.
|
||||
//
|
||||
// @Summary Create knowledge document
|
||||
// @Tags Knowledge
|
||||
// @Accept json
|
||||
// @Produce json
|
||||
// @Security BearerAuth
|
||||
// @Param json body models.CreateKnowledgeDocRequest true "Create request"
|
||||
// @Success 201 {object} response.successResp{data=models.KnowledgeDocEntry}
|
||||
// @Failure 400 {object} response.errorResp "invalid request"
|
||||
// @Failure 403 {object} response.errorResp "not permitted"
|
||||
// @Failure 503 {object} response.errorResp "embedding provider not configured"
|
||||
// @Failure 500 {object} response.errorResp "internal error"
|
||||
// @Router /knowledge/ [post]
|
||||
func (s *KnowledgeService) CreateDocument(c *gin.Context) {
|
||||
uid := int64(c.GetUint64("uid"))
|
||||
|
||||
var req models.CreateKnowledgeDocRequest
|
||||
if err := c.ShouldBindJSON(&req); err != nil {
|
||||
logger.FromContext(c).WithError(err).Error("error binding create request")
|
||||
response.Error(c, response.ErrKnowledgeInvalidRequest, err)
|
||||
return
|
||||
}
|
||||
if err := req.Valid(); err != nil {
|
||||
logger.FromContext(c).WithError(err).Error("invalid create request")
|
||||
response.Error(c, response.ErrKnowledgeInvalidRequest, err)
|
||||
return
|
||||
}
|
||||
|
||||
doc, err := s.store.CreateDocument(c.Request.Context(), uid, req.ToGQL())
|
||||
if err != nil {
|
||||
if isKnowledgeStoreUnavailable(err) {
|
||||
response.Error(c, response.ErrKnowledgeStoreUnavail, err)
|
||||
return
|
||||
}
|
||||
logger.FromContext(c).WithError(err).Error("error creating knowledge document")
|
||||
response.Error(c, response.ErrInternal, err)
|
||||
return
|
||||
}
|
||||
|
||||
response.Success(c, http.StatusCreated, models.KnowledgeDocFromGQL(doc))
|
||||
}
|
||||
|
||||
// ---- UpdateDocument ---------------------------------------------------------
|
||||
|
||||
// UpdateDocument updates a knowledge document, re-computing its embedding.
|
||||
// Admin can update any document; regular users can only update their own.
|
||||
//
|
||||
// @Summary Update knowledge document
|
||||
// @Tags Knowledge
|
||||
// @Accept json
|
||||
// @Produce json
|
||||
// @Security BearerAuth
|
||||
// @Param id path string true "Document UUID"
|
||||
// @Param json body models.UpdateKnowledgeDocRequest true "Update request"
|
||||
// @Success 200 {object} response.successResp{data=models.KnowledgeDocEntry}
|
||||
// @Failure 400 {object} response.errorResp "invalid request"
|
||||
// @Failure 403 {object} response.errorResp "not permitted or wrong owner"
|
||||
// @Failure 404 {object} response.errorResp "document not found"
|
||||
// @Failure 503 {object} response.errorResp "embedding provider not configured"
|
||||
// @Failure 500 {object} response.errorResp "internal error"
|
||||
// @Router /knowledge/{id} [put]
|
||||
func (s *KnowledgeService) UpdateDocument(c *gin.Context) {
|
||||
uid := int64(c.GetUint64("uid"))
|
||||
admin := isKnowledgeAdmin(c)
|
||||
id := c.Param("id")
|
||||
|
||||
if id == "" {
|
||||
response.Error(c, response.ErrKnowledgeInvalidRequest, errors.New("document id is required"))
|
||||
return
|
||||
}
|
||||
|
||||
var req models.UpdateKnowledgeDocRequest
|
||||
if err := c.ShouldBindJSON(&req); err != nil {
|
||||
logger.FromContext(c).WithError(err).Error("error binding update request")
|
||||
response.Error(c, response.ErrKnowledgeInvalidRequest, err)
|
||||
return
|
||||
}
|
||||
if err := req.Valid(); err != nil {
|
||||
logger.FromContext(c).WithError(err).Error("invalid update request")
|
||||
response.Error(c, response.ErrKnowledgeInvalidRequest, err)
|
||||
return
|
||||
}
|
||||
|
||||
ctx := c.Request.Context()
|
||||
input := req.ToGQL()
|
||||
|
||||
var (
|
||||
doc *gqlmodel.KnowledgeDocument
|
||||
err error
|
||||
)
|
||||
if admin {
|
||||
doc, err = s.store.UpdateDocument(ctx, uid, id, input)
|
||||
} else {
|
||||
doc, err = s.store.UpdateUserDocument(ctx, uid, id, input)
|
||||
}
|
||||
if err != nil {
|
||||
if isKnowledgeStoreUnavailable(err) {
|
||||
response.Error(c, response.ErrKnowledgeStoreUnavail, err)
|
||||
return
|
||||
}
|
||||
if isKnowledgeNotFound(err) {
|
||||
response.Error(c, response.ErrKnowledgeNotFound, err)
|
||||
return
|
||||
}
|
||||
logger.FromContext(c).WithError(err).Errorf("error updating knowledge document %s", id)
|
||||
response.Error(c, response.ErrInternal, err)
|
||||
return
|
||||
}
|
||||
|
||||
response.Success(c, http.StatusOK, models.KnowledgeDocFromGQL(doc))
|
||||
}
|
||||
|
||||
// ---- DeleteDocument ---------------------------------------------------------
|
||||
|
||||
// DeleteDocument deletes a knowledge document by UUID.
|
||||
// Admin can delete any document; regular users can only delete their own.
|
||||
//
|
||||
// @Summary Delete knowledge document
|
||||
// @Tags Knowledge
|
||||
// @Produce json
|
||||
// @Security BearerAuth
|
||||
// @Param id path string true "Document UUID"
|
||||
// @Success 200 {object} response.successResp "document deleted"
|
||||
// @Failure 403 {object} response.errorResp "not permitted or wrong owner"
|
||||
// @Failure 404 {object} response.errorResp "document not found"
|
||||
// @Failure 500 {object} response.errorResp "internal error"
|
||||
// @Router /knowledge/{id} [delete]
|
||||
func (s *KnowledgeService) DeleteDocument(c *gin.Context) {
|
||||
uid := int64(c.GetUint64("uid"))
|
||||
admin := isKnowledgeAdmin(c)
|
||||
id := c.Param("id")
|
||||
|
||||
if id == "" {
|
||||
response.Error(c, response.ErrKnowledgeInvalidRequest, errors.New("document id is required"))
|
||||
return
|
||||
}
|
||||
|
||||
ctx := c.Request.Context()
|
||||
|
||||
var err error
|
||||
if admin {
|
||||
err = s.store.DeleteDocument(ctx, uid, id)
|
||||
} else {
|
||||
err = s.store.DeleteUserDocument(ctx, uid, id)
|
||||
}
|
||||
if err != nil {
|
||||
if isKnowledgeNotFound(err) {
|
||||
response.Error(c, response.ErrKnowledgeNotFound, err)
|
||||
return
|
||||
}
|
||||
logger.FromContext(c).WithError(err).Errorf("error deleting knowledge document %s", id)
|
||||
response.Error(c, response.ErrInternal, err)
|
||||
return
|
||||
}
|
||||
|
||||
response.Success(c, http.StatusOK, gin.H{"message": "knowledge document deleted successfully"})
|
||||
}
|
||||
|
||||
// ---- error classification ---------------------------------------------------
|
||||
|
||||
// isKnowledgeStoreUnavailable returns true when the error originates from a
|
||||
// missing embedding provider (store or embedder not configured).
|
||||
func isKnowledgeStoreUnavailable(err error) bool {
|
||||
return err != nil && strings.Contains(err.Error(), "knowledge: embedding provider")
|
||||
}
|
||||
|
||||
// isKnowledgeNotFound returns true when the underlying DB returned no rows,
|
||||
// which means the document does not exist or does not belong to the caller.
|
||||
func isKnowledgeNotFound(err error) bool {
|
||||
return err != nil && (strings.Contains(err.Error(), "no rows") ||
|
||||
strings.Contains(err.Error(), "not found"))
|
||||
}
|
||||
@@ -0,0 +1,221 @@
|
||||
package services
|
||||
|
||||
import (
|
||||
"errors"
|
||||
"net/http"
|
||||
"slices"
|
||||
"strconv"
|
||||
|
||||
"pentagi/pkg/server/logger"
|
||||
"pentagi/pkg/server/models"
|
||||
"pentagi/pkg/server/rdb"
|
||||
"pentagi/pkg/server/response"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/jinzhu/gorm"
|
||||
)
|
||||
|
||||
type msglogs struct {
|
||||
MsgLogs []models.Msglog `json:"msglogs"`
|
||||
Total uint64 `json:"total"`
|
||||
}
|
||||
|
||||
type msglogsGrouped struct {
|
||||
Grouped []string `json:"grouped"`
|
||||
Total uint64 `json:"total"`
|
||||
}
|
||||
|
||||
var msglogsSQLMappers = map[string]any{
|
||||
"id": "{{table}}.id",
|
||||
"type": "{{table}}.type",
|
||||
"message": "{{table}}.message",
|
||||
"thinking": "{{table}}.thinking",
|
||||
"result": "{{table}}.result",
|
||||
"result_format": "{{table}}.result_format",
|
||||
"flow_id": "{{table}}.flow_id",
|
||||
"task_id": "{{table}}.task_id",
|
||||
"subtask_id": "{{table}}.subtask_id",
|
||||
"created_at": "{{table}}.created_at",
|
||||
"data": "({{table}}.type || ' ' || {{table}}.message || ' ' || {{table}}.thinking || ' ' || {{table}}.result)",
|
||||
}
|
||||
|
||||
type MsglogService struct {
|
||||
db *gorm.DB
|
||||
}
|
||||
|
||||
func NewMsglogService(db *gorm.DB) *MsglogService {
|
||||
return &MsglogService{
|
||||
db: db,
|
||||
}
|
||||
}
|
||||
|
||||
// GetMsglogs is a function to return msglogs list
|
||||
// @Summary Retrieve msglogs list
|
||||
// @Tags Msglogs
|
||||
// @Produce json
|
||||
// @Security BearerAuth
|
||||
// @Param request query rdb.TableQuery true "query table params"
|
||||
// @Success 200 {object} response.successResp{data=msglogs} "msglogs list received successful"
|
||||
// @Failure 400 {object} response.errorResp "invalid query request data"
|
||||
// @Failure 403 {object} response.errorResp "getting msglogs not permitted"
|
||||
// @Failure 500 {object} response.errorResp "internal error on getting msglogs"
|
||||
// @Router /msglogs/ [get]
|
||||
func (s *MsglogService) GetMsglogs(c *gin.Context) {
|
||||
var (
|
||||
err error
|
||||
query rdb.TableQuery
|
||||
resp msglogs
|
||||
)
|
||||
|
||||
if err = c.ShouldBindQuery(&query); err != nil {
|
||||
logger.FromContext(c).WithError(err).Errorf("error binding query")
|
||||
response.Error(c, response.ErrMsglogsInvalidRequest, err)
|
||||
return
|
||||
}
|
||||
|
||||
uid := c.GetUint64("uid")
|
||||
privs := c.GetStringSlice("prm")
|
||||
var scope func(db *gorm.DB) *gorm.DB
|
||||
if slices.Contains(privs, "msglogs.admin") {
|
||||
scope = func(db *gorm.DB) *gorm.DB {
|
||||
return db.
|
||||
Joins("INNER JOIN flows f ON f.id = msglogs.flow_id")
|
||||
}
|
||||
} else if slices.Contains(privs, "msglogs.view") {
|
||||
scope = func(db *gorm.DB) *gorm.DB {
|
||||
return db.
|
||||
Joins("INNER JOIN flows f ON f.id = msglogs.flow_id").
|
||||
Where("f.user_id = ?", uid)
|
||||
}
|
||||
} else {
|
||||
logger.FromContext(c).Errorf("error filtering user role permissions: permission not found")
|
||||
response.Error(c, response.ErrNotPermitted, nil)
|
||||
return
|
||||
}
|
||||
|
||||
query.Init("msglogs", msglogsSQLMappers)
|
||||
|
||||
if query.Group != "" {
|
||||
if _, ok := msglogsSQLMappers[query.Group]; !ok {
|
||||
logger.FromContext(c).Errorf("error finding msglogs grouped: group field not found")
|
||||
response.Error(c, response.ErrMsglogsInvalidRequest, errors.New("group field not found"))
|
||||
return
|
||||
}
|
||||
|
||||
var respGrouped msglogsGrouped
|
||||
if respGrouped.Total, err = query.QueryGrouped(s.db, &respGrouped.Grouped, scope); err != nil {
|
||||
logger.FromContext(c).WithError(err).Errorf("error finding msglogs grouped")
|
||||
response.Error(c, response.ErrInternal, err)
|
||||
return
|
||||
}
|
||||
|
||||
response.Success(c, http.StatusOK, respGrouped)
|
||||
return
|
||||
}
|
||||
|
||||
if resp.Total, err = query.Query(s.db, &resp.MsgLogs, scope); err != nil {
|
||||
logger.FromContext(c).WithError(err).Errorf("error finding msglogs")
|
||||
response.Error(c, response.ErrInternal, err)
|
||||
return
|
||||
}
|
||||
|
||||
for i := 0; i < len(resp.MsgLogs); i++ {
|
||||
if err = resp.MsgLogs[i].Valid(); err != nil {
|
||||
logger.FromContext(c).WithError(err).Errorf("error validating msglog data '%d'", resp.MsgLogs[i].ID)
|
||||
response.Error(c, response.ErrMsglogsInvalidData, err)
|
||||
return
|
||||
}
|
||||
}
|
||||
|
||||
response.Success(c, http.StatusOK, resp)
|
||||
}
|
||||
|
||||
// GetFlowMsglogs is a function to return msglogs list by flow id
|
||||
// @Summary Retrieve msglogs list by flow id
|
||||
// @Tags Msglogs
|
||||
// @Produce json
|
||||
// @Security BearerAuth
|
||||
// @Param flowID path int true "flow id" minimum(0)
|
||||
// @Param request query rdb.TableQuery true "query table params"
|
||||
// @Success 200 {object} response.successResp{data=msglogs} "msglogs list received successful"
|
||||
// @Failure 400 {object} response.errorResp "invalid query request data"
|
||||
// @Failure 403 {object} response.errorResp "getting msglogs not permitted"
|
||||
// @Failure 500 {object} response.errorResp "internal error on getting msglogs"
|
||||
// @Router /flows/{flowID}/msglogs/ [get]
|
||||
func (s *MsglogService) GetFlowMsglogs(c *gin.Context) {
|
||||
var (
|
||||
err error
|
||||
flowID uint64
|
||||
query rdb.TableQuery
|
||||
resp msglogs
|
||||
)
|
||||
|
||||
if flowID, err = strconv.ParseUint(c.Param("flowID"), 10, 64); err != nil {
|
||||
logger.FromContext(c).WithError(err).Errorf("error parsing flow id")
|
||||
response.Error(c, response.ErrMsglogsInvalidRequest, err)
|
||||
return
|
||||
}
|
||||
|
||||
if err = c.ShouldBindQuery(&query); err != nil {
|
||||
logger.FromContext(c).WithError(err).Errorf("error binding query")
|
||||
response.Error(c, response.ErrMsglogsInvalidRequest, err)
|
||||
return
|
||||
}
|
||||
|
||||
uid := c.GetUint64("uid")
|
||||
privs := c.GetStringSlice("prm")
|
||||
var scope func(db *gorm.DB) *gorm.DB
|
||||
if slices.Contains(privs, "msglogs.admin") {
|
||||
scope = func(db *gorm.DB) *gorm.DB {
|
||||
return db.
|
||||
Joins("INNER JOIN flows f ON f.id = msglogs.flow_id").
|
||||
Where("f.id = ?", flowID)
|
||||
}
|
||||
} else if slices.Contains(privs, "msglogs.view") {
|
||||
scope = func(db *gorm.DB) *gorm.DB {
|
||||
return db.
|
||||
Joins("INNER JOIN flows f ON f.id = msglogs.flow_id").
|
||||
Where("f.id = ? AND f.user_id = ?", flowID, uid)
|
||||
}
|
||||
} else {
|
||||
logger.FromContext(c).Errorf("error filtering user role permissions: permission not found")
|
||||
response.Error(c, response.ErrNotPermitted, nil)
|
||||
return
|
||||
}
|
||||
|
||||
query.Init("msglogs", msglogsSQLMappers)
|
||||
|
||||
if query.Group != "" {
|
||||
if _, ok := msglogsSQLMappers[query.Group]; !ok {
|
||||
logger.FromContext(c).Errorf("error finding msglogs grouped: group field not found")
|
||||
response.Error(c, response.ErrMsglogsInvalidRequest, errors.New("group field not found"))
|
||||
return
|
||||
}
|
||||
|
||||
var respGrouped msglogsGrouped
|
||||
if respGrouped.Total, err = query.QueryGrouped(s.db, &respGrouped.Grouped, scope); err != nil {
|
||||
logger.FromContext(c).WithError(err).Errorf("error finding msglogs grouped")
|
||||
response.Error(c, response.ErrInternal, err)
|
||||
return
|
||||
}
|
||||
|
||||
response.Success(c, http.StatusOK, respGrouped)
|
||||
return
|
||||
}
|
||||
|
||||
if resp.Total, err = query.Query(s.db, &resp.MsgLogs, scope); err != nil {
|
||||
logger.FromContext(c).WithError(err).Errorf("error finding msglogs")
|
||||
response.Error(c, response.ErrInternal, err)
|
||||
return
|
||||
}
|
||||
|
||||
for i := 0; i < len(resp.MsgLogs); i++ {
|
||||
if err = resp.MsgLogs[i].Valid(); err != nil {
|
||||
logger.FromContext(c).WithError(err).Errorf("error validating msglog data '%d'", resp.MsgLogs[i].ID)
|
||||
response.Error(c, response.ErrMsglogsInvalidData, err)
|
||||
return
|
||||
}
|
||||
}
|
||||
|
||||
response.Success(c, http.StatusOK, resp)
|
||||
}
|
||||
@@ -0,0 +1,391 @@
|
||||
package services
|
||||
|
||||
import (
|
||||
"errors"
|
||||
"net/http"
|
||||
"slices"
|
||||
|
||||
"pentagi/pkg/server/logger"
|
||||
"pentagi/pkg/server/models"
|
||||
"pentagi/pkg/server/rdb"
|
||||
"pentagi/pkg/server/response"
|
||||
"pentagi/pkg/templates"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/jinzhu/gorm"
|
||||
)
|
||||
|
||||
type prompts struct {
|
||||
Prompts []models.Prompt `json:"prompts"`
|
||||
Total uint64 `json:"total"`
|
||||
}
|
||||
|
||||
type promptsGrouped struct {
|
||||
Grouped []string `json:"grouped"`
|
||||
Total uint64 `json:"total"`
|
||||
}
|
||||
|
||||
var promptsSQLMappers = map[string]any{
|
||||
"type": "{{table}}.type",
|
||||
"prompt": "{{table}}.prompt",
|
||||
"created_at": "{{table}}.created_at",
|
||||
"updated_at": "{{table}}.updated_at",
|
||||
"data": "({{table}}.type || ' ' || {{table}}.prompt)",
|
||||
}
|
||||
|
||||
type PromptService struct {
|
||||
db *gorm.DB
|
||||
prompter templates.Prompter
|
||||
}
|
||||
|
||||
func NewPromptService(db *gorm.DB) *PromptService {
|
||||
return &PromptService{
|
||||
db: db,
|
||||
prompter: templates.NewDefaultPrompter(),
|
||||
}
|
||||
}
|
||||
|
||||
// GetPrompts is a function to return prompts list
|
||||
// @Summary Retrieve prompts list
|
||||
// @Tags Prompts
|
||||
// @Produce json
|
||||
// @Security BearerAuth
|
||||
// @Param request query rdb.TableQuery true "query table params"
|
||||
// @Success 200 {object} response.successResp{data=prompts} "prompts list received successful"
|
||||
// @Failure 400 {object} response.errorResp "invalid query request data"
|
||||
// @Failure 403 {object} response.errorResp "getting prompts not permitted"
|
||||
// @Failure 500 {object} response.errorResp "internal error on getting prompts"
|
||||
// @Router /prompts/ [get]
|
||||
func (s *PromptService) GetPrompts(c *gin.Context) {
|
||||
var (
|
||||
err error
|
||||
query rdb.TableQuery
|
||||
resp prompts
|
||||
)
|
||||
|
||||
if err = c.ShouldBindQuery(&query); err != nil {
|
||||
logger.FromContext(c).WithError(err).Errorf("error binding query")
|
||||
response.Error(c, response.ErrPromptsInvalidRequest, err)
|
||||
return
|
||||
}
|
||||
|
||||
privs := c.GetStringSlice("prm")
|
||||
if !slices.Contains(privs, "settings.prompts.view") {
|
||||
logger.FromContext(c).Errorf("error filtering user role permissions: permission not found")
|
||||
response.Error(c, response.ErrNotPermitted, nil)
|
||||
return
|
||||
}
|
||||
|
||||
uid := c.GetUint64("uid")
|
||||
scope := func(db *gorm.DB) *gorm.DB {
|
||||
return db.Where("user_id = ?", uid)
|
||||
}
|
||||
|
||||
query.Init("prompts", promptsSQLMappers)
|
||||
|
||||
if query.Group != "" {
|
||||
if _, ok := promptsSQLMappers[query.Group]; !ok {
|
||||
logger.FromContext(c).Errorf("error finding prompts grouped: group field not found")
|
||||
response.Error(c, response.ErrPromptsInvalidRequest, errors.New("group field not found"))
|
||||
return
|
||||
}
|
||||
|
||||
var respGrouped promptsGrouped
|
||||
if respGrouped.Total, err = query.QueryGrouped(s.db, &respGrouped.Grouped, scope); err != nil {
|
||||
logger.FromContext(c).WithError(err).Errorf("error finding prompts grouped")
|
||||
response.Error(c, response.ErrInternal, err)
|
||||
return
|
||||
}
|
||||
|
||||
response.Success(c, http.StatusOK, respGrouped)
|
||||
return
|
||||
}
|
||||
|
||||
if resp.Total, err = query.Query(s.db, &resp.Prompts, scope); err != nil {
|
||||
logger.FromContext(c).WithError(err).Errorf("error finding prompts")
|
||||
response.Error(c, response.ErrInternal, err)
|
||||
return
|
||||
}
|
||||
|
||||
for i := 0; i < len(resp.Prompts); i++ {
|
||||
if err = resp.Prompts[i].Valid(); err != nil {
|
||||
logger.FromContext(c).WithError(err).Errorf("error validating prompt data '%s'", resp.Prompts[i].Type)
|
||||
response.Error(c, response.ErrPromptsInvalidData, err)
|
||||
return
|
||||
}
|
||||
}
|
||||
|
||||
response.Success(c, http.StatusOK, resp)
|
||||
}
|
||||
|
||||
// GetPrompt is a function to return prompt by type
|
||||
// @Summary Retrieve prompt by type
|
||||
// @Tags Prompts
|
||||
// @Produce json
|
||||
// @Security BearerAuth
|
||||
// @Param promptType path string true "prompt type"
|
||||
// @Success 200 {object} response.successResp{data=models.Prompt} "prompt received successful"
|
||||
// @Failure 400 {object} response.errorResp "invalid prompt request data"
|
||||
// @Failure 403 {object} response.errorResp "getting prompt not permitted"
|
||||
// @Failure 404 {object} response.errorResp "prompt not found"
|
||||
// @Failure 500 {object} response.errorResp "internal error on getting prompt"
|
||||
// @Router /prompts/{promptType} [get]
|
||||
func (s *PromptService) GetPrompt(c *gin.Context) {
|
||||
var (
|
||||
err error
|
||||
promptType models.PromptType = models.PromptType(c.Param("promptType"))
|
||||
resp models.Prompt
|
||||
)
|
||||
|
||||
if err = models.PromptType(promptType).Valid(); err != nil {
|
||||
logger.FromContext(c).WithError(err).Errorf("error validating prompt type '%s'", promptType)
|
||||
response.Error(c, response.ErrPromptsInvalidRequest, err)
|
||||
return
|
||||
}
|
||||
|
||||
privs := c.GetStringSlice("prm")
|
||||
if !slices.Contains(privs, "settings.prompts.view") {
|
||||
logger.FromContext(c).Errorf("error filtering user role permissions: permission not found")
|
||||
response.Error(c, response.ErrNotPermitted, nil)
|
||||
return
|
||||
}
|
||||
|
||||
uid := c.GetUint64("uid")
|
||||
scope := func(db *gorm.DB) *gorm.DB {
|
||||
return db.Where("type = ? AND user_id = ?", promptType, uid)
|
||||
}
|
||||
|
||||
if err = s.db.Scopes(scope).Take(&resp).Error; err != nil {
|
||||
logger.FromContext(c).WithError(err).Errorf("error finding prompt by type")
|
||||
if errors.Is(err, gorm.ErrRecordNotFound) {
|
||||
response.Error(c, response.ErrPromptsNotFound, err)
|
||||
} else {
|
||||
response.Error(c, response.ErrInternal, err)
|
||||
}
|
||||
return
|
||||
}
|
||||
if err = resp.Valid(); err != nil {
|
||||
logger.FromContext(c).WithError(err).Errorf("error validating prompt data '%s'", resp.Type)
|
||||
response.Error(c, response.ErrPromptsInvalidData, err)
|
||||
return
|
||||
}
|
||||
|
||||
response.Success(c, http.StatusOK, resp)
|
||||
}
|
||||
|
||||
// PatchPrompt is a function to update prompt by type
|
||||
// @Summary Update prompt
|
||||
// @Tags Prompts
|
||||
// @Accept json
|
||||
// @Produce json
|
||||
// @Security BearerAuth
|
||||
// @Param promptType path string true "prompt type"
|
||||
// @Param json body models.PatchPrompt true "prompt model to update"
|
||||
// @Success 200 {object} response.successResp{data=models.Prompt} "prompt updated successful"
|
||||
// @Success 201 {object} response.successResp{data=models.Prompt} "prompt created successful"
|
||||
// @Failure 400 {object} response.errorResp "invalid prompt request data"
|
||||
// @Failure 403 {object} response.errorResp "updating prompt not permitted"
|
||||
// @Failure 404 {object} response.errorResp "prompt not found"
|
||||
// @Failure 500 {object} response.errorResp "internal error on updating prompt"
|
||||
// @Router /prompts/{promptType} [put]
|
||||
func (s *PromptService) PatchPrompt(c *gin.Context) {
|
||||
var (
|
||||
err error
|
||||
prompt models.PatchPrompt
|
||||
promptType models.PromptType = models.PromptType(c.Param("promptType"))
|
||||
resp models.Prompt
|
||||
)
|
||||
|
||||
if err = c.ShouldBindJSON(&prompt); err != nil {
|
||||
logger.FromContext(c).WithError(err).Errorf("error binding JSON")
|
||||
response.Error(c, response.ErrPromptsInvalidRequest, err)
|
||||
return
|
||||
} else if err = prompt.Valid(); err != nil {
|
||||
logger.FromContext(c).WithError(err).Errorf("error validating prompt JSON")
|
||||
response.Error(c, response.ErrPromptsInvalidRequest, err)
|
||||
return
|
||||
} else if err = promptType.Valid(); err != nil {
|
||||
logger.FromContext(c).WithError(err).Errorf("error validating prompt type '%s'", promptType)
|
||||
response.Error(c, response.ErrPromptsInvalidRequest, err)
|
||||
return
|
||||
}
|
||||
|
||||
privs := c.GetStringSlice("prm")
|
||||
if !slices.Contains(privs, "settings.prompts.edit") {
|
||||
logger.FromContext(c).Errorf("error filtering user role permissions: permission not found")
|
||||
response.Error(c, response.ErrNotPermitted, nil)
|
||||
return
|
||||
}
|
||||
|
||||
uid := c.GetUint64("uid")
|
||||
scope := func(db *gorm.DB) *gorm.DB {
|
||||
return db.Where("type = ? AND user_id = ?", promptType, uid)
|
||||
}
|
||||
|
||||
err = s.db.Scopes(scope).Take(&resp).Error
|
||||
if err != nil && errors.Is(err, gorm.ErrRecordNotFound) {
|
||||
resp = models.Prompt{
|
||||
Type: promptType,
|
||||
UserID: uid,
|
||||
Prompt: prompt.Prompt,
|
||||
}
|
||||
if err = s.db.Create(&resp).Error; err != nil {
|
||||
logger.FromContext(c).WithError(err).Errorf("error creating prompt by type '%s'", promptType)
|
||||
response.Error(c, response.ErrInternal, err)
|
||||
return
|
||||
}
|
||||
|
||||
response.Success(c, http.StatusCreated, resp)
|
||||
return
|
||||
} else if err != nil {
|
||||
logger.FromContext(c).WithError(err).Errorf("error finding updated prompt by type '%s'", promptType)
|
||||
response.Error(c, response.ErrInternal, err)
|
||||
return
|
||||
}
|
||||
|
||||
resp.Prompt = prompt.Prompt
|
||||
|
||||
err = s.db.Scopes(scope).Save(&resp).Error
|
||||
if err != nil {
|
||||
logger.FromContext(c).WithError(err).Errorf("error updating prompt by type '%s'", promptType)
|
||||
response.Error(c, response.ErrInternal, err)
|
||||
return
|
||||
}
|
||||
|
||||
response.Success(c, http.StatusOK, resp)
|
||||
}
|
||||
|
||||
// ResetPrompt is a function to reset prompt by type to default value
|
||||
// @Summary Reset prompt by type to default value
|
||||
// @Tags Prompts
|
||||
// @Accept json
|
||||
// @Produce json
|
||||
// @Security BearerAuth
|
||||
// @Param promptType path string true "prompt type"
|
||||
// @Success 200 {object} response.successResp{data=models.Prompt} "prompt reset successful"
|
||||
// @Success 201 {object} response.successResp{data=models.Prompt} "prompt created with default value successful"
|
||||
// @Failure 400 {object} response.errorResp "invalid prompt request data"
|
||||
// @Failure 403 {object} response.errorResp "updating prompt not permitted"
|
||||
// @Failure 404 {object} response.errorResp "prompt not found"
|
||||
// @Failure 500 {object} response.errorResp "internal error on resetting prompt"
|
||||
// @Router /prompts/{promptType}/default [post]
|
||||
func (s *PromptService) ResetPrompt(c *gin.Context) {
|
||||
var (
|
||||
err error
|
||||
promptType models.PromptType = models.PromptType(c.Param("promptType"))
|
||||
resp models.Prompt
|
||||
)
|
||||
|
||||
if err = promptType.Valid(); err != nil {
|
||||
logger.FromContext(c).WithError(err).Errorf("error validating prompt type '%s'", promptType)
|
||||
response.Error(c, response.ErrPromptsInvalidRequest, err)
|
||||
return
|
||||
}
|
||||
|
||||
privs := c.GetStringSlice("prm")
|
||||
if !slices.Contains(privs, "settings.prompts.edit") {
|
||||
logger.FromContext(c).Errorf("error filtering user role permissions: permission not found")
|
||||
response.Error(c, response.ErrNotPermitted, nil)
|
||||
return
|
||||
}
|
||||
|
||||
uid := c.GetUint64("uid")
|
||||
scope := func(db *gorm.DB) *gorm.DB {
|
||||
return db.Where("type = ? AND user_id = ?", promptType, uid)
|
||||
}
|
||||
|
||||
template, err := s.prompter.GetTemplate(templates.PromptType(promptType))
|
||||
if err != nil {
|
||||
logger.FromContext(c).WithError(err).Errorf("error getting template '%s'", promptType)
|
||||
response.Error(c, response.ErrPromptsInvalidRequest, err)
|
||||
return
|
||||
}
|
||||
|
||||
err = s.db.Scopes(scope).Take(&resp).Error
|
||||
if err != nil && errors.Is(err, gorm.ErrRecordNotFound) {
|
||||
resp = models.Prompt{
|
||||
Type: promptType,
|
||||
UserID: uid,
|
||||
Prompt: template,
|
||||
}
|
||||
err = s.db.Create(&resp).Error
|
||||
if err != nil {
|
||||
logger.FromContext(c).WithError(err).Errorf("error creating default prompt by type '%s'", promptType)
|
||||
response.Error(c, response.ErrInternal, err)
|
||||
}
|
||||
|
||||
response.Success(c, http.StatusCreated, resp)
|
||||
return
|
||||
} else if err != nil {
|
||||
logger.FromContext(c).WithError(err).Errorf("error finding updated prompt by type '%s'", promptType)
|
||||
response.Error(c, response.ErrInternal, err)
|
||||
return
|
||||
}
|
||||
|
||||
resp.Prompt = template
|
||||
|
||||
err = s.db.Scopes(scope).Save(&resp).Error
|
||||
if err != nil {
|
||||
logger.FromContext(c).WithError(err).Errorf("error resetting prompt by type '%s'", promptType)
|
||||
response.Error(c, response.ErrInternal, err)
|
||||
return
|
||||
}
|
||||
|
||||
response.Success(c, http.StatusOK, resp)
|
||||
}
|
||||
|
||||
// DeletePrompt is a function to delete prompt by type
|
||||
// @Summary Delete prompt by type
|
||||
// @Tags Prompts
|
||||
// @Produce json
|
||||
// @Security BearerAuth
|
||||
// @Param promptType path string true "prompt type"
|
||||
// @Success 200 {object} response.successResp "prompt deleted successful"
|
||||
// @Failure 400 {object} response.errorResp "invalid prompt request data"
|
||||
// @Failure 403 {object} response.errorResp "deleting prompt not permitted"
|
||||
// @Failure 404 {object} response.errorResp "prompt not found"
|
||||
// @Failure 500 {object} response.errorResp "internal error on deleting prompt"
|
||||
// @Router /prompts/{promptType} [delete]
|
||||
func (s *PromptService) DeletePrompt(c *gin.Context) {
|
||||
var (
|
||||
err error
|
||||
promptType models.PromptType = models.PromptType(c.Param("promptType"))
|
||||
resp models.Prompt
|
||||
)
|
||||
|
||||
if err = promptType.Valid(); err != nil {
|
||||
logger.FromContext(c).WithError(err).Errorf("error validating prompt type '%s'", promptType)
|
||||
response.Error(c, response.ErrPromptsInvalidRequest, err)
|
||||
return
|
||||
}
|
||||
|
||||
privs := c.GetStringSlice("prm")
|
||||
if !slices.Contains(privs, "settings.prompts.edit") {
|
||||
logger.FromContext(c).Errorf("error filtering user role permissions: permission not found")
|
||||
response.Error(c, response.ErrNotPermitted, nil)
|
||||
return
|
||||
}
|
||||
|
||||
uid := c.GetUint64("uid")
|
||||
scope := func(db *gorm.DB) *gorm.DB {
|
||||
return db.Where("type = ? AND user_id = ?", promptType, uid)
|
||||
}
|
||||
|
||||
if err = s.db.Scopes(scope).Take(&resp).Error; err != nil {
|
||||
logger.FromContext(c).WithError(err).Errorf("error finding prompt by type '%s'", promptType)
|
||||
if errors.Is(err, gorm.ErrRecordNotFound) {
|
||||
response.Error(c, response.ErrPromptsNotFound, err)
|
||||
} else {
|
||||
response.Error(c, response.ErrInternal, err)
|
||||
}
|
||||
return
|
||||
}
|
||||
|
||||
if err = s.db.Scopes(scope).Delete(&resp).Error; err != nil {
|
||||
logger.FromContext(c).WithError(err).Errorf("error deleting prompt by type '%s'", promptType)
|
||||
response.Error(c, response.ErrInternal, err)
|
||||
return
|
||||
}
|
||||
|
||||
response.Success(c, http.StatusOK, nil)
|
||||
}
|
||||
@@ -0,0 +1,123 @@
|
||||
package services
|
||||
|
||||
import (
|
||||
"net/http"
|
||||
"slices"
|
||||
|
||||
"pentagi/pkg/providers"
|
||||
"pentagi/pkg/providers/pconfig"
|
||||
"pentagi/pkg/providers/provider"
|
||||
"pentagi/pkg/server/logger"
|
||||
"pentagi/pkg/server/models"
|
||||
"pentagi/pkg/server/response"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
)
|
||||
|
||||
type ProviderService struct {
|
||||
providers providers.ProviderController
|
||||
}
|
||||
|
||||
func NewProviderService(providers providers.ProviderController) *ProviderService {
|
||||
return &ProviderService{
|
||||
providers: providers,
|
||||
}
|
||||
}
|
||||
|
||||
// GetProviders is a function to return providers list
|
||||
// @Summary Retrieve providers list
|
||||
// @Tags Providers
|
||||
// @Produce json
|
||||
// @Security BearerAuth
|
||||
// @Success 200 {object} response.successResp{data=[]models.ProviderInfo} "providers list received successful"
|
||||
// @Failure 403 {object} response.errorResp "getting providers not permitted"
|
||||
// @Router /providers/ [get]
|
||||
func (s *ProviderService) GetProviders(c *gin.Context) {
|
||||
privs := c.GetStringSlice("prm")
|
||||
if !slices.Contains(privs, "providers.view") {
|
||||
logger.FromContext(c).Errorf("error filtering user role permissions: permission not found")
|
||||
response.Error(c, response.ErrNotPermitted, nil)
|
||||
return
|
||||
}
|
||||
|
||||
providers, err := s.providers.GetProviders(c, int64(c.GetUint64("uid")))
|
||||
if err != nil {
|
||||
logger.FromContext(c).Errorf("error getting providers: %v", err)
|
||||
response.Error(c, response.ErrInternal, nil)
|
||||
return
|
||||
}
|
||||
|
||||
providerInfos := make([]models.ProviderInfo, len(providers))
|
||||
for i, name := range providers.ListNames() {
|
||||
prv := providers[name]
|
||||
providerInfos[i] = models.ProviderInfo{
|
||||
Name: name.String(),
|
||||
Type: models.ProviderType(prv.Type()),
|
||||
DefaultModel: prv.Model(pconfig.OptionsTypePrimaryAgent),
|
||||
Models: buildModelInfos(prv),
|
||||
}
|
||||
}
|
||||
|
||||
response.Success(c, http.StatusOK, providerInfos)
|
||||
}
|
||||
|
||||
func buildModelInfos(prv provider.Provider) []models.ModelInfo {
|
||||
modelsConfig := prv.GetModels()
|
||||
|
||||
// Build lookup: model name -> price from ModelsConfig (models.yml)
|
||||
modelConfigPrice := make(map[string]*pconfig.PriceInfo, len(modelsConfig))
|
||||
for _, mc := range modelsConfig {
|
||||
if mc.Price != nil {
|
||||
modelConfigPrice[mc.Name] = mc.Price
|
||||
}
|
||||
}
|
||||
|
||||
// Collect unique models actually in use across all agent types.
|
||||
// For price priority: AgentConfig.Price > ModelsConfig price.
|
||||
type entry struct {
|
||||
price *pconfig.PriceInfo
|
||||
agentTypes []string
|
||||
}
|
||||
seen := make(map[string]*entry)
|
||||
order := make([]string, 0)
|
||||
|
||||
for _, optType := range pconfig.AllAgentTypes {
|
||||
modelName := prv.Model(optType)
|
||||
if modelName == "" {
|
||||
continue
|
||||
}
|
||||
|
||||
e, exists := seen[modelName]
|
||||
if !exists {
|
||||
price := prv.GetPriceInfo(optType)
|
||||
if price == nil {
|
||||
price = modelConfigPrice[modelName]
|
||||
}
|
||||
e = &entry{price: price}
|
||||
seen[modelName] = e
|
||||
order = append(order, modelName)
|
||||
}
|
||||
|
||||
e.agentTypes = append(e.agentTypes, string(optType))
|
||||
}
|
||||
|
||||
modelInfos := make([]models.ModelInfo, 0, len(order))
|
||||
for _, name := range order {
|
||||
e := seen[name]
|
||||
mi := models.ModelInfo{
|
||||
Name: name,
|
||||
AgentTypes: e.agentTypes,
|
||||
}
|
||||
if e.price != nil {
|
||||
mi.PriceInfo = &models.ModelPriceInfo{
|
||||
Input: e.price.Input,
|
||||
Output: e.price.Output,
|
||||
CacheRead: e.price.CacheRead,
|
||||
CacheWrite: e.price.CacheWrite,
|
||||
}
|
||||
}
|
||||
modelInfos = append(modelInfos, mi)
|
||||
}
|
||||
|
||||
return modelInfos
|
||||
}
|
||||
File diff suppressed because it is too large
Load Diff
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,164 @@
|
||||
package services
|
||||
|
||||
import (
|
||||
"errors"
|
||||
"net/http"
|
||||
"slices"
|
||||
"strconv"
|
||||
|
||||
"pentagi/pkg/server/logger"
|
||||
"pentagi/pkg/server/models"
|
||||
"pentagi/pkg/server/rdb"
|
||||
"pentagi/pkg/server/response"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/jinzhu/gorm"
|
||||
)
|
||||
|
||||
type roles struct {
|
||||
Roles []models.RolePrivileges `json:"roles"`
|
||||
Total uint64 `json:"total"`
|
||||
}
|
||||
|
||||
var rolesSQLMappers = map[string]any{
|
||||
"id": "{{table}}.id",
|
||||
"name": "{{table}}.name",
|
||||
"data": "{{table}}.name",
|
||||
}
|
||||
|
||||
type RoleService struct {
|
||||
db *gorm.DB
|
||||
}
|
||||
|
||||
func NewRoleService(db *gorm.DB) *RoleService {
|
||||
return &RoleService{
|
||||
db: db,
|
||||
}
|
||||
}
|
||||
|
||||
// GetRoles is a function to return roles list
|
||||
// @Summary Retrieve roles list
|
||||
// @Tags Roles
|
||||
// @Produce json
|
||||
// @Param request query rdb.TableQuery true "query table params"
|
||||
// @Success 200 {object} response.successResp{data=roles} "roles list received successful"
|
||||
// @Failure 400 {object} response.errorResp "invalid query request data"
|
||||
// @Failure 403 {object} response.errorResp "getting roles not permitted"
|
||||
// @Failure 500 {object} response.errorResp "internal error on getting roles"
|
||||
// @Router /roles/ [get]
|
||||
func (s *RoleService) GetRoles(c *gin.Context) {
|
||||
var (
|
||||
err error
|
||||
query rdb.TableQuery
|
||||
resp roles
|
||||
rids []uint64
|
||||
)
|
||||
|
||||
if err = c.ShouldBindQuery(&query); err != nil {
|
||||
logger.FromContext(c).WithError(err).Errorf("error binding query")
|
||||
response.Error(c, response.ErrRolesInvalidRequest, err)
|
||||
return
|
||||
}
|
||||
|
||||
rid := c.GetUint64("rid")
|
||||
privs := c.GetStringSlice("prm")
|
||||
scope := func(db *gorm.DB) *gorm.DB {
|
||||
if !slices.Contains(privs, "roles.view") {
|
||||
return db.Where("role_id = ?", rid)
|
||||
}
|
||||
return db
|
||||
}
|
||||
|
||||
query.Init("roles", rolesSQLMappers)
|
||||
|
||||
if resp.Total, err = query.Query(s.db, &resp.Roles, scope); err != nil {
|
||||
logger.FromContext(c).WithError(err).Errorf("error finding roles")
|
||||
response.Error(c, response.ErrInternal, err)
|
||||
return
|
||||
}
|
||||
|
||||
for _, role := range resp.Roles {
|
||||
rids = append(rids, role.ID)
|
||||
}
|
||||
|
||||
var privsObjs []models.Privilege
|
||||
if err = s.db.Find(&privsObjs, "role_id IN (?)", rids).Error; err != nil {
|
||||
logger.FromContext(c).WithError(err).Errorf("error finding linked roles")
|
||||
response.Error(c, response.ErrInternal, err)
|
||||
return
|
||||
}
|
||||
|
||||
privsRoles := make(map[uint64][]models.Privilege)
|
||||
for i := range privsObjs {
|
||||
privsRoles[privsObjs[i].RoleID] = append(privsRoles[privsObjs[i].RoleID], privsObjs[i])
|
||||
}
|
||||
|
||||
for i := range resp.Roles {
|
||||
resp.Roles[i].Privileges = privsRoles[resp.Roles[i].ID]
|
||||
}
|
||||
|
||||
for i := 0; i < len(resp.Roles); i++ {
|
||||
if err = resp.Roles[i].Valid(); err != nil {
|
||||
logger.FromContext(c).WithError(err).Errorf("error validating role data '%d'", resp.Roles[i].ID)
|
||||
response.Error(c, response.ErrRolesInvalidData, err)
|
||||
return
|
||||
}
|
||||
}
|
||||
|
||||
response.Success(c, http.StatusOK, resp)
|
||||
}
|
||||
|
||||
// GetRole is a function to return role by id
|
||||
// @Summary Retrieve role by id
|
||||
// @Tags Roles
|
||||
// @Produce json
|
||||
// @Param roleID path uint64 true "role id"
|
||||
// @Success 200 {object} response.successResp{data=models.RolePrivileges} "role received successful"
|
||||
// @Failure 400 {object} response.errorResp "invalid query request data"
|
||||
// @Failure 403 {object} response.errorResp "getting role not permitted"
|
||||
// @Failure 500 {object} response.errorResp "internal error on getting role"
|
||||
// @Router /roles/{roleID} [get]
|
||||
func (s *RoleService) GetRole(c *gin.Context) {
|
||||
var (
|
||||
err error
|
||||
resp models.RolePrivileges
|
||||
roleID uint64
|
||||
)
|
||||
|
||||
rid := c.GetUint64("rid")
|
||||
privs := c.GetStringSlice("prm")
|
||||
scope := func(db *gorm.DB) *gorm.DB {
|
||||
if !slices.Contains(privs, "roles.view") {
|
||||
return db.Where("role_id = ?", rid)
|
||||
}
|
||||
return db
|
||||
}
|
||||
|
||||
if roleID, err = strconv.ParseUint(c.Param("roleID"), 10, 64); err != nil {
|
||||
logger.FromContext(c).WithError(err).Errorf("error parsing role id")
|
||||
response.Error(c, response.ErrRolesInvalidRequest, err)
|
||||
return
|
||||
}
|
||||
|
||||
if err := s.db.Scopes(scope).Take(&resp, "id = ?", roleID).Error; err != nil {
|
||||
logger.FromContext(c).WithError(err).Errorf("error finding role by id")
|
||||
if errors.Is(err, gorm.ErrRecordNotFound) {
|
||||
response.Error(c, response.ErrRolesNotFound, err)
|
||||
} else {
|
||||
response.Error(c, response.ErrInternal, err)
|
||||
}
|
||||
return
|
||||
}
|
||||
if err := s.db.Model(&resp).Association("privileges").Find(&resp.Privileges).Error; err != nil {
|
||||
logger.FromContext(c).WithError(err).Errorf("error finding role privileges by role model")
|
||||
response.Error(c, response.ErrInternal, err)
|
||||
return
|
||||
}
|
||||
if err := resp.Valid(); err != nil {
|
||||
logger.FromContext(c).WithError(err).Errorf("error validating role data '%d'", resp.ID)
|
||||
response.Error(c, response.ErrRolesInvalidData, err)
|
||||
return
|
||||
}
|
||||
|
||||
response.Success(c, http.StatusOK, resp)
|
||||
}
|
||||
@@ -0,0 +1,354 @@
|
||||
package services
|
||||
|
||||
import (
|
||||
"errors"
|
||||
"fmt"
|
||||
"net/http"
|
||||
"path/filepath"
|
||||
"slices"
|
||||
"strconv"
|
||||
|
||||
"pentagi/pkg/server/logger"
|
||||
"pentagi/pkg/server/models"
|
||||
"pentagi/pkg/server/rdb"
|
||||
"pentagi/pkg/server/response"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/jinzhu/gorm"
|
||||
)
|
||||
|
||||
type screenshots struct {
|
||||
Screenshots []models.Screenshot `json:"screenshots"`
|
||||
Total uint64 `json:"total"`
|
||||
}
|
||||
|
||||
type screenshotsGrouped struct {
|
||||
Grouped []string `json:"grouped"`
|
||||
Total uint64 `json:"total"`
|
||||
}
|
||||
|
||||
var screenshotsSQLMappers = map[string]any{
|
||||
"id": "{{table}}.id",
|
||||
"name": "{{table}}.name",
|
||||
"url": "{{table}}.url",
|
||||
"flow_id": "{{table}}.flow_id",
|
||||
"task_id": "{{table}}.task_id",
|
||||
"subtask_id": "{{table}}.subtask_id",
|
||||
"created_at": "{{table}}.created_at",
|
||||
"data": "({{table}}.name || ' ' || {{table}}.url)",
|
||||
}
|
||||
|
||||
type ScreenshotService struct {
|
||||
db *gorm.DB
|
||||
dataDir string
|
||||
}
|
||||
|
||||
func NewScreenshotService(db *gorm.DB, dataDir string) *ScreenshotService {
|
||||
return &ScreenshotService{
|
||||
db: db,
|
||||
dataDir: dataDir,
|
||||
}
|
||||
}
|
||||
|
||||
// GetScreenshots is a function to return screenshots list
|
||||
// @Summary Retrieve screenshots list
|
||||
// @Tags Screenshots
|
||||
// @Produce json
|
||||
// @Security BearerAuth
|
||||
// @Param request query rdb.TableQuery true "query table params"
|
||||
// @Success 200 {object} response.successResp{data=screenshots} "screenshots list received successful"
|
||||
// @Failure 400 {object} response.errorResp "invalid query request data"
|
||||
// @Failure 403 {object} response.errorResp "getting screenshots not permitted"
|
||||
// @Failure 500 {object} response.errorResp "internal error on getting screenshots"
|
||||
// @Router /screenshots/ [get]
|
||||
func (s *ScreenshotService) GetScreenshots(c *gin.Context) {
|
||||
var (
|
||||
err error
|
||||
query rdb.TableQuery
|
||||
resp screenshots
|
||||
)
|
||||
|
||||
if err = c.ShouldBindQuery(&query); err != nil {
|
||||
logger.FromContext(c).WithError(err).Errorf("error binding query")
|
||||
response.Error(c, response.ErrScreenshotsInvalidRequest, err)
|
||||
return
|
||||
}
|
||||
|
||||
uid := c.GetUint64("uid")
|
||||
privs := c.GetStringSlice("prm")
|
||||
var scope func(db *gorm.DB) *gorm.DB
|
||||
if slices.Contains(privs, "screenshots.admin") {
|
||||
scope = func(db *gorm.DB) *gorm.DB {
|
||||
return db.
|
||||
Joins("INNER JOIN flows f ON f.id = screenshots.flow_id")
|
||||
}
|
||||
} else if slices.Contains(privs, "screenshots.view") {
|
||||
scope = func(db *gorm.DB) *gorm.DB {
|
||||
return db.
|
||||
Joins("INNER JOIN flows f ON f.id = screenshots.flow_id").
|
||||
Where("f.user_id = ?", uid)
|
||||
}
|
||||
} else {
|
||||
logger.FromContext(c).Errorf("error filtering user role permissions: permission not found")
|
||||
response.Error(c, response.ErrNotPermitted, nil)
|
||||
return
|
||||
}
|
||||
|
||||
query.Init("screenshots", screenshotsSQLMappers)
|
||||
|
||||
if query.Group != "" {
|
||||
if _, ok := screenshotsSQLMappers[query.Group]; !ok {
|
||||
logger.FromContext(c).Errorf("error finding screenshots grouped: group field not found")
|
||||
response.Error(c, response.ErrScreenshotsInvalidRequest, errors.New("group field not found"))
|
||||
return
|
||||
}
|
||||
|
||||
var respGrouped screenshotsGrouped
|
||||
if respGrouped.Total, err = query.QueryGrouped(s.db, &respGrouped.Grouped, scope); err != nil {
|
||||
logger.FromContext(c).WithError(err).Errorf("error finding screenshots grouped")
|
||||
response.Error(c, response.ErrInternal, err)
|
||||
return
|
||||
}
|
||||
|
||||
response.Success(c, http.StatusOK, respGrouped)
|
||||
return
|
||||
}
|
||||
|
||||
if resp.Total, err = query.Query(s.db, &resp.Screenshots, scope); err != nil {
|
||||
logger.FromContext(c).WithError(err).Errorf("error finding screenshots")
|
||||
response.Error(c, response.ErrInternal, err)
|
||||
return
|
||||
}
|
||||
|
||||
for i := 0; i < len(resp.Screenshots); i++ {
|
||||
if err = resp.Screenshots[i].Valid(); err != nil {
|
||||
logger.FromContext(c).WithError(err).Errorf("error validating screenshot data '%d'", resp.Screenshots[i].ID)
|
||||
response.Error(c, response.ErrScreenshotsInvalidData, err)
|
||||
return
|
||||
}
|
||||
}
|
||||
|
||||
response.Success(c, http.StatusOK, resp)
|
||||
}
|
||||
|
||||
// GetFlowScreenshots is a function to return screenshots list by flow id
|
||||
// @Summary Retrieve screenshots list by flow id
|
||||
// @Tags Screenshots
|
||||
// @Produce json
|
||||
// @Security BearerAuth
|
||||
// @Param flowID path int true "flow id" minimum(0)
|
||||
// @Param request query rdb.TableQuery true "query table params"
|
||||
// @Success 200 {object} response.successResp{data=screenshots} "screenshots list received successful"
|
||||
// @Failure 400 {object} response.errorResp "invalid query request data"
|
||||
// @Failure 403 {object} response.errorResp "getting screenshots not permitted"
|
||||
// @Failure 500 {object} response.errorResp "internal error on getting screenshots"
|
||||
// @Router /flows/{flowID}/screenshots/ [get]
|
||||
func (s *ScreenshotService) GetFlowScreenshots(c *gin.Context) {
|
||||
var (
|
||||
err error
|
||||
flowID uint64
|
||||
query rdb.TableQuery
|
||||
resp screenshots
|
||||
)
|
||||
|
||||
if flowID, err = strconv.ParseUint(c.Param("flowID"), 10, 64); err != nil {
|
||||
logger.FromContext(c).WithError(err).Errorf("error parsing flow id")
|
||||
response.Error(c, response.ErrScreenshotsInvalidRequest, err)
|
||||
return
|
||||
}
|
||||
|
||||
if err = c.ShouldBindQuery(&query); err != nil {
|
||||
logger.FromContext(c).WithError(err).Errorf("error binding query")
|
||||
response.Error(c, response.ErrScreenshotsInvalidRequest, err)
|
||||
return
|
||||
}
|
||||
|
||||
uid := c.GetUint64("uid")
|
||||
privs := c.GetStringSlice("prm")
|
||||
var scope func(db *gorm.DB) *gorm.DB
|
||||
if slices.Contains(privs, "screenshots.admin") {
|
||||
scope = func(db *gorm.DB) *gorm.DB {
|
||||
return db.
|
||||
Joins("INNER JOIN flows f ON f.id = screenshots.flow_id").
|
||||
Where("f.id = ?", flowID)
|
||||
}
|
||||
} else if slices.Contains(privs, "screenshots.view") {
|
||||
scope = func(db *gorm.DB) *gorm.DB {
|
||||
return db.
|
||||
Joins("INNER JOIN flows f ON f.id = screenshots.flow_id").
|
||||
Where("f.id = ? AND f.user_id = ?", flowID, uid)
|
||||
}
|
||||
} else {
|
||||
logger.FromContext(c).Errorf("error filtering user role permissions: permission not found")
|
||||
response.Error(c, response.ErrNotPermitted, nil)
|
||||
return
|
||||
}
|
||||
|
||||
query.Init("screenshots", screenshotsSQLMappers)
|
||||
|
||||
if query.Group != "" {
|
||||
if _, ok := screenshotsSQLMappers[query.Group]; !ok {
|
||||
logger.FromContext(c).Errorf("error finding screenshots grouped: group field not found")
|
||||
response.Error(c, response.ErrScreenshotsInvalidRequest, errors.New("group field not found"))
|
||||
return
|
||||
}
|
||||
|
||||
var respGrouped screenshotsGrouped
|
||||
if respGrouped.Total, err = query.QueryGrouped(s.db, &respGrouped.Grouped, scope); err != nil {
|
||||
logger.FromContext(c).WithError(err).Errorf("error finding screenshots grouped")
|
||||
response.Error(c, response.ErrInternal, err)
|
||||
return
|
||||
}
|
||||
|
||||
response.Success(c, http.StatusOK, respGrouped)
|
||||
return
|
||||
}
|
||||
|
||||
if resp.Total, err = query.Query(s.db, &resp.Screenshots, scope); err != nil {
|
||||
logger.FromContext(c).WithError(err).Errorf("error finding screenshots")
|
||||
response.Error(c, response.ErrInternal, err)
|
||||
return
|
||||
}
|
||||
|
||||
for i := 0; i < len(resp.Screenshots); i++ {
|
||||
if err = resp.Screenshots[i].Valid(); err != nil {
|
||||
logger.FromContext(c).WithError(err).Errorf("error validating screenshot data '%d'", resp.Screenshots[i].ID)
|
||||
response.Error(c, response.ErrScreenshotsInvalidData, err)
|
||||
return
|
||||
}
|
||||
}
|
||||
|
||||
response.Success(c, http.StatusOK, resp)
|
||||
}
|
||||
|
||||
// GetFlowScreenshot is a function to return screenshot info by id and flow id
|
||||
// @Summary Retrieve screenshot info by id and flow id
|
||||
// @Tags Screenshots
|
||||
// @Produce json
|
||||
// @Security BearerAuth
|
||||
// @Param flowID path int true "flow id" minimum(0)
|
||||
// @Param screenshotID path int true "screenshot id" minimum(0)
|
||||
// @Success 200 {object} response.successResp{data=models.Screenshot} "screenshot info received successful"
|
||||
// @Failure 403 {object} response.errorResp "getting screenshot not permitted"
|
||||
// @Failure 404 {object} response.errorResp "screenshot not found"
|
||||
// @Failure 500 {object} response.errorResp "internal error on getting screenshot"
|
||||
// @Router /flows/{flowID}/screenshots/{screenshotID} [get]
|
||||
func (s *ScreenshotService) GetFlowScreenshot(c *gin.Context) {
|
||||
var (
|
||||
err error
|
||||
flowID uint64
|
||||
screenshotID uint64
|
||||
resp models.Screenshot
|
||||
)
|
||||
|
||||
if flowID, err = strconv.ParseUint(c.Param("flowID"), 10, 64); err != nil {
|
||||
logger.FromContext(c).WithError(err).Errorf("error parsing flow id")
|
||||
response.Error(c, response.ErrScreenshotsInvalidRequest, err)
|
||||
return
|
||||
}
|
||||
if screenshotID, err = strconv.ParseUint(c.Param("screenshotID"), 10, 64); err != nil {
|
||||
logger.FromContext(c).WithError(err).Errorf("error parsing screenshot id")
|
||||
response.Error(c, response.ErrScreenshotsInvalidRequest, err)
|
||||
return
|
||||
}
|
||||
|
||||
uid := c.GetUint64("uid")
|
||||
privs := c.GetStringSlice("prm")
|
||||
var scope func(db *gorm.DB) *gorm.DB
|
||||
if slices.Contains(privs, "screenshots.admin") {
|
||||
scope = func(db *gorm.DB) *gorm.DB {
|
||||
return db.Where("f.id = ?", flowID)
|
||||
}
|
||||
} else if slices.Contains(privs, "screenshots.view") {
|
||||
scope = func(db *gorm.DB) *gorm.DB {
|
||||
return db.Where("f.id = ? AND f.user_id = ?", flowID, uid)
|
||||
}
|
||||
} else {
|
||||
logger.FromContext(c).Errorf("error filtering user role permissions: permission not found")
|
||||
response.Error(c, response.ErrNotPermitted, nil)
|
||||
return
|
||||
}
|
||||
|
||||
err = s.db.Model(&resp).
|
||||
Joins("INNER JOIN flows f ON f.id = flow_id").
|
||||
Scopes(scope).
|
||||
Where("screenshots.id = ?", screenshotID).
|
||||
Take(&resp).Error
|
||||
if err != nil {
|
||||
logger.FromContext(c).WithError(err).Errorf("error on getting screenshot by id")
|
||||
if gorm.IsRecordNotFoundError(err) {
|
||||
response.Error(c, response.ErrScreenshotsNotFound, err)
|
||||
} else {
|
||||
response.Error(c, response.ErrInternal, err)
|
||||
}
|
||||
return
|
||||
}
|
||||
|
||||
response.Success(c, http.StatusOK, resp)
|
||||
}
|
||||
|
||||
// GetFlowScreenshotFile is a function to return screenshot file by id and flow id
|
||||
// @Summary Retrieve screenshot file by id and flow id
|
||||
// @Tags Screenshots
|
||||
// @Produce png,json
|
||||
// @Security BearerAuth
|
||||
// @Param flowID path int true "flow id" minimum(0)
|
||||
// @Param screenshotID path int true "screenshot id" minimum(0)
|
||||
// @Success 200 {file} file "screenshot file"
|
||||
// @Failure 403 {object} response.errorResp "getting screenshot not permitted"
|
||||
// @Failure 500 {object} response.errorResp "internal error on getting screenshot"
|
||||
// @Router /flows/{flowID}/screenshots/{screenshotID}/file [get]
|
||||
func (s *ScreenshotService) GetFlowScreenshotFile(c *gin.Context) {
|
||||
var (
|
||||
err error
|
||||
flowID uint64
|
||||
screenshotID uint64
|
||||
resp models.Screenshot
|
||||
)
|
||||
|
||||
if flowID, err = strconv.ParseUint(c.Param("flowID"), 10, 64); err != nil {
|
||||
logger.FromContext(c).WithError(err).Errorf("error parsing flow id")
|
||||
response.Error(c, response.ErrScreenshotsInvalidRequest, err)
|
||||
return
|
||||
}
|
||||
if screenshotID, err = strconv.ParseUint(c.Param("screenshotID"), 10, 64); err != nil {
|
||||
logger.FromContext(c).WithError(err).Errorf("error parsing screenshot id")
|
||||
response.Error(c, response.ErrScreenshotsInvalidRequest, err)
|
||||
return
|
||||
}
|
||||
|
||||
uid := c.GetUint64("uid")
|
||||
privs := c.GetStringSlice("prm")
|
||||
var scope func(db *gorm.DB) *gorm.DB
|
||||
if slices.Contains(privs, "screenshots.admin") {
|
||||
scope = func(db *gorm.DB) *gorm.DB {
|
||||
return db.Where("f.id = ?", flowID)
|
||||
}
|
||||
} else if slices.Contains(privs, "screenshots.download") {
|
||||
scope = func(db *gorm.DB) *gorm.DB {
|
||||
return db.Where("f.id = ? AND f.user_id = ?", flowID, uid)
|
||||
}
|
||||
} else {
|
||||
logger.FromContext(c).Errorf("error filtering user role permissions: permission not found")
|
||||
response.Error(c, response.ErrNotPermitted, nil)
|
||||
return
|
||||
}
|
||||
|
||||
err = s.db.Model(&resp).
|
||||
Joins("INNER JOIN flows f ON f.id = flow_id").
|
||||
Scopes(scope).
|
||||
Where("screenshots.id = ?", screenshotID).
|
||||
Take(&resp).Error
|
||||
if err != nil {
|
||||
logger.FromContext(c).WithError(err).Errorf("error on getting screenshot by id")
|
||||
if gorm.IsRecordNotFoundError(err) {
|
||||
response.Error(c, response.ErrScreenshotsNotFound, err)
|
||||
} else {
|
||||
response.Error(c, response.ErrInternal, err)
|
||||
}
|
||||
return
|
||||
}
|
||||
|
||||
flowDirName := fmt.Sprintf("flow-%d", resp.FlowID)
|
||||
c.File(filepath.Join(s.dataDir, "screenshots", flowDirName, resp.Name))
|
||||
}
|
||||
@@ -0,0 +1,221 @@
|
||||
package services
|
||||
|
||||
import (
|
||||
"errors"
|
||||
"net/http"
|
||||
"slices"
|
||||
"strconv"
|
||||
|
||||
"pentagi/pkg/server/logger"
|
||||
"pentagi/pkg/server/models"
|
||||
"pentagi/pkg/server/rdb"
|
||||
"pentagi/pkg/server/response"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/jinzhu/gorm"
|
||||
)
|
||||
|
||||
type searchlogs struct {
|
||||
SearchLogs []models.Searchlog `json:"searchlogs"`
|
||||
Total uint64 `json:"total"`
|
||||
}
|
||||
|
||||
type searchlogsGrouped struct {
|
||||
Grouped []string `json:"grouped"`
|
||||
Total uint64 `json:"total"`
|
||||
}
|
||||
|
||||
var searchlogsSQLMappers = map[string]any{
|
||||
"id": "{{table}}.id",
|
||||
"initiator": "{{table}}.initiator",
|
||||
"executor": "{{table}}.executor",
|
||||
"engine": "{{table}}.engine",
|
||||
"query": "{{table}}.query",
|
||||
"result": "{{table}}.result",
|
||||
"flow_id": "{{table}}.flow_id",
|
||||
"task_id": "{{table}}.task_id",
|
||||
"subtask_id": "{{table}}.subtask_id",
|
||||
"created_at": "{{table}}.created_at",
|
||||
"data": "({{table}}.query || ' ' || {{table}}.result)",
|
||||
}
|
||||
|
||||
type SearchlogService struct {
|
||||
db *gorm.DB
|
||||
}
|
||||
|
||||
func NewSearchlogService(db *gorm.DB) *SearchlogService {
|
||||
return &SearchlogService{
|
||||
db: db,
|
||||
}
|
||||
}
|
||||
|
||||
// GetSearchlogs is a function to return searchlogs list
|
||||
// @Summary Retrieve searchlogs list
|
||||
// @Tags Searchlogs
|
||||
// @Produce json
|
||||
// @Security BearerAuth
|
||||
// @Param request query rdb.TableQuery true "query table params"
|
||||
// @Success 200 {object} response.successResp{data=searchlogs} "searchlogs list received successful"
|
||||
// @Failure 400 {object} response.errorResp "invalid query request data"
|
||||
// @Failure 403 {object} response.errorResp "getting searchlogs not permitted"
|
||||
// @Failure 500 {object} response.errorResp "internal error on getting searchlogs"
|
||||
// @Router /searchlogs/ [get]
|
||||
func (s *SearchlogService) GetSearchlogs(c *gin.Context) {
|
||||
var (
|
||||
err error
|
||||
query rdb.TableQuery
|
||||
resp searchlogs
|
||||
)
|
||||
|
||||
if err = c.ShouldBindQuery(&query); err != nil {
|
||||
logger.FromContext(c).WithError(err).Errorf("error binding query")
|
||||
response.Error(c, response.ErrSearchlogsInvalidRequest, err)
|
||||
return
|
||||
}
|
||||
|
||||
uid := c.GetUint64("uid")
|
||||
privs := c.GetStringSlice("prm")
|
||||
var scope func(db *gorm.DB) *gorm.DB
|
||||
if slices.Contains(privs, "searchlogs.admin") {
|
||||
scope = func(db *gorm.DB) *gorm.DB {
|
||||
return db.
|
||||
Joins("INNER JOIN flows f ON f.id = flow_id")
|
||||
}
|
||||
} else if slices.Contains(privs, "searchlogs.view") {
|
||||
scope = func(db *gorm.DB) *gorm.DB {
|
||||
return db.
|
||||
Joins("INNER JOIN flows f ON f.id = flow_id").
|
||||
Where("f.user_id = ?", uid)
|
||||
}
|
||||
} else {
|
||||
logger.FromContext(c).Errorf("error filtering user role permissions: permission not found")
|
||||
response.Error(c, response.ErrNotPermitted, nil)
|
||||
return
|
||||
}
|
||||
|
||||
query.Init("searchlogs", searchlogsSQLMappers)
|
||||
|
||||
if query.Group != "" {
|
||||
if _, ok := searchlogsSQLMappers[query.Group]; !ok {
|
||||
logger.FromContext(c).Errorf("error finding searchlogs grouped: group field not found")
|
||||
response.Error(c, response.ErrSearchlogsInvalidRequest, errors.New("group field not found"))
|
||||
return
|
||||
}
|
||||
|
||||
var respGrouped searchlogsGrouped
|
||||
if respGrouped.Total, err = query.QueryGrouped(s.db, &respGrouped.Grouped, scope); err != nil {
|
||||
logger.FromContext(c).WithError(err).Errorf("error finding searchlogs grouped")
|
||||
response.Error(c, response.ErrInternal, err)
|
||||
return
|
||||
}
|
||||
|
||||
response.Success(c, http.StatusOK, respGrouped)
|
||||
return
|
||||
}
|
||||
|
||||
if resp.Total, err = query.Query(s.db, &resp.SearchLogs, scope); err != nil {
|
||||
logger.FromContext(c).WithError(err).Errorf("error finding searchlogs")
|
||||
response.Error(c, response.ErrInternal, err)
|
||||
return
|
||||
}
|
||||
|
||||
for i := 0; i < len(resp.SearchLogs); i++ {
|
||||
if err = resp.SearchLogs[i].Valid(); err != nil {
|
||||
logger.FromContext(c).WithError(err).Errorf("error validating searchlog data '%d'", resp.SearchLogs[i].ID)
|
||||
response.Error(c, response.ErrSearchlogsInvalidData, err)
|
||||
return
|
||||
}
|
||||
}
|
||||
|
||||
response.Success(c, http.StatusOK, resp)
|
||||
}
|
||||
|
||||
// GetFlowSearchlogs is a function to return searchlogs list by flow id
|
||||
// @Summary Retrieve searchlogs list by flow id
|
||||
// @Tags Searchlogs
|
||||
// @Produce json
|
||||
// @Security BearerAuth
|
||||
// @Param flowID path int true "flow id" minimum(0)
|
||||
// @Param request query rdb.TableQuery true "query table params"
|
||||
// @Success 200 {object} response.successResp{data=searchlogs} "searchlogs list received successful"
|
||||
// @Failure 400 {object} response.errorResp "invalid query request data"
|
||||
// @Failure 403 {object} response.errorResp "getting searchlogs not permitted"
|
||||
// @Failure 500 {object} response.errorResp "internal error on getting searchlogs"
|
||||
// @Router /flows/{flowID}/searchlogs/ [get]
|
||||
func (s *SearchlogService) GetFlowSearchlogs(c *gin.Context) {
|
||||
var (
|
||||
err error
|
||||
flowID uint64
|
||||
query rdb.TableQuery
|
||||
resp searchlogs
|
||||
)
|
||||
|
||||
if flowID, err = strconv.ParseUint(c.Param("flowID"), 10, 64); err != nil {
|
||||
logger.FromContext(c).WithError(err).Errorf("error parsing flow id")
|
||||
response.Error(c, response.ErrSearchlogsInvalidRequest, err)
|
||||
return
|
||||
}
|
||||
|
||||
if err = c.ShouldBindQuery(&query); err != nil {
|
||||
logger.FromContext(c).WithError(err).Errorf("error binding query")
|
||||
response.Error(c, response.ErrSearchlogsInvalidRequest, err)
|
||||
return
|
||||
}
|
||||
|
||||
uid := c.GetUint64("uid")
|
||||
privs := c.GetStringSlice("prm")
|
||||
var scope func(db *gorm.DB) *gorm.DB
|
||||
if slices.Contains(privs, "searchlogs.admin") {
|
||||
scope = func(db *gorm.DB) *gorm.DB {
|
||||
return db.
|
||||
Joins("INNER JOIN flows f ON f.id = flow_id").
|
||||
Where("f.id = ?", flowID)
|
||||
}
|
||||
} else if slices.Contains(privs, "searchlogs.view") {
|
||||
scope = func(db *gorm.DB) *gorm.DB {
|
||||
return db.
|
||||
Joins("INNER JOIN flows f ON f.id = flow_id").
|
||||
Where("f.id = ? AND f.user_id = ?", flowID, uid)
|
||||
}
|
||||
} else {
|
||||
logger.FromContext(c).Errorf("error filtering user role permissions: permission not found")
|
||||
response.Error(c, response.ErrNotPermitted, nil)
|
||||
return
|
||||
}
|
||||
|
||||
query.Init("searchlogs", searchlogsSQLMappers)
|
||||
|
||||
if query.Group != "" {
|
||||
if _, ok := searchlogsSQLMappers[query.Group]; !ok {
|
||||
logger.FromContext(c).Errorf("error finding searchlogs grouped: group field not found")
|
||||
response.Error(c, response.ErrSearchlogsInvalidRequest, errors.New("group field not found"))
|
||||
return
|
||||
}
|
||||
|
||||
var respGrouped searchlogsGrouped
|
||||
if respGrouped.Total, err = query.QueryGrouped(s.db, &respGrouped.Grouped, scope); err != nil {
|
||||
logger.FromContext(c).WithError(err).Errorf("error finding searchlogs grouped")
|
||||
response.Error(c, response.ErrInternal, err)
|
||||
return
|
||||
}
|
||||
|
||||
response.Success(c, http.StatusOK, respGrouped)
|
||||
return
|
||||
}
|
||||
|
||||
if resp.Total, err = query.Query(s.db, &resp.SearchLogs, scope); err != nil {
|
||||
logger.FromContext(c).WithError(err).Errorf("error finding searchlogs")
|
||||
response.Error(c, response.ErrInternal, err)
|
||||
return
|
||||
}
|
||||
|
||||
for i := 0; i < len(resp.SearchLogs); i++ {
|
||||
if err = resp.SearchLogs[i].Valid(); err != nil {
|
||||
logger.FromContext(c).WithError(err).Errorf("error validating searchlog data '%d'", resp.SearchLogs[i].ID)
|
||||
response.Error(c, response.ErrSearchlogsInvalidData, err)
|
||||
return
|
||||
}
|
||||
}
|
||||
|
||||
response.Success(c, http.StatusOK, resp)
|
||||
}
|
||||
@@ -0,0 +1,50 @@
|
||||
package services
|
||||
|
||||
import (
|
||||
"net/http"
|
||||
"slices"
|
||||
|
||||
"pentagi/pkg/config"
|
||||
"pentagi/pkg/server/logger"
|
||||
"pentagi/pkg/server/models"
|
||||
"pentagi/pkg/server/response"
|
||||
"pentagi/pkg/version"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
)
|
||||
|
||||
type SettingsService struct {
|
||||
cfg *config.Config
|
||||
}
|
||||
|
||||
func NewSettingsService(cfg *config.Config) *SettingsService {
|
||||
return &SettingsService{cfg: cfg}
|
||||
}
|
||||
|
||||
// GetSettings is a function to return settings
|
||||
// @Summary Retrieve settings
|
||||
// @Tags Settings
|
||||
// @Produce json
|
||||
// @Security BearerAuth
|
||||
// @Success 200 {object} response.successResp{data=models.Settings} "settings received successful"
|
||||
// @Failure 403 {object} response.errorResp "getting settings not permitted"
|
||||
// @Router /settings/ [get]
|
||||
func (s *SettingsService) GetSettings(c *gin.Context) {
|
||||
privs := c.GetStringSlice("prm")
|
||||
if !slices.Contains(privs, "settings.view") {
|
||||
logger.FromContext(c).Errorf("error filtering user role permissions: permission not found")
|
||||
response.Error(c, response.ErrNotPermitted, nil)
|
||||
return
|
||||
}
|
||||
|
||||
settings := models.Settings{
|
||||
Debug: s.cfg.Debug,
|
||||
AskUser: s.cfg.AskUser,
|
||||
Version: version.GetBinaryVersion(),
|
||||
DockerInside: s.cfg.DockerInside,
|
||||
IsDevelopMode: version.IsDevelopMode(),
|
||||
AssistantUseAgents: s.cfg.AssistantUseAgents,
|
||||
}
|
||||
|
||||
response.Success(c, http.StatusOK, settings)
|
||||
}
|
||||
@@ -0,0 +1,321 @@
|
||||
package services
|
||||
|
||||
import (
|
||||
"errors"
|
||||
"net/http"
|
||||
"slices"
|
||||
"strconv"
|
||||
|
||||
"pentagi/pkg/server/logger"
|
||||
"pentagi/pkg/server/models"
|
||||
"pentagi/pkg/server/rdb"
|
||||
"pentagi/pkg/server/response"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/jinzhu/gorm"
|
||||
)
|
||||
|
||||
type subtasks struct {
|
||||
Subtasks []models.Subtask `json:"subtasks"`
|
||||
Total uint64 `json:"total"`
|
||||
}
|
||||
|
||||
type subtasksGrouped struct {
|
||||
Grouped []string `json:"grouped"`
|
||||
Total uint64 `json:"total"`
|
||||
}
|
||||
|
||||
var subtasksSQLMappers = map[string]any{
|
||||
"id": "{{table}}.id",
|
||||
"status": "{{table}}.status",
|
||||
"title": "{{table}}.title",
|
||||
"description": "{{table}}.description",
|
||||
"context": "{{table}}.context",
|
||||
"result": "{{table}}.result",
|
||||
"task_id": "{{table}}.task_id",
|
||||
"created_at": "{{table}}.created_at",
|
||||
"updated_at": "{{table}}.updated_at",
|
||||
"data": "({{table}}.status || ' ' || {{table}}.title || ' ' || {{table}}.description || ' ' || {{table}}.context || ' ' || {{table}}.result)",
|
||||
}
|
||||
|
||||
type SubtaskService struct {
|
||||
db *gorm.DB
|
||||
}
|
||||
|
||||
func NewSubtaskService(db *gorm.DB) *SubtaskService {
|
||||
return &SubtaskService{
|
||||
db: db,
|
||||
}
|
||||
}
|
||||
|
||||
// GetFlowSubtasks is a function to return flow subtasks list
|
||||
// @Summary Retrieve flow subtasks list
|
||||
// @Tags Subtasks
|
||||
// @Produce json
|
||||
// @Security BearerAuth
|
||||
// @Param flowID path int true "flow id" minimum(0)
|
||||
// @Param request query rdb.TableQuery true "query table params"
|
||||
// @Success 200 {object} response.successResp{data=subtasks} "flow subtasks list received successful"
|
||||
// @Failure 400 {object} response.errorResp "invalid query request data"
|
||||
// @Failure 403 {object} response.errorResp "getting flow subtasks not permitted"
|
||||
// @Failure 500 {object} response.errorResp "internal error on getting flow subtasks"
|
||||
// @Router /flows/{flowID}/subtasks/ [get]
|
||||
func (s *SubtaskService) GetFlowSubtasks(c *gin.Context) {
|
||||
var (
|
||||
err error
|
||||
flowID uint64
|
||||
query rdb.TableQuery
|
||||
resp subtasks
|
||||
)
|
||||
|
||||
if err = c.ShouldBindQuery(&query); err != nil {
|
||||
logger.FromContext(c).WithError(err).Errorf("error binding query")
|
||||
response.Error(c, response.ErrSubtasksInvalidRequest, err)
|
||||
return
|
||||
}
|
||||
|
||||
if flowID, err = strconv.ParseUint(c.Param("flowID"), 10, 64); err != nil {
|
||||
logger.FromContext(c).WithError(err).Errorf("error parsing flow id")
|
||||
response.Error(c, response.ErrSubtasksInvalidRequest, err)
|
||||
return
|
||||
}
|
||||
|
||||
uid := c.GetUint64("uid")
|
||||
privs := c.GetStringSlice("prm")
|
||||
var scope func(db *gorm.DB) *gorm.DB
|
||||
if slices.Contains(privs, "subtasks.admin") {
|
||||
scope = func(db *gorm.DB) *gorm.DB {
|
||||
return db.
|
||||
Joins("INNER JOIN tasks t ON t.id = subtasks.task_id").
|
||||
Joins("INNER JOIN flows f ON f.id = t.flow_id").
|
||||
Where("f.id = ?", flowID)
|
||||
}
|
||||
} else if slices.Contains(privs, "subtasks.view") {
|
||||
scope = func(db *gorm.DB) *gorm.DB {
|
||||
return db.
|
||||
Joins("INNER JOIN tasks t ON t.id = subtasks.task_id").
|
||||
Joins("INNER JOIN flows f ON f.id = t.flow_id").
|
||||
Where("f.id = ? AND f.user_id = ?", flowID, uid)
|
||||
}
|
||||
} else {
|
||||
logger.FromContext(c).Errorf("error filtering user role permissions: permission not found")
|
||||
response.Error(c, response.ErrNotPermitted, nil)
|
||||
return
|
||||
}
|
||||
|
||||
query.Init("subtasks", subtasksSQLMappers)
|
||||
|
||||
if query.Group != "" {
|
||||
if _, ok := subtasksSQLMappers[query.Group]; !ok {
|
||||
logger.FromContext(c).Errorf("error finding subtasks grouped: group field not found")
|
||||
response.Error(c, response.ErrSubtasksInvalidRequest, errors.New("group field not found"))
|
||||
return
|
||||
}
|
||||
|
||||
var respGrouped subtasksGrouped
|
||||
if respGrouped.Total, err = query.QueryGrouped(s.db, &respGrouped.Grouped, scope); err != nil {
|
||||
logger.FromContext(c).WithError(err).Errorf("error finding subtasks grouped")
|
||||
response.Error(c, response.ErrInternal, err)
|
||||
return
|
||||
}
|
||||
|
||||
response.Success(c, http.StatusOK, respGrouped)
|
||||
return
|
||||
}
|
||||
|
||||
if resp.Total, err = query.Query(s.db, &resp.Subtasks, scope); err != nil {
|
||||
logger.FromContext(c).WithError(err).Errorf("error finding subtasks")
|
||||
response.Error(c, response.ErrInternal, err)
|
||||
return
|
||||
}
|
||||
|
||||
for i := 0; i < len(resp.Subtasks); i++ {
|
||||
if err = resp.Subtasks[i].Valid(); err != nil {
|
||||
logger.FromContext(c).WithError(err).Errorf("error validating subtask data '%d'", resp.Subtasks[i].ID)
|
||||
response.Error(c, response.ErrSubtasksInvalidData, err)
|
||||
return
|
||||
}
|
||||
}
|
||||
|
||||
response.Success(c, http.StatusOK, resp)
|
||||
}
|
||||
|
||||
// GetFlowTaskSubtasks is a function to return flow task subtasks list
|
||||
// @Summary Retrieve flow task subtasks list
|
||||
// @Tags Subtasks
|
||||
// @Produce json
|
||||
// @Security BearerAuth
|
||||
// @Param flowID path int true "flow id" minimum(0)
|
||||
// @Param taskID path int true "task id" minimum(0)
|
||||
// @Param request query rdb.TableQuery true "query table params"
|
||||
// @Success 200 {object} response.successResp{data=subtasks} "flow task subtasks list received successful"
|
||||
// @Failure 400 {object} response.errorResp "invalid query request data"
|
||||
// @Failure 403 {object} response.errorResp "getting flow task subtasks not permitted"
|
||||
// @Failure 500 {object} response.errorResp "internal error on getting flow subtasks"
|
||||
// @Router /flows/{flowID}/tasks/{taskID}/subtasks/ [get]
|
||||
func (s *SubtaskService) GetFlowTaskSubtasks(c *gin.Context) {
|
||||
var (
|
||||
err error
|
||||
flowID uint64
|
||||
taskID uint64
|
||||
query rdb.TableQuery
|
||||
resp subtasks
|
||||
)
|
||||
|
||||
if err = c.ShouldBindQuery(&query); err != nil {
|
||||
logger.FromContext(c).WithError(err).Errorf("error binding query")
|
||||
response.Error(c, response.ErrSubtasksInvalidRequest, err)
|
||||
return
|
||||
}
|
||||
|
||||
if flowID, err = strconv.ParseUint(c.Param("flowID"), 10, 64); err != nil {
|
||||
logger.FromContext(c).WithError(err).Errorf("error parsing flow id")
|
||||
response.Error(c, response.ErrSubtasksInvalidRequest, err)
|
||||
return
|
||||
}
|
||||
|
||||
if taskID, err = strconv.ParseUint(c.Param("taskID"), 10, 64); err != nil {
|
||||
logger.FromContext(c).WithError(err).Errorf("error parsing task id")
|
||||
response.Error(c, response.ErrSubtasksInvalidRequest, err)
|
||||
return
|
||||
}
|
||||
|
||||
uid := c.GetUint64("uid")
|
||||
privs := c.GetStringSlice("prm")
|
||||
var scope func(db *gorm.DB) *gorm.DB
|
||||
if slices.Contains(privs, "subtasks.admin") {
|
||||
scope = func(db *gorm.DB) *gorm.DB {
|
||||
return db.
|
||||
Joins("INNER JOIN tasks t ON t.id = subtasks.task_id").
|
||||
Joins("INNER JOIN flows f ON f.id = t.flow_id").
|
||||
Where("f.id = ? AND t.id = ?", flowID, taskID)
|
||||
}
|
||||
} else if slices.Contains(privs, "subtasks.view") {
|
||||
scope = func(db *gorm.DB) *gorm.DB {
|
||||
return db.
|
||||
Joins("INNER JOIN tasks t ON t.id = subtasks.task_id").
|
||||
Joins("INNER JOIN flows f ON f.id = t.flow_id").
|
||||
Where("f.id = ? AND f.user_id = ? AND t.id = ?", flowID, uid, taskID)
|
||||
}
|
||||
} else {
|
||||
logger.FromContext(c).Errorf("error filtering user role permissions: permission not found")
|
||||
response.Error(c, response.ErrNotPermitted, nil)
|
||||
return
|
||||
}
|
||||
|
||||
query.Init("subtasks", subtasksSQLMappers)
|
||||
|
||||
if query.Group != "" {
|
||||
if _, ok := subtasksSQLMappers[query.Group]; !ok {
|
||||
logger.FromContext(c).Errorf("error finding subtasks grouped: group field not found")
|
||||
response.Error(c, response.ErrSubtasksInvalidRequest, errors.New("group field not found"))
|
||||
return
|
||||
}
|
||||
|
||||
var respGrouped subtasksGrouped
|
||||
if respGrouped.Total, err = query.QueryGrouped(s.db, &respGrouped.Grouped, scope); err != nil {
|
||||
logger.FromContext(c).WithError(err).Errorf("error finding subtasks grouped")
|
||||
response.Error(c, response.ErrInternal, err)
|
||||
return
|
||||
}
|
||||
|
||||
response.Success(c, http.StatusOK, respGrouped)
|
||||
return
|
||||
}
|
||||
|
||||
if resp.Total, err = query.Query(s.db, &resp.Subtasks, scope); err != nil {
|
||||
logger.FromContext(c).WithError(err).Errorf("error finding subtasks")
|
||||
response.Error(c, response.ErrInternal, err)
|
||||
return
|
||||
}
|
||||
|
||||
for i := 0; i < len(resp.Subtasks); i++ {
|
||||
if err = resp.Subtasks[i].Valid(); err != nil {
|
||||
logger.FromContext(c).WithError(err).Errorf("error validating subtask data '%d'", resp.Subtasks[i].ID)
|
||||
response.Error(c, response.ErrSubtasksInvalidData, err)
|
||||
return
|
||||
}
|
||||
}
|
||||
|
||||
response.Success(c, http.StatusOK, resp)
|
||||
}
|
||||
|
||||
// GetFlowTaskSubtask is a function to return flow task subtask by id
|
||||
// @Summary Retrieve flow task subtask by id
|
||||
// @Tags Subtasks
|
||||
// @Produce json
|
||||
// @Security BearerAuth
|
||||
// @Param flowID path int true "flow id" minimum(0)
|
||||
// @Param taskID path int true "task id" minimum(0)
|
||||
// @Param subtaskID path int true "subtask id" minimum(0)
|
||||
// @Success 200 {object} response.successResp{data=models.Subtask} "flow task subtask received successful"
|
||||
// @Failure 403 {object} response.errorResp "getting flow task subtask not permitted"
|
||||
// @Failure 404 {object} response.errorResp "flow task subtask not found"
|
||||
// @Failure 500 {object} response.errorResp "internal error on getting flow task subtask"
|
||||
// @Router /flows/{flowID}/tasks/{taskID}/subtasks/{subtaskID} [get]
|
||||
func (s *SubtaskService) GetFlowTaskSubtask(c *gin.Context) {
|
||||
var (
|
||||
err error
|
||||
flowID uint64
|
||||
taskID uint64
|
||||
subtaskID uint64
|
||||
resp models.Subtask
|
||||
)
|
||||
|
||||
if flowID, err = strconv.ParseUint(c.Param("flowID"), 10, 64); err != nil {
|
||||
logger.FromContext(c).WithError(err).Errorf("error parsing flow id")
|
||||
response.Error(c, response.ErrSubtasksInvalidRequest, err)
|
||||
return
|
||||
}
|
||||
|
||||
if taskID, err = strconv.ParseUint(c.Param("taskID"), 10, 64); err != nil {
|
||||
logger.FromContext(c).WithError(err).Errorf("error parsing task id")
|
||||
response.Error(c, response.ErrSubtasksInvalidRequest, err)
|
||||
return
|
||||
}
|
||||
|
||||
if subtaskID, err = strconv.ParseUint(c.Param("subtaskID"), 10, 64); err != nil {
|
||||
logger.FromContext(c).WithError(err).Errorf("error parsing subtask id")
|
||||
response.Error(c, response.ErrSubtasksInvalidRequest, err)
|
||||
return
|
||||
}
|
||||
|
||||
uid := c.GetUint64("uid")
|
||||
privs := c.GetStringSlice("prm")
|
||||
var scope func(db *gorm.DB) *gorm.DB
|
||||
if slices.Contains(privs, "subtasks.admin") {
|
||||
scope = func(db *gorm.DB) *gorm.DB {
|
||||
return db.
|
||||
Joins("INNER JOIN tasks t ON t.id = subtasks.task_id").
|
||||
Joins("INNER JOIN flows f ON f.id = t.flow_id").
|
||||
Where("f.id = ? AND t.id = ?", flowID, taskID)
|
||||
}
|
||||
} else if slices.Contains(privs, "subtasks.view") {
|
||||
scope = func(db *gorm.DB) *gorm.DB {
|
||||
return db.
|
||||
Joins("INNER JOIN tasks t ON t.id = subtasks.task_id").
|
||||
Joins("INNER JOIN flows f ON f.id = t.flow_id").
|
||||
Where("f.id = ? AND f.user_id = ? AND t.id = ?", flowID, uid, taskID)
|
||||
}
|
||||
} else {
|
||||
logger.FromContext(c).Errorf("error filtering user role permissions: permission not found")
|
||||
response.Error(c, response.ErrNotPermitted, nil)
|
||||
return
|
||||
}
|
||||
|
||||
err = s.db.Model(&resp).
|
||||
Scopes(scope).
|
||||
Where("subtasks.id = ?", subtaskID).
|
||||
Take(&resp).Error
|
||||
if err != nil {
|
||||
logger.FromContext(c).WithError(err).Errorf("error on getting flow task subtask by id")
|
||||
if gorm.IsRecordNotFoundError(err) {
|
||||
response.Error(c, response.ErrSubtasksNotFound, err)
|
||||
} else {
|
||||
response.Error(c, response.ErrInternal, err)
|
||||
}
|
||||
return
|
||||
}
|
||||
|
||||
response.Success(c, http.StatusOK, resp)
|
||||
}
|
||||
@@ -0,0 +1,307 @@
|
||||
package services
|
||||
|
||||
import (
|
||||
"errors"
|
||||
"net/http"
|
||||
"slices"
|
||||
"strconv"
|
||||
|
||||
"pentagi/pkg/server/logger"
|
||||
"pentagi/pkg/server/models"
|
||||
"pentagi/pkg/server/rdb"
|
||||
"pentagi/pkg/server/response"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/jinzhu/gorm"
|
||||
)
|
||||
|
||||
type tasks struct {
|
||||
Tasks []models.Task `json:"tasks"`
|
||||
Total uint64 `json:"total"`
|
||||
}
|
||||
|
||||
type tasksGrouped struct {
|
||||
Grouped []string `json:"grouped"`
|
||||
Total uint64 `json:"total"`
|
||||
}
|
||||
|
||||
var tasksSQLMappers = map[string]any{
|
||||
"id": "{{table}}.id",
|
||||
"status": "{{table}}.status",
|
||||
"title": "{{table}}.title",
|
||||
"input": "{{table}}.input",
|
||||
"result": "{{table}}.result",
|
||||
"flow_id": "{{table}}.flow_id",
|
||||
"created_at": "{{table}}.created_at",
|
||||
"updated_at": "{{table}}.updated_at",
|
||||
"data": "({{table}}.status || ' ' || {{table}}.title || ' ' || {{table}}.input || ' ' || {{table}}.result)",
|
||||
}
|
||||
|
||||
type TaskService struct {
|
||||
db *gorm.DB
|
||||
}
|
||||
|
||||
func NewTaskService(db *gorm.DB) *TaskService {
|
||||
return &TaskService{
|
||||
db: db,
|
||||
}
|
||||
}
|
||||
|
||||
// GetFlowTasks is a function to return flow tasks list
|
||||
// @Summary Retrieve flow tasks list
|
||||
// @Tags Tasks
|
||||
// @Produce json
|
||||
// @Security BearerAuth
|
||||
// @Param flowID path int true "flow id" minimum(0)
|
||||
// @Param request query rdb.TableQuery true "query table params"
|
||||
// @Success 200 {object} response.successResp{data=tasks} "flow tasks list received successful"
|
||||
// @Failure 400 {object} response.errorResp "invalid query request data"
|
||||
// @Failure 403 {object} response.errorResp "getting flow tasks not permitted"
|
||||
// @Failure 500 {object} response.errorResp "internal error on getting flow tasks"
|
||||
// @Router /flows/{flowID}/tasks/ [get]
|
||||
func (s *TaskService) GetFlowTasks(c *gin.Context) {
|
||||
var (
|
||||
err error
|
||||
flowID uint64
|
||||
query rdb.TableQuery
|
||||
resp tasks
|
||||
)
|
||||
|
||||
if err = c.ShouldBindQuery(&query); err != nil {
|
||||
logger.FromContext(c).WithError(err).Errorf("error binding query")
|
||||
response.Error(c, response.ErrTasksInvalidRequest, err)
|
||||
return
|
||||
}
|
||||
|
||||
if flowID, err = strconv.ParseUint(c.Param("flowID"), 10, 64); err != nil {
|
||||
logger.FromContext(c).WithError(err).Errorf("error parsing flow id")
|
||||
response.Error(c, response.ErrTasksInvalidRequest, err)
|
||||
return
|
||||
}
|
||||
|
||||
uid := c.GetUint64("uid")
|
||||
privs := c.GetStringSlice("prm")
|
||||
var scope func(db *gorm.DB) *gorm.DB
|
||||
if slices.Contains(privs, "tasks.admin") {
|
||||
scope = func(db *gorm.DB) *gorm.DB {
|
||||
return db.
|
||||
Joins("INNER JOIN flows f ON f.id = tasks.flow_id").
|
||||
Where("f.id = ?", flowID)
|
||||
}
|
||||
} else if slices.Contains(privs, "tasks.view") {
|
||||
scope = func(db *gorm.DB) *gorm.DB {
|
||||
return db.
|
||||
Joins("INNER JOIN flows f ON f.id = tasks.flow_id").
|
||||
Where("f.id = ? AND f.user_id = ?", flowID, uid)
|
||||
}
|
||||
} else {
|
||||
logger.FromContext(c).Errorf("error filtering user role permissions: permission not found")
|
||||
response.Error(c, response.ErrNotPermitted, nil)
|
||||
return
|
||||
}
|
||||
|
||||
query.Init("tasks", tasksSQLMappers)
|
||||
|
||||
if query.Group != "" {
|
||||
if _, ok := tasksSQLMappers[query.Group]; !ok {
|
||||
logger.FromContext(c).Errorf("error finding tasks grouped: group field not found")
|
||||
response.Error(c, response.ErrTasksInvalidRequest, errors.New("group field not found"))
|
||||
return
|
||||
}
|
||||
|
||||
var respGrouped tasksGrouped
|
||||
if respGrouped.Total, err = query.QueryGrouped(s.db, &respGrouped.Grouped, scope); err != nil {
|
||||
logger.FromContext(c).WithError(err).Errorf("error finding tasks grouped")
|
||||
response.Error(c, response.ErrInternal, err)
|
||||
return
|
||||
}
|
||||
|
||||
response.Success(c, http.StatusOK, respGrouped)
|
||||
return
|
||||
}
|
||||
|
||||
if resp.Total, err = query.Query(s.db, &resp.Tasks, scope); err != nil {
|
||||
logger.FromContext(c).WithError(err).Errorf("error finding tasks")
|
||||
response.Error(c, response.ErrInternal, err)
|
||||
return
|
||||
}
|
||||
|
||||
for i := 0; i < len(resp.Tasks); i++ {
|
||||
if err = resp.Tasks[i].Valid(); err != nil {
|
||||
logger.FromContext(c).WithError(err).Errorf("error validating task data '%d'", resp.Tasks[i].ID)
|
||||
response.Error(c, response.ErrTasksInvalidData, err)
|
||||
return
|
||||
}
|
||||
}
|
||||
|
||||
response.Success(c, http.StatusOK, resp)
|
||||
}
|
||||
|
||||
// GetFlowTask is a function to return flow task by id
|
||||
// @Summary Retrieve flow task by id
|
||||
// @Tags Tasks
|
||||
// @Produce json
|
||||
// @Security BearerAuth
|
||||
// @Param flowID path int true "flow id" minimum(0)
|
||||
// @Param taskID path int true "task id" minimum(0)
|
||||
// @Success 200 {object} response.successResp{data=models.Task} "flow task received successful"
|
||||
// @Failure 403 {object} response.errorResp "getting flow task not permitted"
|
||||
// @Failure 404 {object} response.errorResp "flow task not found"
|
||||
// @Failure 500 {object} response.errorResp "internal error on getting flow task"
|
||||
// @Router /flows/{flowID}/tasks/{taskID} [get]
|
||||
func (s *TaskService) GetFlowTask(c *gin.Context) {
|
||||
var (
|
||||
err error
|
||||
flowID uint64
|
||||
taskID uint64
|
||||
resp models.Task
|
||||
)
|
||||
|
||||
if flowID, err = strconv.ParseUint(c.Param("flowID"), 10, 64); err != nil {
|
||||
logger.FromContext(c).WithError(err).Errorf("error parsing flow id")
|
||||
response.Error(c, response.ErrTasksInvalidRequest, err)
|
||||
return
|
||||
}
|
||||
|
||||
if taskID, err = strconv.ParseUint(c.Param("taskID"), 10, 64); err != nil {
|
||||
logger.FromContext(c).WithError(err).Errorf("error parsing task id")
|
||||
response.Error(c, response.ErrTasksInvalidRequest, err)
|
||||
return
|
||||
}
|
||||
|
||||
uid := c.GetUint64("uid")
|
||||
privs := c.GetStringSlice("prm")
|
||||
var scope func(db *gorm.DB) *gorm.DB
|
||||
if slices.Contains(privs, "tasks.admin") {
|
||||
scope = func(db *gorm.DB) *gorm.DB {
|
||||
return db.
|
||||
Joins("INNER JOIN flows f ON f.id = tasks.flow_id").
|
||||
Where("f.id = ?", flowID)
|
||||
}
|
||||
} else if slices.Contains(privs, "tasks.view") {
|
||||
scope = func(db *gorm.DB) *gorm.DB {
|
||||
return db.
|
||||
Joins("INNER JOIN flows f ON f.id = tasks.flow_id").
|
||||
Where("f.id = ? AND f.user_id = ?", flowID, uid)
|
||||
}
|
||||
} else {
|
||||
logger.FromContext(c).Errorf("error filtering user role permissions: permission not found")
|
||||
response.Error(c, response.ErrNotPermitted, nil)
|
||||
return
|
||||
}
|
||||
|
||||
err = s.db.Model(&resp).
|
||||
Scopes(scope).
|
||||
Where("tasks.id = ?", taskID).
|
||||
Take(&resp).Error
|
||||
if err != nil {
|
||||
logger.FromContext(c).WithError(err).Errorf("error on getting flow task by id")
|
||||
if gorm.IsRecordNotFoundError(err) {
|
||||
response.Error(c, response.ErrTasksNotFound, err)
|
||||
} else {
|
||||
response.Error(c, response.ErrInternal, err)
|
||||
}
|
||||
return
|
||||
}
|
||||
|
||||
response.Success(c, http.StatusOK, resp)
|
||||
}
|
||||
|
||||
// GetFlowTaskGraph is a function to return flow task graph by id
|
||||
// @Summary Retrieve flow task graph by id
|
||||
// @Tags Tasks
|
||||
// @Produce json
|
||||
// @Security BearerAuth
|
||||
// @Param flowID path int true "flow id" minimum(0)
|
||||
// @Param taskID path int true "task id" minimum(0)
|
||||
// @Success 200 {object} response.successResp{data=models.FlowTasksSubtasks} "flow task graph received successful"
|
||||
// @Failure 403 {object} response.errorResp "getting flow task graph not permitted"
|
||||
// @Failure 404 {object} response.errorResp "flow task graph not found"
|
||||
// @Failure 500 {object} response.errorResp "internal error on getting flow task graph"
|
||||
// @Router /flows/{flowID}/tasks/{taskID}/graph [get]
|
||||
func (s *TaskService) GetFlowTaskGraph(c *gin.Context) {
|
||||
var (
|
||||
err error
|
||||
flow models.Flow
|
||||
flowID uint64
|
||||
taskID uint64
|
||||
resp models.TaskSubtasks
|
||||
)
|
||||
|
||||
if flowID, err = strconv.ParseUint(c.Param("flowID"), 10, 64); err != nil {
|
||||
logger.FromContext(c).WithError(err).Errorf("error parsing flow id")
|
||||
response.Error(c, response.ErrTasksInvalidRequest, err)
|
||||
return
|
||||
}
|
||||
|
||||
if taskID, err = strconv.ParseUint(c.Param("taskID"), 10, 64); err != nil {
|
||||
logger.FromContext(c).WithError(err).Errorf("error parsing task id")
|
||||
response.Error(c, response.ErrTasksInvalidRequest, err)
|
||||
return
|
||||
}
|
||||
|
||||
uid := c.GetUint64("uid")
|
||||
privs := c.GetStringSlice("prm")
|
||||
var scope func(db *gorm.DB) *gorm.DB
|
||||
if slices.Contains(privs, "tasks.admin") {
|
||||
scope = func(db *gorm.DB) *gorm.DB {
|
||||
return db.Where("f.id = ?", flowID)
|
||||
}
|
||||
} else if slices.Contains(privs, "tasks.view") {
|
||||
scope = func(db *gorm.DB) *gorm.DB {
|
||||
return db.Where("f.id = ? AND f.user_id = ?", flowID, uid)
|
||||
}
|
||||
} else {
|
||||
logger.FromContext(c).Errorf("error filtering user role permissions: permission not found")
|
||||
response.Error(c, response.ErrNotPermitted, nil)
|
||||
return
|
||||
}
|
||||
|
||||
err = s.db.Model(&resp).
|
||||
Joins("INNER JOIN flows f ON f.id = tasks.flow_id").
|
||||
Scopes(scope).
|
||||
Where("tasks.id = ?", taskID).
|
||||
Take(&resp).Error
|
||||
if err != nil {
|
||||
logger.FromContext(c).WithError(err).Errorf("error on getting flow task by id")
|
||||
if gorm.IsRecordNotFoundError(err) {
|
||||
response.Error(c, response.ErrTasksNotFound, err)
|
||||
} else {
|
||||
response.Error(c, response.ErrInternal, err)
|
||||
}
|
||||
return
|
||||
}
|
||||
|
||||
err = s.db.Where("id = ?", resp.FlowID).Take(&flow).Error
|
||||
if err != nil {
|
||||
logger.FromContext(c).WithError(err).Errorf("error on getting flow by id")
|
||||
if gorm.IsRecordNotFoundError(err) {
|
||||
response.Error(c, response.ErrTasksNotFound, err)
|
||||
} else {
|
||||
response.Error(c, response.ErrInternal, err)
|
||||
}
|
||||
return
|
||||
}
|
||||
|
||||
isSubtasksAdmin := slices.Contains(privs, "subtasks.admin")
|
||||
isSubtasksView := slices.Contains(privs, "subtasks.view")
|
||||
if !(flow.UserID == uid && isSubtasksView) && !(flow.UserID != uid && isSubtasksAdmin) {
|
||||
response.Success(c, http.StatusOK, resp)
|
||||
return
|
||||
}
|
||||
|
||||
err = s.db.Model(&resp).Association("subtasks").Find(&resp.Subtasks).Error
|
||||
if err != nil {
|
||||
logger.FromContext(c).WithError(err).Errorf("error on getting task subtasks")
|
||||
response.Error(c, response.ErrInternal, err)
|
||||
return
|
||||
}
|
||||
|
||||
if err = resp.Valid(); err != nil {
|
||||
logger.FromContext(c).WithError(err).Errorf("error validating task data '%d'", taskID)
|
||||
response.Error(c, response.ErrTasksInvalidData, err)
|
||||
return
|
||||
}
|
||||
|
||||
response.Success(c, http.StatusOK, resp)
|
||||
}
|
||||
@@ -0,0 +1,219 @@
|
||||
package services
|
||||
|
||||
import (
|
||||
"errors"
|
||||
"net/http"
|
||||
"slices"
|
||||
"strconv"
|
||||
|
||||
"pentagi/pkg/server/logger"
|
||||
"pentagi/pkg/server/models"
|
||||
"pentagi/pkg/server/rdb"
|
||||
"pentagi/pkg/server/response"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/jinzhu/gorm"
|
||||
)
|
||||
|
||||
type termlogs struct {
|
||||
TermLogs []models.Termlog `json:"termlogs"`
|
||||
Total uint64 `json:"total"`
|
||||
}
|
||||
|
||||
type termlogsGrouped struct {
|
||||
Grouped []string `json:"grouped"`
|
||||
Total uint64 `json:"total"`
|
||||
}
|
||||
|
||||
var termlogsSQLMappers = map[string]any{
|
||||
"id": "{{table}}.id",
|
||||
"type": "{{table}}.type",
|
||||
"text": "{{table}}.text",
|
||||
"container_id": "{{table}}.container_id",
|
||||
"flow_id": "{{table}}.flow_id",
|
||||
"task_id": "{{table}}.task_id",
|
||||
"subtask_id": "{{table}}.subtask_id",
|
||||
"created_at": "{{table}}.created_at",
|
||||
"data": "({{table}}.type || ' ' || {{table}}.text)",
|
||||
}
|
||||
|
||||
type TermlogService struct {
|
||||
db *gorm.DB
|
||||
}
|
||||
|
||||
func NewTermlogService(db *gorm.DB) *TermlogService {
|
||||
return &TermlogService{
|
||||
db: db,
|
||||
}
|
||||
}
|
||||
|
||||
// GetTermlogs is a function to return termlogs list
|
||||
// @Summary Retrieve termlogs list
|
||||
// @Tags Termlogs
|
||||
// @Produce json
|
||||
// @Security BearerAuth
|
||||
// @Param request query rdb.TableQuery true "query table params"
|
||||
// @Success 200 {object} response.successResp{data=termlogs} "termlogs list received successful"
|
||||
// @Failure 400 {object} response.errorResp "invalid query request data"
|
||||
// @Failure 403 {object} response.errorResp "getting termlogs not permitted"
|
||||
// @Failure 500 {object} response.errorResp "internal error on getting termlogs"
|
||||
// @Router /termlogs/ [get]
|
||||
func (s *TermlogService) GetTermlogs(c *gin.Context) {
|
||||
var (
|
||||
err error
|
||||
query rdb.TableQuery
|
||||
resp termlogs
|
||||
)
|
||||
|
||||
if err = c.ShouldBindQuery(&query); err != nil {
|
||||
logger.FromContext(c).WithError(err).Errorf("error binding query")
|
||||
response.Error(c, response.ErrTermlogsInvalidRequest, err)
|
||||
return
|
||||
}
|
||||
|
||||
uid := c.GetUint64("uid")
|
||||
privs := c.GetStringSlice("prm")
|
||||
var scope func(db *gorm.DB) *gorm.DB
|
||||
if slices.Contains(privs, "termlogs.admin") {
|
||||
scope = func(db *gorm.DB) *gorm.DB {
|
||||
return db.
|
||||
Joins("INNER JOIN flows f ON f.id = flow_id")
|
||||
}
|
||||
} else if slices.Contains(privs, "termlogs.view") {
|
||||
scope = func(db *gorm.DB) *gorm.DB {
|
||||
return db.
|
||||
Joins("INNER JOIN flows f ON f.id = flow_id").
|
||||
Where("f.user_id = ?", uid)
|
||||
}
|
||||
} else {
|
||||
logger.FromContext(c).Errorf("error filtering user role permissions: permission not found")
|
||||
response.Error(c, response.ErrNotPermitted, nil)
|
||||
return
|
||||
}
|
||||
|
||||
query.Init("termlogs", termlogsSQLMappers)
|
||||
|
||||
if query.Group != "" {
|
||||
if _, ok := termlogsSQLMappers[query.Group]; !ok {
|
||||
logger.FromContext(c).Errorf("error finding termlogs grouped: group field not found")
|
||||
response.Error(c, response.ErrTermlogsInvalidRequest, errors.New("group field not found"))
|
||||
return
|
||||
}
|
||||
|
||||
var respGrouped termlogsGrouped
|
||||
if respGrouped.Total, err = query.QueryGrouped(s.db, &respGrouped.Grouped, scope); err != nil {
|
||||
logger.FromContext(c).WithError(err).Errorf("error finding termlogs grouped")
|
||||
response.Error(c, response.ErrInternal, err)
|
||||
return
|
||||
}
|
||||
|
||||
response.Success(c, http.StatusOK, respGrouped)
|
||||
return
|
||||
}
|
||||
|
||||
if resp.Total, err = query.Query(s.db, &resp.TermLogs, scope); err != nil {
|
||||
logger.FromContext(c).WithError(err).Errorf("error finding termlogs")
|
||||
response.Error(c, response.ErrInternal, err)
|
||||
return
|
||||
}
|
||||
|
||||
for i := 0; i < len(resp.TermLogs); i++ {
|
||||
if err = resp.TermLogs[i].Valid(); err != nil {
|
||||
logger.FromContext(c).WithError(err).Errorf("error validating termlog data '%d'", resp.TermLogs[i].ID)
|
||||
response.Error(c, response.ErrTermlogsInvalidData, err)
|
||||
return
|
||||
}
|
||||
}
|
||||
|
||||
response.Success(c, http.StatusOK, resp)
|
||||
}
|
||||
|
||||
// GetFlowTermlogs is a function to return termlogs list by flow id
|
||||
// @Summary Retrieve termlogs list by flow id
|
||||
// @Tags Termlogs
|
||||
// @Produce json
|
||||
// @Security BearerAuth
|
||||
// @Param flowID path int true "flow id" minimum(0)
|
||||
// @Param request query rdb.TableQuery true "query table params"
|
||||
// @Success 200 {object} response.successResp{data=termlogs} "termlogs list received successful"
|
||||
// @Failure 400 {object} response.errorResp "invalid query request data"
|
||||
// @Failure 403 {object} response.errorResp "getting termlogs not permitted"
|
||||
// @Failure 500 {object} response.errorResp "internal error on getting termlogs"
|
||||
// @Router /flows/{flowID}/termlogs/ [get]
|
||||
func (s *TermlogService) GetFlowTermlogs(c *gin.Context) {
|
||||
var (
|
||||
err error
|
||||
flowID uint64
|
||||
query rdb.TableQuery
|
||||
resp termlogs
|
||||
)
|
||||
|
||||
if flowID, err = strconv.ParseUint(c.Param("flowID"), 10, 64); err != nil {
|
||||
logger.FromContext(c).WithError(err).Errorf("error parsing flow id")
|
||||
response.Error(c, response.ErrTermlogsInvalidRequest, err)
|
||||
return
|
||||
}
|
||||
|
||||
if err = c.ShouldBindQuery(&query); err != nil {
|
||||
logger.FromContext(c).WithError(err).Errorf("error binding query")
|
||||
response.Error(c, response.ErrTermlogsInvalidRequest, err)
|
||||
return
|
||||
}
|
||||
|
||||
uid := c.GetUint64("uid")
|
||||
privs := c.GetStringSlice("prm")
|
||||
var scope func(db *gorm.DB) *gorm.DB
|
||||
if slices.Contains(privs, "termlogs.admin") {
|
||||
scope = func(db *gorm.DB) *gorm.DB {
|
||||
return db.
|
||||
Joins("INNER JOIN flows f ON f.id = flow_id").
|
||||
Where("f.id = ?", flowID)
|
||||
}
|
||||
} else if slices.Contains(privs, "termlogs.view") {
|
||||
scope = func(db *gorm.DB) *gorm.DB {
|
||||
return db.
|
||||
Joins("INNER JOIN flows f ON f.id = flow_id").
|
||||
Where("f.id = ? AND f.user_id = ?", flowID, uid)
|
||||
}
|
||||
} else {
|
||||
logger.FromContext(c).Errorf("error filtering user role permissions: permission not found")
|
||||
response.Error(c, response.ErrNotPermitted, nil)
|
||||
return
|
||||
}
|
||||
|
||||
query.Init("termlogs", termlogsSQLMappers)
|
||||
|
||||
if query.Group != "" {
|
||||
if _, ok := termlogsSQLMappers[query.Group]; !ok {
|
||||
logger.FromContext(c).Errorf("error finding termlogs grouped: group field not found")
|
||||
response.Error(c, response.ErrTermlogsInvalidRequest, errors.New("group field not found"))
|
||||
return
|
||||
}
|
||||
|
||||
var respGrouped termlogsGrouped
|
||||
if respGrouped.Total, err = query.QueryGrouped(s.db, &respGrouped.Grouped, scope); err != nil {
|
||||
logger.FromContext(c).WithError(err).Errorf("error finding termlogs grouped")
|
||||
response.Error(c, response.ErrInternal, err)
|
||||
return
|
||||
}
|
||||
|
||||
response.Success(c, http.StatusOK, respGrouped)
|
||||
return
|
||||
}
|
||||
|
||||
if resp.Total, err = query.Query(s.db, &resp.TermLogs, scope); err != nil {
|
||||
logger.FromContext(c).WithError(err).Errorf("error finding termlogs")
|
||||
response.Error(c, response.ErrInternal, err)
|
||||
return
|
||||
}
|
||||
|
||||
for i := 0; i < len(resp.TermLogs); i++ {
|
||||
if err = resp.TermLogs[i].Valid(); err != nil {
|
||||
logger.FromContext(c).WithError(err).Errorf("error validating termlog data '%d'", resp.TermLogs[i].ID)
|
||||
response.Error(c, response.ErrTermlogsInvalidData, err)
|
||||
return
|
||||
}
|
||||
}
|
||||
|
||||
response.Success(c, http.StatusOK, resp)
|
||||
}
|
||||
@@ -0,0 +1,290 @@
|
||||
package services
|
||||
|
||||
import (
|
||||
"errors"
|
||||
"net/http"
|
||||
"slices"
|
||||
"strconv"
|
||||
|
||||
"pentagi/pkg/server/logger"
|
||||
"pentagi/pkg/server/models"
|
||||
"pentagi/pkg/server/rdb"
|
||||
"pentagi/pkg/server/response"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/jinzhu/gorm"
|
||||
)
|
||||
|
||||
type toolcalls struct {
|
||||
Toolcalls []models.Toolcall `json:"toolcalls"`
|
||||
Total uint64 `json:"total"`
|
||||
}
|
||||
|
||||
type toolcallsGrouped struct {
|
||||
Grouped []string `json:"grouped"`
|
||||
Total uint64 `json:"total"`
|
||||
}
|
||||
|
||||
var toolcallsSQLMappers = map[string]any{
|
||||
"id": "{{table}}.id",
|
||||
"call_id": "{{table}}.call_id",
|
||||
"status": "{{table}}.status",
|
||||
"name": "{{table}}.name",
|
||||
"args": "{{table}}.args",
|
||||
"result": "{{table}}.result",
|
||||
"duration_seconds": "{{table}}.duration_seconds",
|
||||
"flow_id": "{{table}}.flow_id",
|
||||
"task_id": "{{table}}.task_id",
|
||||
"subtask_id": "{{table}}.subtask_id",
|
||||
"created_at": "{{table}}.created_at",
|
||||
"updated_at": "{{table}}.updated_at",
|
||||
"data": "({{table}}.name || ' ' || {{table}}.call_id || ' ' || {{table}}.args::text || ' ' || {{table}}.result)",
|
||||
}
|
||||
|
||||
type ToolcallService struct {
|
||||
db *gorm.DB
|
||||
}
|
||||
|
||||
func NewToolcallService(db *gorm.DB) *ToolcallService {
|
||||
return &ToolcallService{
|
||||
db: db,
|
||||
}
|
||||
}
|
||||
|
||||
// GetToolcalls is a function to return toolcalls list
|
||||
// @Summary Retrieve toolcalls list
|
||||
// @Tags Toolcalls
|
||||
// @Produce json
|
||||
// @Security BearerAuth
|
||||
// @Param request query rdb.TableQuery true "query table params"
|
||||
// @Success 200 {object} response.successResp{data=toolcalls} "toolcalls list received successful"
|
||||
// @Failure 400 {object} response.errorResp "invalid query request data"
|
||||
// @Failure 403 {object} response.errorResp "getting toolcalls not permitted"
|
||||
// @Failure 500 {object} response.errorResp "internal error on getting toolcalls"
|
||||
// @Router /toolcalls/ [get]
|
||||
func (s *ToolcallService) GetToolcalls(c *gin.Context) {
|
||||
var (
|
||||
err error
|
||||
query rdb.TableQuery
|
||||
resp toolcalls
|
||||
)
|
||||
|
||||
if err = c.ShouldBindQuery(&query); err != nil {
|
||||
logger.FromContext(c).WithError(err).Errorf("error binding query")
|
||||
response.Error(c, response.ErrToolcallsInvalidRequest, err)
|
||||
return
|
||||
}
|
||||
|
||||
uid := c.GetUint64("uid")
|
||||
privs := c.GetStringSlice("prm")
|
||||
var scope func(db *gorm.DB) *gorm.DB
|
||||
if slices.Contains(privs, "toolcalls.admin") {
|
||||
scope = func(db *gorm.DB) *gorm.DB {
|
||||
return db.
|
||||
Joins("INNER JOIN flows f ON f.id = toolcalls.flow_id")
|
||||
}
|
||||
} else if slices.Contains(privs, "toolcalls.view") {
|
||||
scope = func(db *gorm.DB) *gorm.DB {
|
||||
return db.
|
||||
Joins("INNER JOIN flows f ON f.id = toolcalls.flow_id").
|
||||
Where("f.user_id = ?", uid)
|
||||
}
|
||||
} else {
|
||||
logger.FromContext(c).Errorf("error filtering user role permissions: permission not found")
|
||||
response.Error(c, response.ErrNotPermitted, nil)
|
||||
return
|
||||
}
|
||||
|
||||
query.Init("toolcalls", toolcallsSQLMappers)
|
||||
|
||||
if query.Group != "" {
|
||||
if _, ok := toolcallsSQLMappers[query.Group]; !ok {
|
||||
logger.FromContext(c).Errorf("error finding toolcalls grouped: group field not found")
|
||||
response.Error(c, response.ErrToolcallsInvalidRequest, errors.New("group field not found"))
|
||||
return
|
||||
}
|
||||
|
||||
var respGrouped toolcallsGrouped
|
||||
if respGrouped.Total, err = query.QueryGrouped(s.db, &respGrouped.Grouped, scope); err != nil {
|
||||
logger.FromContext(c).WithError(err).Errorf("error finding toolcalls grouped")
|
||||
response.Error(c, response.ErrInternal, err)
|
||||
return
|
||||
}
|
||||
|
||||
response.Success(c, http.StatusOK, respGrouped)
|
||||
return
|
||||
}
|
||||
|
||||
if resp.Total, err = query.Query(s.db, &resp.Toolcalls, scope); err != nil {
|
||||
logger.FromContext(c).WithError(err).Errorf("error finding toolcalls")
|
||||
response.Error(c, response.ErrInternal, err)
|
||||
return
|
||||
}
|
||||
|
||||
for i := range resp.Toolcalls {
|
||||
if err = resp.Toolcalls[i].Valid(); err != nil {
|
||||
logger.FromContext(c).WithError(err).Errorf("error validating toolcall data '%d'", resp.Toolcalls[i].ID)
|
||||
response.Error(c, response.ErrToolcallsInvalidData, err)
|
||||
return
|
||||
}
|
||||
}
|
||||
|
||||
response.Success(c, http.StatusOK, resp)
|
||||
}
|
||||
|
||||
// GetFlowToolcalls is a function to return toolcalls list by flow id
|
||||
// @Summary Retrieve toolcalls list by flow id
|
||||
// @Tags Toolcalls
|
||||
// @Produce json
|
||||
// @Security BearerAuth
|
||||
// @Param flowID path int true "flow id" minimum(0)
|
||||
// @Param request query rdb.TableQuery true "query table params"
|
||||
// @Success 200 {object} response.successResp{data=toolcalls} "toolcalls list received successful"
|
||||
// @Failure 400 {object} response.errorResp "invalid query request data"
|
||||
// @Failure 403 {object} response.errorResp "getting toolcalls not permitted"
|
||||
// @Failure 500 {object} response.errorResp "internal error on getting toolcalls"
|
||||
// @Router /flows/{flowID}/toolcalls/ [get]
|
||||
func (s *ToolcallService) GetFlowToolcalls(c *gin.Context) {
|
||||
var (
|
||||
err error
|
||||
flowID uint64
|
||||
query rdb.TableQuery
|
||||
resp toolcalls
|
||||
)
|
||||
|
||||
if flowID, err = strconv.ParseUint(c.Param("flowID"), 10, 64); err != nil {
|
||||
logger.FromContext(c).WithError(err).Errorf("error parsing flow id")
|
||||
response.Error(c, response.ErrToolcallsInvalidRequest, err)
|
||||
return
|
||||
}
|
||||
|
||||
if err = c.ShouldBindQuery(&query); err != nil {
|
||||
logger.FromContext(c).WithError(err).Errorf("error binding query")
|
||||
response.Error(c, response.ErrToolcallsInvalidRequest, err)
|
||||
return
|
||||
}
|
||||
|
||||
uid := c.GetUint64("uid")
|
||||
privs := c.GetStringSlice("prm")
|
||||
var scope func(db *gorm.DB) *gorm.DB
|
||||
if slices.Contains(privs, "toolcalls.admin") {
|
||||
scope = func(db *gorm.DB) *gorm.DB {
|
||||
return db.
|
||||
Joins("INNER JOIN flows f ON f.id = toolcalls.flow_id").
|
||||
Where("f.id = ?", flowID)
|
||||
}
|
||||
} else if slices.Contains(privs, "toolcalls.view") {
|
||||
scope = func(db *gorm.DB) *gorm.DB {
|
||||
return db.
|
||||
Joins("INNER JOIN flows f ON f.id = toolcalls.flow_id").
|
||||
Where("f.id = ? AND f.user_id = ?", flowID, uid)
|
||||
}
|
||||
} else {
|
||||
logger.FromContext(c).Errorf("error filtering user role permissions: permission not found")
|
||||
response.Error(c, response.ErrNotPermitted, nil)
|
||||
return
|
||||
}
|
||||
|
||||
query.Init("toolcalls", toolcallsSQLMappers)
|
||||
|
||||
if query.Group != "" {
|
||||
if _, ok := toolcallsSQLMappers[query.Group]; !ok {
|
||||
logger.FromContext(c).Errorf("error finding toolcalls grouped: group field not found")
|
||||
response.Error(c, response.ErrToolcallsInvalidRequest, errors.New("group field not found"))
|
||||
return
|
||||
}
|
||||
|
||||
var respGrouped toolcallsGrouped
|
||||
if respGrouped.Total, err = query.QueryGrouped(s.db, &respGrouped.Grouped, scope); err != nil {
|
||||
logger.FromContext(c).WithError(err).Errorf("error finding toolcalls grouped")
|
||||
response.Error(c, response.ErrInternal, err)
|
||||
return
|
||||
}
|
||||
|
||||
response.Success(c, http.StatusOK, respGrouped)
|
||||
return
|
||||
}
|
||||
|
||||
if resp.Total, err = query.Query(s.db, &resp.Toolcalls, scope); err != nil {
|
||||
logger.FromContext(c).WithError(err).Errorf("error finding toolcalls")
|
||||
response.Error(c, response.ErrInternal, err)
|
||||
return
|
||||
}
|
||||
|
||||
for i := range resp.Toolcalls {
|
||||
if err = resp.Toolcalls[i].Valid(); err != nil {
|
||||
logger.FromContext(c).WithError(err).Errorf("error validating toolcall data '%d'", resp.Toolcalls[i].ID)
|
||||
response.Error(c, response.ErrToolcallsInvalidData, err)
|
||||
return
|
||||
}
|
||||
}
|
||||
|
||||
response.Success(c, http.StatusOK, resp)
|
||||
}
|
||||
|
||||
// GetFlowToolcall is a function to return toolcall info by id and flow id
|
||||
// @Summary Retrieve toolcall info by id and flow id
|
||||
// @Tags Toolcalls
|
||||
// @Produce json
|
||||
// @Security BearerAuth
|
||||
// @Param flowID path int true "flow id" minimum(0)
|
||||
// @Param toolcallID path int true "toolcall id" minimum(0)
|
||||
// @Success 200 {object} response.successResp{data=models.Toolcall} "toolcall info received successful"
|
||||
// @Failure 400 {object} response.errorResp "invalid request data"
|
||||
// @Failure 403 {object} response.errorResp "getting toolcall not permitted"
|
||||
// @Failure 404 {object} response.errorResp "toolcall not found"
|
||||
// @Failure 500 {object} response.errorResp "internal error on getting toolcall"
|
||||
// @Router /flows/{flowID}/toolcalls/{toolcallID} [get]
|
||||
func (s *ToolcallService) GetFlowToolcall(c *gin.Context) {
|
||||
var (
|
||||
err error
|
||||
flowID uint64
|
||||
toolcallID uint64
|
||||
resp models.Toolcall
|
||||
)
|
||||
|
||||
if flowID, err = strconv.ParseUint(c.Param("flowID"), 10, 64); err != nil {
|
||||
logger.FromContext(c).WithError(err).Errorf("error parsing flow id")
|
||||
response.Error(c, response.ErrToolcallsInvalidRequest, err)
|
||||
return
|
||||
}
|
||||
if toolcallID, err = strconv.ParseUint(c.Param("toolcallID"), 10, 64); err != nil {
|
||||
logger.FromContext(c).WithError(err).Errorf("error parsing toolcall id")
|
||||
response.Error(c, response.ErrToolcallsInvalidRequest, err)
|
||||
return
|
||||
}
|
||||
|
||||
uid := c.GetUint64("uid")
|
||||
privs := c.GetStringSlice("prm")
|
||||
var scope func(db *gorm.DB) *gorm.DB
|
||||
if slices.Contains(privs, "toolcalls.admin") {
|
||||
scope = func(db *gorm.DB) *gorm.DB {
|
||||
return db.Where("f.id = ?", flowID)
|
||||
}
|
||||
} else if slices.Contains(privs, "toolcalls.view") {
|
||||
scope = func(db *gorm.DB) *gorm.DB {
|
||||
return db.Where("f.id = ? AND f.user_id = ?", flowID, uid)
|
||||
}
|
||||
} else {
|
||||
logger.FromContext(c).Errorf("error filtering user role permissions: permission not found")
|
||||
response.Error(c, response.ErrNotPermitted, nil)
|
||||
return
|
||||
}
|
||||
|
||||
err = s.db.Model(&resp).
|
||||
Joins("INNER JOIN flows f ON f.id = toolcalls.flow_id").
|
||||
Scopes(scope).
|
||||
Where("toolcalls.id = ?", toolcallID).
|
||||
Take(&resp).Error
|
||||
if err != nil {
|
||||
logger.FromContext(c).WithError(err).Errorf("error on getting toolcall by id")
|
||||
if gorm.IsRecordNotFoundError(err) {
|
||||
response.Error(c, response.ErrToolcallsNotFound, err)
|
||||
} else {
|
||||
response.Error(c, response.ErrInternal, err)
|
||||
}
|
||||
return
|
||||
}
|
||||
|
||||
response.Success(c, http.StatusOK, resp)
|
||||
}
|
||||
@@ -0,0 +1,684 @@
|
||||
package services
|
||||
|
||||
import (
|
||||
"errors"
|
||||
"net/http"
|
||||
"slices"
|
||||
|
||||
"pentagi/pkg/server/auth"
|
||||
"pentagi/pkg/server/logger"
|
||||
"pentagi/pkg/server/models"
|
||||
"pentagi/pkg/server/rdb"
|
||||
"pentagi/pkg/server/response"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/jinzhu/gorm"
|
||||
"golang.org/x/crypto/bcrypt"
|
||||
)
|
||||
|
||||
type users struct {
|
||||
Users []models.UserRole `json:"users"`
|
||||
Total uint64 `json:"total"`
|
||||
}
|
||||
|
||||
type usersGrouped struct {
|
||||
Grouped []string `json:"grouped"`
|
||||
Total uint64 `json:"total"`
|
||||
}
|
||||
|
||||
var usersSQLMappers = map[string]any{
|
||||
"id": "{{table}}.id",
|
||||
"hash": "{{table}}.hash",
|
||||
"type": "{{table}}.type",
|
||||
"mail": "{{table}}.mail",
|
||||
"name": "{{table}}.name",
|
||||
"role_id": "{{table}}.role_id",
|
||||
"status": "{{table}}.status",
|
||||
"created_at": "{{table}}.created_at",
|
||||
"data": "({{table}}.hash || ' ' || {{table}}.mail || ' ' || {{table}}.name || ' ' || {{table}}.status)",
|
||||
}
|
||||
|
||||
type UserService struct {
|
||||
db *gorm.DB
|
||||
userCache *auth.UserCache
|
||||
}
|
||||
|
||||
func NewUserService(db *gorm.DB, userCache *auth.UserCache) *UserService {
|
||||
return &UserService{
|
||||
db: db,
|
||||
userCache: userCache,
|
||||
}
|
||||
}
|
||||
|
||||
// GetCurrentUser is a function to return account information
|
||||
// @Summary Retrieve current user information
|
||||
// @Tags Users
|
||||
// @Produce json
|
||||
// @Success 200 {object} response.successResp{data=models.UserRolePrivileges} "user info received successful"
|
||||
// @Failure 403 {object} response.errorResp "getting current user not permitted"
|
||||
// @Failure 404 {object} response.errorResp "current user not found"
|
||||
// @Failure 500 {object} response.errorResp "internal error on getting current user"
|
||||
// @Router /user/ [get]
|
||||
func (s *UserService) GetCurrentUser(c *gin.Context) {
|
||||
var (
|
||||
err error
|
||||
resp models.UserRolePrivileges
|
||||
)
|
||||
|
||||
uid := c.GetUint64("uid")
|
||||
|
||||
if err = s.db.Take(&resp.User, "id = ?", uid).Error; err != nil {
|
||||
logger.FromContext(c).WithError(err).Errorf("error finding current user")
|
||||
if errors.Is(err, gorm.ErrRecordNotFound) {
|
||||
response.Error(c, response.ErrUsersNotFound, err)
|
||||
} else {
|
||||
response.Error(c, response.ErrInternal, err)
|
||||
}
|
||||
return
|
||||
}
|
||||
|
||||
if err = s.db.Take(&resp.Role, "id = ?", resp.User.RoleID).Error; err != nil {
|
||||
logger.FromContext(c).WithError(err).Errorf("error finding role by role id")
|
||||
if errors.Is(err, gorm.ErrRecordNotFound) {
|
||||
response.Error(c, response.ErrGetUserModelsNotFound, err)
|
||||
} else {
|
||||
response.Error(c, response.ErrInternal, err)
|
||||
}
|
||||
return
|
||||
}
|
||||
|
||||
if err = s.db.Model(&resp.Role).Association("privileges").Find(&resp.Role.Privileges).Error; err != nil {
|
||||
logger.FromContext(c).WithError(err).Errorf("error finding privileges by role id")
|
||||
if errors.Is(err, gorm.ErrRecordNotFound) {
|
||||
response.Error(c, response.ErrGetUserModelsNotFound, err)
|
||||
} else {
|
||||
response.Error(c, response.ErrInternal, err)
|
||||
}
|
||||
return
|
||||
}
|
||||
|
||||
if err = resp.Valid(); err != nil {
|
||||
logger.FromContext(c).WithError(err).Errorf("error validating user data '%s'", resp.Hash)
|
||||
response.Error(c, response.ErrUsersInvalidData, err)
|
||||
return
|
||||
}
|
||||
|
||||
response.Success(c, http.StatusOK, resp)
|
||||
}
|
||||
|
||||
// ChangePasswordCurrentUser is a function to update account password
|
||||
// @Summary Update password for current user (account)
|
||||
// @Tags Users
|
||||
// @Accept json
|
||||
// @Produce json
|
||||
// @Param json body models.Password true "container to validate and update account password"
|
||||
// @Success 200 {object} response.successResp "account password updated successful"
|
||||
// @Failure 400 {object} response.errorResp "invalid account password form data"
|
||||
// @Failure 403 {object} response.errorResp "updating account password not permitted"
|
||||
// @Failure 404 {object} response.errorResp "current user not found"
|
||||
// @Failure 500 {object} response.errorResp "internal error on updating account password"
|
||||
// @Router /user/password [put]
|
||||
func (s *UserService) ChangePasswordCurrentUser(c *gin.Context) {
|
||||
var (
|
||||
encPass []byte
|
||||
err error
|
||||
form models.Password
|
||||
user models.UserPassword
|
||||
)
|
||||
|
||||
if err = c.ShouldBindJSON(&form); err != nil || form.Valid() != nil {
|
||||
if err == nil {
|
||||
err = form.Valid()
|
||||
}
|
||||
logger.FromContext(c).WithError(err).Errorf("error binding JSON")
|
||||
response.Error(c, response.ErrChangePasswordCurrentUserInvalidPassword, err)
|
||||
return
|
||||
}
|
||||
|
||||
uid := c.GetUint64("uid")
|
||||
scope := func(db *gorm.DB) *gorm.DB {
|
||||
return db.Where("id = ?", uid)
|
||||
}
|
||||
|
||||
if err = s.db.Scopes(scope).Take(&user).Error; err != nil {
|
||||
logger.FromContext(c).WithError(err).Errorf("error finding current user")
|
||||
if errors.Is(err, gorm.ErrRecordNotFound) {
|
||||
response.Error(c, response.ErrUsersNotFound, err)
|
||||
} else {
|
||||
response.Error(c, response.ErrInternal, err)
|
||||
}
|
||||
return
|
||||
} else if err = user.Valid(); err != nil {
|
||||
logger.FromContext(c).WithError(err).Errorf("error validating user data '%s'", user.Hash)
|
||||
response.Error(c, response.ErrUsersInvalidData, err)
|
||||
return
|
||||
}
|
||||
|
||||
if err = bcrypt.CompareHashAndPassword([]byte(user.Password), []byte(form.CurrentPassword)); err != nil {
|
||||
logger.FromContext(c).WithError(err).Errorf("error checking password for current user")
|
||||
response.Error(c, response.ErrChangePasswordCurrentUserInvalidCurrentPassword, err)
|
||||
return
|
||||
}
|
||||
|
||||
if encPass, err = rdb.EncryptPassword(form.Password); err != nil {
|
||||
logger.FromContext(c).WithError(err).Errorf("error making new password for current user")
|
||||
response.Error(c, response.ErrChangePasswordCurrentUserInvalidNewPassword, err)
|
||||
return
|
||||
}
|
||||
|
||||
// Use map to update fields to avoid GORM ignoring zero values (false for bool)
|
||||
updates := map[string]any{
|
||||
"password": string(encPass),
|
||||
"password_change_required": false,
|
||||
}
|
||||
|
||||
if err = s.db.Model(&user).Scopes(scope).Updates(updates).Error; err != nil {
|
||||
logger.FromContext(c).WithError(err).Errorf("error updating password for current user")
|
||||
response.Error(c, response.ErrInternal, err)
|
||||
return
|
||||
}
|
||||
|
||||
response.Success(c, http.StatusOK, struct{}{})
|
||||
}
|
||||
|
||||
// GetUsers returns users list
|
||||
// @Summary Retrieve users list by filters
|
||||
// @Tags Users
|
||||
// @Produce json
|
||||
// @Param request query rdb.TableQuery true "query table params"
|
||||
// @Success 200 {object} response.successResp{data=users} "users list received successful"
|
||||
// @Failure 400 {object} response.errorResp "invalid query request data"
|
||||
// @Failure 403 {object} response.errorResp "getting users not permitted"
|
||||
// @Failure 500 {object} response.errorResp "internal error on getting users"
|
||||
// @Router /users/ [get]
|
||||
func (s *UserService) GetUsers(c *gin.Context) {
|
||||
var (
|
||||
err error
|
||||
query rdb.TableQuery
|
||||
resp users
|
||||
rids []uint64
|
||||
roles []models.Role
|
||||
)
|
||||
|
||||
if err = c.ShouldBindQuery(&query); err != nil {
|
||||
logger.FromContext(c).WithError(err).Errorf("error binding query")
|
||||
response.Error(c, response.ErrUsersInvalidRequest, err)
|
||||
return
|
||||
}
|
||||
|
||||
uid := c.GetUint64("uid")
|
||||
privs := c.GetStringSlice("prm")
|
||||
scope := func(db *gorm.DB) *gorm.DB {
|
||||
if !slices.Contains(privs, "users.view") {
|
||||
return db.Where("id = ?", uid)
|
||||
}
|
||||
return db
|
||||
}
|
||||
|
||||
query.Init("users", usersSQLMappers)
|
||||
|
||||
if query.Group != "" {
|
||||
if _, ok := usersSQLMappers[query.Group]; !ok {
|
||||
logger.FromContext(c).Errorf("error finding users grouped: group field not found")
|
||||
response.Error(c, response.ErrUsersInvalidRequest, errors.New("group field not found"))
|
||||
return
|
||||
}
|
||||
|
||||
var respGrouped usersGrouped
|
||||
if respGrouped.Total, err = query.QueryGrouped(s.db, &respGrouped.Grouped, scope); err != nil {
|
||||
logger.FromContext(c).WithError(err).Errorf("error finding users grouped")
|
||||
response.Error(c, response.ErrInternal, err)
|
||||
return
|
||||
}
|
||||
|
||||
response.Success(c, http.StatusOK, respGrouped)
|
||||
return
|
||||
}
|
||||
|
||||
if resp.Total, err = query.Query(s.db, &resp.Users, scope); err != nil {
|
||||
logger.FromContext(c).WithError(err).Errorf("error finding users")
|
||||
response.Error(c, response.ErrInternal, err)
|
||||
return
|
||||
}
|
||||
|
||||
for _, user := range resp.Users {
|
||||
rids = append(rids, user.RoleID)
|
||||
}
|
||||
|
||||
if err = s.db.Find(&roles, "id IN (?)", rids).Error; err != nil {
|
||||
logger.FromContext(c).WithError(err).Errorf("error finding linked roles")
|
||||
response.Error(c, response.ErrInternal, err)
|
||||
return
|
||||
}
|
||||
|
||||
for i := range resp.Users {
|
||||
roleID := resp.Users[i].RoleID
|
||||
for _, role := range roles {
|
||||
if roleID == role.ID {
|
||||
resp.Users[i].Role = role
|
||||
break
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
for i := 0; i < len(resp.Users); i++ {
|
||||
if err = resp.Users[i].Valid(); err != nil {
|
||||
logger.FromContext(c).WithError(err).Errorf("error validating user data '%s'", resp.Users[i].Hash)
|
||||
response.Error(c, response.ErrUsersInvalidData, err)
|
||||
return
|
||||
}
|
||||
}
|
||||
|
||||
response.Success(c, http.StatusOK, resp)
|
||||
}
|
||||
|
||||
// GetUser is a function to return user by hash
|
||||
// @Summary Retrieve user by hash
|
||||
// @Tags Users
|
||||
// @Produce json
|
||||
// @Param hash path string true "hash in hex format (md5)" minlength(32) maxlength(32)
|
||||
// @Success 200 {object} response.successResp{data=models.UserRolePrivileges} "user received successful"
|
||||
// @Failure 403 {object} response.errorResp "getting user not permitted"
|
||||
// @Failure 404 {object} response.errorResp "user not found"
|
||||
// @Failure 500 {object} response.errorResp "internal error on getting user"
|
||||
// @Router /users/{hash} [get]
|
||||
func (s *UserService) GetUser(c *gin.Context) {
|
||||
var (
|
||||
err error
|
||||
hash string = c.Param("hash")
|
||||
resp models.UserRolePrivileges
|
||||
)
|
||||
|
||||
uhash := c.GetString("uhash")
|
||||
privs := c.GetStringSlice("prm")
|
||||
if !slices.Contains(privs, "users.view") && uhash != hash {
|
||||
logger.FromContext(c).Errorf("error filtering user role permissions: permission not found")
|
||||
response.Error(c, response.ErrNotPermitted, nil)
|
||||
return
|
||||
}
|
||||
|
||||
if err = s.db.Take(&resp.User, "hash = ?", hash).Error; err != nil {
|
||||
logger.FromContext(c).WithError(err).Errorf("error finding user by hash")
|
||||
if errors.Is(err, gorm.ErrRecordNotFound) {
|
||||
response.Error(c, response.ErrUsersNotFound, err)
|
||||
} else {
|
||||
response.Error(c, response.ErrInternal, err)
|
||||
}
|
||||
return
|
||||
}
|
||||
|
||||
if err = s.db.Take(&resp.Role, "id = ?", resp.User.RoleID).Error; err != nil {
|
||||
logger.FromContext(c).WithError(err).Errorf("error finding role by role id")
|
||||
if errors.Is(err, gorm.ErrRecordNotFound) {
|
||||
response.Error(c, response.ErrGetUserModelsNotFound, err)
|
||||
} else {
|
||||
response.Error(c, response.ErrInternal, err)
|
||||
}
|
||||
return
|
||||
}
|
||||
|
||||
if err = s.db.Model(&resp.Role).Association("privileges").Find(&resp.Role.Privileges).Error; err != nil {
|
||||
logger.FromContext(c).WithError(err).Errorf("error finding privileges by role id")
|
||||
if errors.Is(err, gorm.ErrRecordNotFound) {
|
||||
response.Error(c, response.ErrGetUserModelsNotFound, err)
|
||||
} else {
|
||||
response.Error(c, response.ErrInternal, err)
|
||||
}
|
||||
return
|
||||
}
|
||||
if err = resp.Valid(); err != nil {
|
||||
logger.FromContext(c).WithError(err).Errorf("error validating user data '%s'", resp.Hash)
|
||||
response.Error(c, response.ErrUsersInvalidData, err)
|
||||
return
|
||||
}
|
||||
|
||||
response.Success(c, http.StatusOK, resp)
|
||||
}
|
||||
|
||||
// CreateUser is a function to create new user
|
||||
// @Summary Create new user
|
||||
// @Tags Users
|
||||
// @Accept json
|
||||
// @Produce json
|
||||
// @Param json body models.UserPassword true "user model to create from"
|
||||
// @Success 201 {object} response.successResp{data=models.UserRole} "user created successful"
|
||||
// @Failure 400 {object} response.errorResp "invalid user request data"
|
||||
// @Failure 403 {object} response.errorResp "creating user not permitted"
|
||||
// @Failure 500 {object} response.errorResp "internal error on creating user"
|
||||
// @Router /users/ [post]
|
||||
func (s *UserService) CreateUser(c *gin.Context) {
|
||||
var (
|
||||
encPassword []byte
|
||||
err error
|
||||
resp models.UserRole
|
||||
user models.UserPassword
|
||||
)
|
||||
|
||||
if err = c.ShouldBindJSON(&user); err != nil {
|
||||
logger.FromContext(c).WithError(err).Errorf("error binding JSON")
|
||||
response.Error(c, response.ErrUsersInvalidRequest, err)
|
||||
return
|
||||
}
|
||||
|
||||
rid := c.GetUint64("rid")
|
||||
privs := c.GetStringSlice("prm")
|
||||
if !slices.Contains(privs, "users.create") {
|
||||
logger.FromContext(c).Errorf("error filtering user role permissions: permission not found")
|
||||
response.Error(c, response.ErrNotPermitted, nil)
|
||||
return
|
||||
}
|
||||
|
||||
privsCurrentUser, err := s.GetUserPrivileges(c, rid)
|
||||
if err != nil {
|
||||
logger.FromContext(c).WithError(err).Errorf("error getting current user privileges")
|
||||
response.Error(c, response.ErrInternal, err)
|
||||
return
|
||||
}
|
||||
|
||||
privsNewUser, err := s.GetUserPrivileges(c, user.RoleID)
|
||||
if err != nil {
|
||||
logger.FromContext(c).WithError(err).Errorf("error getting new user privileges")
|
||||
response.Error(c, response.ErrInternal, err)
|
||||
return
|
||||
}
|
||||
|
||||
if !s.CheckPrivilege(c, privsCurrentUser, privsNewUser) {
|
||||
logger.FromContext(c).Errorf("error checking new user privileges")
|
||||
response.Error(c, response.ErrNotPermitted, nil)
|
||||
return
|
||||
}
|
||||
|
||||
user.ID = 0
|
||||
user.Hash = rdb.MakeUserHash(user.Name)
|
||||
if err = user.Valid(); err != nil {
|
||||
logger.FromContext(c).WithError(err).Errorf("error validating user")
|
||||
response.Error(c, response.ErrCreateUserInvalidUser, err)
|
||||
return
|
||||
}
|
||||
|
||||
if encPassword, err = rdb.EncryptPassword(user.Password); err != nil {
|
||||
logger.FromContext(c).WithError(err).Errorf("error encoding password")
|
||||
response.Error(c, response.ErrInternal, err)
|
||||
return
|
||||
} else {
|
||||
user.Password = string(encPassword)
|
||||
}
|
||||
|
||||
tx := s.db.Begin()
|
||||
if tx.Error != nil {
|
||||
logger.FromContext(c).WithError(tx.Error).Errorf("error starting transaction")
|
||||
response.Error(c, response.ErrInternal, tx.Error)
|
||||
return
|
||||
}
|
||||
|
||||
if err = tx.Create(&user).Error; err != nil {
|
||||
tx.Rollback()
|
||||
logger.FromContext(c).WithError(err).Errorf("error creating user")
|
||||
response.Error(c, response.ErrInternal, err)
|
||||
return
|
||||
}
|
||||
|
||||
preferences := models.NewUserPreferences(user.ID)
|
||||
if err = tx.Create(preferences).Error; err != nil {
|
||||
tx.Rollback()
|
||||
logger.FromContext(c).WithError(err).Errorf("error creating user preferences")
|
||||
response.Error(c, response.ErrInternal, err)
|
||||
return
|
||||
}
|
||||
|
||||
if err = tx.Commit().Error; err != nil {
|
||||
logger.FromContext(c).WithError(err).Errorf("error committing transaction")
|
||||
response.Error(c, response.ErrInternal, err)
|
||||
return
|
||||
}
|
||||
|
||||
if err = s.db.Take(&resp.User, "hash = ?", user.Hash).Error; err != nil {
|
||||
logger.FromContext(c).WithError(err).Errorf("error finding user by hash")
|
||||
response.Error(c, response.ErrInternal, err)
|
||||
return
|
||||
}
|
||||
|
||||
if err = s.db.Take(&resp.Role, "id = ?", resp.User.RoleID).Error; err != nil {
|
||||
logger.FromContext(c).WithError(err).Errorf("error finding role by role id")
|
||||
response.Error(c, response.ErrInternal, err)
|
||||
return
|
||||
}
|
||||
|
||||
if err = resp.Valid(); err != nil {
|
||||
logger.FromContext(c).WithError(err).Errorf("error validating user data '%s'", resp.Hash)
|
||||
response.Error(c, response.ErrUsersInvalidData, err)
|
||||
return
|
||||
}
|
||||
|
||||
s.userCache.Invalidate(resp.User.ID)
|
||||
|
||||
response.Success(c, http.StatusCreated, resp)
|
||||
}
|
||||
|
||||
// PatchUser is a function to update user by hash
|
||||
// @Summary Update user
|
||||
// @Tags Users
|
||||
// @Accept json
|
||||
// @Produce json
|
||||
// @Param hash path string true "user hash in hex format (md5)" minlength(32) maxlength(32)
|
||||
// @Param json body models.UserPassword true "user model to update"
|
||||
// @Success 200 {object} response.successResp{data=models.UserRole} "user updated successful"
|
||||
// @Failure 400 {object} response.errorResp "invalid user request data"
|
||||
// @Failure 403 {object} response.errorResp "updating user not permitted"
|
||||
// @Failure 404 {object} response.errorResp "user not found"
|
||||
// @Failure 500 {object} response.errorResp "internal error on updating user"
|
||||
// @Router /users/{hash} [put]
|
||||
func (s *UserService) PatchUser(c *gin.Context) {
|
||||
var (
|
||||
err error
|
||||
hash = c.Param("hash")
|
||||
resp models.UserRole
|
||||
user models.UserPassword
|
||||
)
|
||||
|
||||
if err = c.ShouldBindJSON(&user); err != nil {
|
||||
logger.FromContext(c).WithError(err).Errorf("error binding JSON")
|
||||
response.Error(c, response.ErrUsersInvalidRequest, err)
|
||||
return
|
||||
} else if hash != user.Hash {
|
||||
logger.FromContext(c).Errorf("mismatch user hash to requested one")
|
||||
response.Error(c, response.ErrUsersInvalidRequest, nil)
|
||||
return
|
||||
} else if err = user.User.Valid(); err != nil {
|
||||
logger.FromContext(c).WithError(err).Errorf("error validating user JSON")
|
||||
response.Error(c, response.ErrUsersInvalidRequest, err)
|
||||
return
|
||||
} else if err = user.Valid(); user.Password != "" && err != nil {
|
||||
logger.FromContext(c).WithError(err).Errorf("error validating user password")
|
||||
response.Error(c, response.ErrUsersInvalidRequest, err)
|
||||
return
|
||||
}
|
||||
|
||||
uid := c.GetUint64("uid")
|
||||
uhash := c.GetString("uhash")
|
||||
privs := c.GetStringSlice("prm")
|
||||
scope := func(db *gorm.DB) *gorm.DB {
|
||||
if slices.Contains(privs, "users.edit") {
|
||||
return db.Where("hash = ?", hash)
|
||||
} else {
|
||||
return db.Where("hash = ? AND id = ?", hash, uid)
|
||||
}
|
||||
}
|
||||
if !slices.Contains(privs, "users.edit") && uhash != hash {
|
||||
logger.FromContext(c).Errorf("error filtering user role permissions: permission not found")
|
||||
response.Error(c, response.ErrNotPermitted, nil)
|
||||
return
|
||||
}
|
||||
|
||||
// Check if user exists before updating
|
||||
var existingUser models.User
|
||||
if err = s.db.Scopes(scope).Take(&existingUser).Error; err != nil {
|
||||
logger.FromContext(c).WithError(err).Errorf("error finding user by hash")
|
||||
if errors.Is(err, gorm.ErrRecordNotFound) {
|
||||
response.Error(c, response.ErrUsersNotFound, err)
|
||||
} else {
|
||||
response.Error(c, response.ErrInternal, err)
|
||||
}
|
||||
return
|
||||
}
|
||||
|
||||
if user.Password != "" {
|
||||
var encPassword []byte
|
||||
encPassword, err = rdb.EncryptPassword(user.Password)
|
||||
if err != nil {
|
||||
logger.FromContext(c).WithError(err).Errorf("error encoding password")
|
||||
response.Error(c, response.ErrInternal, err)
|
||||
return
|
||||
}
|
||||
// Use map to update fields to avoid GORM ignoring zero values (false for bool)
|
||||
updates := map[string]any{
|
||||
"name": user.Name,
|
||||
"status": user.Status,
|
||||
"password": string(encPassword),
|
||||
"password_change_required": false,
|
||||
}
|
||||
err = s.db.Model(&existingUser).Updates(updates).Error
|
||||
} else {
|
||||
updates := map[string]any{
|
||||
"name": user.Name,
|
||||
"status": user.Status,
|
||||
}
|
||||
err = s.db.Model(&existingUser).Updates(updates).Error
|
||||
}
|
||||
|
||||
if err != nil {
|
||||
logger.FromContext(c).WithError(err).Errorf("error updating user by hash '%s'", hash)
|
||||
response.Error(c, response.ErrInternal, err)
|
||||
return
|
||||
}
|
||||
|
||||
if err = s.db.Scopes(scope).Take(&resp.User).Error; err != nil {
|
||||
logger.FromContext(c).WithError(err).Errorf("error finding user by hash")
|
||||
if errors.Is(err, gorm.ErrRecordNotFound) {
|
||||
response.Error(c, response.ErrUsersNotFound, err)
|
||||
} else {
|
||||
response.Error(c, response.ErrInternal, err)
|
||||
}
|
||||
return
|
||||
}
|
||||
|
||||
if err = s.db.Take(&resp.Role, "id = ?", resp.User.RoleID).Error; err != nil {
|
||||
logger.FromContext(c).WithError(err).Errorf("error finding role by role id")
|
||||
if errors.Is(err, gorm.ErrRecordNotFound) {
|
||||
response.Error(c, response.ErrPatchUserModelsNotFound, err)
|
||||
} else {
|
||||
response.Error(c, response.ErrInternal, err)
|
||||
}
|
||||
return
|
||||
}
|
||||
if err = resp.Valid(); err != nil {
|
||||
logger.FromContext(c).WithError(err).Errorf("error validating user data '%s'", resp.Hash)
|
||||
response.Error(c, response.ErrInternal, err)
|
||||
return
|
||||
}
|
||||
|
||||
s.userCache.Invalidate(resp.User.ID)
|
||||
|
||||
response.Success(c, http.StatusOK, resp)
|
||||
}
|
||||
|
||||
// DeleteUser is a function to delete user by hash
|
||||
// @Summary Delete user by hash
|
||||
// @Tags Users
|
||||
// @Produce json
|
||||
// @Param hash path string true "hash in hex format (md5)" minlength(32) maxlength(32)
|
||||
// @Success 200 {object} response.successResp "user deleted successful"
|
||||
// @Failure 403 {object} response.errorResp "deleting user not permitted"
|
||||
// @Failure 404 {object} response.errorResp "user not found"
|
||||
// @Failure 500 {object} response.errorResp "internal error on deleting user"
|
||||
// @Router /users/{hash} [delete]
|
||||
func (s *UserService) DeleteUser(c *gin.Context) {
|
||||
var (
|
||||
err error
|
||||
hash string = c.Param("hash")
|
||||
user models.UserRole
|
||||
)
|
||||
|
||||
uid := c.GetUint64("uid")
|
||||
uhash := c.GetString("uhash")
|
||||
privs := c.GetStringSlice("prm")
|
||||
scope := func(db *gorm.DB) *gorm.DB {
|
||||
if slices.Contains(privs, "users.delete") {
|
||||
return db.Where("hash = ?", hash)
|
||||
} else {
|
||||
return db.Where("hash = ? AND id = ?", hash, uid)
|
||||
}
|
||||
}
|
||||
if !slices.Contains(privs, "users.delete") && uhash != hash {
|
||||
logger.FromContext(c).Errorf("error filtering user role permissions: permission not found")
|
||||
response.Error(c, response.ErrNotPermitted, nil)
|
||||
return
|
||||
}
|
||||
|
||||
if err = s.db.Scopes(scope).Take(&user.User).Error; err != nil {
|
||||
logger.FromContext(c).WithError(err).Errorf("error finding user by hash")
|
||||
if errors.Is(err, gorm.ErrRecordNotFound) {
|
||||
response.Error(c, response.ErrUsersNotFound, err)
|
||||
} else {
|
||||
response.Error(c, response.ErrInternal, err)
|
||||
}
|
||||
return
|
||||
}
|
||||
|
||||
if err = s.db.Take(&user.Role, "id = ?", user.User.RoleID).Error; err != nil {
|
||||
logger.FromContext(c).WithError(err).Errorf("error finding role by role id")
|
||||
if errors.Is(err, gorm.ErrRecordNotFound) {
|
||||
response.Error(c, response.ErrDeleteUserModelsNotFound, err)
|
||||
} else {
|
||||
response.Error(c, response.ErrInternal, err)
|
||||
}
|
||||
return
|
||||
}
|
||||
if err = user.Valid(); err != nil {
|
||||
logger.FromContext(c).WithError(err).Errorf("error validating user data '%s'", user.Hash)
|
||||
response.Error(c, response.ErrUsersInvalidData, err)
|
||||
return
|
||||
}
|
||||
|
||||
if err = s.db.Delete(&user.User).Error; err != nil {
|
||||
logger.FromContext(c).WithError(err).Errorf("error deleting user by hash '%s'", hash)
|
||||
response.Error(c, response.ErrInternal, err)
|
||||
return
|
||||
}
|
||||
|
||||
s.userCache.Invalidate(user.ID)
|
||||
|
||||
response.Success(c, http.StatusOK, struct{}{})
|
||||
}
|
||||
|
||||
// GetUserPrivileges is a function to return user privileges
|
||||
func (s *UserService) GetUserPrivileges(c *gin.Context, rid uint64) ([]string, error) {
|
||||
var (
|
||||
err error
|
||||
privs []string
|
||||
resp []models.Privilege
|
||||
)
|
||||
|
||||
if err = s.db.Model(&models.Privilege{}).Where("role_id = ?", rid).Find(&resp).Error; err != nil {
|
||||
return nil, err
|
||||
}
|
||||
for _, p := range resp {
|
||||
if err = p.Valid(); err != nil {
|
||||
logger.FromContext(c).WithError(err).Errorf("error validating privilege data '%s'", p.Name)
|
||||
return nil, err
|
||||
}
|
||||
privs = append(privs, p.Name)
|
||||
}
|
||||
|
||||
return privs, nil
|
||||
}
|
||||
|
||||
// CheckPrivilege is a function to check if user has privilege
|
||||
func (s *UserService) CheckPrivilege(c *gin.Context, privsCurrentUser, privsNewUser []string) bool {
|
||||
for _, priv := range privsNewUser {
|
||||
if !slices.Contains(privsCurrentUser, priv) {
|
||||
return false
|
||||
}
|
||||
}
|
||||
return true
|
||||
}
|
||||
@@ -0,0 +1,354 @@
|
||||
package services
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"encoding/json"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"testing"
|
||||
|
||||
"pentagi/pkg/server/auth"
|
||||
"pentagi/pkg/server/models"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/jinzhu/gorm"
|
||||
_ "github.com/jinzhu/gorm/dialects/sqlite"
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
func TestCreateUser_CreatesUserPreferences(t *testing.T) {
|
||||
db := setupTestDB(t)
|
||||
defer db.Close()
|
||||
|
||||
userCache := auth.NewUserCache(db)
|
||||
service := NewUserService(db, userCache)
|
||||
|
||||
gin.SetMode(gin.TestMode)
|
||||
w := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(w)
|
||||
|
||||
// Set up context with admin permissions
|
||||
c.Set("uid", uint64(1))
|
||||
c.Set("rid", uint64(1))
|
||||
c.Set("uhash", "testhash1")
|
||||
c.Set("prm", []string{"users.create"})
|
||||
|
||||
// Create request body
|
||||
userRequest := models.UserPassword{
|
||||
User: models.User{
|
||||
Mail: "newuser@test.com",
|
||||
Name: "New User",
|
||||
RoleID: 2,
|
||||
Status: models.UserStatusActive,
|
||||
Type: models.UserTypeLocal,
|
||||
},
|
||||
Password: "SecurePass123!",
|
||||
}
|
||||
|
||||
body, err := json.Marshal(userRequest)
|
||||
require.NoError(t, err)
|
||||
|
||||
c.Request, _ = http.NewRequest("POST", "/users/", bytes.NewBuffer(body))
|
||||
c.Request.Header.Set("Content-Type", "application/json")
|
||||
|
||||
// Call the handler
|
||||
service.CreateUser(c)
|
||||
|
||||
// Check response status
|
||||
assert.Equal(t, http.StatusCreated, w.Code, "Expected HTTP 201 Created")
|
||||
|
||||
// Verify user was created
|
||||
var createdUser models.User
|
||||
err = db.Where("mail = ?", "newuser@test.com").First(&createdUser).Error
|
||||
require.NoError(t, err, "User should be created in database")
|
||||
assert.Equal(t, "New User", createdUser.Name)
|
||||
assert.Equal(t, uint64(2), createdUser.RoleID)
|
||||
|
||||
// Verify user_preferences was created
|
||||
var userPrefs models.UserPreferences
|
||||
err = db.Where("user_id = ?", createdUser.ID).First(&userPrefs).Error
|
||||
require.NoError(t, err, "User preferences should be created in database")
|
||||
assert.Equal(t, createdUser.ID, userPrefs.UserID)
|
||||
assert.NotNil(t, userPrefs.Preferences.FavoriteFlows)
|
||||
assert.Equal(t, 0, len(userPrefs.Preferences.FavoriteFlows), "FavoriteFlows should be empty array")
|
||||
}
|
||||
|
||||
func TestCreateUser_RollbackOnPreferencesError(t *testing.T) {
|
||||
db := setupTestDB(t)
|
||||
defer db.Close()
|
||||
|
||||
// Drop user_preferences table to simulate error
|
||||
db.Exec("DROP TABLE user_preferences")
|
||||
|
||||
userCache := auth.NewUserCache(db)
|
||||
service := NewUserService(db, userCache)
|
||||
|
||||
gin.SetMode(gin.TestMode)
|
||||
w := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(w)
|
||||
|
||||
c.Set("uid", uint64(1))
|
||||
c.Set("rid", uint64(1))
|
||||
c.Set("uhash", "testhash1")
|
||||
c.Set("prm", []string{"users.create"})
|
||||
|
||||
userRequest := models.UserPassword{
|
||||
User: models.User{
|
||||
Mail: "failuser@test.com",
|
||||
Name: "Fail User",
|
||||
RoleID: 2,
|
||||
Status: models.UserStatusActive,
|
||||
Type: models.UserTypeLocal,
|
||||
},
|
||||
Password: "SecurePass123!",
|
||||
}
|
||||
|
||||
body, err := json.Marshal(userRequest)
|
||||
require.NoError(t, err)
|
||||
|
||||
c.Request, _ = http.NewRequest("POST", "/users/", bytes.NewBuffer(body))
|
||||
c.Request.Header.Set("Content-Type", "application/json")
|
||||
|
||||
service.CreateUser(c)
|
||||
|
||||
// Should return error
|
||||
assert.Equal(t, http.StatusInternalServerError, w.Code, "Expected HTTP 500 on preferences creation error")
|
||||
|
||||
// Verify user was NOT created (transaction rolled back)
|
||||
var user models.User
|
||||
err = db.Where("mail = ?", "failuser@test.com").First(&user).Error
|
||||
assert.Error(t, err, "User should not exist due to transaction rollback")
|
||||
assert.Equal(t, gorm.ErrRecordNotFound, err)
|
||||
}
|
||||
|
||||
func TestCreateUser_InvalidPermissions(t *testing.T) {
|
||||
db := setupTestDB(t)
|
||||
defer db.Close()
|
||||
|
||||
userCache := auth.NewUserCache(db)
|
||||
service := NewUserService(db, userCache)
|
||||
|
||||
gin.SetMode(gin.TestMode)
|
||||
w := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(w)
|
||||
|
||||
// Set up context WITHOUT users.create permission
|
||||
c.Set("uid", uint64(2))
|
||||
c.Set("rid", uint64(2))
|
||||
c.Set("uhash", "testhash2")
|
||||
c.Set("prm", []string{"flows.view"})
|
||||
|
||||
userRequest := models.UserPassword{
|
||||
User: models.User{
|
||||
Mail: "unauthorized@test.com",
|
||||
Name: "Unauthorized User",
|
||||
RoleID: 2,
|
||||
Status: models.UserStatusActive,
|
||||
Type: models.UserTypeLocal,
|
||||
},
|
||||
Password: "SecurePass123!",
|
||||
}
|
||||
|
||||
body, err := json.Marshal(userRequest)
|
||||
require.NoError(t, err)
|
||||
|
||||
c.Request, _ = http.NewRequest("POST", "/users/", bytes.NewBuffer(body))
|
||||
c.Request.Header.Set("Content-Type", "application/json")
|
||||
|
||||
service.CreateUser(c)
|
||||
|
||||
// Should return forbidden
|
||||
assert.Equal(t, http.StatusForbidden, w.Code, "Expected HTTP 403 Forbidden")
|
||||
|
||||
// Verify user was NOT created
|
||||
var user models.User
|
||||
err = db.Where("mail = ?", "unauthorized@test.com").First(&user).Error
|
||||
assert.Error(t, err, "User should not be created")
|
||||
}
|
||||
|
||||
func TestCreateUser_MultipleUsers(t *testing.T) {
|
||||
db := setupTestDB(t)
|
||||
defer db.Close()
|
||||
|
||||
userCache := auth.NewUserCache(db)
|
||||
service := NewUserService(db, userCache)
|
||||
|
||||
testCases := []struct {
|
||||
name string
|
||||
mail string
|
||||
username string
|
||||
roleID uint64
|
||||
}{
|
||||
{
|
||||
name: "create first user",
|
||||
mail: "newuser1@test.com",
|
||||
username: "User One",
|
||||
roleID: 2,
|
||||
},
|
||||
{
|
||||
name: "create second user",
|
||||
mail: "newuser2@test.com",
|
||||
username: "User Two",
|
||||
roleID: 2,
|
||||
},
|
||||
{
|
||||
name: "create third user",
|
||||
mail: "newuser3@test.com",
|
||||
username: "User Three",
|
||||
roleID: 2,
|
||||
},
|
||||
}
|
||||
|
||||
for _, tc := range testCases {
|
||||
t.Run(tc.name, func(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
w := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(w)
|
||||
|
||||
c.Set("uid", uint64(1))
|
||||
c.Set("rid", uint64(1))
|
||||
c.Set("uhash", "testhash1")
|
||||
c.Set("prm", []string{"users.create"})
|
||||
|
||||
userRequest := models.UserPassword{
|
||||
User: models.User{
|
||||
Mail: tc.mail,
|
||||
Name: tc.username,
|
||||
RoleID: tc.roleID,
|
||||
Status: models.UserStatusActive,
|
||||
Type: models.UserTypeLocal,
|
||||
},
|
||||
Password: "SecurePass123!",
|
||||
}
|
||||
|
||||
body, err := json.Marshal(userRequest)
|
||||
require.NoError(t, err)
|
||||
|
||||
c.Request, _ = http.NewRequest("POST", "/users/", bytes.NewBuffer(body))
|
||||
c.Request.Header.Set("Content-Type", "application/json")
|
||||
|
||||
service.CreateUser(c)
|
||||
|
||||
assert.Equal(t, http.StatusCreated, w.Code, "Expected HTTP 201 Created")
|
||||
|
||||
// Verify both user and preferences were created
|
||||
var user models.User
|
||||
err = db.Where("mail = ?", tc.mail).First(&user).Error
|
||||
require.NoError(t, err)
|
||||
|
||||
var prefs models.UserPreferences
|
||||
err = db.Where("user_id = ?", user.ID).First(&prefs).Error
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, user.ID, prefs.UserID)
|
||||
})
|
||||
}
|
||||
|
||||
// Verify all users and preferences exist
|
||||
var userCount int
|
||||
db.Model(&models.User{}).Where("mail LIKE ?", "newuser%@test.com").Count(&userCount)
|
||||
assert.Equal(t, 3, userCount, "Should have 3 newly created users")
|
||||
|
||||
var prefsCount int
|
||||
db.Model(&models.UserPreferences{}).Count(&prefsCount)
|
||||
assert.Equal(t, 5, prefsCount, "Should have 5 user preferences total (2 initial + 3 created)")
|
||||
}
|
||||
|
||||
func TestCreateUser_InvalidJSON(t *testing.T) {
|
||||
db := setupTestDB(t)
|
||||
defer db.Close()
|
||||
|
||||
userCache := auth.NewUserCache(db)
|
||||
service := NewUserService(db, userCache)
|
||||
|
||||
gin.SetMode(gin.TestMode)
|
||||
w := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(w)
|
||||
|
||||
c.Set("uid", uint64(1))
|
||||
c.Set("rid", uint64(1))
|
||||
c.Set("uhash", "testhash1")
|
||||
c.Set("prm", []string{"users.create"})
|
||||
|
||||
// Invalid JSON
|
||||
c.Request, _ = http.NewRequest("POST", "/users/", bytes.NewBufferString("{invalid json"))
|
||||
c.Request.Header.Set("Content-Type", "application/json")
|
||||
|
||||
service.CreateUser(c)
|
||||
|
||||
assert.Equal(t, http.StatusBadRequest, w.Code, "Expected HTTP 400 Bad Request")
|
||||
}
|
||||
|
||||
func TestCreateUser_DuplicateEmail(t *testing.T) {
|
||||
db := setupTestDB(t)
|
||||
defer db.Close()
|
||||
|
||||
userCache := auth.NewUserCache(db)
|
||||
service := NewUserService(db, userCache)
|
||||
|
||||
// Create first user
|
||||
gin.SetMode(gin.TestMode)
|
||||
w1 := httptest.NewRecorder()
|
||||
c1, _ := gin.CreateTestContext(w1)
|
||||
|
||||
c1.Set("uid", uint64(1))
|
||||
c1.Set("rid", uint64(1))
|
||||
c1.Set("uhash", "testhash1")
|
||||
c1.Set("prm", []string{"users.create"})
|
||||
|
||||
userRequest := models.UserPassword{
|
||||
User: models.User{
|
||||
Mail: "duplicate@test.com",
|
||||
Name: "First User",
|
||||
RoleID: 2,
|
||||
Status: models.UserStatusActive,
|
||||
Type: models.UserTypeLocal,
|
||||
},
|
||||
Password: "SecurePass123!",
|
||||
}
|
||||
|
||||
body, err := json.Marshal(userRequest)
|
||||
require.NoError(t, err)
|
||||
|
||||
c1.Request, _ = http.NewRequest("POST", "/users/", bytes.NewBuffer(body))
|
||||
c1.Request.Header.Set("Content-Type", "application/json")
|
||||
|
||||
service.CreateUser(c1)
|
||||
assert.Equal(t, http.StatusCreated, w1.Code)
|
||||
|
||||
// Try to create second user with same email
|
||||
w2 := httptest.NewRecorder()
|
||||
c2, _ := gin.CreateTestContext(w2)
|
||||
|
||||
c2.Set("uid", uint64(1))
|
||||
c2.Set("rid", uint64(1))
|
||||
c2.Set("uhash", "testhash1")
|
||||
c2.Set("prm", []string{"users.create"})
|
||||
|
||||
userRequest2 := models.UserPassword{
|
||||
User: models.User{
|
||||
Mail: "duplicate@test.com", // Same email
|
||||
Name: "Second User",
|
||||
RoleID: 2,
|
||||
Status: models.UserStatusActive,
|
||||
Type: models.UserTypeLocal,
|
||||
},
|
||||
Password: "AnotherPass456!",
|
||||
}
|
||||
|
||||
body2, err := json.Marshal(userRequest2)
|
||||
require.NoError(t, err)
|
||||
|
||||
c2.Request, _ = http.NewRequest("POST", "/users/", bytes.NewBuffer(body2))
|
||||
c2.Request.Header.Set("Content-Type", "application/json")
|
||||
|
||||
service.CreateUser(c2)
|
||||
|
||||
// Should fail due to unique constraint
|
||||
assert.Equal(t, http.StatusInternalServerError, w2.Code, "Expected error on duplicate email")
|
||||
|
||||
// Verify only one user exists
|
||||
var count int
|
||||
db.Model(&models.User{}).Where("mail = ?", "duplicate@test.com").Count(&count)
|
||||
assert.Equal(t, 1, count, "Should have only one user with this email")
|
||||
}
|
||||
@@ -0,0 +1,222 @@
|
||||
package services
|
||||
|
||||
import (
|
||||
"errors"
|
||||
"net/http"
|
||||
"slices"
|
||||
"strconv"
|
||||
|
||||
"pentagi/pkg/server/logger"
|
||||
"pentagi/pkg/server/models"
|
||||
"pentagi/pkg/server/rdb"
|
||||
"pentagi/pkg/server/response"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/jinzhu/gorm"
|
||||
)
|
||||
|
||||
type vecstorelogs struct {
|
||||
VecstoreLogs []models.Vecstorelog `json:"vecstorelogs"`
|
||||
Total uint64 `json:"total"`
|
||||
}
|
||||
|
||||
type vecstorelogsGrouped struct {
|
||||
Grouped []string `json:"grouped"`
|
||||
Total uint64 `json:"total"`
|
||||
}
|
||||
|
||||
var vecstorelogsSQLMappers = map[string]any{
|
||||
"id": "{{table}}.id",
|
||||
"initiator": "{{table}}.initiator",
|
||||
"executor": "{{table}}.executor",
|
||||
"filter": "{{table}}.filter",
|
||||
"query": "{{table}}.query",
|
||||
"action": "{{table}}.action",
|
||||
"result": "{{table}}.result",
|
||||
"flow_id": "{{table}}.flow_id",
|
||||
"task_id": "{{table}}.task_id",
|
||||
"subtask_id": "{{table}}.subtask_id",
|
||||
"created_at": "{{table}}.created_at",
|
||||
"data": "({{table}}.filter || ' ' || {{table}}.query || ' ' || {{table}}.result)",
|
||||
}
|
||||
|
||||
type VecstorelogService struct {
|
||||
db *gorm.DB
|
||||
}
|
||||
|
||||
func NewVecstorelogService(db *gorm.DB) *VecstorelogService {
|
||||
return &VecstorelogService{
|
||||
db: db,
|
||||
}
|
||||
}
|
||||
|
||||
// GetVecstorelogs is a function to return vecstorelogs list
|
||||
// @Summary Retrieve vecstorelogs list
|
||||
// @Tags Vecstorelogs
|
||||
// @Produce json
|
||||
// @Security BearerAuth
|
||||
// @Param request query rdb.TableQuery true "query table params"
|
||||
// @Success 200 {object} response.successResp{data=vecstorelogs} "vecstorelogs list received successful"
|
||||
// @Failure 400 {object} response.errorResp "invalid query request data"
|
||||
// @Failure 403 {object} response.errorResp "getting vecstorelogs not permitted"
|
||||
// @Failure 500 {object} response.errorResp "internal error on getting vecstorelogs"
|
||||
// @Router /vecstorelogs/ [get]
|
||||
func (s *VecstorelogService) GetVecstorelogs(c *gin.Context) {
|
||||
var (
|
||||
err error
|
||||
query rdb.TableQuery
|
||||
resp vecstorelogs
|
||||
)
|
||||
|
||||
if err = c.ShouldBindQuery(&query); err != nil {
|
||||
logger.FromContext(c).WithError(err).Errorf("error binding query")
|
||||
response.Error(c, response.ErrVecstorelogsInvalidRequest, err)
|
||||
return
|
||||
}
|
||||
|
||||
uid := c.GetUint64("uid")
|
||||
privs := c.GetStringSlice("prm")
|
||||
var scope func(db *gorm.DB) *gorm.DB
|
||||
if slices.Contains(privs, "vecstorelogs.admin") {
|
||||
scope = func(db *gorm.DB) *gorm.DB {
|
||||
return db.
|
||||
Joins("INNER JOIN flows f ON f.id = flow_id")
|
||||
}
|
||||
} else if slices.Contains(privs, "vecstorelogs.view") {
|
||||
scope = func(db *gorm.DB) *gorm.DB {
|
||||
return db.
|
||||
Joins("INNER JOIN flows f ON f.id = flow_id").
|
||||
Where("f.user_id = ?", uid)
|
||||
}
|
||||
} else {
|
||||
logger.FromContext(c).Errorf("error filtering user role permissions: permission not found")
|
||||
response.Error(c, response.ErrNotPermitted, nil)
|
||||
return
|
||||
}
|
||||
|
||||
query.Init("vecstorelogs", vecstorelogsSQLMappers)
|
||||
|
||||
if query.Group != "" {
|
||||
if _, ok := vecstorelogsSQLMappers[query.Group]; !ok {
|
||||
logger.FromContext(c).Errorf("error finding vecstorelogs grouped: group field not found")
|
||||
response.Error(c, response.ErrVecstorelogsInvalidRequest, errors.New("group field not found"))
|
||||
return
|
||||
}
|
||||
|
||||
var respGrouped vecstorelogsGrouped
|
||||
if respGrouped.Total, err = query.QueryGrouped(s.db, &respGrouped.Grouped, scope); err != nil {
|
||||
logger.FromContext(c).WithError(err).Errorf("error finding vecstorelogs grouped")
|
||||
response.Error(c, response.ErrInternal, err)
|
||||
return
|
||||
}
|
||||
|
||||
response.Success(c, http.StatusOK, respGrouped)
|
||||
return
|
||||
}
|
||||
|
||||
if resp.Total, err = query.Query(s.db, &resp.VecstoreLogs, scope); err != nil {
|
||||
logger.FromContext(c).WithError(err).Errorf("error finding vecstorelogs")
|
||||
response.Error(c, response.ErrInternal, err)
|
||||
return
|
||||
}
|
||||
|
||||
for i := 0; i < len(resp.VecstoreLogs); i++ {
|
||||
if err = resp.VecstoreLogs[i].Valid(); err != nil {
|
||||
logger.FromContext(c).WithError(err).Errorf("error validating vecstorelog data '%d'", resp.VecstoreLogs[i].ID)
|
||||
response.Error(c, response.ErrVecstorelogsInvalidData, err)
|
||||
return
|
||||
}
|
||||
}
|
||||
|
||||
response.Success(c, http.StatusOK, resp)
|
||||
}
|
||||
|
||||
// GetFlowVecstorelogs is a function to return vecstorelogs list by flow id
|
||||
// @Summary Retrieve vecstorelogs list by flow id
|
||||
// @Tags Vecstorelogs
|
||||
// @Produce json
|
||||
// @Security BearerAuth
|
||||
// @Param flowID path int true "flow id" minimum(0)
|
||||
// @Param request query rdb.TableQuery true "query table params"
|
||||
// @Success 200 {object} response.successResp{data=vecstorelogs} "vecstorelogs list received successful"
|
||||
// @Failure 400 {object} response.errorResp "invalid query request data"
|
||||
// @Failure 403 {object} response.errorResp "getting vecstorelogs not permitted"
|
||||
// @Failure 500 {object} response.errorResp "internal error on getting vecstorelogs"
|
||||
// @Router /flows/{flowID}/vecstorelogs/ [get]
|
||||
func (s *VecstorelogService) GetFlowVecstorelogs(c *gin.Context) {
|
||||
var (
|
||||
err error
|
||||
flowID uint64
|
||||
query rdb.TableQuery
|
||||
resp vecstorelogs
|
||||
)
|
||||
|
||||
if flowID, err = strconv.ParseUint(c.Param("flowID"), 10, 64); err != nil {
|
||||
logger.FromContext(c).WithError(err).Errorf("error parsing flow id")
|
||||
response.Error(c, response.ErrVecstorelogsInvalidRequest, err)
|
||||
return
|
||||
}
|
||||
|
||||
if err = c.ShouldBindQuery(&query); err != nil {
|
||||
logger.FromContext(c).WithError(err).Errorf("error binding query")
|
||||
response.Error(c, response.ErrVecstorelogsInvalidRequest, err)
|
||||
return
|
||||
}
|
||||
|
||||
uid := c.GetUint64("uid")
|
||||
privs := c.GetStringSlice("prm")
|
||||
var scope func(db *gorm.DB) *gorm.DB
|
||||
if slices.Contains(privs, "vecstorelogs.admin") {
|
||||
scope = func(db *gorm.DB) *gorm.DB {
|
||||
return db.
|
||||
Joins("INNER JOIN flows f ON f.id = flow_id").
|
||||
Where("f.id = ?", flowID)
|
||||
}
|
||||
} else if slices.Contains(privs, "vecstorelogs.view") {
|
||||
scope = func(db *gorm.DB) *gorm.DB {
|
||||
return db.
|
||||
Joins("INNER JOIN flows f ON f.id = flow_id").
|
||||
Where("f.id = ? AND f.user_id = ?", flowID, uid)
|
||||
}
|
||||
} else {
|
||||
logger.FromContext(c).Errorf("error filtering user role permissions: permission not found")
|
||||
response.Error(c, response.ErrNotPermitted, nil)
|
||||
return
|
||||
}
|
||||
|
||||
query.Init("vecstorelogs", vecstorelogsSQLMappers)
|
||||
|
||||
if query.Group != "" {
|
||||
if _, ok := vecstorelogsSQLMappers[query.Group]; !ok {
|
||||
logger.FromContext(c).Errorf("error finding vecstorelogs grouped: group field not found")
|
||||
response.Error(c, response.ErrVecstorelogsInvalidRequest, errors.New("group field not found"))
|
||||
return
|
||||
}
|
||||
|
||||
var respGrouped vecstorelogsGrouped
|
||||
if respGrouped.Total, err = query.QueryGrouped(s.db, &respGrouped.Grouped, scope); err != nil {
|
||||
logger.FromContext(c).WithError(err).Errorf("error finding vecstorelogs grouped")
|
||||
response.Error(c, response.ErrInternal, err)
|
||||
return
|
||||
}
|
||||
|
||||
response.Success(c, http.StatusOK, respGrouped)
|
||||
return
|
||||
}
|
||||
|
||||
if resp.Total, err = query.Query(s.db, &resp.VecstoreLogs, scope); err != nil {
|
||||
logger.FromContext(c).WithError(err).Errorf("error finding vecstorelogs")
|
||||
response.Error(c, response.ErrInternal, err)
|
||||
return
|
||||
}
|
||||
|
||||
for i := 0; i < len(resp.VecstoreLogs); i++ {
|
||||
if err = resp.VecstoreLogs[i].Valid(); err != nil {
|
||||
logger.FromContext(c).WithError(err).Errorf("error validating vecstorelog data '%d'", resp.VecstoreLogs[i].ID)
|
||||
response.Error(c, response.ErrVecstorelogsInvalidData, err)
|
||||
return
|
||||
}
|
||||
}
|
||||
|
||||
response.Success(c, http.StatusOK, resp)
|
||||
}
|
||||
Reference in New Issue
Block a user