Bellhop/internal/matrix/client.go

package matrix

import (
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"log/slog"
	"net/http"
	"os"
	"path/filepath"
	"sync"
	"time"

	"bellhop/internal/config"

	"maunium.net/go/mautrix"
	"maunium.net/go/mautrix/crypto/cryptohelper"
	"maunium.net/go/mautrix/event"
	"maunium.net/go/mautrix/id"
)

// MessageHandler is called when a text message is received in an allowlisted room.
type MessageHandler func(roomID id.RoomID, eventID id.EventID, sender id.UserID, body string)

type deviceInfo struct {
	AccessToken string `json:"access_token"`
	DeviceID    string `json:"device_id"`
	UserID      string `json:"user_id"`
}

type Client struct {
	mx           *mautrix.Client
	allowedRooms map[id.RoomID]struct{}
	userID       id.UserID
	cfg          config.MatrixConfig
	crypto       *cryptohelper.CryptoHelper

	mu         sync.RWMutex
	onMessage  MessageHandler
	cancelSync context.CancelFunc
}

// New creates a Matrix client with password login, device persistence, and E2EE.
func New(cfg config.MatrixConfig) (*Client, error) {
	if err := os.MkdirAll(cfg.DataDir, 0o755); err != nil {
		return nil, fmt.Errorf("create data dir: %w", err)
	}

	devicePath := filepath.Join(cfg.DataDir, "device.json")

	device, err := loadDevice(devicePath)
	if err != nil {
		// File exists but is unparseable — refuse to overwrite. A silent fresh
		// login would orphan the prior device's megolm sessions and break
		// decryption in every encrypted room the bot is already in.
		return nil, fmt.Errorf("load device file %s: %w (refusing to overwrite — fix or delete the file)", devicePath, err)
	}

	var mx *mautrix.Client

	if device != nil {
		if isTokenValid(cfg.Homeserver, device.AccessToken, cfg.UserID) {
			slog.Info("existing device credentials valid", "device_id", device.DeviceID)
			mx, err = mautrix.NewClient(cfg.Homeserver, id.UserID(device.UserID), device.AccessToken)
			if err != nil {
				return nil, fmt.Errorf("create client with existing token: %w", err)
			}
			mx.DeviceID = id.DeviceID(device.DeviceID)
		} else {
			slog.Warn("existing device credentials invalid, logging in again")
			device = nil
		}
	}

	if device == nil {
		mx, err = mautrix.NewClient(cfg.Homeserver, "", "")
		if err != nil {
			return nil, fmt.Errorf("create client: %w", err)
		}

		resp, err := mx.Login(context.Background(), &mautrix.ReqLogin{
			Type: mautrix.AuthTypePassword,
			Identifier: mautrix.UserIdentifier{
				Type: mautrix.IdentifierTypeUser,
				User: cfg.UserID,
			},
			Password:                 cfg.Password,
			InitialDeviceDisplayName: cfg.DisplayName,
		})
		if err != nil {
			return nil, fmt.Errorf("login: %w", err)
		}

		mx.AccessToken = resp.AccessToken
		mx.UserID = resp.UserID
		mx.DeviceID = resp.DeviceID

		if err := saveDevice(devicePath, &deviceInfo{
			AccessToken: resp.AccessToken,
			DeviceID:    string(resp.DeviceID),
			UserID:      string(resp.UserID),
		}); err != nil {
			slog.Warn("failed to save device info", "err", err)
		}

		slog.Info("logged in successfully", "user_id", resp.UserID, "device_id", resp.DeviceID)
	}

	cryptoDBPath := filepath.Join(cfg.DataDir, "crypto.db")
	ch, err := cryptohelper.NewCryptoHelper(mx, []byte(cfg.PickleKey), cryptoDBPath)
	if err != nil {
		return nil, fmt.Errorf("init crypto helper: %w", err)
	}

	// Surface decrypt failures — the default silent no-op hides missing-megolm-session
	// errors, which look like the bot ignoring commands.
	ch.DecryptErrorCallback = func(evt *event.Event, err error) {
		slog.Warn("matrix: failed to decrypt incoming event",
			"room", evt.RoomID, "event_id", evt.ID, "sender", evt.Sender, "err", err)
	}

	// Intentionally NOT setting ch.LoginAs: cryptohelper.Init will Login() every
	// time LoginAs is non-nil (see mautrix cryptohelper.go), which would mint a
	// new access_token on every startup and silently invalidate the device.json
	// we just wrote. We already have valid credentials in mx by this point —
	// Init will reuse them.
	if err := ch.Init(context.Background()); err != nil {
		return nil, fmt.Errorf("crypto helper init: %w", err)
	}

	mx.Crypto = ch

	mach := ch.Machine()
	existingKeys, err := mach.GetOwnCrossSigningPublicKeys(context.Background())
	if err != nil {
		slog.Warn("cross-signing: failed to fetch existing keys", "err", err)
	}
	if existingKeys == nil || existingKeys.MasterKey == "" {
		_, _, 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": cfg.UserID,
				},
				"password": cfg.Password,
				"session":  ui.Session,
			}
		}, "")
		if err != nil {
			slog.Warn("cross-signing: key upload failed", "err", err)
		} else {
			slog.Info("cross-signing: keys uploaded")
		}
		if err := mach.SignOwnDevice(context.Background(), mach.OwnIdentity()); err != nil {
			slog.Warn("cross-signing: sign own device failed", "err", err)
		}
		if err := mach.SignOwnMasterKey(context.Background()); err != nil {
			slog.Warn("cross-signing: sign master key failed", "err", err)
		}
	} else {
		slog.Info("cross-signing: already configured, skipping bootstrap")
	}

	slog.Info("E2EE initialized", "device_id", mx.DeviceID)

	allowed := make(map[id.RoomID]struct{}, len(cfg.AllowedRooms))
	for _, r := range cfg.AllowedRooms {
		allowed[id.RoomID(r)] = struct{}{}
	}

	return &Client{
		mx:           mx,
		allowedRooms: allowed,
		userID:       mx.UserID,
		cfg:          cfg,
		crypto:       ch,
	}, nil
}

// SetMessageHandler registers a callback for text messages in allowlisted rooms.
func (c *Client) SetMessageHandler(fn MessageHandler) {
	c.mu.Lock()
	defer c.mu.Unlock()
	c.onMessage = fn
}

// Start begins the Matrix sync loop in the background.
func (c *Client) Start(ctx context.Context) error {
	syncer, ok := c.mx.Syncer.(*mautrix.DefaultSyncer)
	if !ok {
		return fmt.Errorf("matrix syncer is not *mautrix.DefaultSyncer (got %T)", c.mx.Syncer)
	}

	syncer.OnEventType(event.EventMessage, func(ctx context.Context, evt *event.Event) {
		if evt.Sender == c.userID {
			return
		}
		if _, ok := c.allowedRooms[evt.RoomID]; !ok {
			return
		}
		msg := evt.Content.AsMessage()
		if msg == nil || msg.MsgType != event.MsgText {
			return
		}
		c.mu.RLock()
		handler := c.onMessage
		c.mu.RUnlock()
		if handler != nil {
			handler(evt.RoomID, evt.ID, evt.Sender, msg.Body)
		}
	})

	// Auto-join on invite to any room — allowlist still gates message handling.
	syncer.OnEventType(event.StateMember, func(ctx context.Context, evt *event.Event) {
		if evt.GetStateKey() == string(c.userID) && evt.Content.AsMember().Membership == event.MembershipInvite {
			if _, err := c.mx.JoinRoomByID(ctx, evt.RoomID); err != nil {
				slog.Error("failed to auto-join room", "room", evt.RoomID, "err", err)
			} else {
				slog.Info("auto-joined room", "room", evt.RoomID)
			}
		}
	})

	syncCtx, cancel := context.WithCancel(ctx)
	c.cancelSync = cancel

	go func() {
		for {
			err := c.mx.SyncWithContext(syncCtx)
			if syncCtx.Err() != nil {
				return
			}
			if err != nil {
				slog.Error("matrix sync stopped, restarting in 5s", "err", err)
			} else {
				slog.Warn("matrix sync returned without error, restarting in 5s")
			}
			select {
			case <-time.After(5 * time.Second):
			case <-syncCtx.Done():
				return
			}
		}
	}()
	return nil
}

// Stop halts the sync loop and closes the crypto store.
func (c *Client) Stop() {
	if c.cancelSync != nil {
		c.cancelSync()
	}
	if c.crypto != nil {
		if err := c.crypto.Close(); err != nil {
			slog.Error("failed to close crypto helper", "err", err)
		}
	}
}

// PostThreadedReply sends a plain/HTML message into the thread rooted at rootEventID.
// IsFallingBack + m.in_reply_to gives non-thread-aware clients a sensible quoted reply.
// The caller's ctx governs cancellation and deadline.
func (c *Client) PostThreadedReply(ctx context.Context, roomID id.RoomID, rootEventID id.EventID, plain, htmlBody string) error {
	content := &event.MessageEventContent{
		MsgType:       event.MsgText,
		Body:          plain,
		Format:        event.FormatHTML,
		FormattedBody: htmlBody,
		RelatesTo: &event.RelatesTo{
			Type:          event.RelThread,
			EventID:       rootEventID,
			InReplyTo:     &event.InReplyTo{EventID: rootEventID},
			IsFallingBack: true,
		},
	}
	_, err := c.mx.SendMessageEvent(ctx, roomID, event.EventMessage, content)
	return err
}

// loadDevice returns (nil, nil) when no device file exists, the parsed device
// when it does, or (nil, err) when the file is present but unreadable/corrupt.
// Distinguishing the corrupt case lets the caller refuse to silently overwrite
// it with a brand-new device registration.
func loadDevice(path string) (*deviceInfo, error) {
	data, err := os.ReadFile(path)
	if err != nil {
		if errors.Is(err, os.ErrNotExist) {
			return nil, nil
		}
		return nil, err
	}
	var info deviceInfo
	if err := json.Unmarshal(data, &info); err != nil {
		return nil, fmt.Errorf("parse: %w", err)
	}
	if info.AccessToken == "" || info.DeviceID == "" || info.UserID == "" {
		return nil, fmt.Errorf("missing required fields (access_token/device_id/user_id)")
	}
	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)
}

// isTokenValid hits /whoami and confirms both that the token is accepted AND
// that it identifies the user we expect. A bare status-code check would accept
// a token belonging to a different account if the operator swapped user_id
// without clearing data_dir.
func isTokenValid(homeserver, accessToken, expectedUserID 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)

	client := &http.Client{Timeout: 10 * time.Second}
	resp, err := client.Do(req)
	if err != nil {
		return false
	}
	defer resp.Body.Close()
	if resp.StatusCode != http.StatusOK {
		return false
	}
	var body struct {
		UserID string `json:"user_id"`
	}
	if err := json.NewDecoder(resp.Body).Decode(&body); err != nil {
		return false
	}
	if body.UserID != expectedUserID {
		slog.Warn("saved token belongs to a different user, will re-login",
			"saved", body.UserID, "configured", expectedUserID)
		return false
	}
	return true
}