In-Memory DB 구현 10편 - Decoder와 Checksum으로 영속화 된 데이터 불러오기

import "hash/crc32"

// 방법 1: 한 번에 계산
checksum := crc32.ChecksumIEEE(data) // []byte → uint32

// 방법 2: 스트리밍 방식 (데이터를 나눠서 넣을 수 있음)
h := crc32.NewIEEE()
h.Write(part1)
h.Write(part2)
checksum := h.Sum32()

[Header]   MINIDB + Version          7 bytes
[Entry]    타입별 엔트리 반복             가변
[EOF]      0xFF                      1 byte
[CRC32]    체크섬 (Big Endian)         4 bytes   <- 추가

package persistence

var MagicBytes = [6]byte{'M', 'I', 'N', 'I', 'D', 'B'}

const (
	Version byte = 0x01

	TypeString byte = 0x00
	TypeList   byte = 0x01

	NoExpiry  byte = 0x00
	HasExpiry byte = 0x01

	EOF byte = 0xFF

	ChecksumSize = 4 // 추가
)

package persistence

import (
	"bufio"
	"encoding/binary"
	"hash"
	"hash/crc32"
	"io"
	"time"
)

type Encoder struct {
	w    *bufio.Writer
	hash hash.Hash32
}

func NewEncoder(w io.Writer) *Encoder {
	return &Encoder{
		w:    bufio.NewWriter(w),
		hash: crc32.NewIEEE(),
	}
}

func (e *Encoder) WriteHeader() error {
	if err := e.writeBytes(MagicBytes[:]); err != nil {
		return err
	}
	return e.writeBytes([]byte{Version})
}

func (e *Encoder) WriteStringEntry(key, value string, expireAt *time.Time) error {
	if err := e.writeBytes([]byte{TypeString}); err != nil {
		return err
	}
	if err := e.writeString(key); err != nil {
		return err
	}
	if err := e.writeString(value); err != nil {
		return err
	}
	return e.writeExpiry(expireAt)
}

func (e *Encoder) WriteListEntry(key string, values []string, expireAt *time.Time) error {
	if err := e.writeBytes([]byte{TypeList}); err != nil {
		return err
	}
	if err := e.writeString(key); err != nil {
		return err
	}
	if err := e.writeUint32(uint32(len(values))); err != nil {
		return err
	}
	for _, v := range values {
		if err := e.writeString(v); err != nil {
			return err
		}
	}
	return e.writeExpiry(expireAt)
}

func (e *Encoder) WriteEOF() error {
	return e.writeBytes([]byte{EOF})
}

func (e *Encoder) Flush() error {
	return e.w.Flush()
}

func (e *Encoder) WriteChecksum() error {
	checksum := e.hash.Sum32()
	buf := make([]byte, 4)
	binary.BigEndian.PutUint32(buf, checksum)
	_, err := e.w.Write(buf)
	return err
}

// ========== 헬퍼 메서드 ==========

func (e *Encoder) writeBytes(data []byte) error {
	_, err := e.w.Write(data)
	if err == nil {
		e.hash.Write(data)
	}
	return err
}

func (e *Encoder) writeUint32(n uint32) error {
	buf := make([]byte, 4)
	binary.BigEndian.PutUint32(buf, n)
	return e.writeBytes(buf)
}

func (e *Encoder) writeInt64(n int64) error {
	buf := make([]byte, 8)
	binary.BigEndian.PutUint64(buf, uint64(n))
	return e.writeBytes(buf)
}

func (e *Encoder) writeString(s string) error {
	if err := e.writeUint32(uint32(len(s))); err != nil {
		return err
	}
	return e.writeBytes([]byte(s))
}

func (e *Encoder) writeExpiry(expireAt *time.Time) error {
	if expireAt == nil {
		return e.writeBytes([]byte{NoExpiry})
	}
	if err := e.writeBytes([]byte{HasExpiry}); err != nil {
		return err
	}
	return e.writeInt64(expireAt.UnixMilli())
}

package persistence

import (
	"bufio"
	"bytes"
	"encoding/binary"
	"fmt"
	"hash/crc32"
	"io"
	"os"
	"time"
)

type Decoder struct {
	r *bufio.Reader
}

type DecodedEntry struct {
	Type     byte
	Key      string
	Value    string
	Values   []string
	ExpireAt *time.Time
}

func NewDecoder(r io.Reader) *Decoder {
	return &Decoder{bufio.NewReader(r)}
}

func (d *Decoder) ReadHeader() error {
	magicByte, err := d.readBytes(6)
	if err != nil {
		return err
	}

	if !bytes.Equal(magicByte, MagicBytes[:]) {
		return fmt.Errorf("invalid magic bytes")
	}

	version, err := d.readBytes(1)
	if err != nil {
		return err
	}

	if version[0] != Version {
		return fmt.Errorf("버전이 불일치합니다.")
	}
	return nil
}

func (d *Decoder) ReadEntry() (*DecodedEntry, error) {
	typeBuf, err := d.readBytes(1)
	if err != nil {
		return nil, err
	}

	if typeBuf[0] == EOF {
		return nil, nil
	}

	key, err := d.readString()
	if err != nil {
		return nil, err
	}

	entry := &DecodedEntry{Type: typeBuf[0], Key: key}

	switch typeBuf[0] {
	case TypeString:
		value, err := d.readString()
		if err != nil {
			return nil, err
		}
		entry.Value = value

	case TypeList:
		count, err := d.readUint32()
		if err != nil {
			return nil, err
		}
		values := make([]string, 0, count)
		for i := 0; i < int(count); i++ {
			v, err := d.readString()
			if err != nil {
				return nil, err
			}
			values = append(values, v)
		}
		entry.Values = values

	default:
		return nil, fmt.Errorf("unknown entry type: 0x%02x", typeBuf[0])
	}

	expireAt, err := d.readExpiry()
	if err != nil {
		return nil, err
	}
	entry.ExpireAt = expireAt

	return entry, nil
}

func VerifyChecksum(path string) error {
	file, err := os.ReadFile(path)
	if err != nil {
		return err
	}

	if len(file) < ChecksumSize {
		return fmt.Errorf("파일의 크기가 Checksum 사이즈보다 작습니다.")
	}

	content := file[:len(file)-ChecksumSize]
	stored := binary.BigEndian.Uint32(file[len(file)-ChecksumSize:])
	computed := crc32.ChecksumIEEE(content)

	if stored != computed {
		return fmt.Errorf(
			"Checksum이 일치하지 않습니다: stored=0x%08x, computed=0x%08x",
			stored,
			computed,
		)
	}
	return nil
}

// ========== 헬퍼 메서드 ==========
func (d *Decoder) readBytes(n int) ([]byte, error) {
	buf := make([]byte, n)
	_, err := io.ReadFull(d.r, buf)
	return buf, err
}

func (d *Decoder) readUint32() (uint32, error) {
	buf, err := d.readBytes(4)
	if err != nil {
		return 0, err
	}
	return binary.BigEndian.Uint32(buf), nil
}

func (d *Decoder) readInt64() (int64, error) {
	buf, err := d.readBytes(8)
	if err != nil {
		return 0, err
	}
	return int64(binary.BigEndian.Uint64(buf)), nil
}

func (d *Decoder) readString() (string, error) {
	length, err := d.readUint32()
	if err != nil {
		return "", err
	}
	buf, err := d.readBytes(int(length))
	if err != nil {
		return "", err
	}

	return string(buf), nil
}

func (d *Decoder) readExpiry() (*time.Time, error) {
	buf, err := d.readBytes(1)
	if err != nil {
		return nil, err
	}

	if buf[0] == NoExpiry {
		return nil, nil
	}

	ms, err := d.readInt64()
	if err != nil {
		return nil, err
	}

	t := time.UnixMilli(ms)
	return &t, nil
}

func (s *Store) Save(path string) error {
	s.mu.RLock()
	defer s.mu.RUnlock()

	file, err := os.Create(path)
	if err != nil {
		return err
	}
	defer file.Close()

	encoder := persistence.NewEncoder(file)
	encoder.WriteHeader()

	for key, entry := range s.data {
		if entry.ExpireAt != nil && entry.ExpireAt.Before(time.Now()) {
			continue
		}

		switch entry.Type {
		case TypeString:
			encoder.WriteStringEntry(key, entry.Str, entry.ExpireAt)

		case TypeList:
			values := entry.List.Range(0, entry.List.Length-1)
			encoder.WriteListEntry(key, values, entry.ExpireAt)

		}
	}

	encoder.WriteEOF()
	encoder.WriteChecksum()
	return encoder.Flush()
}

func (s *Store) Load(path string) error {
	data, err := os.ReadFile(path)
	if err != nil {
		if os.IsNotExist(err) {
			return nil
		}
		return err
	}

	if err := persistence.VerifyChecksum(path); err != nil {
		return err
	}

	content := data[:len(data)-persistence.ChecksumSize]
	decoder := persistence.NewDecoder(bytes.NewReader(content))

	if err := decoder.ReadHeader(); err != nil {
		return err
	}

	for {
		entry, err := decoder.ReadEntry()
		if err != nil {
			return err
		}
		if entry == nil {
			break
		}

		if entry.ExpireAt != nil && entry.ExpireAt.Before(time.Now()) {
			continue
		}

		switch entry.Type {
		case persistence.TypeString:
			s.data[entry.Key] = &Entry{
				Type:     TypeString,
				Str:      entry.Value,
				ExpireAt: entry.ExpireAt,
			}

		case persistence.TypeList:
			list := NewList()
			for _, v := range entry.Values {
				list.RPush(v)
			}
			s.data[entry.Key] = &Entry{
				Type:     TypeList,
				List:     list,
				ExpireAt: entry.ExpireAt,
			}
		}

		if entry.ExpireAt != nil {
			s.heap.Push(&HeapItem{
				Key:      entry.Key,
				ExpireAt: *entry.ExpireAt,
			})
		}
	}
	return nil
}

func (s *Server) Start() error {
	listener, err := net.Listen("tcp", s.addr)
	if err != nil {
		return err
	}

	s.listener = listener
	defer s.listener.Close()
	log.Printf("현재 서버가 [%s] 에서 리스닝중입니다.", s.addr)

	if err := s.store.Load("dump.rdb"); err != nil {
		log.Printf("RDB 로딩 실패: %v", err)
	} else {
		log.Println("RDB 파일 로딩 완료")
	}

	s.store.StartExpiry()
	defer s.store.StopExpiry()

	for {
		conn, err := s.listener.Accept()
		if err != nil {
			log.Print("연결 중 오류: ", err)
			continue
		} else {
			go s.handleConnection(conn)
		}
	}
}

func (e *Encoder) writeBytes(data []byte) error {
	_, err := e.w.Write(data)
	if err == nil {
		e.hash.Write(data)
	}
	return err
}

func (e *Encoder) WriteChecksum() error {
	checksum := e.hash.Sum32()
	buf := make([]byte, 4)
	binary.BigEndian.PutUint32(buf, checksum)
	_, err := e.w.Write(buf)  // writeBytes가 아닌 e.w.Write를 직접 사용
	return err
}

func (d *Decoder) ReadEntry() (*DecodedEntry, error) {
	typeBuf, err := d.readBytes(1)
	if err != nil {
		return nil, err
	}

	if typeBuf[0] == EOF {
		return nil, nil  // 정상 종료
	}

	// ...

	default:
		return nil, fmt.Errorf("unknown entry type: 0x%02x", typeBuf[0])
	}

func VerifyChecksum(path string) error {
	file, err := os.ReadFile(path)
	// ...

	content := file[:len(file)-ChecksumSize]
	stored := binary.BigEndian.Uint32(file[len(file)-ChecksumSize:])
	computed := crc32.ChecksumIEEE(content)

	if stored != computed {
		return fmt.Errorf(
			"Checksum이 일치하지 않습니다: stored=0x%08x, computed=0x%08x",
			stored, computed,
		)
	}
	return nil
}

content := data[:len(data)-persistence.ChecksumSize]
decoder := persistence.NewDecoder(bytes.NewReader(content))

if os.IsNotExist(err) {
	return nil
}

func TestReadStringEntry(t *testing.T) {
	// given: Header + String("hello", "world") + EOF
	data := encodeToBytes(t, func(enc *Encoder) {
		enc.WriteHeader()
		enc.WriteStringEntry("hello", "world", nil)
		enc.WriteEOF()
	})
	decoder := NewDecoder(bytes.NewReader(data))
	decoder.ReadHeader()

	// when
	entry, err := decoder.ReadEntry()

	// then
	if err != nil {
		t.Fatalf("에러 발생: %v", err)
	}
	if entry.Type != TypeString {
		t.Fatalf("Type: 0x%02x, expected: 0x%02x", entry.Type, TypeString)
	}
	if entry.Key != "hello" {
		t.Fatalf("Key: %s, expected: hello", entry.Key)
	}
	if entry.Value != "world" {
		t.Fatalf("Value: %s, expected: world", entry.Value)
	}
	if entry.ExpireAt != nil {
		t.Fatalf("ExpireAt이 nil이어야 합니다")
	}
}

func TestVerifyChecksum_Corrupted(t *testing.T) {
	// given: 정상 파일 작성 후 중간 바이트 변조
	path := filepath.Join(t.TempDir(), "corrupted.rdb")
	file, _ := os.Create(path)
	enc := NewEncoder(file)
	enc.WriteHeader()
	enc.WriteStringEntry("key", "value", nil)
	enc.WriteEOF()
	enc.WriteChecksum()
	enc.Flush()
	file.Close()

	// 파일 변조: 10번째 바이트를 뒤집는다
	data, _ := os.ReadFile(path)
	data[10] ^= 0xFF
	os.WriteFile(path, data, 0644)

	// when
	err := VerifyChecksum(path)

	// then
	if err == nil {
		t.Fatal("변조된 파일에 대해 에러가 발생해야 합니다")
	}
}

func TestLoad_StringEntries(t *testing.T) {
	// given: Save로 String 엔트리 저장
	store := New()
	store.Set("name", "gopher")
	store.Set("lang", "go")
	path := filepath.Join(t.TempDir(), "string.rdb")
	store.Save(path)

	// when: 새 Store에 Load
	loaded := New()
	err := loaded.Load(path)

	// then
	if err != nil {
		t.Fatalf("에러 발생: %v", err)
	}
	v1, ok1 := loaded.Get("name")
	if !ok1 || v1 != "gopher" {
		t.Fatalf("name: %s, exist: %v", v1, ok1)
	}
	v2, ok2 := loaded.Get("lang")
	if !ok2 || v2 != "go" {
		t.Fatalf("lang: %s, exist: %v", v2, ok2)
	}
}

func TestSaveAndLoad(t *testing.T) {
	// given: 서버1에서 SET + SAVE
	server1 := New(":16379")
	go server1.Start()
	time.Sleep(time.Second)

	conn1, _ := net.Dial("tcp", "localhost:16379")
	reader1 := bufio.NewReader(conn1)

	conn1.Write([]byte("*3\r\n$3\r\nSET\r\n$8\r\nload-key\r\n$5\r\nhello\r\n"))
	reader1.ReadString('\n')

	conn1.Write([]byte("*1\r\n$4\r\nSAVE\r\n"))
	saveResp, _ := reader1.ReadString('\n')
	if saveResp != "+OK\r\n" {
		t.Fatalf("SAVE 응답: %s", saveResp)
	}
	conn1.Close()

	// when: 새 서버2를 같은 dump.rdb로 시작
	server2 := New(":16380")
	go server2.Start()
	time.Sleep(time.Second)

	conn2, _ := net.Dial("tcp", "localhost:16380")
	defer conn2.Close()
	reader2 := bufio.NewReader(conn2)

	conn2.Write([]byte("*2\r\n$3\r\nGET\r\n$8\r\nload-key\r\n"))

	// then: 이전 서버에서 저장한 값이 복원됨
	lenLine, _ := reader2.ReadString('\n')
	if lenLine != "$5\r\n" {
		t.Fatalf("GET 길이: %s", lenLine)
	}
	valLine, _ := reader2.ReadString('\n')
	if valLine != "hello\r\n" {
		t.Fatalf("GET 값: %s", valLine)
	}

	os.Remove("dump.rdb")
}