package store import ( "fmt" "os" "path/filepath" "time" "github.com/google/uuid" "golang.org/x/crypto/bcrypt" "gorm.io/driver/sqlite" "gorm.io/gorm" "gorm.io/gorm/logger" ) type Store struct { db *gorm.DB } func Open(path string) (*Store, error) { dir := filepath.Dir(path) if err := os.MkdirAll(dir, 0755); err != nil { return nil, err } db, err := gorm.Open(sqlite.Open(path), &gorm.Config{ Logger: logger.Default.LogMode(logger.Warn), }) if err != nil { return nil, err } sqlDB, err := db.DB() if err != nil { return nil, err } sqlDB.SetMaxOpenConns(1) if _, err := sqlDB.Exec("PRAGMA journal_mode=WAL"); err != nil { return nil, fmt.Errorf("enable WAL: %w", err) } if _, err := sqlDB.Exec("PRAGMA busy_timeout=5000"); err != nil { return nil, fmt.Errorf("set busy_timeout: %w", err) } s := &Store{db: db} if err := s.migrate(); err != nil { return nil, err } if err := s.seed(); err != nil { return nil, err } return s, nil } func (s *Store) DB() *gorm.DB { return s.db } func (s *Store) migrate() error { return s.db.AutoMigrate( &User{}, &Client{}, &Tunnel{}, &Host{}, &IPRule{}, &AuthPolicy{}, &FlowStat{}, &RequestIP{}, &Session{}, &SystemSettings{}, &APIKey{}, &ResourceVisitor{}, &VisitorAccessLog{}, ) } func (s *Store) seed() error { var count int64 s.db.Model(&User{}).Count(&count) if count > 0 { return nil } hash, err := bcrypt.GenerateFromPassword([]byte("admin"), bcrypt.DefaultCost) if err != nil { return err } admin := User{ Username: "admin", PasswordHash: string(hash), Role: "admin", Status: true, } return s.db.Create(&admin).Error } func GenerateVerifyKey() string { return uuid.New().String() } func HashPassword(password string) (string, error) { hash, err := bcrypt.GenerateFromPassword([]byte(password), bcrypt.DefaultCost) if err != nil { return "", err } return string(hash), nil } func CheckPassword(hash, password string) bool { return bcrypt.CompareHashAndPassword([]byte(hash), []byte(password)) == nil } func (s *Store) ListClients(ownerID uint, isAdmin bool) ([]Client, error) { var clients []Client q := s.db.Order("id asc") if !isAdmin { q = q.Where("owner_user_id = ?", ownerID) } return clients, q.Find(&clients).Error } func (s *Store) GetClientByID(id uint) (*Client, error) { var c Client if err := s.db.First(&c, id).Error; err != nil { return nil, err } return &c, nil } func (s *Store) GetClientByVerifyKey(key string) (*Client, error) { var c Client if err := s.db.Where("verify_key = ? AND status = ?", key, true).First(&c).Error; err != nil { return nil, err } return &c, nil } func (s *Store) CreateClient(c *Client) error { if c.VerifyKey == "" { c.VerifyKey = GenerateVerifyKey() } return s.db.Create(c).Error } func (s *Store) UpdateClient(c *Client) error { return s.db.Save(c).Error } func (s *Store) DeleteClient(id uint) error { return s.db.Transaction(func(tx *gorm.DB) error { var tunnelIDs, hostIDs []uint tx.Model(&Tunnel{}).Where("client_id = ?", id).Pluck("id", &tunnelIDs) tx.Model(&Host{}).Where("client_id = ?", id).Pluck("id", &hostIDs) for _, tid := range tunnelIDs { if err := tx.Where("resource_type = ? AND resource_id = ?", "tunnel", tid).Delete(&AuthPolicy{}).Error; err != nil { return err } if err := tx.Where("resource_type = ? AND resource_id = ?", "tunnel", tid).Delete(&ResourceVisitor{}).Error; err != nil { return err } if err := tx.Where("scope = ? AND resource_id = ?", "tunnel", tid).Delete(&IPRule{}).Error; err != nil { return err } if err := tx.Where("resource_type = ? AND resource_id = ?", "tunnel", tid).Delete(&FlowStat{}).Error; err != nil { return err } if err := tx.Where("resource_type = ? AND resource_id = ?", "tunnel", tid).Delete(&RequestIP{}).Error; err != nil { return err } } for _, hid := range hostIDs { if err := tx.Where("resource_type = ? AND resource_id = ?", "host", hid).Delete(&AuthPolicy{}).Error; err != nil { return err } if err := tx.Where("resource_type = ? AND resource_id = ?", "host", hid).Delete(&ResourceVisitor{}).Error; err != nil { return err } if err := tx.Where("scope = ? AND resource_id = ?", "host", hid).Delete(&IPRule{}).Error; err != nil { return err } if err := tx.Where("resource_type = ? AND resource_id = ?", "host", hid).Delete(&FlowStat{}).Error; err != nil { return err } if err := tx.Where("resource_type = ? AND resource_id = ?", "host", hid).Delete(&RequestIP{}).Error; err != nil { return err } } if err := tx.Where("scope = ? AND resource_id = ?", "client", id).Delete(&IPRule{}).Error; err != nil { return err } if err := tx.Where("client_id = ?", id).Delete(&Tunnel{}).Error; err != nil { return err } if err := tx.Where("client_id = ?", id).Delete(&Host{}).Error; err != nil { return err } return tx.Delete(&Client{}, id).Error }) } func (s *Store) ListTunnels(clientID uint, ownerID uint, isAdmin bool) ([]Tunnel, error) { var tunnels []Tunnel q := s.db.Order("id asc") if clientID > 0 { q = q.Where("client_id = ?", clientID) } else if !isAdmin { var ids []uint s.db.Model(&Client{}).Where("owner_user_id = ?", ownerID).Pluck("id", &ids) if len(ids) == 0 { return tunnels, nil } q = q.Where("client_id IN ?", ids) } return tunnels, q.Find(&tunnels).Error } func (s *Store) GetTunnelByID(id uint) (*Tunnel, error) { var t Tunnel if err := s.db.First(&t, id).Error; err != nil { return nil, err } return &t, nil } func (s *Store) CreateTunnel(t *Tunnel) error { return s.db.Create(t).Error } func (s *Store) UpdateTunnel(t *Tunnel) error { return s.db.Save(t).Error } func (s *Store) DeleteTunnel(id uint) error { _ = s.DeleteResourceVisitors("tunnel", id) return s.db.Delete(&Tunnel{}, id).Error } func (s *Store) ListHosts(clientID uint, ownerID uint, isAdmin bool) ([]Host, error) { var hosts []Host q := s.db.Order("id asc") if clientID > 0 { q = q.Where("client_id = ?", clientID) } else if !isAdmin { var ids []uint s.db.Model(&Client{}).Where("owner_user_id = ?", ownerID).Pluck("id", &ids) if len(ids) == 0 { return hosts, nil } q = q.Where("client_id IN ?", ids) } return hosts, q.Find(&hosts).Error } func (s *Store) GetHostByID(id uint) (*Host, error) { var h Host if err := s.db.First(&h, id).Error; err != nil { return nil, err } return &h, nil } func (s *Store) CreateHost(h *Host) error { return s.db.Create(h).Error } func (s *Store) UpdateHost(h *Host) error { return s.db.Save(h).Error } func (s *Store) DeleteHost(id uint) error { _ = s.DeleteResourceVisitors("host", id) return s.db.Delete(&Host{}, id).Error } func (s *Store) ListIPRules(scope string, resourceID uint) ([]IPRule, error) { var rules []IPRule q := s.db.Order("priority desc, id asc") if scope != "" { q = q.Where("scope = ?", scope) } if resourceID > 0 { q = q.Where("resource_id = ?", resourceID) } return rules, q.Find(&rules).Error } func (s *Store) CreateIPRule(r *IPRule) error { return s.db.Create(r).Error } func (s *Store) DeleteIPRule(id uint) error { return s.db.Delete(&IPRule{}, id).Error } func (s *Store) ListAuthPolicies(resourceType string, resourceID uint) ([]AuthPolicy, error) { var policies []AuthPolicy q := s.db.Order("id desc") if resourceType != "" { q = q.Where("resource_type = ?", resourceType) } if resourceID > 0 { q = q.Where("resource_id = ?", resourceID) } return policies, q.Find(&policies).Error } func (s *Store) GetAuthPolicy(id uint) (*AuthPolicy, error) { var p AuthPolicy if err := s.db.First(&p, id).Error; err != nil { return nil, err } return &p, nil } func (s *Store) CreateAuthPolicy(p *AuthPolicy) error { return s.db.Create(p).Error } func (s *Store) UpdateAuthPolicy(p *AuthPolicy) error { return s.db.Save(p).Error } func (s *Store) DeleteAuthPolicy(id uint) error { return s.db.Delete(&AuthPolicy{}, id).Error } func (s *Store) ListUsers() ([]User, error) { var users []User return users, s.db.Order("id asc").Find(&users).Error } func (s *Store) GetUserByUsername(username string) (*User, error) { var u User if err := s.db.Where("username = ?", username).First(&u).Error; err != nil { return nil, err } return &u, nil } func (s *Store) GetUserByID(id uint) (*User, error) { var u User if err := s.db.First(&u, id).Error; err != nil { return nil, err } return &u, nil } func (s *Store) CreateUser(u *User) error { return s.db.Create(u).Error } func (s *Store) UpdateUser(u *User) error { return s.db.Save(u).Error } func (s *Store) DeleteUser(id uint, currentUserID uint) error { if id == currentUserID { return fmt.Errorf("cannot delete yourself") } user, err := s.GetUserByID(id) if err != nil { return err } if user.Role == "admin" { var adminCount int64 s.db.Model(&User{}).Where("role = ?", "admin").Count(&adminCount) if adminCount <= 1 { return fmt.Errorf("cannot delete the last admin") } } _ = s.DeleteVisitorAccessLogs(id) return s.db.Delete(&User{}, id).Error } func (s *Store) RecordRequestIP(resourceType string, resourceID uint, visitorIP, directIP string) error { var rec RequestIP err := s.db.Where("resource_type = ? AND resource_id = ? AND ip = ?", resourceType, resourceID, visitorIP).First(&rec).Error now := time.Now() if err == gorm.ErrRecordNotFound { rec = RequestIP{ ResourceType: resourceType, ResourceID: resourceID, IP: visitorIP, DirectIP: directIP, HitCount: 1, LastSeenAt: now, } return s.db.Create(&rec).Error } if err != nil { return err } return s.db.Model(&rec).Updates(map[string]interface{}{ "hit_count": rec.HitCount + 1, "last_seen_at": now, "direct_ip": directIP, }).Error } func (s *Store) ListRequestIPs(resourceType string, resourceID uint, limit, offset int) ([]RequestIP, int64, error) { var items []RequestIP var total int64 q := s.db.Model(&RequestIP{}) if resourceType != "" { q = q.Where("resource_type = ?", resourceType) } if resourceID > 0 { q = q.Where("resource_id = ?", resourceID) } q.Count(&total) err := q.Order("last_seen_at desc").Limit(limit).Offset(offset).Find(&items).Error return items, total, err } func (s *Store) UpsertFlowStat(resourceType string, resourceID uint, inlet, export int64) error { var stat FlowStat err := s.db.Where("resource_type = ? AND resource_id = ?", resourceType, resourceID).First(&stat).Error if err == gorm.ErrRecordNotFound { stat = FlowStat{ ResourceType: resourceType, ResourceID: resourceID, InletBytes: inlet, ExportBytes: export, } return s.db.Create(&stat).Error } if err != nil { return err } return s.db.Model(&stat).Updates(map[string]interface{}{ "inlet_bytes": gorm.Expr("inlet_bytes + ?", inlet), "export_bytes": gorm.Expr("export_bytes + ?", export), }).Error } func (s *Store) AddClientFlow(id uint, inlet, export int64) error { return s.db.Model(&Client{}).Where("id = ?", id).Updates(map[string]interface{}{ "inlet_flow": gorm.Expr("inlet_flow + ?", inlet), "export_flow": gorm.Expr("export_flow + ?", export), }).Error } func (s *Store) AddTunnelFlow(id uint, inlet, export int64) error { return s.db.Model(&Tunnel{}).Where("id = ?", id).Updates(map[string]interface{}{ "inlet_flow": gorm.Expr("inlet_flow + ?", inlet), "export_flow": gorm.Expr("export_flow + ?", export), }).Error } func (s *Store) AddHostFlow(id uint, inlet, export int64) error { return s.db.Model(&Host{}).Where("id = ?", id).Updates(map[string]interface{}{ "inlet_flow": gorm.Expr("inlet_flow + ?", inlet), "export_flow": gorm.Expr("export_flow + ?", export), }).Error } func (s *Store) DashboardStats() (map[string]interface{}, error) { var clientTotal, clientOnline int64 var tunnelTotal, tunnelRunning int64 var hostTotal int64 var inletFlow, exportFlow int64 s.db.Model(&Client{}).Count(&clientTotal) s.db.Model(&Tunnel{}).Count(&tunnelTotal) s.db.Model(&Tunnel{}).Where("run_status = ?", true).Count(&tunnelRunning) s.db.Model(&Host{}).Count(&hostTotal) s.db.Model(&Client{}).Select("COALESCE(SUM(inlet_flow),0)").Scan(&inletFlow) s.db.Model(&Client{}).Select("COALESCE(SUM(export_flow),0)").Scan(&exportFlow) return map[string]interface{}{ "client_total": clientTotal, "client_online": clientOnline, "tunnel_total": tunnelTotal, "tunnel_running": tunnelRunning, "host_total": hostTotal, "inlet_flow": inletFlow, "export_flow": exportFlow, }, nil } func (s *Store) TopClientsByFlow(limit int) ([]Client, error) { var clients []Client err := s.db.Order("(inlet_flow + export_flow) desc").Limit(limit).Find(&clients).Error return clients, err } func (s *Store) TopTunnelsByFlow(limit int) ([]Tunnel, error) { var tunnels []Tunnel err := s.db.Order("(inlet_flow + export_flow) desc").Limit(limit).Find(&tunnels).Error return tunnels, err } func (s *Store) TopHostsByFlow(limit int) ([]Host, error) { var hosts []Host err := s.db.Order("(inlet_flow + export_flow) desc").Limit(limit).Find(&hosts).Error return hosts, err } func (s *Store) TopRequestIPs(limit int) ([]RequestIP, error) { var items []RequestIP err := s.db.Order("hit_count desc").Limit(limit).Find(&items).Error return items, err } func ClientOwnerCheck(s *Store, clientID, userID uint, isAdmin bool) error { if isAdmin { return nil } c, err := s.GetClientByID(clientID) if err != nil { return err } if c.OwnerUserID != userID { return fmt.Errorf("permission denied") } return nil }