fix(storage): make session, peer cache and config files safe to share between processes

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

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 "<file>.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

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])
+	}
+}

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