before merge to main

2026-03-04 16:29:24 +06:00
parent 0481cde1c3
commit 491ccbea89
21 changed files with 6776 additions and 47 deletions
--- a/pkg/linedb/jsonl_file.go
+++ b/pkg/linedb/jsonl_file.go
@@ -2,11 +2,17 @@ package linedb

 import (
 	"bufio"
+	"bytes"
+	"crypto/aes"
+	"crypto/cipher"
+	"crypto/rand"
 	"crypto/sha256"
 	"encoding/base64"
 	"fmt"
+	"io"
 	"os"
 	"path/filepath"
+	"sort"
 	"strings"
 	"sync"
 	"time"
@@ -23,6 +29,151 @@ func defaultJSONUnmarshal(data []byte, v any) error {
 	return goccyjson.Unmarshal(data, v)
 }

+// makeKeyOrderMarshal создаёт JSONMarshal с заданным порядком ключей
+func makeKeyOrderMarshal(keyOrder []KeyOrder) func(any) ([]byte, error) {
+	return func(v any) ([]byte, error) {
+		if m, ok := v.(map[string]any); ok {
+			return marshalMapSorted(m, keyOrder)
+		}
+		return goccyjson.Marshal(v)
+	}
+}
+
+func marshalMapSorted(m map[string]any, keyOrder []KeyOrder) ([]byte, error) {
+	orderMap := make(map[string]int)
+	for _, ko := range keyOrder {
+		orderMap[ko.Key] = ko.Order
+	}
+
+	keys := make([]string, 0, len(m))
+	for k := range m {
+		keys = append(keys, k)
+	}
+	sort.Slice(keys, func(i, j int) bool {
+		oi, hasI := orderMap[keys[i]]
+		oj, hasJ := orderMap[keys[j]]
+		group := func(o int, has bool) int {
+			if !has {
+				return 1
+			}
+			if o >= 0 {
+				return 0
+			}
+			return 2
+		}
+		gi, gj := group(oi, hasI), group(oj, hasJ)
+		if gi != gj {
+			return gi < gj
+		}
+		switch gi {
+		case 0:
+			return oi < oj
+		case 1:
+			return keys[i] < keys[j]
+		default:
+			return oi < oj
+		}
+	})
+
+	var buf bytes.Buffer
+	buf.WriteByte('{')
+	for i, k := range keys {
+		if i > 0 {
+			buf.WriteByte(',')
+		}
+		keyEscaped, _ := goccyjson.Marshal(k)
+		buf.Write(keyEscaped)
+		buf.WriteByte(':')
+		val := m[k]
+		if nested, ok := val.(map[string]any); ok {
+			valBytes, err := marshalMapSorted(nested, keyOrder)
+			if err != nil {
+				return nil, err
+			}
+			buf.Write(valBytes)
+		} else {
+			valBytes, err := goccyjson.Marshal(val)
+			if err != nil {
+				return nil, err
+			}
+			buf.Write(valBytes)
+		}
+	}
+	buf.WriteByte('}')
+	return buf.Bytes(), nil
+}
+
+// encodeKeyBytes превращает строку ключа в 32 байта (SHA256) для AES-256
+func encodeKeyBytes(keyStr string) []byte {
+	h := sha256.Sum256([]byte(keyStr))
+	return h[:]
+}
+
+// aesGCMEncrypt — встроенное AES-256-GCM шифрование
+func aesGCMEncrypt(plaintext, key []byte) ([]byte, error) {
+	block, err := aes.NewCipher(key)
+	if err != nil {
+		return nil, err
+	}
+	aesgcm, err := cipher.NewGCM(block)
+	if err != nil {
+		return nil, err
+	}
+	nonce := make([]byte, aesgcm.NonceSize())
+	if _, err := io.ReadFull(rand.Reader, nonce); err != nil {
+		return nil, err
+	}
+	return aesgcm.Seal(nonce, nonce, plaintext, nil), nil
+}
+
+// aesGCMDecrypt — встроенная расшифровка AES-256-GCM
+func aesGCMDecrypt(ciphertext, key []byte) ([]byte, error) {
+	block, err := aes.NewCipher(key)
+	if err != nil {
+		return nil, err
+	}
+	aesgcm, err := cipher.NewGCM(block)
+	if err != nil {
+		return nil, err
+	}
+	nonceSize := aesgcm.NonceSize()
+	if len(ciphertext) < nonceSize {
+		return nil, fmt.Errorf("ciphertext too short")
+	}
+	nonce, ciphertext := ciphertext[:nonceSize], ciphertext[nonceSize:]
+	return aesgcm.Open(nil, nonce, ciphertext, nil)
+}
+
+// makeEncodedMarshal оборачивает marshal в шифрование: marshal -> encrypt -> base64
+func makeEncodedMarshal(marshal func(any) ([]byte, error), encFn func([]byte, []byte) ([]byte, error), key []byte) func(any) ([]byte, error) {
+	return func(v any) ([]byte, error) {
+		jsonData, err := marshal(v)
+		if err != nil {
+			return nil, err
+		}
+		encrypted, err := encFn(jsonData, key)
+		if err != nil {
+			return nil, err
+		}
+		return []byte(base64.StdEncoding.EncodeToString(encrypted)), nil
+	}
+}
+
+// makeEncodedUnmarshal оборачивает unmarshal в дешифрование: base64 -> decrypt -> unmarshal
+func makeEncodedUnmarshal(unmarshal func([]byte, any) error, decFn func([]byte, []byte) ([]byte, error), key []byte) func([]byte, any) error {
+	return func(data []byte, v any) error {
+		encrypted, err := base64.StdEncoding.DecodeString(string(data))
+		if err != nil {
+			return err
+		}
+		plaintext, err := decFn(encrypted, key)
+		if err != nil {
+			return err
+		}
+		return unmarshal(plaintext, v)
+	}
+}
+
 // JSONLFile представляет адаптер для работы с JSONL файлами
 type JSONLFile struct {
 	filename       string
@@ -61,14 +212,31 @@ func NewJSONLFile(filename string, cypherKey string, options JSONLFileOptions) *
 	jsonMarshal := defaultJSONMarshal
 	jsonUnmarshal := defaultJSONUnmarshal

-	// Используем пользовательские функции если они предоставлены
 	if options.JSONMarshal != nil {
 		jsonMarshal = options.JSONMarshal
+	} else if len(options.KeyOrder) > 0 {
+		jsonMarshal = makeKeyOrderMarshal(options.KeyOrder)
 	}
+
 	if options.JSONUnmarshal != nil {
 		jsonUnmarshal = options.JSONUnmarshal
 	}

+	// Encode: оборачиваем marshal/unmarshal в шифрование
+	if options.Encode && options.EncodeKey != "" {
+		key := encodeKeyBytes(options.EncodeKey)
+		encFn := options.EncryptFn
+		decFn := options.DecryptFn
+		if encFn == nil {
+			encFn = aesGCMEncrypt
+		}
+		if decFn == nil {
+			decFn = aesGCMDecrypt
+		}
+		jsonMarshal = makeEncodedMarshal(jsonMarshal, encFn, key)
+		jsonUnmarshal = makeEncodedUnmarshal(jsonUnmarshal, decFn, key)
+	}
+
 	return &JSONLFile{
 		filename:       filename,
 		cypherKey:      cypherKey,
@@ -99,18 +267,87 @@ func (j *JSONLFile) Init(force bool, options LineDbAdapterOptions) error {
 	}

 	// Создаем файл если не существует
-	if _, err := os.Stat(j.filename); os.IsNotExist(err) {
+	if st, err := os.Stat(j.filename); os.IsNotExist(err) {
 		file, err := os.Create(j.filename)
 		if err != nil {
 			return fmt.Errorf("failed to create file: %w", err)
 		}
 		file.Close()
+	} else if err == nil && st.Size() > 0 {
+		// Файл существует и не пустой — проверяем и нормализуем записи
+		if err := j.normalizeExistingFile(); err != nil {
+			return err
+		}
 	}

 	j.initialized = true
 	return nil
 }

+// normalizeExistingFile проверяет размеры записей и приводит к allocSize-1
+func (j *JSONLFile) normalizeExistingFile() error {
+	data, err := os.ReadFile(j.filename)
+	if err != nil {
+		return fmt.Errorf("failed to read file: %w", err)
+	}
+	lines := strings.Split(strings.TrimSuffix(string(data), "\n"), "\n")
+	var nonEmpty []string
+	for _, l := range lines {
+		if l == "" {
+			continue
+		}
+		nonEmpty = append(nonEmpty, l)
+	}
+	if len(nonEmpty) == 0 {
+		return nil
+	}
+
+	// Максимальная длина среди записей
+	maxLen := 0
+	for _, l := range nonEmpty {
+		if len(l) > maxLen {
+			maxLen = len(l)
+		}
+	}
+
+	// Приводим короткие к maxLen (добавляем пробелы)
+	for i, l := range nonEmpty {
+		if len(l) < maxLen {
+			nonEmpty[i] = l + strings.Repeat(" ", maxLen-len(l))
+		}
+	}
+
+	targetLen := j.allocSize - 1
+	if targetLen < 1 {
+		targetLen = 1
+	}
+
+	if targetLen < maxLen {
+		// Уменьшаем: можно только если данные (без trailing spaces) помещаются
+		for i := range nonEmpty {
+			trimmed := strings.TrimRight(nonEmpty[i], " ")
+			if len(trimmed) > targetLen {
+				return fmt.Errorf("init failed: record data size %d exceeds configured allocSize-1 (%d), cannot reduce without data loss",
+					len(trimmed), targetLen)
+			}
+			nonEmpty[i] = trimmed + strings.Repeat(" ", targetLen-len(trimmed))
+		}
+	} else if targetLen > maxLen {
+		// Расширяем
+		for i := range nonEmpty {
+			nonEmpty[i] = nonEmpty[i] + strings.Repeat(" ", targetLen-len(nonEmpty[i]))
+		}
+	}
+
+	// Перезаписываем файл
+	var buf bytes.Buffer
+	for _, l := range nonEmpty {
+		buf.WriteString(l)
+		buf.WriteByte('\n')
+	}
+	return os.WriteFile(j.filename, buf.Bytes(), 0644)
+}
+
 // Read читает все записи из файла
 func (j *JSONLFile) Read(options LineDbAdapterOptions) ([]any, error) {
 	j.mutex.RLock()
@@ -135,8 +372,8 @@ func (j *JSONLFile) Read(options LineDbAdapterOptions) ([]any, error) {
 			continue
 		}

-		// Расшифровываем если нужно
-		if j.cypherKey != "" {
+		// Расшифровываем если нужно (только cypherKey; Encode обрабатывается в jsonUnmarshal)
+		if j.cypherKey != "" && !j.options.Encode {
 			decoded, err := base64.StdEncoding.DecodeString(line)
 			if err != nil {
 				if j.options.SkipInvalidLines {
@@ -189,14 +426,18 @@ func (j *JSONLFile) Write(data any, options LineDbAdapterOptions) error {

 		line := string(jsonData)

-		// Шифруем если нужно
-		if j.cypherKey != "" {
+		// Шифруем если нужно (только cypherKey; Encode уже в jsonMarshal)
+		if j.cypherKey != "" && !j.options.Encode {
 			line = base64.StdEncoding.EncodeToString([]byte(line))
 		}

-		// Дополняем до allocSize
-		if len(line) < j.allocSize {
-			line += strings.Repeat(" ", j.allocSize-len(line)-1)
+		maxLineLen := j.allocSize - 1
+		if len(line) > maxLineLen {
+			return fmt.Errorf("record size %d exceeds allocSize-1 (%d)", len(line), maxLineLen)
+		}
+		// Дополняем до allocSize-1
+		if len(line) < maxLineLen {
+			line += strings.Repeat(" ", maxLineLen-len(line))
 		}

 		if _, err := file.WriteString(line + "\n"); err != nil {
@@ -406,14 +647,18 @@ func (j *JSONLFile) rewriteFile(records []any) error {

 		line := string(jsonData)

-		// Шифруем если нужно
-		if j.cypherKey != "" {
+		// Шифруем если нужно (только cypherKey; Encode уже в jsonMarshal)
+		if j.cypherKey != "" && !j.options.Encode {
 			line = base64.StdEncoding.EncodeToString([]byte(line))
 		}

-		// Дополняем до allocSize
-		if len(line) < j.allocSize {
-			line += strings.Repeat(" ", j.allocSize-len(line)-1)
+		maxLineLen := j.allocSize - 1
+		if len(line) > maxLineLen {
+			return fmt.Errorf("record size %d exceeds allocSize-1 (%d)", len(line), maxLineLen)
+		}
+		// Дополняем до allocSize-1
+		if len(line) < maxLineLen {
+			line += strings.Repeat(" ", maxLineLen-len(line))
 		}

 		if _, err := file.WriteString(line + "\n"); err != nil {