diff --git a/cmd/tg/config.go b/cmd/tg/config.go index 1e5f87c..79aef24 100644 --- a/cmd/tg/config.go +++ b/cmd/tg/config.go @@ -10,6 +10,8 @@ import ( "github.com/go-faster/errors" "gopkg.in/yaml.v3" + + "github.com/gotd/cli/internal/atomicfile" ) // defaultAccount is the label of the top-level (legacy) account. @@ -108,39 +110,47 @@ func loadConfig(path string) (Config, error) { return cfg, nil } -// writeConfig writes cfg to path, refusing to overwrite an existing file. -func writeConfig(path string, cfg Config) error { +// encodeConfig renders cfg as YAML. +func encodeConfig(cfg Config) ([]byte, error) { buf := new(bytes.Buffer) e := yaml.NewEncoder(buf) e.SetIndent(2) if err := e.Encode(cfg); err != nil { - return errors.Wrap(err, "encode") + return nil, errors.Wrap(err, "encode") } + return buf.Bytes(), nil +} - if _, err := os.Stat(path); err == nil { - return errors.Errorf("file %s exists", path) +// writeConfig writes cfg to path, refusing to overwrite an existing file. +func writeConfig(path string, cfg Config) error { + data, err := encodeConfig(cfg) + if err != nil { + return err } if err := os.MkdirAll(filepath.Dir(path), 0o700); err != nil { return err } - if err := os.WriteFile(path, buf.Bytes(), 0o600); err != nil { + // Created atomically: a crash cannot leave a partial config behind that a + // retried `tg init` would then refuse to replace. + if err := atomicfile.Create(path, data, 0o600); err != nil { + if errors.Is(err, os.ErrExist) { + return errors.Errorf("file %s exists", path) + } return errors.Wrap(err, "write") } return nil } -// saveConfig writes cfg to path, overwriting any existing file. +// saveConfig atomically writes cfg to path, overwriting any existing file. func saveConfig(path string, cfg Config) error { - buf := new(bytes.Buffer) - e := yaml.NewEncoder(buf) - e.SetIndent(2) - if err := e.Encode(cfg); err != nil { - return errors.Wrap(err, "encode") + data, err := encodeConfig(cfg) + if err != nil { + return err } if err := os.MkdirAll(filepath.Dir(path), 0o700); err != nil { return err } - if err := os.WriteFile(path, buf.Bytes(), 0o600); err != nil { + if err := atomicfile.WriteFile(path, data, 0o600); err != nil { return errors.Wrap(err, "write") } return nil diff --git a/cmd/tg/session_store.go b/cmd/tg/session_store.go index 283a922..baed9c9 100644 --- a/cmd/tg/session_store.go +++ b/cmd/tg/session_store.go @@ -2,10 +2,17 @@ package main import ( "context" + "encoding/json" "os" "path/filepath" + "sync" + + "github.com/go-faster/errors" "github.com/gotd/td/session" + + "github.com/gotd/cli/internal/atomicfile" + "github.com/gotd/cli/internal/filelock" ) // sessionService is the Keychain service name under which sessions are stored. @@ -22,14 +29,96 @@ type sessionStore interface { Delete(ctx context.Context) error } -// fileSessionStore is the cross-platform backend: gotd's JSON FileStorage plus -// existence/deletion over the same path. +// fileSessionStore is the cross-platform backend: the same JSON file format as +// gotd's session.FileStorage, but safe to share between tg processes, which +// agents routinely run in parallel on one account. +// +// Guarantees: +// - Readers never see a torn file: writes replace it atomically. +// - Writes are serialized across processes by a lock on ".lock". +// - Processes using the same auth key overwrite each other's copy (last +// writer wins). That is harmless: the copies differ only in volatile state +// such as the server salt, which Telegram corrects on the next request. +// - A process never writes its auth key over a session that another process +// re-keyed (tg login) or removed (tg logout) after this one loaded it. +// gotd re-saves the session on every reconnect, so without this a +// long-running `tg watch` would undo a login made elsewhere. type fileSessionStore struct { - *session.FileStorage + path string + mu sync.Mutex + // keyID is the auth key id this process uses: from the first session it + // loaded, then from each session it stored ("" until known). + keyID string +} + +func (s *fileSessionStore) LoadSession(context.Context) ([]byte, error) { + s.mu.Lock() + defer s.mu.Unlock() + + data, err := os.ReadFile(s.path) + if os.IsNotExist(err) { + return nil, session.ErrNotFound + } + if err != nil { + return nil, errors.Wrap(err, "read session") + } + if s.keyID == "" { + s.keyID = sessionKeyID(data) + } + return data, nil +} + +func (s *fileSessionStore) StoreSession(_ context.Context, data []byte) (rerr error) { + s.mu.Lock() + defer s.mu.Unlock() + + unlock, err := filelock.Lock(s.path + ".lock") + if err != nil { + return errors.Wrap(err, "lock session") + } + defer func() { + if err := unlock(); err != nil && rerr == nil { + rerr = errors.Wrap(err, "unlock session") + } + }() + + newKey := sessionKeyID(data) + if s.keyID != "" && newKey == s.keyID { + current, err := os.ReadFile(s.path) + if err != nil && !os.IsNotExist(err) { + return errors.Wrap(err, "read session") + } + if sessionKeyID(current) != s.keyID { + // Another process logged in again or out since this one loaded + // the session: its state is newer, keep it. + return nil + } + } + if err := atomicfile.WriteFile(s.path, data, 0o600); err != nil { + return errors.Wrap(err, "store session") + } + if newKey != "" { + s.keyID = newKey + } + return nil +} + +// sessionKeyID extracts the auth key id from a gotd session file; "" when the +// data is empty or not in gotd's format. +func sessionKeyID(data []byte) string { + var v struct { + Data struct { + AuthKeyID []byte + } + } + if len(data) == 0 || json.Unmarshal(data, &v) != nil { + return "" + } + return string(v.Data.AuthKeyID) } func (s *fileSessionStore) Exists(context.Context) (bool, error) { - switch _, err := os.Stat(s.Path); { + switch _, err := os.Stat(s.path); { case err == nil: return true, nil case os.IsNotExist(err): @@ -40,7 +129,7 @@ func (s *fileSessionStore) Exists(context.Context) (bool, error) { } func (s *fileSessionStore) Delete(context.Context) error { - if err := os.Remove(s.Path); err != nil && !os.IsNotExist(err) { + if err := os.Remove(s.path); err != nil && !os.IsNotExist(err) { return err } return nil @@ -65,7 +154,7 @@ func newSessionStore(dir, label string, acc Account, kind string, useKeychain bo return ks } } - return &fileSessionStore{FileStorage: &session.FileStorage{Path: path}} + return &fileSessionStore{path: path} } // useKeychain reports whether the Keychain backend is enabled. It defaults to diff --git a/cmd/tg/session_store_test.go b/cmd/tg/session_store_test.go index bd7e827..ab1be37 100644 --- a/cmd/tg/session_store_test.go +++ b/cmd/tg/session_store_test.go @@ -5,6 +5,8 @@ import ( "os" "path/filepath" "testing" + + "github.com/gotd/td/session" ) func TestUseKeychain(t *testing.T) { @@ -76,3 +78,88 @@ func TestFileSessionStore(t *testing.T) { t.Fatal("session file should be gone after delete") } } + +// saveTestSession writes a gotd-format session with the given auth key id +// through store, the way the gotd client does. +func saveTestSession(t *testing.T, store session.Storage, keyID, salt int64) { + t.Helper() + data := &session.Data{AuthKeyID: []byte{byte(keyID)}, AuthKey: []byte{1, 2, 3}, Salt: salt, DC: 2} + if err := (&session.Loader{Storage: store}).Save(context.Background(), data); err != nil { + t.Fatal(err) + } +} + +func loadTestSession(t *testing.T, store session.Storage) *session.Data { + t.Helper() + data, err := (&session.Loader{Storage: store}).Load(context.Background()) + if err != nil { + t.Fatal(err) + } + return data +} + +func TestFileSessionStoreKeepsNewerLogin(t *testing.T) { + dir := t.TempDir() + open := func() sessionStore { return newSessionStore(dir, "work", Account{}, kindUser, false) } + + saveTestSession(t, open(), 1, 100) // logged in with key 1 + watcher := open() // a long-running tg process + loadTestSession(t, watcher) + + saveTestSession(t, open(), 2, 200) // tg login elsewhere: key 2 + + // The watcher reconnects; gotd re-saves the session with its own key. + saveTestSession(t, watcher, 1, 101) + + if got := loadTestSession(t, open()); got.AuthKeyID[0] != 2 { + t.Errorf("auth key = %d, want the newer login (2) kept", got.AuthKeyID[0]) + } +} + +func TestFileSessionStoreDoesNotResurrectLogout(t *testing.T) { + ctx := context.Background() + dir := t.TempDir() + open := func() sessionStore { return newSessionStore(dir, "work", Account{}, kindUser, false) } + + saveTestSession(t, open(), 1, 100) + watcher := open() + loadTestSession(t, watcher) + + if err := open().Delete(ctx); err != nil { // tg logout elsewhere + t.Fatal(err) + } + saveTestSession(t, watcher, 1, 101) + + if ok, err := open().Exists(ctx); err != nil || ok { + t.Errorf("Exists() = %v, %v; a logged-out session must stay removed", ok, err) + } +} + +func TestFileSessionStoreSameKeyLastWriterWins(t *testing.T) { + dir := t.TempDir() + open := func() sessionStore { return newSessionStore(dir, "work", Account{}, kindUser, false) } + + saveTestSession(t, open(), 1, 100) + a, b := open(), open() + loadTestSession(t, a) + loadTestSession(t, b) + + saveTestSession(t, b, 1, 200) + saveTestSession(t, a, 1, 300) // same key: a's newer state wins + + if got := loadTestSession(t, open()); got.Salt != 300 { + t.Errorf("salt = %d, want 300 (last writer)", got.Salt) + } +} + +func TestFileSessionStoreOwnRekeyIsWritten(t *testing.T) { + dir := t.TempDir() + store := newSessionStore(dir, "work", Account{}, kindUser, false) + + saveTestSession(t, store, 1, 100) // fresh login in this process + saveTestSession(t, store, 2, 200) // this process re-keys (e.g. tg login again) + + if got := loadTestSession(t, newSessionStore(dir, "work", Account{}, kindUser, false)); got.AuthKeyID[0] != 2 { + t.Errorf("auth key = %d, want this process's new key (2)", got.AuthKeyID[0]) + } +} diff --git a/go.mod b/go.mod index 53e85eb..69af781 100644 --- a/go.mod +++ b/go.mod @@ -15,6 +15,7 @@ require ( go.uber.org/zap v1.28.0 golang.org/x/net v0.56.0 golang.org/x/sync v0.22.0 + golang.org/x/sys v0.46.0 gopkg.in/yaml.v3 v3.0.1 ) @@ -58,7 +59,6 @@ require ( golang.org/x/crypto v0.53.0 // indirect golang.org/x/exp v0.0.0-20230725093048-515e97ebf090 // indirect golang.org/x/mod v0.37.0 // indirect - golang.org/x/sys v0.46.0 // indirect golang.org/x/term v0.44.0 // indirect golang.org/x/text v0.38.0 // indirect golang.org/x/tools v0.47.0 // indirect diff --git a/internal/atomicfile/atomicfile.go b/internal/atomicfile/atomicfile.go new file mode 100644 index 0000000..9168021 --- /dev/null +++ b/internal/atomicfile/atomicfile.go @@ -0,0 +1,126 @@ +// Package atomicfile replaces files so that concurrent readers never observe a +// partially written file. +// +// Several tg processes commonly share one session directory (an agent running +// commands in parallel, or two agents driving the same account). A plain +// os.WriteFile truncates the file before writing it, so a reader that opens it +// mid-write sees an empty or torn file and fails to parse it. Writing to a +// temporary file in the same directory and renaming it over the target makes +// the replacement atomic: readers see either the old or the new contents. +package atomicfile + +import ( + "os" + "path/filepath" + + "github.com/go-faster/errors" +) + +// WriteFile atomically replaces path with data. The data is synced before the +// rename and the directory after it, so after a crash the file holds either the +// old or the new contents, never a truncated mix, and a completed write +// survives power loss. Use it for state such as sessions and config. +func WriteFile(path string, data []byte, perm os.FileMode) error { + return write(path, data, perm, true) +} + +// Create atomically creates path with data, failing with an error satisfying +// errors.Is(err, os.ErrExist) if path already exists. Like WriteFile, the file +// appears complete or not at all, even after a crash. +func Create(path string, data []byte, perm os.FileMode) error { + tmp, err := writeTemp(path, data, perm, true) + if err != nil { + return err + } + defer func() { _ = os.Remove(tmp) }() + + // A hard link is the portable "rename unless the target exists". + if err := os.Link(tmp, path); err != nil { + if errors.Is(err, os.ErrExist) { + return errors.Wrapf(os.ErrExist, "create %s", path) + } + return createExclusive(path, data, perm) // no hard links on this filesystem + } + return syncDir(filepath.Dir(path)) +} + +// createExclusive is Create without atomicity, for filesystems that do not +// support hard links: the file is created exclusively and synced. +func createExclusive(path string, data []byte, perm os.FileMode) error { + f, err := os.OpenFile(path, os.O_WRONLY|os.O_CREATE|os.O_EXCL, perm) // #nosec G304 // caller-provided path + if err != nil { + return errors.Wrap(err, "create") + } + if _, err := f.Write(data); err != nil { + _ = f.Close() + return errors.Wrap(err, "write") + } + if err := f.Sync(); err != nil { + _ = f.Close() + return errors.Wrap(err, "sync") + } + if err := f.Close(); err != nil { + return errors.Wrap(err, "close") + } + return syncDir(filepath.Dir(path)) +} + +// Replace is WriteFile without the fsync: readers still never observe a +// partial file, but a crash may lose the most recent write. Use it for caches +// that are rewritten often and can be rebuilt. +func Replace(path string, data []byte, perm os.FileMode) error { + return write(path, data, perm, false) +} + +func write(path string, data []byte, perm os.FileMode, durable bool) error { + // Replace the file a symlink points at rather than the symlink itself, so + // a config managed as a dotfile symlink stays a symlink. + if resolved, err := filepath.EvalSymlinks(path); err == nil { + path = resolved + } + + tmp, err := writeTemp(path, data, perm, durable) + if err != nil { + return err + } + if err := os.Rename(tmp, path); err != nil { + _ = os.Remove(tmp) + return errors.Wrap(err, "rename temp file") + } + if durable { + return syncDir(filepath.Dir(path)) + } + return nil +} + +// writeTemp writes data to a new temporary file next to path and returns its +// name; with durable set the data is synced to disk. +func writeTemp(path string, data []byte, perm os.FileMode, durable bool) (_ string, rerr error) { + f, err := os.CreateTemp(filepath.Dir(path), "."+filepath.Base(path)+".tmp-*") + if err != nil { + return "", errors.Wrap(err, "create temp file") + } + tmp := f.Name() + defer func() { + if rerr != nil { + _ = f.Close() + _ = os.Remove(tmp) + } + }() + + if err := f.Chmod(perm); err != nil { + return "", errors.Wrap(err, "chmod temp file") + } + if _, err := f.Write(data); err != nil { + return "", errors.Wrap(err, "write temp file") + } + if durable { + if err := f.Sync(); err != nil { + return "", errors.Wrap(err, "sync temp file") + } + } + if err := f.Close(); err != nil { + return "", errors.Wrap(err, "close temp file") + } + return tmp, nil +} diff --git a/internal/atomicfile/atomicfile_test.go b/internal/atomicfile/atomicfile_test.go new file mode 100644 index 0000000..7e9a4aa --- /dev/null +++ b/internal/atomicfile/atomicfile_test.go @@ -0,0 +1,116 @@ +package atomicfile + +import ( + "errors" + "os" + "path/filepath" + "runtime" + "testing" +) + +func TestWriteFileReplaces(t *testing.T) { + for name, write := range map[string]func(string, []byte, os.FileMode) error{ + "WriteFile": WriteFile, + "Replace": Replace, + } { + t.Run(name, func(t *testing.T) { + dir := t.TempDir() + path := filepath.Join(dir, "state.json") + if err := os.WriteFile(path, []byte("old contents"), 0o600); err != nil { + t.Fatal(err) + } + + if err := write(path, []byte("new"), 0o600); err != nil { + t.Fatal(err) + } + + got, err := os.ReadFile(path) + if err != nil { + t.Fatal(err) + } + if string(got) != "new" { + t.Errorf("contents = %q, want %q", got, "new") + } + entries, err := os.ReadDir(dir) + if err != nil { + t.Fatal(err) + } + if len(entries) != 1 { + t.Errorf("leftover files: %v", entries) + } + }) + } +} + +func TestWriteFilePerm(t *testing.T) { + if runtime.GOOS == "windows" { + t.Skip("POSIX permissions") + } + path := filepath.Join(t.TempDir(), "session.json") + if err := WriteFile(path, []byte("{}"), 0o600); err != nil { + t.Fatal(err) + } + info, err := os.Stat(path) + if err != nil { + t.Fatal(err) + } + if perm := info.Mode().Perm(); perm != 0o600 { + t.Errorf("perm = %o, want 600", perm) + } +} + +func TestWriteFileFollowsSymlink(t *testing.T) { + if runtime.GOOS == "windows" { + t.Skip("symlinks need privileges on windows") + } + dir := t.TempDir() + target := filepath.Join(dir, "real.yaml") + link := filepath.Join(dir, "link.yaml") + if err := os.WriteFile(target, []byte("a"), 0o600); err != nil { + t.Fatal(err) + } + if err := os.Symlink(target, link); err != nil { + t.Fatal(err) + } + + if err := WriteFile(link, []byte("b"), 0o600); err != nil { + t.Fatal(err) + } + + if fi, err := os.Lstat(link); err != nil || fi.Mode()&os.ModeSymlink == 0 { + t.Fatalf("link is no longer a symlink (err=%v)", err) + } + if got, _ := os.ReadFile(target); string(got) != "b" { + t.Errorf("target = %q, want %q", got, "b") + } +} + +func TestWriteFileMissingDir(t *testing.T) { + path := filepath.Join(t.TempDir(), "missing", "file") + if err := WriteFile(path, []byte("x"), 0o600); err == nil { + t.Fatal("expected an error for a missing directory") + } +} + +func TestCreate(t *testing.T) { + dir := t.TempDir() + path := filepath.Join(dir, "config.yaml") + + if err := Create(path, []byte("first"), 0o600); err != nil { + t.Fatal(err) + } + err := Create(path, []byte("second"), 0o600) + if !errors.Is(err, os.ErrExist) { + t.Fatalf("second Create err = %v, want os.ErrExist", err) + } + if got, _ := os.ReadFile(path); string(got) != "first" { + t.Errorf("contents = %q, want the first write kept", got) + } + entries, err := os.ReadDir(dir) + if err != nil { + t.Fatal(err) + } + if len(entries) != 1 { + t.Errorf("leftover files: %v", entries) + } +} diff --git a/internal/atomicfile/syncdir.go b/internal/atomicfile/syncdir.go new file mode 100644 index 0000000..c004616 --- /dev/null +++ b/internal/atomicfile/syncdir.go @@ -0,0 +1,22 @@ +//go:build !windows + +package atomicfile + +import ( + "os" + + "github.com/go-faster/errors" +) + +// syncDir makes a rename or link in dir durable. +func syncDir(dir string) error { + d, err := os.Open(dir) // #nosec G304 // directory of a caller-provided path + if err != nil { + return errors.Wrap(err, "open directory") + } + defer func() { _ = d.Close() }() + if err := d.Sync(); err != nil { + return errors.Wrap(err, "sync directory") + } + return nil +} diff --git a/internal/atomicfile/syncdir_windows.go b/internal/atomicfile/syncdir_windows.go new file mode 100644 index 0000000..c2a2074 --- /dev/null +++ b/internal/atomicfile/syncdir_windows.go @@ -0,0 +1,7 @@ +//go:build windows + +package atomicfile + +// syncDir is a no-op: Windows cannot sync a directory handle, and NTFS +// journals metadata changes such as renames itself. +func syncDir(string) error { return nil } diff --git a/internal/filelock/filelock.go b/internal/filelock/filelock.go new file mode 100644 index 0000000..c08cad2 --- /dev/null +++ b/internal/filelock/filelock.go @@ -0,0 +1,39 @@ +// Package filelock serializes read-modify-write cycles on files shared by +// several tg processes. +// +// Replacing a file atomically (see internal/atomicfile) keeps readers from +// seeing a torn file, but two processes that both read, modify and replace the +// same file still lose one another's changes. Holding an exclusive lock on a +// sidecar lock file for the whole cycle makes the cycles run one at a time. +// +// The lock is advisory (only cooperating tg processes honor it) and is released +// by the operating system if the holder dies, so a crash never leaves it stuck. +package filelock + +import ( + "os" + + "github.com/go-faster/errors" +) + +// Lock blocks until it holds an exclusive lock on path, creating the file if +// needed, and returns the function that releases it. Lock files are left in +// place: removing one while another process waits on it would break mutual +// exclusion. +func Lock(path string) (unlock func() error, err error) { + f, err := os.OpenFile(path, os.O_RDWR|os.O_CREATE, 0o600) // #nosec G304 // path derived from the config dir + if err != nil { + return nil, errors.Wrap(err, "open lock file") + } + if err := lockFile(f); err != nil { + _ = f.Close() + return nil, errors.Wrap(err, "lock") + } + return func() error { + uerr := unlockFile(f) + if cerr := f.Close(); uerr == nil { + uerr = cerr + } + return uerr + }, nil +} diff --git a/internal/filelock/filelock_other.go b/internal/filelock/filelock_other.go new file mode 100644 index 0000000..8e0f95e --- /dev/null +++ b/internal/filelock/filelock_other.go @@ -0,0 +1,11 @@ +//go:build !unix && !windows + +package filelock + +import "os" + +// Platforms without file locking (js, wasip1, plan9) run tg as a single +// process; locking is a no-op there. +func lockFile(*os.File) error { return nil } + +func unlockFile(*os.File) error { return nil } diff --git a/internal/filelock/filelock_test.go b/internal/filelock/filelock_test.go new file mode 100644 index 0000000..34aadcc --- /dev/null +++ b/internal/filelock/filelock_test.go @@ -0,0 +1,62 @@ +package filelock + +import ( + "path/filepath" + "sync" + "testing" + "time" +) + +func TestLockIsExclusive(t *testing.T) { + path := filepath.Join(t.TempDir(), "state.lock") + + // Separate Lock calls open separate file descriptions, so they exclude + // each other just as separate processes do. + var ( + mu sync.Mutex + inside int + maxInner int + ) + var wg sync.WaitGroup + for range 8 { + wg.Add(1) + go func() { + defer wg.Done() + unlock, err := Lock(path) + if err != nil { + t.Error(err) + return + } + mu.Lock() + inside++ + maxInner = max(maxInner, inside) + mu.Unlock() + + time.Sleep(5 * time.Millisecond) + + mu.Lock() + inside-- + mu.Unlock() + if err := unlock(); err != nil { + t.Error(err) + } + }() + } + wg.Wait() + if maxInner != 1 { + t.Errorf("%d holders at once, want 1", maxInner) + } +} + +func TestLockReusable(t *testing.T) { + path := filepath.Join(t.TempDir(), "state.lock") + for range 3 { + unlock, err := Lock(path) + if err != nil { + t.Fatal(err) + } + if err := unlock(); err != nil { + t.Fatal(err) + } + } +} diff --git a/internal/filelock/filelock_unix.go b/internal/filelock/filelock_unix.go new file mode 100644 index 0000000..63ddd0d --- /dev/null +++ b/internal/filelock/filelock_unix.go @@ -0,0 +1,21 @@ +//go:build unix + +package filelock + +import ( + "os" + "syscall" +) + +func lockFile(f *os.File) error { + for { + err := syscall.Flock(int(f.Fd()), syscall.LOCK_EX) // #nosec G115 // fd fits in int + if err != syscall.EINTR { + return err + } + } +} + +func unlockFile(f *os.File) error { + return syscall.Flock(int(f.Fd()), syscall.LOCK_UN) // #nosec G115 // fd fits in int +} diff --git a/internal/filelock/filelock_windows.go b/internal/filelock/filelock_windows.go new file mode 100644 index 0000000..0374da3 --- /dev/null +++ b/internal/filelock/filelock_windows.go @@ -0,0 +1,18 @@ +//go:build windows + +package filelock + +import ( + "os" + + "golang.org/x/sys/windows" +) + +// Lock the first byte; the lock file holds no data. +func lockFile(f *os.File) error { + return windows.LockFileEx(windows.Handle(f.Fd()), windows.LOCKFILE_EXCLUSIVE_LOCK, 0, 1, 0, new(windows.Overlapped)) +} + +func unlockFile(f *os.File) error { + return windows.UnlockFileEx(windows.Handle(f.Fd()), 0, 1, 0, new(windows.Overlapped)) +} diff --git a/internal/peercache/peercache.go b/internal/peercache/peercache.go index 98ac5f9..34aa116 100644 --- a/internal/peercache/peercache.go +++ b/internal/peercache/peercache.go @@ -4,6 +4,12 @@ // access-hashes would be lost between CLI invocations. This storage persists // them as JSON in the session directory, keyed per account, so peer resolution // is cheap and survives restarts. +// +// Several tg processes may share one cache file. Each write re-reads the file, +// merges in entries other processes added, and replaces it atomically, all +// under an exclusive lock on a sidecar ".lock" (internal/filelock), so +// concurrent commands neither corrupt the file nor drop each other's access +// hashes (re-resolving a username costs a rate-limited RPC). package peercache import ( @@ -16,6 +22,9 @@ import ( "github.com/go-faster/errors" "github.com/gotd/td/telegram/peers" + + "github.com/gotd/cli/internal/atomicfile" + "github.com/gotd/cli/internal/filelock" ) // Storage is a JSON-file-backed peers.Storage. @@ -37,48 +46,84 @@ type data struct { var _ peers.Storage = (*Storage)(nil) -// Open loads the cache at path, creating an empty one if it does not exist. +// Open loads the cache at path, creating an empty one if it does not exist. An +// unreadable (corrupt) cache is treated as empty and rewritten on the next +// save: it only holds data that can be fetched again. func Open(path string) (*Storage, error) { - s := &Storage{ - path: path, - data: data{ - Peers: map[string]int64{}, - Phones: map[string]string{}, - }, + d, err := load(path) + if err != nil { + return nil, err } + return &Storage{path: path, data: d}, nil +} +// load reads the cache file at path. A missing or corrupt file yields an empty +// cache; only I/O errors are returned. +func load(path string) (data, error) { + d := data{ + Peers: map[string]int64{}, + Phones: map[string]string{}, + } raw, err := os.ReadFile(path) // #nosec G304 // path derived from config dir if err != nil { if os.IsNotExist(err) { - return s, nil + return d, nil } - return nil, errors.Wrap(err, "read peer cache") + return d, errors.Wrap(err, "read peer cache") } - if len(raw) > 0 { - if err := json.Unmarshal(raw, &s.data); err != nil { - return nil, errors.Wrap(err, "parse peer cache") - } + var onDisk data + if json.Unmarshal(raw, &onDisk) != nil { + return d, nil } - if s.data.Peers == nil { - s.data.Peers = map[string]int64{} + for k, v := range onDisk.Peers { + d.Peers[k] = v } - if s.data.Phones == nil { - s.data.Phones = map[string]string{} + for k, v := range onDisk.Phones { + d.Phones[k] = v } - return s, nil + d.ContactsHash = onDisk.ContactsHash + return d, nil } func keyString(k peers.Key) string { return k.Prefix + ":" + strconv.FormatInt(k.ID, 10) } -// flush persists the cache; caller must hold mu. -func (s *Storage) flush() error { +// flush persists the cache; caller must hold mu. Under the cross-process file +// lock, entries another process saved since this cache was loaded are merged +// in first (this process's values win for keys both know), so concurrent +// writers do not drop each other's entries. +func (s *Storage) flush() (rerr error) { + unlock, err := filelock.Lock(s.path + ".lock") + if err != nil { + return errors.Wrap(err, "lock peer cache") + } + defer func() { + if err := unlock(); err != nil && rerr == nil { + rerr = errors.Wrap(err, "unlock peer cache") + } + }() + + onDisk, err := load(s.path) + if err != nil { + return err + } + for k, v := range onDisk.Peers { + if _, ok := s.data.Peers[k]; !ok { + s.data.Peers[k] = v + } + } + for k, v := range onDisk.Phones { + if _, ok := s.data.Phones[k]; !ok { + s.data.Phones[k] = v + } + } + raw, err := json.Marshal(s.data) if err != nil { return errors.Wrap(err, "marshal peer cache") } - if err := os.WriteFile(s.path, raw, 0o600); err != nil { + if err := atomicfile.Replace(s.path, raw, 0o600); err != nil { return errors.Wrap(err, "write peer cache") } return nil @@ -100,6 +145,14 @@ func (s *Storage) Find(_ context.Context, key peers.Key) (peers.Value, bool, err return peers.Value{AccessHash: hash}, ok, nil } +// Key prefixes of gotd's peers.Storage entries (":"); they are +// unexported in telegram/peers/storage.go, so they are mirrored here. +const ( + prefixUsers = "users_" + prefixChannel = "channel_" + prefixChats = "chats_" +) + // Peer kinds returned by Kind, matching the type names used in CLI output. const ( KindUser = "user" @@ -110,17 +163,14 @@ const ( // Kind reports the cached peer kind ("user", "chat" or "channel") for a numeric // id, if the peer has been seen before (e.g. via `tg chats list`). Users and // channels are checked before chats so a populated access hash wins. -// -// gotd's peers.Storage keys entries as ":"; the prefixes are -// unexported there (telegram/peers/storage.go), so we mirror them here. func (s *Storage) Kind(id int64) (string, bool) { s.mu.Lock() defer s.mu.Unlock() idStr := strconv.FormatInt(id, 10) for _, kp := range []struct{ kind, prefix string }{ - {KindUser, "users_"}, - {KindChannel, "channel_"}, - {KindChat, "chats_"}, + {KindUser, prefixUsers}, + {KindChannel, prefixChannel}, + {KindChat, prefixChats}, } { if _, ok := s.data.Peers[kp.prefix+":"+idStr]; ok { return kp.kind, true diff --git a/internal/peercache/peercache_test.go b/internal/peercache/peercache_test.go index 01a8a93..4c85d1a 100644 --- a/internal/peercache/peercache_test.go +++ b/internal/peercache/peercache_test.go @@ -2,7 +2,9 @@ package peercache import ( "context" + "os" "path/filepath" + "sync" "testing" "github.com/gotd/td/telegram/peers" @@ -66,9 +68,9 @@ func TestKind(t *testing.T) { t.Fatal(err) } } - save("users_", 42, 111) - save("channel_", 7, 222) - save("chats_", 9, 0) // basic groups have no access hash + save(prefixUsers, 42, 111) + save(prefixChannel, 7, 222) + save(prefixChats, 9, 0) // basic groups have no access hash cases := []struct { id int64 @@ -98,3 +100,102 @@ func TestFindMissing(t *testing.T) { t.Error("expected not found") } } + +func TestConcurrentWritersMerge(t *testing.T) { + ctx := context.Background() + path := filepath.Join(t.TempDir(), "peers.json") + + // Two processes open the same (empty) cache, then each saves a different peer. + a, err := Open(path) + if err != nil { + t.Fatal(err) + } + b, err := Open(path) + if err != nil { + t.Fatal(err) + } + if err := a.Save(ctx, peers.Key{Prefix: prefixUsers, ID: 1}, peers.Value{AccessHash: 11}); err != nil { + t.Fatal(err) + } + if err := b.Save(ctx, peers.Key{Prefix: prefixUsers, ID: 2}, peers.Value{AccessHash: 22}); err != nil { + t.Fatal(err) + } + + c, err := Open(path) + if err != nil { + t.Fatal(err) + } + for id, want := range map[int64]int64{1: 11, 2: 22} { + v, found, err := c.Find(ctx, peers.Key{Prefix: prefixUsers, ID: id}) + if err != nil || !found || v.AccessHash != want { + t.Errorf("peer %d: Find = %+v, %v, %v; want hash %d", id, v, found, err, want) + } + } +} + +func TestCorruptCacheIsRebuilt(t *testing.T) { + ctx := context.Background() + path := filepath.Join(t.TempDir(), "peers.json") + // A torn write from an older version: truncated JSON. + if err := os.WriteFile(path, []byte(`{"peers":{"users_:1":`), 0o600); err != nil { + t.Fatal(err) + } + + s, err := Open(path) + if err != nil { + t.Fatalf("Open should tolerate a corrupt cache: %v", err) + } + key := peers.Key{Prefix: prefixUsers, ID: 5} + if err := s.Save(ctx, key, peers.Value{AccessHash: 55}); err != nil { + t.Fatal(err) + } + s2, err := Open(path) + if err != nil { + t.Fatal(err) + } + if v, found, _ := s2.Find(ctx, key); !found || v.AccessHash != 55 { + t.Errorf("Find after rebuild = %+v, %v", v, found) + } +} + +func TestOverlappingWritersKeepAllEntries(t *testing.T) { + // Each Storage stands in for a separate tg process sharing the cache file: + // separate instances share no in-memory lock, only the file. + const writers, saves = 8, 25 + path := filepath.Join(t.TempDir(), "peers.json") + ctx := context.Background() + + var wg sync.WaitGroup + for w := range writers { + s, err := Open(path) + if err != nil { + t.Fatal(err) + } + wg.Add(1) + go func() { + defer wg.Done() + for i := range saves { + id := int64(w*saves + i + 1) + if err := s.Save(ctx, peers.Key{Prefix: prefixUsers, ID: id}, peers.Value{AccessHash: id}); err != nil { + t.Error(err) + return + } + } + }() + } + wg.Wait() + + final, err := Open(path) + if err != nil { + t.Fatal(err) + } + missing := 0 + for id := int64(1); id <= writers*saves; id++ { + if _, found, _ := final.Find(ctx, peers.Key{Prefix: prefixUsers, ID: id}); !found { + missing++ + } + } + if missing > 0 { + t.Errorf("%d of %d entries lost by overlapping writers", missing, writers*saves) + } +}