package model import ( "encoding/json" "testing" "with_you/internal/pkg/crypto" ) const testKey = "12345678901234567890123456789012" // initEnc 初始化(幂等,因 crypto 用 sync.Once)加密器。 func initEnc(t *testing.T) { t.Helper() if err := crypto.InitMessageEncryptor(testKey, 1); err != nil { t.Fatalf("InitMessageEncryptor: %v", err) } if crypto.GetMessageEncryptor() == nil { t.Fatal("encryptor not initialized") } } // TestMessage_EncryptDecryptRoundTrip 验证:BeforeCreate 加密落库 → // 清空内存 Segments → Decrypt 还原 → 内容与原文一致。 func TestMessage_EncryptDecryptRoundTrip(t *testing.T) { initEnc(t) original := MessageSegments{ {Type: "text", Data: MessageSegmentData{"content": "你好世界"}}, {Type: "image", Data: MessageSegmentData{"url": "https://example.com/a.png"}}, } msg := &Message{Segments: original} // 模拟 GORM 创建钩子:加密 if err := msg.BeforeCreate(nil); err != nil { t.Fatalf("BeforeCreate (encrypt) failed: %v", err) } // 加密后应写入 SegmentsEncrypted,且不应是明文 JSON(除非降级) if msg.SegmentsEncrypted == "" { t.Fatal("SegmentsEncrypted should be populated after BeforeCreate") } // 校验确实加密了:明文 JSON 中应包含 "content",密文 base64 中不应直接出现该明文 if contains(msg.SegmentsEncrypted, "你好世界") { t.Errorf("SegmentsEncrypted appears to contain plaintext: %q", msg.SegmentsEncrypted) } if msg.SegmentsKeyVersion != crypto.GetMessageEncryptor().GetKeyVersion() { t.Errorf("SegmentsKeyVersion = %d, want %d", msg.SegmentsKeyVersion, crypto.GetMessageEncryptor().GetKeyVersion()) } // 模拟从 DB 取回后清空内存段(仅 SegmentsEncrypted 落库) msg.Segments = nil msg.decrypted = false // 解密还原 if err := msg.Decrypt(); err != nil { t.Fatalf("Decrypt failed: %v", err) } if len(msg.Segments) != len(original) { t.Fatalf("after Decrypt, len(Segments) = %d, want %d", len(msg.Segments), len(original)) } for i, seg := range msg.Segments { if seg.Type != original[i].Type { t.Errorf("segment[%d].Type = %q, want %q", i, seg.Type, original[i].Type) } if seg.Data["content"] != original[i].Data["content"] { t.Errorf("segment[%d].content mismatch: got %v, want %v", i, seg.Data["content"], original[i].Data["content"]) } } } // TestMessage_EncryptEmptySegments 验证:空 Segments 时加密为 no-op。 func TestMessage_EncryptEmptySegments(t *testing.T) { initEnc(t) msg := &Message{} if err := msg.BeforeCreate(nil); err != nil { t.Fatalf("BeforeCreate failed: %v", err) } if msg.SegmentsEncrypted != "" { t.Errorf("empty Segments should not produce ciphertext, got %q", msg.SegmentsEncrypted) } } // TestBatchDecryptMessagesParallel_RoundTrip 验证批量解密路径与单条一致。 func TestBatchDecryptMessagesParallel_RoundTrip(t *testing.T) { initEnc(t) contents := []string{"消息A", "消息B", "消息C"} msgs := make([]*Message, len(contents)) for i, c := range contents { m := &Message{ Segments: MessageSegments{ {Type: "text", Data: MessageSegmentData{"content": c}}, }, } if err := m.BeforeCreate(nil); err != nil { t.Fatalf("BeforeCreate[%d]: %v", i, err) } m.Segments = nil // 模拟从 DB 取回 msgs[i] = m } BatchDecryptMessagesParallel(msgs) for i, m := range msgs { if !m.decrypted { t.Errorf("msg[%d] not marked decrypted", i) } if len(m.Segments) != 1 { t.Errorf("msg[%d] segments len = %d, want 1", i, len(m.Segments)) continue } if got := m.Segments[0].Data["content"]; got != contents[i] { t.Errorf("msg[%d].content = %v, want %q", i, got, contents[i]) } } } // contains 简单子串判断(避免引入 strings 仅为此)。 func contains(s, sub string) bool { return len(s) >= len(sub) && (indexOf(s, sub) >= 0) } func indexOf(s, sub string) int { bs := []byte(s) for i := 0; i+len(sub) <= len(bs); i++ { if string(bs[i:i+len(sub)]) == sub { return i } } return -1 } // 编译期保证 json 被使用(兼容路径用到)。 var _ = json.Marshal