im_server/server/store/local_message_crypto_integration_test.go

//go:build mysql

package store_test

import (
	"database/sql"
	"encoding/base64"
	"encoding/json"
	"fmt"
	"os"
	"strings"
	"testing"
	"time"

	_ "github.com/go-sql-driver/mysql"
	_ "github.com/tinode/chat/server/db/mysql"
	"github.com/tinode/chat/server/store"
	"github.com/tinode/chat/server/store/types"
)

type localCryptoProvider struct {
	key string
}

func (p localCryptoProvider) GetTenantByCode(string) (*types.Tenant, error) {
	return nil, types.ErrUnsupported
}
func (p localCryptoProvider) GetTenantByID(types.TenantID) (*types.Tenant, error) {
	return nil, types.ErrUnsupported
}
func (p localCryptoProvider) ListActiveTenants(int) ([]types.Tenant, error) {
	return nil, types.ErrUnsupported
}
func (p localCryptoProvider) RegistrationSMSEnabled() (bool, error) {
	return false, types.ErrUnsupported
}
func (p localCryptoProvider) ThirdPartyConfig(string, ...string) (map[string]string, error) {
	return nil, types.ErrUnsupported
}
func (p localCryptoProvider) LatestAppRelease(string, string) (*types.AppRelease, error) {
	return nil, types.ErrUnsupported
}
func (p localCryptoProvider) ModerationProvider() (*types.ModerationProviderConfig, error) {
	return nil, types.ErrUnsupported
}
func (p localCryptoProvider) ModerationPolicy(string) (*types.ModerationPolicy, error) {
	return nil, types.ErrUnsupported
}
func (p localCryptoProvider) SensitiveWords(string) (*types.SensitiveWordSet, error) {
	return nil, types.ErrUnsupported
}
func (p localCryptoProvider) SecurityPolicy(string) (*types.SecurityPolicy, error) {
	return nil, types.ErrUnsupported
}
func (p localCryptoProvider) RealNamePolicy(types.TenantID) (*types.RealNamePolicy, error) {
	return nil, types.ErrUnsupported
}
func (p localCryptoProvider) MessageKey(tenantCode string, version int) (*types.MessageKey, error) {
	if tenantCode != "default" || (version != 0 && version != 7) {
		return nil, types.ErrNotFound
	}
	return &types.MessageKey{TenantID: 1, TenantCode: tenantCode, Enabled: true, Version: 7, Algorithm: "AES-256-GCM", Status: "active", Key: p.key}, nil
}
func (p localCryptoProvider) MessageKeyAck(string, int) error { return nil }

func TestLocalMySQLMessageCryptoFlow(t *testing.T) {
	if os.Getenv("IM_LOCAL_CRYPTO_TEST") != "1" {
		t.Skip("set IM_LOCAL_CRYPTO_TEST=1 to run against local MySQL")
	}
	config := json.RawMessage(`{
		"uid_key":"dkXxWLwqUNRfYmdXRzOkZw==",
		"max_results":1024,
		"use_adapter":"mysql",
		"adapters":{"mysql":{
			"dsn":"mysql://root:root@tcp(127.0.0.1:3306)/tinode?parseTime=true&collation=utf8mb4_0900_ai_ci",
			"tenant_registry_db":"hotgo",
			"tenant_db_routing":true,
			"sql_timeout":10
		}}
	}`)
	if err := store.Store.Open(991, config); err != nil {
		t.Fatal(err)
	}
	defer store.Store.Close()
	defer store.SetAdminInternalConfigProvider(nil)
	store.SetAdminInternalConfigProvider(localCryptoProvider{key: base64.StdEncoding.EncodeToString([]byte("0123456789abcdef0123456789abcdef"))})
	store.InvalidateMessageEncryptionKeys(1)

	db, err := sql.Open("mysql", "root:root@tcp(127.0.0.1:3306)/tinode?parseTime=true&collation=utf8mb4_0900_ai_ci")
	if err != nil {
		t.Fatal(err)
	}
	defer db.Close()

	var topic string
	var senderID, originalTopicSeq int64
	err = db.QueryRow(`SELECT s.topic,s.userid,t.seqid FROM im_subscriptions s JOIN im_topics t ON t.name=s.topic
		WHERE s.topic LIKE 'p2p%' AND s.deletedat IS NULL ORDER BY t.id DESC LIMIT 1`).Scan(&topic, &senderID, &originalTopicSeq)
	if err != nil {
		t.Fatal(err)
	}
	sender := store.EncodeUid(senderID)
	marker := fmt.Sprintf("LOCAL-CRYPTO-%d", time.Now().UnixNano())
	msg := &types.Message{
		TenantID: 1,
		SeqId:    originalTopicSeq + 1,
		Topic:    topic,
		From:     sender.String(),
		Head:     types.KVMap{"x-im-client-mid": marker, "x-im-type": "text"},
		Content:  map[string]any{"txt": marker},
	}
	msg.InitTimes()
	msg.SetUid(store.Store.GetUid())
	fallbackSeq := int64(0)
	defer func() {
		_, _ = db.Exec("DELETE FROM im_message_revisions WHERE topic=? AND seqid=?", topic, msg.SeqId)
		_, _ = db.Exec("DELETE FROM im_messages WHERE topic=? AND seqid=?", topic, msg.SeqId)
		if fallbackSeq > 0 {
			_, _ = db.Exec("DELETE FROM im_messages WHERE topic=? AND seqid=?", topic, fallbackSeq)
		}
		_, _ = db.Exec("UPDATE im_topics SET seqid=? WHERE name=? AND seqid IN (?,?)", originalTopicSeq, topic, msg.SeqId, fallbackSeq)
	}()

	if _, err = store.PersistPreparedMessage(1, msg, nil, false); err != nil {
		t.Fatal(err)
	}
	var rawContent string
	var nonce []byte
	var keyVersion int
	if err = db.QueryRow("SELECT JSON_UNQUOTE(content),content_nonce,key_version FROM im_messages WHERE topic=? AND seqid=?", topic, msg.SeqId).Scan(&rawContent, &nonce, &keyVersion); err != nil {
		t.Fatal(err)
	}
	if keyVersion != 7 || len(nonce) != 12 || strings.Contains(rawContent, marker) {
		t.Fatalf("message not encrypted at rest: key_version=%d nonce=%d content=%q", keyVersion, len(nonce), rawContent)
	}
	if _, err = base64.StdEncoding.DecodeString(rawContent); err != nil {
		t.Fatalf("stored ciphertext is not base64: %v", err)
	}

	loaded, err := store.Messages.GetOne(1, topic, msg.SeqId)
	if err != nil || fmt.Sprint(loaded.Content.(map[string]any)["txt"]) != marker {
		t.Fatalf("GetOne plaintext mismatch: content=%#v err=%v", loaded, err)
	}
	editedMarker := marker + "-EDITED"
	edited, err := store.Messages.Edit(1, topic, msg.SeqId, sender, loaded.Version, loaded.Head, map[string]any{"txt": editedMarker}, "local integration test", types.TimeNow())
	if err != nil || fmt.Sprint(edited.Content.(map[string]any)["txt"]) != editedMarker {
		t.Fatalf("Edit plaintext mismatch: content=%#v err=%v", edited, err)
	}
	var revisionVersion, currentVersion int
	if err = db.QueryRow("SELECT key_version FROM im_message_revisions WHERE topic=? AND seqid=? ORDER BY version DESC LIMIT 1", topic, msg.SeqId).Scan(&revisionVersion); err != nil {
		t.Fatal(err)
	}
	if err = db.QueryRow("SELECT key_version FROM im_messages WHERE topic=? AND seqid=?", topic, msg.SeqId).Scan(&currentVersion); err != nil {
		t.Fatal(err)
	}
	if revisionVersion != 7 || currentVersion != 7 {
		t.Fatalf("unexpected key versions: revision=%d current=%d", revisionVersion, currentVersion)
	}

	history, err := store.Messages.GetAll(1, topic, sender, &types.QueryOpt{Since: msg.SeqId, Before: msg.SeqId + 1, Limit: 10})
	if err != nil || len(history) != 1 || fmt.Sprint(history[0].Content.(map[string]any)["txt"]) != editedMarker {
		t.Fatalf("history plaintext mismatch: messages=%#v err=%v", history, err)
	}
	found, _, err := store.Messages.Search(1, topic, sender, editedMarker, "text", "", 10)
	if err != nil || len(found) != 1 || found[0].SeqId != msg.SeqId {
		t.Fatalf("search plaintext mismatch: messages=%#v err=%v", found, err)
	}

	store.SetAdminInternalConfigProvider(localCryptoProvider{})
	store.InvalidateMessageEncryptionKeys(1)
	fallback := &types.Message{
		TenantID: 1,
		SeqId:    msg.SeqId + 1,
		Topic:    topic,
		From:     sender.String(),
		Head:     types.KVMap{"x-im-client-mid": marker + "-FALLBACK", "x-im-type": "text"},
		Content:  map[string]any{"txt": marker + "-PLAINTEXT"},
	}
	fallback.InitTimes()
	fallback.SetUid(store.Store.GetUid())
	fallbackSeq = fallback.SeqId
	if _, err = store.PersistPreparedMessage(1, fallback, nil, false); err != nil {
		t.Fatal(err)
	}
	if err = db.QueryRow("SELECT JSON_UNQUOTE(content),key_version FROM im_messages WHERE topic=? AND seqid=?", topic, fallback.SeqId).Scan(&rawContent, &keyVersion); err != nil {
		t.Fatal(err)
	}
	if keyVersion != 0 || !strings.Contains(rawContent, marker+"-PLAINTEXT") {
		t.Fatalf("missing-key fallback was not plaintext: key_version=%d content=%q", keyVersion, rawContent)
	}
}