Fix daily token budget file permissions and race

- Add internal/flock package with portable advisory file locking
  (syscall.Flock on Unix, LockFileEx on Windows).
- Serialize CheckDailyBudget and DailyTokenUsage with a lock file.
- Write the budget file with 0600 instead of 0644 permissions.
- Add tests for flock serialization, budget file permissions, and
  concurrent budget billing.

Author: jkyberneees

go.mod

@@ -9,4 +9,4 @@ require (
 	golang.org/x/term v0.43.0
 )
 
-require golang.org/x/sys v0.44.0 // indirect
+require golang.org/x/sys v0.44.0

internal/flock/flock.go

@@ -0,0 +1,30 @@
+// Package flock provides a portable advisory file lock.
+//
+// Lock opens or creates a lock file and acquires an exclusive lock on it.
+// The returned release function must be called to unlock and close the file.
+// The lock is advisory: it only serializes callers that also use this package
+// (or otherwise cooperate on the same lock file).
+package flock
+
+import (
+	"fmt"
+	"os"
+)
+
+// Lock acquires an exclusive advisory lock on path. It creates the lock file
+// with 0600 permissions if it does not exist. The returned release function
+// unlocks and closes the lock file; callers should defer it.
+func Lock(path string) (func(), error) {
+	f, err := os.OpenFile(path, os.O_CREATE|os.O_RDWR, 0600)
+	if err != nil {
+		return nil, fmt.Errorf("flock: open: %w", err)
+	}
+	if err := lockFile(int(f.Fd())); err != nil {
+		f.Close()
+		return nil, fmt.Errorf("flock: lock: %w", err)
+	}
+	return func() {
+		unlockFile(int(f.Fd()))
+		f.Close()
+	}, nil
+}

internal/flock/flock_test.go

@@ -0,0 +1,91 @@
+package flock
+
+import (
+	"os"
+	"path/filepath"
+	"strconv"
+	"sync"
+	"testing"
+)
+
+func TestLock_AcquireAndRelease(t *testing.T) {
+	dir := t.TempDir()
+	lockPath := filepath.Join(dir, "test.lock")
+
+	release, err := Lock(lockPath)
+	if err != nil {
+		t.Fatalf("Lock: %v", err)
+	}
+
+	// Lock file should exist with restricted permissions.
+	info, err := os.Stat(lockPath)
+	if err != nil {
+		t.Fatalf("stat lock file: %v", err)
+	}
+	if info.Mode().Perm()&0077 != 0 {
+		t.Errorf("lock file is world/group accessible: %o", info.Mode().Perm())
+	}
+
+	release()
+
+	// After release, the lock file may be left behind; that's fine.
+}
+
+func TestLock_SerializesConcurrentWriters(t *testing.T) {
+	dir := t.TempDir()
+	counterPath := filepath.Join(dir, "counter")
+	lockPath := filepath.Join(dir, "counter.lock")
+
+	if err := os.WriteFile(counterPath, []byte("0"), 0600); err != nil {
+		t.Fatalf("write counter: %v", err)
+	}
+
+	var wg sync.WaitGroup
+	workers := 20
+	increments := 50
+	for i := 0; i < workers; i++ {
+		wg.Add(1)
+		go func() {
+			defer wg.Done()
+			for j := 0; j < increments; j++ {
+				release, err := Lock(lockPath)
+				if err != nil {
+					t.Errorf("Lock: %v", err)
+					return
+				}
+				data, err := os.ReadFile(counterPath)
+				if err != nil {
+					release()
+					t.Errorf("read counter: %v", err)
+					return
+				}
+				n, err := strconv.Atoi(string(data))
+				if err != nil {
+					release()
+					t.Errorf("parse counter: %v", err)
+					return
+				}
+				if err := os.WriteFile(counterPath, []byte(strconv.Itoa(n+1)), 0600); err != nil {
+					release()
+					t.Errorf("write counter: %v", err)
+					return
+				}
+				release()
+			}
+		}()
+	}
+	wg.Wait()
+
+	data, err := os.ReadFile(counterPath)
+	if err != nil {
+		t.Fatalf("read final counter: %v", err)
+	}
+	got, err := strconv.Atoi(string(data))
+	if err != nil {
+		t.Fatalf("parse final counter: %v", err)
+	}
+	want := workers * increments
+	if got != want {
+		t.Errorf("counter = %d, want %d (race detected)", got, want)
+	}
+}

internal/flock/flock_unix.go

@@ -0,0 +1,13 @@
+//go:build !windows
+
+package flock
+
+import "syscall"
+
+func lockFile(fd int) error {
+	return syscall.Flock(fd, syscall.LOCK_EX)
+}
+
+func unlockFile(fd int) error {
+	return syscall.Flock(fd, syscall.LOCK_UN)
+}

internal/flock/flock_windows.go

@@ -0,0 +1,26 @@
+//go:build windows
+
+package flock
+
+import (
+	"golang.org/x/sys/windows"
+)
+
+func lockFile(fd int) error {
+	h := windows.Handle(fd)
+	var overlapped windows.Overlapped
+	return windows.LockFileEx(
+		h,
+		windows.LOCKFILE_EXCLUSIVE_LOCK,
+		0,
+		1,
+		0,
+		&overlapped,
+	)
+}
+
+func unlockFile(fd int) error {
+	h := windows.Handle(fd)
+	var overlapped windows.Overlapped
+	return windows.UnlockFileEx(h, 0, 1, 0, &overlapped)
+}

internal/telegram/bot.go

@@ -16,6 +16,7 @@ import (
 	"sync"
 	"time"
 
+	"github.com/BackendStack21/odek/internal/flock"
 	"github.com/BackendStack21/odek/internal/transport"
 )
 
@@ -680,6 +681,9 @@ func budgetFilePath() string {
 // adds the given number of tokens, and returns an error if the total
 // exceeds the configured DailyTokenBudget. If the budget is zero (unset),
 // no check is performed and nil is returned.
+//
+// The read-modify-write cycle is protected by an advisory file lock so
+// concurrent odek processes and goroutines cannot clobber the counter.
 func (b *Bot) CheckDailyBudget(tokens int64) error {
 	if b.DailyTokenBudget <= 0 {
 		return nil // budget not configured
@@ -696,6 +700,13 @@ func (b *Bot) CheckDailyBudget(tokens int64) error {
 		return fmt.Errorf("telegram: create budget dir: %w", err)
 	}
 
+	// Serialize read-modify-write across processes.
+	release, err := flock.Lock(path + ".lock")
+	if err != nil {
+		return fmt.Errorf("telegram: lock budget file: %w", err)
+	}
+	defer release()
+
 	// Read current usage (file may not exist yet — that's fine).
 	var current int64
 	data, err := os.ReadFile(path)
@@ -715,8 +726,8 @@ func (b *Bot) CheckDailyBudget(tokens int64) error {
 		)
 	}
 
-	// Write the updated count.
-	if err := os.WriteFile(path, []byte(strconv.FormatInt(total, 10)), 0644); err != nil {
+	// Write the updated count with owner-only permissions.
+	if err := os.WriteFile(path, []byte(strconv.FormatInt(total, 10)), 0600); err != nil {
 		return fmt.Errorf("telegram: write budget file: %w", err)
 	}
 
@@ -730,6 +741,13 @@ func (b *Bot) DailyTokenUsage() (used int64, limit int64) {
 		return 0, 0
 	}
 	path := budgetFilePath()
+
+	release, err := flock.Lock(path + ".lock")
+	if err != nil {
+		return 0, b.DailyTokenBudget
+	}
+	defer release()
+
 	data, err := os.ReadFile(path)
 	if err == nil {
 		if parsed, err := strconv.ParseInt(string(data), 10, 64); err == nil {

internal/telegram/bot_test.go

@@ -11,6 +11,7 @@ import (
 	"os"
 	"path/filepath"
 	"strings"
+	"sync"
 	"testing"
 	"time"
 )
@@ -1315,6 +1316,55 @@ func TestBot_CheckDailyBudget_SequentialBillings(t *testing.T) {
 	}
 }
 
+func TestBot_CheckDailyBudget_FilePermissionsAreRestricted(t *testing.T) {
+	tmpDir := t.TempDir()
+	t.Setenv("HOME", tmpDir)
+
+	bot := NewBot("testtoken")
+	bot.SetDailyTokenBudget(10_000)
+
+	if err := bot.CheckDailyBudget(100); err != nil {
+		t.Fatalf("CheckDailyBudget: %v", err)
+	}
+
+	date := time.Now().Format("2006-01-02")
+	budgetPath := filepath.Join(tmpDir, ".odek", "telegram_token_usage_"+date)
+	info, err := os.Stat(budgetPath)
+	if err != nil {
+		t.Fatalf("stat budget file: %v", err)
+	}
+	if info.Mode().Perm()&0077 != 0 {
+		t.Errorf("budget file is world/group accessible: %o", info.Mode().Perm())
+	}
+}
+
+func TestBot_CheckDailyBudget_ConcurrentBillingsAreSafe(t *testing.T) {
+	tmpDir := t.TempDir()
+	t.Setenv("HOME", tmpDir)
+
+	bot := NewBot("testtoken")
+	bot.SetDailyTokenBudget(1_000_000)
+
+	var wg sync.WaitGroup
+	workers := 20
+	billEach := 1000
+	for i := 0; i < workers; i++ {
+		wg.Add(1)
+		go func() {
+			defer wg.Done()
+			if err := bot.CheckDailyBudget(int64(billEach)); err != nil {
+				t.Errorf("CheckDailyBudget: %v", err)
+			}
+		}()
+	}
+	wg.Wait()
+
+	used, _ := bot.DailyTokenUsage()
+	want := int64(workers * billEach)
+	if used != want {
+		t.Errorf("DailyTokenUsage = %d, want %d (race detected)", used, want)
+	}
+}
 // ---------------------------------------------------------------------------
 // DailyTokenUsage
 // ---------------------------------------------------------------------------