Skip to content

Commit 91c98cb

Browse files
authored
bugfix/menu fixes (#8)
- Bug fix for Quarantine Servers (0) after app start. Improvements in state management for tray UI updates. - Update autostart menu item icons for macOS to improve visual clarity and user experience.
1 parent afcbe29 commit 91c98cb

4 files changed

Lines changed: 131 additions & 43 deletions

File tree

internal/server/server.go

Lines changed: 26 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -651,34 +651,50 @@ func (s *Server) GetAllServers() ([]map[string]interface{}, error) {
651651

652652
// GetQuarantinedServers returns information about quarantined servers for tray UI
653653
func (s *Server) GetQuarantinedServers() ([]map[string]interface{}, error) {
654+
s.logger.Debug("GetQuarantinedServers called")
655+
654656
// Check if storage manager is available
655657
if s.storageManager == nil {
658+
s.logger.Warn("Storage manager is nil in GetQuarantinedServers")
656659
return []map[string]interface{}{}, nil
657660
}
658661

662+
s.logger.Debug("Calling storage manager ListQuarantinedUpstreamServers")
659663
quarantinedServers, err := s.storageManager.ListQuarantinedUpstreamServers()
660664
if err != nil {
661665
// Handle database closed gracefully
662666
if strings.Contains(err.Error(), "database not open") || strings.Contains(err.Error(), "closed") {
663667
s.logger.Debug("Database not available for GetQuarantinedServers, returning empty list")
664668
return []map[string]interface{}{}, nil
665669
}
670+
s.logger.Error("Failed to get quarantined servers from storage", zap.Error(err))
666671
return nil, err
667672
}
668673

674+
s.logger.Debug("Retrieved quarantined servers from storage",
675+
zap.Int("count", len(quarantinedServers)))
676+
669677
var result []map[string]interface{}
670678
for _, server := range quarantinedServers {
671-
result = append(result, map[string]interface{}{
679+
serverMap := map[string]interface{}{
672680
"name": server.Name,
673681
"url": server.URL,
674682
"command": server.Command,
675683
"protocol": server.Protocol,
676684
"enabled": server.Enabled,
677685
"quarantined": server.Quarantined,
678686
"created": server.Created,
679-
})
687+
}
688+
result = append(result, serverMap)
689+
690+
s.logger.Debug("Added quarantined server to result",
691+
zap.String("server", server.Name),
692+
zap.Bool("quarantined", server.Quarantined))
680693
}
681694

695+
s.logger.Debug("GetQuarantinedServers completed",
696+
zap.Int("total_result_count", len(result)))
697+
682698
return result, nil
683699
}
684700

@@ -722,11 +738,19 @@ func (s *Server) QuarantineServer(serverName string, quarantined bool) error {
722738
zap.String("server", serverName),
723739
zap.Bool("quarantined", quarantined))
724740

741+
s.logger.Debug("Calling storage manager QuarantineUpstreamServer",
742+
zap.String("server", serverName),
743+
zap.Bool("quarantined", quarantined))
744+
725745
if err := s.storageManager.QuarantineUpstreamServer(serverName, quarantined); err != nil {
726746
s.logger.Error("Failed to update server quarantine state in storage", zap.Error(err))
727747
return fmt.Errorf("failed to update quarantine state for server '%s' in storage: %w", serverName, err)
728748
}
729749

750+
s.logger.Debug("Successfully updated quarantine state in storage, saving configuration",
751+
zap.String("server", serverName),
752+
zap.Bool("quarantined", quarantined))
753+
730754
if err := s.SaveConfiguration(); err != nil {
731755
s.logger.Error("Failed to save configuration after quarantine state change", zap.Error(err))
732756
}

internal/storage/manager.go

Lines changed: 45 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -139,13 +139,25 @@ func (m *Manager) ListQuarantinedUpstreamServers() ([]*config.ServerConfig, erro
139139
m.mu.RLock()
140140
defer m.mu.RUnlock()
141141

142+
m.logger.Debug("ListQuarantinedUpstreamServers called")
143+
142144
records, err := m.db.ListUpstreams()
143145
if err != nil {
146+
m.logger.Errorw("Failed to list upstreams for quarantine filtering",
147+
"error", err)
144148
return nil, err
145149
}
146150

151+
m.logger.Debugw("Retrieved all upstream records for quarantine filtering",
152+
"total_records", len(records))
153+
147154
var quarantinedServers []*config.ServerConfig
148155
for _, record := range records {
156+
m.logger.Debugw("Checking server quarantine status",
157+
"server", record.Name,
158+
"quarantined", record.Quarantined,
159+
"enabled", record.Enabled)
160+
149161
if record.Quarantined {
150162
quarantinedServers = append(quarantinedServers, &config.ServerConfig{
151163
Name: record.Name,
@@ -160,9 +172,16 @@ func (m *Manager) ListQuarantinedUpstreamServers() ([]*config.ServerConfig, erro
160172
Created: record.Created,
161173
Updated: record.Updated,
162174
})
175+
176+
m.logger.Debugw("Added server to quarantined list",
177+
"server", record.Name,
178+
"total_quarantined_so_far", len(quarantinedServers))
163179
}
164180
}
165181

182+
m.logger.Debugw("ListQuarantinedUpstreamServers completed",
183+
"total_quarantined", len(quarantinedServers))
184+
166185
return quarantinedServers, nil
167186
}
168187

@@ -226,14 +245,39 @@ func (m *Manager) QuarantineUpstreamServer(name string, quarantined bool) error
226245
m.mu.Lock()
227246
defer m.mu.Unlock()
228247

248+
m.logger.Debugw("QuarantineUpstreamServer called",
249+
"server", name,
250+
"quarantined", quarantined)
251+
229252
record, err := m.db.GetUpstream(name)
230253
if err != nil {
254+
m.logger.Errorw("Failed to get upstream record for quarantine operation",
255+
"server", name,
256+
"error", err)
231257
return err
232258
}
233259

260+
m.logger.Debugw("Retrieved upstream record for quarantine",
261+
"server", name,
262+
"current_quarantined", record.Quarantined,
263+
"new_quarantined", quarantined)
264+
234265
record.Quarantined = quarantined
235266
record.Updated = time.Now()
236-
return m.db.SaveUpstream(record)
267+
268+
if err := m.db.SaveUpstream(record); err != nil {
269+
m.logger.Errorw("Failed to save quarantine status to database",
270+
"server", name,
271+
"quarantined", quarantined,
272+
"error", err)
273+
return err
274+
}
275+
276+
m.logger.Debugw("Successfully saved quarantine status to database",
277+
"server", name,
278+
"quarantined", quarantined)
279+
280+
return nil
237281
}
238282

239283
// Tool statistics operations

internal/tray/managers.go

Lines changed: 51 additions & 36 deletions
Original file line numberDiff line numberDiff line change
@@ -27,9 +27,10 @@ type ServerStateManager struct {
2727
mu sync.RWMutex
2828

2929
// Current state tracking
30-
allServers []map[string]interface{}
31-
quarantinedServers []map[string]interface{}
32-
lastUpdate time.Time
30+
allServers []map[string]interface{}
31+
quarantinedServers []map[string]interface{}
32+
lastUpdate time.Time
33+
lastQuarantineUpdate time.Time // Separate timestamp for quarantine data
3334
}
3435

3536
// NewServerStateManager creates a new server state manager
@@ -40,26 +41,29 @@ func NewServerStateManager(server ServerInterface, logger *zap.SugaredLogger) *S
4041
}
4142
}
4243

43-
// RefreshState loads the current state from the server
44+
// RefreshState forces a refresh of server state from the server
4445
func (m *ServerStateManager) RefreshState() error {
4546
m.mu.Lock()
4647
defer m.mu.Unlock()
4748

4849
// Get all servers
4950
allServers, err := m.server.GetAllServers()
5051
if err != nil {
52+
m.logger.Error("RefreshState failed to get all servers", zap.Error(err))
5153
return fmt.Errorf("failed to get all servers: %w", err)
5254
}
5355

5456
// Get quarantined servers
5557
quarantinedServers, err := m.server.GetQuarantinedServers()
5658
if err != nil {
59+
m.logger.Error("RefreshState failed to get quarantined servers", zap.Error(err))
5760
return fmt.Errorf("failed to get quarantined servers: %w", err)
5861
}
5962

6063
m.allServers = allServers
6164
m.quarantinedServers = quarantinedServers
6265
m.lastUpdate = time.Now()
66+
m.lastQuarantineUpdate = time.Now()
6367

6468
m.logger.Debug("Server state refreshed",
6569
zap.Int("all_servers", len(allServers)),
@@ -73,8 +77,8 @@ func (m *ServerStateManager) GetAllServers() ([]map[string]interface{}, error) {
7377
m.mu.RLock()
7478
defer m.mu.RUnlock()
7579

76-
// Return cached data if available and recent
77-
if time.Since(m.lastUpdate) < 2*time.Second && len(m.allServers) > 0 {
80+
// Return cached data if available and recent (only if THIS data has been loaded before)
81+
if time.Since(m.lastUpdate) < 2*time.Second && !m.lastUpdate.IsZero() && m.allServers != nil {
7882
return m.allServers, nil
7983
}
8084

@@ -91,6 +95,7 @@ func (m *ServerStateManager) GetAllServers() ([]map[string]interface{}, error) {
9195
m.logger.Debug("Database not available and no cached data, returning empty server list")
9296
return []map[string]interface{}{}, nil
9397
}
98+
m.logger.Error("Failed to get fresh all servers data", zap.Error(err))
9499
return nil, err
95100
}
96101

@@ -105,8 +110,8 @@ func (m *ServerStateManager) GetQuarantinedServers() ([]map[string]interface{},
105110
m.mu.RLock()
106111
defer m.mu.RUnlock()
107112

108-
// Return cached data if available and recent
109-
if time.Since(m.lastUpdate) < 2*time.Second && len(m.quarantinedServers) >= 0 {
113+
// Return cached data if available and recent (only if data has been loaded before)
114+
if time.Since(m.lastQuarantineUpdate) < 2*time.Second && !m.lastQuarantineUpdate.IsZero() {
110115
return m.quarantinedServers, nil
111116
}
112117

@@ -115,20 +120,21 @@ func (m *ServerStateManager) GetQuarantinedServers() ([]map[string]interface{},
115120
if err != nil {
116121
// If database is closed, return cached data if available
117122
if strings.Contains(err.Error(), "database not open") || strings.Contains(err.Error(), "closed") {
118-
if len(m.quarantinedServers) >= 0 {
123+
if len(m.quarantinedServers) > 0 {
119124
m.logger.Debug("Database not available, returning cached quarantine data")
120125
return m.quarantinedServers, nil
121126
}
122127
// No cached data available, return empty slice
123128
m.logger.Debug("Database not available and no cached data, returning empty quarantine list")
124129
return []map[string]interface{}{}, nil
125130
}
131+
m.logger.Error("Failed to get fresh quarantined servers data", zap.Error(err))
126132
return nil, err
127133
}
128134

129135
// Cache the fresh data
130136
m.quarantinedServers = servers
131-
m.lastUpdate = time.Now()
137+
m.lastQuarantineUpdate = time.Now()
132138
return servers, nil
133139
}
134140

@@ -326,17 +332,16 @@ func (m *MenuManager) UpdateQuarantineMenu(quarantinedServers []map[string]inter
326332
menuTitle := fmt.Sprintf("Security Quarantine (%d)", quarantineCount)
327333
if m.quarantineMenu != nil {
328334
m.quarantineMenu.SetTitle(menuTitle)
335+
} else {
336+
m.logger.Error("Quarantine menu is nil, cannot update title")
337+
return
329338
}
330339

331-
// --- Initialize Info Items on First Run ---
332-
if m.quarantineInfoEmpty == nil && m.quarantineMenu != nil {
333-
m.quarantineInfoEmpty = m.quarantineMenu.AddSubMenuItem("No servers quarantined", "All servers are approved")
334-
m.quarantineInfoEmpty.Disable()
335-
336-
m.quarantineInfoHelp = m.quarantineMenu.AddSubMenuItem("ℹ️ Click server to unquarantine", "Click on a quarantined server to remove it from quarantine")
337-
m.quarantineInfoHelp.Disable()
338-
339-
// Add a separator that is always visible
340+
// --- Create Info Items if Needed ---
341+
if m.quarantineInfoEmpty == nil || m.quarantineInfoHelp == nil {
342+
m.quarantineInfoEmpty = m.quarantineMenu.AddSubMenuItem("(No servers quarantined)", "No servers are currently quarantined")
343+
m.quarantineInfoHelp = m.quarantineMenu.AddSubMenuItem("Click to unquarantine", "Click on a quarantined server to remove it from quarantine")
344+
// Add empty separator for visual separation
340345
m.quarantineMenu.AddSubMenuItem("", "")
341346
}
342347

@@ -356,6 +361,8 @@ func (m *MenuManager) UpdateQuarantineMenu(quarantinedServers []map[string]inter
356361
for _, server := range quarantinedServers {
357362
if name, ok := server["name"].(string); ok {
358363
currentQuarantineMap[name] = true
364+
} else {
365+
m.logger.Warn("Quarantined server missing name field", zap.Any("server", server))
359366
}
360367
}
361368

@@ -366,7 +373,6 @@ func (m *MenuManager) UpdateQuarantineMenu(quarantinedServers []map[string]inter
366373
menuItem.Show()
367374
} else {
368375
// Server is no longer quarantined, hide it
369-
m.logger.Info("Hiding menu item for unquarantined server", zap.String("server", serverName))
370376
menuItem.Hide()
371377
}
372378
}
@@ -375,11 +381,21 @@ func (m *MenuManager) UpdateQuarantineMenu(quarantinedServers []map[string]inter
375381
for serverName := range currentQuarantineMap {
376382
if _, exists := m.quarantineMenuItems[serverName]; !exists {
377383
// This is a newly quarantined server, create its menu item
378-
m.logger.Info("Creating quarantine menu item for server", zap.String("server", serverName))
384+
if m.quarantineMenu == nil {
385+
m.logger.Error("Cannot create quarantine menu item - quarantineMenu is nil!", zap.String("server", serverName))
386+
continue
387+
}
388+
379389
quarantineMenuItem := m.quarantineMenu.AddSubMenuItem(
380390
fmt.Sprintf("🔒 %s", serverName),
381391
fmt.Sprintf("Click to unquarantine %s", serverName),
382392
)
393+
394+
if quarantineMenuItem == nil {
395+
m.logger.Error("Failed to create quarantine menu item", zap.String("server", serverName))
396+
continue
397+
}
398+
383399
m.quarantineMenuItems[serverName] = quarantineMenuItem
384400

385401
// Set up the one-time click handler
@@ -558,12 +574,12 @@ func (m *SynchronizationManager) Stop() {
558574
}
559575
}
560576

561-
// SyncNow forces an immediate synchronization
577+
// SyncNow performs immediate synchronization
562578
func (m *SynchronizationManager) SyncNow() error {
563579
return m.performSync()
564580
}
565581

566-
// SyncDelayed triggers a delayed synchronization (debounces rapid changes)
582+
// SyncDelayed schedules a delayed synchronization to batch updates
567583
func (m *SynchronizationManager) SyncDelayed() {
568584
if m.syncTimer != nil {
569585
m.syncTimer.Stop()
@@ -594,44 +610,43 @@ func (m *SynchronizationManager) syncLoop() {
594610

595611
// performSync performs the actual synchronization
596612
func (m *SynchronizationManager) performSync() error {
597-
m.logger.Debug("Performing synchronization")
598-
599613
// Check if the state manager's server is available and running
600614
// If not, skip the sync to avoid database errors
601-
if m.stateManager.server != nil && !m.stateManager.server.IsRunning() {
602-
m.logger.Debug("Server is stopped, skipping synchronization")
603-
return nil
604-
}
615+
//
616+
// FIXME: remove this if no issue with DB connection
617+
//
618+
// if m.stateManager.server != nil && !m.stateManager.server.IsRunning() {
619+
// m.logger.Debug("Server is stopped, skipping synchronization")
620+
// return nil
621+
// }
605622

606623
// Get current state with error handling for database issues
607624
allServers, err := m.stateManager.GetAllServers()
608625
if err != nil {
609626
// Check if it's a database closed error and handle gracefully
610627
if strings.Contains(err.Error(), "database not open") || strings.Contains(err.Error(), "closed") {
611-
m.logger.Debug("Database not available, skipping synchronization")
628+
m.logger.Warn("DATABASE NOT AVAILABLE - skipping synchronization", zap.Error(err))
612629
return nil
613630
}
631+
m.logger.Error("Failed to get all servers", zap.Error(err))
614632
return fmt.Errorf("failed to get all servers: %w", err)
615633
}
616634

617635
quarantinedServers, err := m.stateManager.GetQuarantinedServers()
618636
if err != nil {
619637
// Check if it's a database closed error and handle gracefully
620638
if strings.Contains(err.Error(), "database not open") || strings.Contains(err.Error(), "closed") {
621-
m.logger.Debug("Database not available for quarantine data, skipping synchronization")
639+
m.logger.Warn("DATABASE NOT AVAILABLE FOR QUARANTINE - skipping synchronization", zap.Error(err))
622640
return nil
623641
}
642+
m.logger.Error("Failed to get quarantined servers", zap.Error(err))
624643
return fmt.Errorf("failed to get quarantined servers: %w", err)
625644
}
626645

627-
// Update menus
646+
// Update menus using the manager
628647
m.menuManager.UpdateUpstreamServersMenu(allServers)
629648
m.menuManager.UpdateQuarantineMenu(quarantinedServers)
630649

631-
m.logger.Debug("Synchronization completed",
632-
zap.Int("total_servers", len(allServers)),
633-
zap.Int("quarantined_servers", len(quarantinedServers)))
634-
635650
return nil
636651
}
637652

0 commit comments

Comments
 (0)