Mathis/llamactl
1
0
Fork 0
mirror of https://github.com/lordmathis/llamactl.git synced 2025-12-22 17:14:22 +00:00
Code Issues Packages Projects Releases Wiki Activity

Initial api key store implementation

Browse Source
This commit is contained in:
Matúš Námešný
2025-12-03 20:31:38 +01:00
parent 470f90076f
commit 9eee42c673
18 changed files with 986 additions and 116 deletions
Download Patch File Download Diff File Expand all files Collapse all files

211
pkg/database/apikeys.go Normal file
View File

@@ -0,0 +1,211 @@
package database
import (
"context"
"database/sql"
"fmt"
"llamactl/pkg/auth"
"time"
)
// CreateKey inserts a new API key with permissions (transactional)
func (db *sqliteDB) CreateKey(ctx context.Context, key *auth.APIKey, permissions []auth.KeyPermission) error {
tx, err := db.BeginTx(ctx, nil)
if err != nil {
return fmt.Errorf("failed to begin transaction: %w", err)
}
defer tx.Rollback()
// Insert the API key
query := `
INSERT INTO api_keys (key_hash, name, user_id, permission_mode, expires_at, enabled, created_at, updated_at)
VALUES (?, ?, ?, ?, ?, ?, ?, ?)
`
var expiresAt sql.NullInt64
if key.ExpiresAt != nil {
expiresAt = sql.NullInt64{Int64: *key.ExpiresAt, Valid: true}
}
result, err := tx.ExecContext(ctx, query,
key.KeyHash, key.Name, key.UserID, key.PermissionMode,
expiresAt, key.Enabled, key.CreatedAt, key.UpdatedAt,
)
if err != nil {
return fmt.Errorf("failed to insert API key: %w", err)
}
keyID, err := result.LastInsertId()
if err != nil {
return fmt.Errorf("failed to get last insert ID: %w", err)
}
key.ID = int(keyID)
// Insert permissions if per-instance mode
if key.PermissionMode == auth.PermissionModePerInstance {
for _, perm := range permissions {
query := `
INSERT INTO key_permissions (key_id, instance_id, can_infer, can_view_logs)
VALUES (?, ?, ?, ?)
`
_, err := tx.ExecContext(ctx, query, perm.KeyID, perm.InstanceID, perm.CanInfer, perm.CanViewLogs)
if err != nil {
return fmt.Errorf("failed to insert permission for instance %d: %w", perm.InstanceID, err)
}
}
}
return tx.Commit()
}
// GetKeyByID retrieves an API key by ID
func (db *sqliteDB) GetKeyByID(ctx context.Context, id int) (*auth.APIKey, error) {
query := `
SELECT id, key_hash, name, user_id, permission_mode, expires_at, enabled, created_at, updated_at, last_used_at
FROM api_keys
WHERE id = ?
`
var key auth.APIKey
var expiresAt sql.NullInt64
var lastUsedAt sql.NullInt64
err := db.QueryRowContext(ctx, query, id).Scan(
&key.ID, &key.KeyHash, &key.Name, &key.UserID, &key.PermissionMode,
&expiresAt, &key.Enabled, &key.CreatedAt, &key.UpdatedAt, &lastUsedAt,
)
if err != nil {
if err == sql.ErrNoRows {
return nil, fmt.Errorf("API key not found")
}
return nil, fmt.Errorf("failed to query API key: %w", err)
}
if expiresAt.Valid {
key.ExpiresAt = &expiresAt.Int64
}
if lastUsedAt.Valid {
key.LastUsedAt = &lastUsedAt.Int64
}
return &key, nil
}
// GetUserKeys retrieves all API keys for a user
func (db *sqliteDB) GetUserKeys(ctx context.Context, userID string) ([]*auth.APIKey, error) {
query := `
SELECT id, key_hash, name, user_id, permission_mode, expires_at, enabled, created_at, updated_at, last_used_at
FROM api_keys
WHERE user_id = ?
ORDER BY created_at DESC
`
rows, err := db.QueryContext(ctx, query, userID)
if err != nil {
return nil, fmt.Errorf("failed to query API keys: %w", err)
}
defer rows.Close()
var keys []*auth.APIKey
for rows.Next() {
var key auth.APIKey
var expiresAt sql.NullInt64
var lastUsedAt sql.NullInt64
err := rows.Scan(
&key.ID, &key.KeyHash, &key.Name, &key.UserID, &key.PermissionMode,
&expiresAt, &key.Enabled, &key.CreatedAt, &key.UpdatedAt, &lastUsedAt,
)
if err != nil {
return nil, fmt.Errorf("failed to scan API key: %w", err)
}
if expiresAt.Valid {
key.ExpiresAt = &expiresAt.Int64
}
if lastUsedAt.Valid {
key.LastUsedAt = &lastUsedAt.Int64
}
keys = append(keys, &key)
}
return keys, nil
}
// GetActiveKeys retrieves all enabled, non-expired API keys
func (db *sqliteDB) GetActiveKeys(ctx context.Context) ([]*auth.APIKey, error) {
query := `
SELECT id, key_hash, name, user_id, permission_mode, expires_at, enabled, created_at, updated_at, last_used_at
FROM api_keys
WHERE enabled = 1 AND (expires_at IS NULL OR expires_at > ?)
ORDER BY created_at DESC
`
now := time.Now().Unix()
rows, err := db.QueryContext(ctx, query, now)
if err != nil {
return nil, fmt.Errorf("failed to query active API keys: %w", err)
}
defer rows.Close()
var keys []*auth.APIKey
for rows.Next() {
var key auth.APIKey
var expiresAt sql.NullInt64
var lastUsedAt sql.NullInt64
err := rows.Scan(
&key.ID, &key.KeyHash, &key.Name, &key.UserID, &key.PermissionMode,
&expiresAt, &key.Enabled, &key.CreatedAt, &key.UpdatedAt, &lastUsedAt,
)
if err != nil {
return nil, fmt.Errorf("failed to scan API key: %w", err)
}
if expiresAt.Valid {
key.ExpiresAt = &expiresAt.Int64
}
if lastUsedAt.Valid {
key.LastUsedAt = &lastUsedAt.Int64
}
keys = append(keys, &key)
}
return keys, nil
}
// DeleteKey removes an API key (cascades to permissions)
func (db *sqliteDB) DeleteKey(ctx context.Context, id int) error {
query := `DELETE FROM api_keys WHERE id = ?`
result, err := db.ExecContext(ctx, query, id)
if err != nil {
return fmt.Errorf("failed to delete API key: %w", err)
}
rowsAffected, err := result.RowsAffected()
if err != nil {
return fmt.Errorf("failed to get rows affected: %w", err)
}
if rowsAffected == 0 {
return fmt.Errorf("API key not found")
}
return nil
}
// TouchKey updates the last_used_at timestamp
func (db *sqliteDB) TouchKey(ctx context.Context, id int) error {
query := `UPDATE api_keys SET last_used_at = ?, updated_at = ? WHERE id = ?`
now := time.Now().Unix()
_, err := db.ExecContext(ctx, query, now, now, id)
if err != nil {
return fmt.Errorf("failed to update last used timestamp: %w", err)
}
return nil
}

30
pkg/database/database.go
View File

@@ -1,8 +1,10 @@
package database
import (
"context"
"database/sql"
"fmt"
"llamactl/pkg/auth"
"llamactl/pkg/instance"
"log"
"path/filepath"
@@ -11,14 +13,26 @@ import (
_ "github.com/mattn/go-sqlite3"
)
// DB defines the interface for instance persistence operations
type DB interface {
// InstanceStore defines interface for instance persistence operations
type InstanceStore interface {
Save(inst *instance.Instance) error
Delete(name string) error
LoadAll() ([]*instance.Instance, error)
Close() error
}
// AuthStore defines the interface for authentication operations
type AuthStore interface {
CreateKey(ctx context.Context, key *auth.APIKey, permissions []auth.KeyPermission) error
GetUserKeys(ctx context.Context, userID string) ([]*auth.APIKey, error)
GetActiveKeys(ctx context.Context) ([]*auth.APIKey, error)
GetKeyByID(ctx context.Context, id int) (*auth.APIKey, error)
DeleteKey(ctx context.Context, id int) error
TouchKey(ctx context.Context, id int) error
GetPermissions(ctx context.Context, keyID int) ([]auth.KeyPermission, error)
HasPermission(ctx context.Context, keyID, instanceID int) (bool, error)
}
// Config contains database configuration settings
type Config struct {
// Database file path (relative to data_dir or absolute)
@@ -30,13 +44,13 @@ type Config struct {
ConnMaxLifetime time.Duration
}
// sqliteDB wraps the database connection with configuration
// sqliteDB wraps database connection with configuration
type sqliteDB struct {
*sql.DB
config *Config
}
// Open creates a new database connection with the provided configuration
// Open creates a new database connection with provided configuration
func Open(config *Config) (*sqliteDB, error) {
if config == nil {
return nil, fmt.Errorf("database config cannot be nil")
@@ -46,10 +60,10 @@ func Open(config *Config) (*sqliteDB, error) {
return nil, fmt.Errorf("database path cannot be empty")
}
// Ensure the database directory exists
// Ensure that database directory exists
dbDir := filepath.Dir(config.Path)
if dbDir != "." && dbDir != "/" {
// Directory will be created by the manager if auto_create_dirs is enabled
// Directory will be created by manager if auto_create_dirs is enabled
log.Printf("Database will be created at: %s", config.Path)
}
@@ -89,7 +103,7 @@ func Open(config *Config) (*sqliteDB, error) {
}, nil
}
// Close closes the database connection
// Close closes database connection
func (db *sqliteDB) Close() error {
if db.DB != nil {
log.Println("Closing database connection")
@@ -98,7 +112,7 @@ func (db *sqliteDB) Close() error {
return nil
}
// HealthCheck verifies the database is accessible
// HealthCheck verifies that database is accessible
func (db *sqliteDB) HealthCheck() error {
if db.DB == nil {
return fmt.Errorf("database connection is nil")

12
pkg/database/migrations/001_initial_schema.down.sql
View File

@@ -1,7 +1,11 @@
-- Drop indexes first
DROP INDEX IF EXISTS idx_instances_backend_type;
-- Drop API key related indexes and tables first
DROP INDEX IF EXISTS idx_key_permissions_instance_id;
DROP INDEX IF EXISTS idx_api_keys_expires_at;
DROP INDEX IF EXISTS idx_api_keys_user_id;
DROP TABLE IF EXISTS key_permissions;
DROP TABLE IF EXISTS api_keys;
-- Drop instance related indexes and tables
DROP INDEX IF EXISTS idx_instances_status;
DROP INDEX IF EXISTS idx_instances_name;
-- Drop tables
DROP TABLE IF EXISTS instances;

36
pkg/database/migrations/001_initial_schema.up.sql
View File

@@ -25,3 +25,39 @@ CREATE TABLE IF NOT EXISTS instances (
-- -----------------------------------------------------------------------------
CREATE UNIQUE INDEX IF NOT EXISTS idx_instances_name ON instances(name);
CREATE INDEX IF NOT EXISTS idx_instances_status ON instances(status);
-- -----------------------------------------------------------------------------
-- API Keys Table: Database-backed inference API keys
-- -----------------------------------------------------------------------------
CREATE TABLE IF NOT EXISTS api_keys (
id INTEGER PRIMARY KEY AUTOINCREMENT,
key_hash TEXT NOT NULL,
name TEXT NOT NULL,
user_id TEXT NOT NULL,
permission_mode TEXT NOT NULL CHECK(permission_mode IN ('allow_all', 'per_instance')) DEFAULT 'per_instance',
expires_at INTEGER NULL,
enabled INTEGER NOT NULL DEFAULT 1,
created_at INTEGER NOT NULL,
updated_at INTEGER NOT NULL,
last_used_at INTEGER NULL
);
-- -----------------------------------------------------------------------------
-- Key Permissions Table: Per-instance permissions for API keys
-- -----------------------------------------------------------------------------
CREATE TABLE IF NOT EXISTS key_permissions (
key_id INTEGER NOT NULL,
instance_id INTEGER NOT NULL,
can_infer INTEGER NOT NULL DEFAULT 0,
can_view_logs INTEGER NOT NULL DEFAULT 0,
PRIMARY KEY (key_id, instance_id),
FOREIGN KEY (key_id) REFERENCES api_keys (id) ON DELETE CASCADE,
FOREIGN KEY (instance_id) REFERENCES instances (id) ON DELETE CASCADE
);
-- -----------------------------------------------------------------------------
-- Indexes for API keys and permissions
-- -----------------------------------------------------------------------------
CREATE INDEX IF NOT EXISTS idx_api_keys_user_id ON api_keys(user_id);
CREATE INDEX IF NOT EXISTS idx_api_keys_expires_at ON api_keys(expires_at);
CREATE INDEX IF NOT EXISTS idx_key_permissions_instance_id ON key_permissions(instance_id);

57
pkg/database/permissions.go Normal file
View File

@@ -0,0 +1,57 @@
package database
import (
"context"
"database/sql"
"fmt"
"llamactl/pkg/auth"
)
// GetPermissions retrieves all permissions for a key
func (db *sqliteDB) GetPermissions(ctx context.Context, keyID int) ([]auth.KeyPermission, error) {
query := `
SELECT key_id, instance_id, can_infer, can_view_logs
FROM key_permissions
WHERE key_id = ?
ORDER BY instance_id
`
rows, err := db.QueryContext(ctx, query, keyID)
if err != nil {
return nil, fmt.Errorf("failed to query key permissions: %w", err)
}
defer rows.Close()
var permissions []auth.KeyPermission
for rows.Next() {
var perm auth.KeyPermission
err := rows.Scan(&perm.KeyID, &perm.InstanceID, &perm.CanInfer, &perm.CanViewLogs)
if err != nil {
return nil, fmt.Errorf("failed to scan key permission: %w", err)
}
permissions = append(permissions, perm)
}
return permissions, nil
}
// HasPermission checks if key has inference permission for instance
func (db *sqliteDB) HasPermission(ctx context.Context, keyID, instanceID int) (bool, error) {
query := `
SELECT can_infer
FROM key_permissions
WHERE key_id = ? AND instance_id = ?
`
var canInfer bool
err := db.QueryRowContext(ctx, query, keyID, instanceID).Scan(&canInfer)
if err != nil {
if err == sql.ErrNoRows {
// No permission record found, deny access
return false, nil
}
return false, fmt.Errorf("failed to check key permission: %w", err)
}
return canInfer, nil
}