introduce rw lock for db, ish...

Signed-off-by: Kristoffer Dalby <[email protected]>
meson800 · Sep 19, 2023 · eff529f · eff529f
1 parent a1a3ff4
commit eff529f
Show file tree

Hide file tree

Showing 12 changed files with 368 additions and 155 deletions.
diff --git a/hscontrol/auth.go b/hscontrol/auth.go
@@ -309,7 +309,7 @@ func (h *Headscale) handleAuthKey(
 
 		machine.NodeKey = nodeKey
 		machine.AuthKeyID = uint(pak.ID)
-		err := h.db.RefreshMachine(machine, registerRequest.Expiry)
+		err := h.db.MachineSetExpiry(machine, registerRequest.Expiry)
 		if err != nil {
 			log.Error().
 				Caller().
@@ -510,7 +510,8 @@ func (h *Headscale) handleMachineLogOut(
 		Str("machine", machine.Hostname).
 		Msg("Client requested logout")
 
-	err := h.db.ExpireMachine(&machine)
+	now := time.Now()
+	err := h.db.MachineSetExpiry(&machine, now)
 	if err != nil {
 		log.Error().
 			Caller().
@@ -552,7 +553,7 @@ func (h *Headscale) handleMachineLogOut(
 	}
 
 	if machine.IsEphemeral() {
-		err = h.db.HardDeleteMachine(&machine)
+		err = h.db.DeleteMachine(&machine)
 		if err != nil {
 			log.Error().
 				Err(err).

diff --git a/hscontrol/db/api_key.go b/hscontrol/db/api_key.go
@@ -22,6 +22,9 @@ var ErrAPIKeyFailedToParse = errors.New("failed to parse ApiKey")
 func (hsdb *HSDatabase) CreateAPIKey(
 	expiration *time.Time,
 ) (string, *types.APIKey, error) {
+	hsdb.mu.Lock()
+	defer hsdb.mu.Unlock()
+
 	prefix, err := util.GenerateRandomStringURLSafe(apiPrefixLength)
 	if err != nil {
 		return "", nil, err
@@ -55,6 +58,9 @@ func (hsdb *HSDatabase) CreateAPIKey(
 
 // ListAPIKeys returns the list of ApiKeys for a user.
 func (hsdb *HSDatabase) ListAPIKeys() ([]types.APIKey, error) {
+	hsdb.mu.RLock()
+	defer hsdb.mu.RUnlock()
+
 	keys := []types.APIKey{}
 	if err := hsdb.db.Find(&keys).Error; err != nil {
 		return nil, err
@@ -65,6 +71,9 @@ func (hsdb *HSDatabase) ListAPIKeys() ([]types.APIKey, error) {
 
 // GetAPIKey returns a ApiKey for a given key.
 func (hsdb *HSDatabase) GetAPIKey(prefix string) (*types.APIKey, error) {
+	hsdb.mu.RLock()
+	defer hsdb.mu.RUnlock()
+
 	key := types.APIKey{}
 	if result := hsdb.db.First(&key, "prefix = ?", prefix); result.Error != nil {
 		return nil, result.Error
@@ -75,6 +84,9 @@ func (hsdb *HSDatabase) GetAPIKey(prefix string) (*types.APIKey, error) {
 
 // GetAPIKeyByID returns a ApiKey for a given id.
 func (hsdb *HSDatabase) GetAPIKeyByID(id uint64) (*types.APIKey, error) {
+	hsdb.mu.RLock()
+	defer hsdb.mu.RUnlock()
+
 	key := types.APIKey{}
 	if result := hsdb.db.Find(&types.APIKey{ID: id}).First(&key); result.Error != nil {
 		return nil, result.Error
@@ -86,6 +98,9 @@ func (hsdb *HSDatabase) GetAPIKeyByID(id uint64) (*types.APIKey, error) {
 // DestroyAPIKey destroys a ApiKey. Returns error if the ApiKey
 // does not exist.
 func (hsdb *HSDatabase) DestroyAPIKey(key types.APIKey) error {
+	hsdb.mu.Lock()
+	defer hsdb.mu.Unlock()
+
 	if result := hsdb.db.Unscoped().Delete(key); result.Error != nil {
 		return result.Error
 	}
@@ -95,6 +110,9 @@ func (hsdb *HSDatabase) DestroyAPIKey(key types.APIKey) error {
 
 // ExpireAPIKey marks a ApiKey as expired.
 func (hsdb *HSDatabase) ExpireAPIKey(key *types.APIKey) error {
+	hsdb.mu.Lock()
+	defer hsdb.mu.Unlock()
+
 	if err := hsdb.db.Model(&key).Update("Expiration", time.Now()).Error; err != nil {
 		return err
 	}
@@ -103,6 +121,9 @@ func (hsdb *HSDatabase) ExpireAPIKey(key *types.APIKey) error {
 }
 
 func (hsdb *HSDatabase) ValidateAPIKey(keyStr string) (bool, error) {
+	hsdb.mu.RLock()
+	defer hsdb.mu.RUnlock()
+
 	prefix, hash, found := strings.Cut(keyStr, ".")
 	if !found {
 		return false, ErrAPIKeyFailedToParse

diff --git a/hscontrol/db/db.go b/hscontrol/db/db.go
@@ -40,6 +40,8 @@ type HSDatabase struct {
 	db       *gorm.DB
 	notifier *notifier.Notifier
 
+	mu sync.RWMutex
+
 	ipAllocationMutex sync.Mutex
 
 	ipPrefixes []netip.Prefix