Add token auto-refresh, device persistence, and cross-signing

Following gogobee's auth pattern: if MATRIX_BOT_PASSWORD is set, the bot persists device credentials to device.json, validates tokens on startup, re-logs in automatically when tokens expire, and bootstraps cross-signing so the device shows as verified without manual emoji verification. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
2026-03-17 18:35:23 -07:00
parent 0cfed504a1
commit d6c8d4f599
6 changed files with 682 additions and 31 deletions
--- a/internal/config/config.go
+++ b/internal/config/config.go
@@ -13,6 +13,7 @@ type Config struct {
 	MatrixHomeserverURL  string
 	MatrixBotUserID      string
 	MatrixBotAccessToken string
+	MatrixBotPassword    string
 	MatrixDealsRoomID    string
 	ITADAPIKey           string
 	DealSources          []string
@@ -42,6 +43,8 @@ func Load() (*Config, error) {
 	if err != nil {
 		return nil, err
 	}
+	c.MatrixBotPassword = os.Getenv("MATRIX_BOT_PASSWORD")
+
 	c.MatrixDealsRoomID, err = require("MATRIX_DEALS_ROOM_ID")
 	if err != nil {
 		return nil, err
--- a/internal/matrix/client.go
+++ b/internal/matrix/client.go
@@ -1,9 +1,14 @@
 package matrix

 import (
+	"bytes"
 	"context"
+	"encoding/json"
 	"fmt"
+	"io"
 	"log/slog"
+	"net/http"
+	"os"
 	"sync"
 	"time"

@@ -13,65 +18,185 @@ import (
 	"maunium.net/go/mautrix/id"
 )

+// DeviceInfo holds persisted device credentials for stable device identity
+// across restarts.
+type DeviceInfo struct {
+	AccessToken string `json:"access_token"`
+	DeviceID    string `json:"device_id"`
+	UserID      string `json:"user_id"`
+}
+
+// ClientConfig holds the parameters needed to create a Matrix client.
+type ClientConfig struct {
+	HomeserverURL string
+	UserID        string
+	AccessToken   string
+	Password      string // optional: enables auto-refresh and cross-signing
+	CryptoDBPath  string // e.g. "crypto.db", empty to skip E2EE
+	DevicePath    string // e.g. "device.json", empty to skip persistence
+}
+
 type Client struct {
 	client     *mautrix.Client
 	crypto     *cryptohelper.CryptoHelper
 	cancel     context.CancelFunc
 	syncCancel context.CancelFunc
-	hasE2EE    bool
 	startTime  time.Time
+	cfg        ClientConfig

 	dmMu    sync.Mutex
-	dmCache map[id.UserID]id.RoomID // userID -> DM room
+	dmCache map[id.UserID]id.RoomID
 }

-// New creates a new Matrix client with E2EE support via cryptohelper.
-// The cryptoDBPath is the path to the SQLite database for the crypto store
-// (e.g. "crypto.db"). Call RegisterMessageHandler before StartSync to
-// receive DM events.
-func New(homeserverURL, userID, accessToken, cryptoDBPath string) (*Client, error) {
-	uid := id.UserID(userID)
-	client, err := mautrix.NewClient(homeserverURL, uid, accessToken)
-	if err != nil {
-		return nil, fmt.Errorf("failed to create matrix client: %w", err)
-	}
+// New creates a new Matrix client. If a password is provided, the client will:
+//   - Persist device credentials to DevicePath for stable device identity
+//   - Validate existing tokens on startup and re-login if expired
+//   - Set LoginAs on the cryptohelper so it can auto-refresh mid-operation
+//   - Bootstrap cross-signing so the bot's device is automatically verified
+//
+// If only an access token is provided (no password), the client works but
+// cannot auto-refresh — the token must be manually replaced if it expires.
+func New(cfg ClientConfig) (*Client, error) {
+	uid := id.UserID(cfg.UserID)
+	accessToken := cfg.AccessToken

 	c := &Client{
-		client:    client,
+		cfg:       cfg,
 		dmCache:   make(map[id.UserID]id.RoomID),
 		startTime: time.Now(),
 	}

-	if cryptoDBPath != "" {
-		ch, err := cryptohelper.NewCryptoHelper(client, []byte("pastel_pickle_key"), cryptoDBPath)
+	// If we have a device file and password, try to load persisted credentials
+	if cfg.DevicePath != "" && cfg.Password != "" {
+		device, err := loadDevice(cfg.DevicePath)
+		if err == nil && device.AccessToken != "" {
+			if isTokenValid(cfg.HomeserverURL, device.AccessToken) {
+				slog.Info("existing device credentials valid", "device_id", device.DeviceID)
+				accessToken = device.AccessToken
+
+				client, err := mautrix.NewClient(cfg.HomeserverURL, id.UserID(device.UserID), device.AccessToken)
+				if err != nil {
+					return nil, fmt.Errorf("create client with existing token: %w", err)
+				}
+				client.DeviceID = id.DeviceID(device.DeviceID)
+				c.client = client
+			} else {
+				slog.Warn("existing device credentials expired, logging in again")
+			}
+		}
+	}
+
+	// If we don't have a valid client yet, try password login or fall back to token
+	if c.client == nil && cfg.Password != "" {
+		loginResp, err := loginWithPassword(cfg.HomeserverURL, cfg.UserID, cfg.Password)
+		if err != nil {
+			return nil, fmt.Errorf("login: %w", err)
+		}
+
+		client, err := mautrix.NewClient(cfg.HomeserverURL, id.UserID(loginResp.UserID), loginResp.AccessToken)
+		if err != nil {
+			return nil, fmt.Errorf("create client after login: %w", err)
+		}
+		client.DeviceID = id.DeviceID(loginResp.DeviceID)
+		c.client = client
+		accessToken = loginResp.AccessToken
+
+		// Save device credentials for next restart
+		if cfg.DevicePath != "" {
+			if err := saveDevice(cfg.DevicePath, &DeviceInfo{
+				AccessToken: loginResp.AccessToken,
+				DeviceID:    loginResp.DeviceID,
+				UserID:      loginResp.UserID,
+			}); err != nil {
+				slog.Warn("failed to save device info", "error", err)
+			}
+		}
+
+		slog.Info("logged in with password", "user_id", loginResp.UserID, "device_id", loginResp.DeviceID)
+	}
+
+	// Fall back to bare access token (no password, no auto-refresh)
+	if c.client == nil {
+		client, err := mautrix.NewClient(cfg.HomeserverURL, uid, accessToken)
+		if err != nil {
+			return nil, fmt.Errorf("create client: %w", err)
+		}
+		c.client = client
+	}
+
+	// Set up E2EE
+	if cfg.CryptoDBPath != "" {
+		ch, err := cryptohelper.NewCryptoHelper(c.client, []byte("pastel_pickle_key"), cfg.CryptoDBPath)
 		if err != nil {
 			return nil, fmt.Errorf("init crypto helper: %w", err)
 		}

+		// If password is available, enable auto-refresh on token expiry
+		if cfg.Password != "" {
+			ch.LoginAs = &mautrix.ReqLogin{
+				Type: mautrix.AuthTypePassword,
+				Identifier: mautrix.UserIdentifier{
+					Type: mautrix.IdentifierTypeUser,
+					User: cfg.UserID,
+				},
+				Password: cfg.Password,
+			}
+		}
+
 		if err := ch.Init(context.Background()); err != nil {
 			return nil, fmt.Errorf("crypto helper init: %w", err)
 		}

-		client.Crypto = ch
+		c.client.Crypto = ch
 		c.crypto = ch

-		slog.Info("E2EE initialized", "device_id", client.DeviceID)
+		// Bootstrap cross-signing if password is available
+		if cfg.Password != "" {
+			c.bootstrapCrossSigning()
+		}
+
+		slog.Info("E2EE initialized", "device_id", c.client.DeviceID)
 	}

 	return c, nil
 }

+func (c *Client) bootstrapCrossSigning() {
+	mach := c.crypto.Machine()
+
+	_, _, err := mach.GenerateAndUploadCrossSigningKeys(context.Background(), func(ui *mautrix.RespUserInteractive) interface{} {
+		return map[string]interface{}{
+			"type": mautrix.AuthTypePassword,
+			"identifier": map[string]interface{}{
+				"type": mautrix.IdentifierTypeUser,
+				"user": c.cfg.UserID,
+			},
+			"password": c.cfg.Password,
+			"session":  ui.Session,
+		}
+	}, "")
+	if err != nil {
+		slog.Debug("cross-signing: key upload skipped (may already exist)", "error", err)
+	}
+
+	if err := mach.SignOwnDevice(context.Background(), mach.OwnIdentity()); err != nil {
+		slog.Debug("cross-signing: sign own device skipped", "error", err)
+	}
+
+	if err := mach.SignOwnMasterKey(context.Background()); err != nil {
+		slog.Debug("cross-signing: sign master key skipped", "error", err)
+	}
+}
+
 // RegisterMessageHandler registers a callback for incoming messages.
 // Must be called before StartSync. Ignores the bot's own messages and
 // messages from before the bot started (avoids replaying history).
 func (c *Client) RegisterMessageHandler(fn func(senderID id.UserID, roomID id.RoomID, body string)) {
 	syncer := c.client.Syncer.(*mautrix.DefaultSyncer)
 	syncer.OnEventType(event.EventMessage, func(ctx context.Context, evt *event.Event) {
-		// Ignore own messages
 		if evt.Sender == c.client.UserID {
 			return
 		}
-		// Ignore messages from before bot startup (history replay)
 		evtTime := time.UnixMilli(evt.Timestamp)
 		if evtTime.Before(c.startTime) {
 			return
@@ -153,9 +278,6 @@ func (c *Client) SendNotice(roomID, text string) error {
 }

 // GetDMRoom returns the DM room for a user, reusing an existing one if available.
-// Checks the in-memory cache first, then m.direct account data, and only creates
-// a new encrypted room as a last resort. Follows gogobee's pattern to avoid
-// opening duplicate DM rooms against the same user.
 func (c *Client) GetDMRoom(userID id.UserID) (id.RoomID, error) {
 	c.dmMu.Lock()
 	if roomID, ok := c.dmCache[userID]; ok {
@@ -164,12 +286,11 @@ func (c *Client) GetDMRoom(userID id.UserID) (id.RoomID, error) {
 	}
 	c.dmMu.Unlock()

-	// Check m.direct account data for existing DM rooms
 	var dmRooms map[id.UserID][]id.RoomID
 	err := c.client.GetAccountData(context.Background(), "m.direct", &dmRooms)
 	if err == nil {
 		if rooms, ok := dmRooms[userID]; ok && len(rooms) > 0 {
-			roomID := rooms[len(rooms)-1] // use most recent
+			roomID := rooms[len(rooms)-1]
 			c.dmMu.Lock()
 			c.dmCache[userID] = roomID
 			c.dmMu.Unlock()
@@ -177,7 +298,6 @@ func (c *Client) GetDMRoom(userID id.UserID) (id.RoomID, error) {
 		}
 	}

-	// No existing DM room — create an encrypted one
 	resp, err := c.client.CreateRoom(context.Background(), &mautrix.ReqCreateRoom{
 		Preset:   "trusted_private_chat",
 		Invite:   []id.UserID{userID},
@@ -228,7 +348,6 @@ func (c *Client) StartPresenceHeartbeat() {
 		ticker := time.NewTicker(60 * time.Second)
 		defer ticker.Stop()

-		// Send immediately
 		c.sendPresence()

 		for {
@@ -261,3 +380,75 @@ func (c *Client) Stop() {
 		c.crypto.Close()
 	}
 }
+
+// --- auth helpers ---
+
+type loginResponse struct {
+	AccessToken string `json:"access_token"`
+	DeviceID    string `json:"device_id"`
+	UserID      string `json:"user_id"`
+}
+
+func loginWithPassword(homeserver, user, password string) (*loginResponse, error) {
+	body, _ := json.Marshal(map[string]any{
+		"type": "m.login.password",
+		"identifier": map[string]string{
+			"type": "m.id.user",
+			"user": user,
+		},
+		"password": password,
+	})
+
+	resp, err := (&http.Client{Timeout: 30 * time.Second}).Post(
+		homeserver+"/_matrix/client/v3/login", "application/json", bytes.NewReader(body))
+	if err != nil {
+		return nil, fmt.Errorf("login request: %w", err)
+	}
+	defer resp.Body.Close()
+
+	respBody, _ := io.ReadAll(resp.Body)
+	if resp.StatusCode != 200 {
+		return nil, fmt.Errorf("login failed (HTTP %d): %s", resp.StatusCode, respBody)
+	}
+
+	var result loginResponse
+	if err := json.Unmarshal(respBody, &result); err != nil {
+		return nil, fmt.Errorf("parse login response: %w", err)
+	}
+	return &result, nil
+}
+
+func isTokenValid(homeserver, accessToken string) bool {
+	req, err := http.NewRequest("GET", homeserver+"/_matrix/client/v3/account/whoami", nil)
+	if err != nil {
+		return false
+	}
+	req.Header.Set("Authorization", "Bearer "+accessToken)
+
+	resp, err := (&http.Client{Timeout: 10 * time.Second}).Do(req)
+	if err != nil {
+		return false
+	}
+	resp.Body.Close()
+	return resp.StatusCode == 200
+}
+
+func loadDevice(path string) (*DeviceInfo, error) {
+	data, err := os.ReadFile(path)
+	if err != nil {
+		return nil, err
+	}
+	var info DeviceInfo
+	if err := json.Unmarshal(data, &info); err != nil {
+		return nil, err
+	}
+	return &info, nil
+}
+
+func saveDevice(path string, info *DeviceInfo) error {
+	data, err := json.MarshalIndent(info, "", "  ")
+	if err != nil {
+		return err
+	}
+	return os.WriteFile(path, data, 0o600)
+}
--- a/internal/preflight/preflight.go
+++ b/internal/preflight/preflight.go
@@ -59,7 +59,11 @@ func PrintResults(results []Result) bool {
 }

 func checkMatrix(cfg *config.Config) Result {
-	client, err := matrix.New(cfg.MatrixHomeserverURL, cfg.MatrixBotUserID, cfg.MatrixBotAccessToken, "")
+	client, err := matrix.New(matrix.ClientConfig{
+		HomeserverURL: cfg.MatrixHomeserverURL,
+		UserID:        cfg.MatrixBotUserID,
+		AccessToken:   cfg.MatrixBotAccessToken,
+	})
 	if err != nil {
 		return Result{"Matrix", "fail", fmt.Sprintf("client error: %v", err)}
 	}