Skip to content

Commit 7eabb0e

Browse files
ToeiReiToeiRei
committed
implement global key assignment restrictions and add cleanup utility for account_keys (possible bun bug)
1 parent e882e48 commit 7eabb0e

File tree

5 files changed

+231
-16
lines changed

5 files changed

+231
-16
lines changed

‎internal/core/db/assign_global_key_test.go‎

Lines changed: 106 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,106 @@
1+
// Copyright (c) 2026 Keymaster Team
2+
// Keymaster - SSH key management system
3+
// This source code is licensed under the MIT license found in the LICENSE file.
4+
5+
package db
6+
7+
import (
8+
"context"
9+
"strings"
10+
"testing"
11+
"time"
12+
)
13+
14+
func TestAssignKeyToAccount_GlobalKey_ReturnsError(t *testing.T) {
15+
bStore, err := New("sqlite", ":memory:")
16+
if err != nil {
17+
t.Fatalf("New failed: %v", err)
18+
}
19+
bdb := bStore.BunDB()
20+
ctx := context.Background()
21+
22+
// Create account
23+
_, err = ExecRaw(ctx, bdb, "INSERT INTO accounts(username, hostname, label) VALUES(?, ?, ?)", "user1", "host1.example.com", "Test1")
24+
if err != nil {
25+
t.Fatalf("insert account failed: %v", err)
26+
}
27+
acc, _ := GetAccountByIDBun(bdb, 1)
28+
if acc == nil {
29+
t.Fatal("account not found after insert")
30+
}
31+
32+
// Create a GLOBAL key
33+
err = AddPublicKeyBun(bdb, "ssh-ed25519", "AAAAC3test", "global-key", true, time.Time{})
34+
if err != nil {
35+
t.Fatalf("AddPublicKeyBun failed: %v", err)
36+
}
37+
pk, _ := GetPublicKeyByCommentBun(bdb, "global-key")
38+
if pk == nil {
39+
t.Fatal("key not found after insert")
40+
}
41+
42+
// Try to assign the global key to an account - should fail
43+
err = AssignKeyToAccountBun(bdb, pk.ID, acc.ID)
44+
if err == nil {
45+
t.Fatal("expected error when assigning global key, got nil")
46+
}
47+
if !strings.Contains(err.Error(), "cannot assign global key") {
48+
t.Fatalf("expected error about global key, got: %v", err)
49+
}
50+
51+
// Verify the key was NOT added to account_keys
52+
keys, err := GetKeysForAccountBun(bdb, acc.ID)
53+
if err != nil {
54+
t.Fatalf("GetKeysForAccountBun failed: %v", err)
55+
}
56+
if len(keys) != 0 {
57+
t.Fatalf("expected 0 keys in account_keys for account %d, got %d", acc.ID, len(keys))
58+
}
59+
}
60+
61+
func TestAssignKeyToAccount_NonGlobalKey_Succeeds(t *testing.T) {
62+
bStore, err := New("sqlite", ":memory:")
63+
if err != nil {
64+
t.Fatalf("New failed: %v", err)
65+
}
66+
bdb := bStore.BunDB()
67+
ctx := context.Background()
68+
69+
// Create account
70+
_, err = ExecRaw(ctx, bdb, "INSERT INTO accounts(username, hostname, label) VALUES(?, ?, ?)", "user2", "host2.example.com", "Test2")
71+
if err != nil {
72+
t.Fatalf("insert account failed: %v", err)
73+
}
74+
acc, _ := GetAccountByIDBun(bdb, 1)
75+
if acc == nil {
76+
t.Fatal("account not found after insert")
77+
}
78+
79+
// Create a NON-GLOBAL key
80+
err = AddPublicKeyBun(bdb, "ssh-ed25519", "AAAAC3test2", "regular-key", false, time.Time{})
81+
if err != nil {
82+
t.Fatalf("AddPublicKeyBun failed: %v", err)
83+
}
84+
pk, _ := GetPublicKeyByCommentBun(bdb, "regular-key")
85+
if pk == nil {
86+
t.Fatal("key not found after insert")
87+
}
88+
89+
// Assign the non-global key to an account - should succeed
90+
err = AssignKeyToAccountBun(bdb, pk.ID, acc.ID)
91+
if err != nil {
92+
t.Fatalf("expected no error when assigning non-global key, got: %v", err)
93+
}
94+
95+
// Verify the key WAS added to account_keys
96+
keys, err := GetKeysForAccountBun(bdb, acc.ID)
97+
if err != nil {
98+
t.Fatalf("GetKeysForAccountBun failed: %v", err)
99+
}
100+
if len(keys) != 1 {
101+
t.Fatalf("expected 1 key in account_keys for account %d, got %d", acc.ID, len(keys))
102+
}
103+
if keys[0].ID != pk.ID {
104+
t.Fatalf("expected key ID %d, got %d", pk.ID, keys[0].ID)
105+
}
106+
}

‎internal/core/db/bun_adapter.go‎

Lines changed: 34 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -291,8 +291,24 @@ func DeleteAccountBun(bdb *bun.DB, id int) error {
291291
}
292292

293293
// AssignKeyToAccountBun creates an association in account_keys.
294+
// Global keys should not be assigned to individual accounts since they're
295+
// automatically deployed everywhere. This function returns an error if attempting
296+
// to assign a global key.
294297
func AssignKeyToAccountBun(bdb *bun.DB, keyID, accountID int) error {
295298
ctx := context.Background()
299+
300+
// Check if the key is global - global keys should not be in account_keys
301+
pk, err := GetPublicKeyByIDBun(bdb, keyID)
302+
if err != nil {
303+
return err
304+
}
305+
if pk == nil {
306+
return fmt.Errorf("key with ID %d not found", keyID)
307+
}
308+
if pk.IsGlobal {
309+
return fmt.Errorf("cannot assign global key '%s' to individual accounts (it's already deployed everywhere)", pk.Comment)
310+
}
311+
296312
// Use raw insert since account_keys likely has no PK model in codebase.
297313
if _, err := ExecRaw(ctx, bdb, "INSERT INTO account_keys(key_id, account_id) VALUES(?, ?)", keyID, accountID); err != nil {
298314
return MapDBError(err)
@@ -321,19 +337,25 @@ func UnassignKeyFromAccountBun(bdb *bun.DB, keyID, accountID int) error {
321337
func GetKeysForAccountBun(bdb *bun.DB, accountID int) ([]model.PublicKey, error) {
322338
ctx := context.Background()
323339
var pks []PublicKeyModel
324-
err := bdb.NewSelect().Model(&pks).
325-
TableExpr("public_keys AS pk").
326-
Join("JOIN account_keys ak ON pk.id = ak.key_id").
327-
Where("ak.account_id = ?", accountID).
328-
OrderExpr("pk.comment").
329-
Scan(ctx)
340+
// Use raw query since BUN's Model + Join was not properly applying the WHERE clause
341+
query := `SELECT pk.* FROM public_keys pk
342+
INNER JOIN account_keys ak ON pk.id = ak.key_id
343+
WHERE ak.account_id = ?
344+
ORDER BY pk.comment`
345+
err := bdb.NewRaw(query, accountID).Scan(ctx, &pks)
330346
if err != nil {
331347
return nil, err
332348
}
333349
out := make([]model.PublicKey, 0, len(pks))
334350
for _, p := range pks {
335351
out = append(out, publicKeyModelToModel(p))
336352
}
353+
if dbDebugEnabled {
354+
dbLogf("GetKeysForAccountBun(accountID=%d): returning %d keys", accountID, len(out))
355+
for _, k := range out {
356+
dbLogf(" - Key ID=%d, Comment=%s, IsGlobal=%v", k.ID, k.Comment, k.IsGlobal)
357+
}
358+
}
337359
return out, nil
338360
}
339361

@@ -732,6 +754,12 @@ func GetGlobalPublicKeysBun(bdb *bun.DB) ([]model.PublicKey, error) {
732754
for _, p := range pks {
733755
out = append(out, publicKeyModelToModel(p))
734756
}
757+
if dbDebugEnabled {
758+
dbLogf("GetGlobalPublicKeysBun(): returning %d global keys", len(out))
759+
for _, k := range out {
760+
dbLogf(" - Key ID=%d, Comment=%s", k.ID, k.Comment)
761+
}
762+
}
735763
return out, nil
736764
}
737765

‎internal/core/db/bun_adapter_extra_test.go‎

Lines changed: 7 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -5,6 +5,7 @@
55
package db
66

77
import (
8+
"strings"
89
"testing"
910
"time"
1011
)
@@ -243,16 +244,13 @@ func TestPublicKeys_List_Toggle_Search_Delete_Assignments(t *testing.T) {
243244
t.Fatalf("expected at least 2 global keys, got %d", len(globals))
244245
}
245246

246-
// Assign pk1 to account and verify accounts for key
247-
if err := AssignKeyToAccountBun(bdb, pk1.ID, accID); err != nil {
248-
t.Fatalf("AssignKeyToAccountBun failed: %v", err)
247+
// Try to assign pk1 (now global) to account - should fail
248+
err = AssignKeyToAccountBun(bdb, pk1.ID, accID)
249+
if err == nil {
250+
t.Fatal("AssignKeyToAccountBun should fail when trying to assign global key")
249251
}
250-
accs, err := GetAccountsForKeyBun(bdb, pk1.ID)
251-
if err != nil {
252-
t.Fatalf("GetAccountsForKeyBun failed: %v", err)
253-
}
254-
if len(accs) == 0 {
255-
t.Fatalf("expected account assignment to be visible for key")
252+
if !strings.Contains(err.Error(), "cannot assign global key") {
253+
t.Fatalf("expected error about global key assignment, got: %v", err)
256254
}
257255

258256
// Delete public key pk2

‎internal/core/db/log.go‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -19,6 +19,6 @@ func SetDebug(enabled bool) {
1919

2020
func dbLogf(format string, v ...any) {
2121
if dbDebugEnabled {
22-
log.Debug(fmt.Sprintf(format, v...))
22+
log.Info(fmt.Sprintf("[DB] "+format, v...))
2323
}
2424
}

‎tools/cleanup_global_keys/main.go‎

Lines changed: 83 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,83 @@
1+
// Copyright (c) 2026 Keymaster Team
2+
// Keymaster - SSH key management system
3+
// This source code is licensed under the MIT license found in the LICENSE file.
4+
5+
// This is a one-time cleanup utility to remove global keys from the account_keys table.
6+
// Global keys should only exist in public_keys with is_global=1, not in account_keys.
7+
8+
package main
9+
10+
import (
11+
"context"
12+
"fmt"
13+
"log"
14+
15+
"github.com/toeirei/keymaster/internal/core/db"
16+
)
17+
18+
func main() {
19+
// Initialize the database
20+
store, err := db.New("sqlite", "keymaster.db")
21+
if err != nil {
22+
log.Fatalf("Failed to open database: %v", err)
23+
}
24+
bdb := store.BunDB()
25+
ctx := context.Background()
26+
27+
// Find all global keys that are incorrectly in account_keys
28+
query := `
29+
SELECT DISTINCT pk.id, pk.comment, COUNT(ak.account_id) as assignment_count
30+
FROM public_keys pk
31+
INNER JOIN account_keys ak ON pk.id = ak.key_id
32+
WHERE pk.is_global = 1
33+
GROUP BY pk.id, pk.comment
34+
`
35+
36+
type result struct {
37+
ID int
38+
Comment string
39+
AssignmentCount int
40+
}
41+
42+
var results []result
43+
err = bdb.NewRaw(query).Scan(ctx, &results)
44+
if err != nil {
45+
log.Fatalf("Failed to query global keys in account_keys: %v", err)
46+
}
47+
48+
if len(results) == 0 {
49+
fmt.Println("✓ No global keys found in account_keys table. Database is clean!")
50+
return
51+
}
52+
53+
fmt.Printf("Found %d global key(s) incorrectly assigned in account_keys:\n", len(results))
54+
for _, r := range results {
55+
fmt.Printf(" - Key ID %d (%s): %d assignment(s)\n", r.ID, r.Comment, r.AssignmentCount)
56+
}
57+
58+
// Delete all global keys from account_keys
59+
deleteQuery := `
60+
DELETE FROM account_keys
61+
WHERE key_id IN (
62+
SELECT id FROM public_keys WHERE is_global = 1
63+
)
64+
`
65+
66+
deleteResult, err := db.ExecRaw(ctx, bdb, deleteQuery)
67+
if err != nil {
68+
log.Fatalf("Failed to delete global keys from account_keys: %v", err)
69+
}
70+
71+
rowsAffected, _ := deleteResult.RowsAffected()
72+
fmt.Printf("\n✓ Removed %d incorrect assignment(s) from account_keys table.\n", rowsAffected)
73+
74+
// Mark all accounts as dirty since key assignments changed
75+
_, err = db.ExecRaw(ctx, bdb, "UPDATE accounts SET is_dirty = 1")
76+
if err != nil {
77+
log.Fatalf("Failed to mark accounts dirty: %v", err)
78+
}
79+
80+
fmt.Println("✓ Marked all accounts as dirty for redeployment.")
81+
fmt.Println("\nCleanup complete! Global keys will now only be deployed via the is_global flag,")
82+
fmt.Println("not through account_keys assignments.")
83+
}

0 commit comments

Comments
 (0)