Zustand localStorage Encryption: Protecting Sensitive Data with partialize and a Custom AES-GCM Adapter

partialize: (state) => ({
  user: state.user,
  accessToken: state.accessToken,
  // UI state like isLoading, modalOpen won't be saved if not returned
})

const myStorage: StateStorage = {
  getItem: (name) => { /* decrypt and return */ },
  setItem: (name, value) => { /* encrypt and save */ },
  removeItem: (name) => { /* delete */ },
}

pnpm add crypto-js
pnpm add -D @types/crypto-js

import { create } from 'zustand'
import { persist, createJSONStorage, StateStorage } from 'zustand/middleware'
import CryptoJS from 'crypto-js'
 
// NEXT_PUBLIC_STORAGE_KEY가 없으면 런타임에 암호화가 빈 키로 동작합니다.
// non-null assertion(!)은 빌드 타임에 환경변수가 반드시 주입되는 경우를 전제합니다.
const SECRET_KEY = process.env.NEXT_PUBLIC_STORAGE_KEY!
 
// 1. 암호화 어댑터
const encryptedStorage: StateStorage = {
  getItem: (name) => {
    const encrypted = localStorage.getItem(name)
    if (!encrypted) return null
    try {
      const bytes = CryptoJS.AES.decrypt(encrypted, SECRET_KEY)
      return bytes.toString(CryptoJS.enc.Utf8)
    } catch {
      return null // 손상된 데이터는 null 반환 → 초기 상태로 안전하게 복귀
    }
  },
  setItem: (name, value) => {
    const encrypted = CryptoJS.AES.encrypt(value, SECRET_KEY).toString()
    localStorage.setItem(name, encrypted)
  },
  removeItem: (name) => localStorage.removeItem(name),
}
 
// 2. 스토어 정의
interface AuthState {
  user: { id: string; name: string } | null
  accessToken: string | null
  isLoading: boolean    // UI 상태
  modalOpen: boolean    // UI 상태
  setUser: (user: AuthState['user'], token: string) => void
  logout: () => void
}
 
const useAuthStore = create<AuthState>()(
  persist(
    (set) => ({
      user: null,
      accessToken: null,
      isLoading: false,
      modalOpen: false,
      setUser: (user, accessToken) => set({ user, accessToken }),
      logout: () => set({ user: null, accessToken: null }),
    }),
    {
      name: 'auth-store',
      storage: createJSONStorage(() => encryptedStorage),
      // isLoading, modalOpen은 저장 대상에서 제외
      partialize: (state) => ({
        user: state.user,
        accessToken: state.accessToken,
      }),
    }
  )
)

import { create } from 'zustand'
import { persist, createJSONStorage, StateStorage } from 'zustand/middleware'
 
const SECRET_KEY = process.env.NEXT_PUBLIC_STORAGE_KEY!
const ALGORITHM = { name: 'AES-GCM', length: 256 } as const
 
// 매번 새로 만들면 비효율적이라 캐싱
// 주의: 키는 모듈이 로드될 때 한 번 계산됩니다.
// 환경변수가 변경되어도(핫리로드 등) 재계산이 일어나지 않습니다.
let cachedKey: CryptoKey | null = null
 
async function deriveKey(secret: string): Promise<CryptoKey> {
  if (cachedKey) return cachedKey
  const enc = new TextEncoder()
  const keyMaterial = await crypto.subtle.importKey(
    'raw', enc.encode(secret), 'PBKDF2', false, ['deriveKey']
  )
  // 솔트를 고정 문자열로 사용하면 PBKDF2의 레인보우 테이블 방어 효과가 약해집니다.
  // 고보안 환경에서는 crypto.getRandomValues()로 생성한 솔트를
  // localStorage에 별도 저장하는 방식을 고려하는 것을 권장합니다.
  cachedKey = await crypto.subtle.deriveKey(
    {
      name: 'PBKDF2',
      salt: enc.encode('zustand-persist-salt-v1'),
      iterations: 100_000,
      hash: 'SHA-256',
    },
    keyMaterial,
    ALGORITHM,
    false,
    ['encrypt', 'decrypt']
  )
  return cachedKey
}
 
async function encrypt(text: string): Promise<string> {
  const key = await deriveKey(SECRET_KEY)
  const iv = crypto.getRandomValues(new Uint8Array(12)) // AES-GCM은 12바이트 IV
  const encoded = new TextEncoder().encode(text)
  const ciphertext = await crypto.subtle.encrypt({ name: 'AES-GCM', iv }, key, encoded)
  // IV + 암호문을 하나의 base64 문자열로 합쳐서 저장
  const combined = new Uint8Array([...iv, ...new Uint8Array(ciphertext)])
  // 스프레드로 String.fromCharCode에 큰 배열을 넘기면 콜스택 초과 위험이 있어 reduce 사용
  const binaryString = combined.reduce((acc, byte) => acc + String.fromCharCode(byte), '')
  return btoa(binaryString)
}
 
async function decrypt(data: string): Promise<string> {
  const key = await deriveKey(SECRET_KEY)
  const combined = Uint8Array.from(atob(data), (c) => c.charCodeAt(0))
  const iv = combined.slice(0, 12)
  const ciphertext = combined.slice(12)
  const decrypted = await crypto.subtle.decrypt({ name: 'AES-GCM', iv }, key, ciphertext)
  return new TextDecoder().decode(decrypted)
}
 
// 비동기 어댑터
const webCryptoStorage: StateStorage = {
  getItem: async (name) => {
    const val = localStorage.getItem(name)
    if (!val) return null
    try {
      return await decrypt(val)
    } catch {
      return null
    }
  },
  setItem: async (name, value) => {
    localStorage.setItem(name, await encrypt(value))
  },
  removeItem: (name) => localStorage.removeItem(name),
}
 
// 하이드레이션 완료 여부 추적을 위해 인터페이스에 명시 필요
interface AuthState {
  user: { id: string; name: string } | null
  accessToken: string | null
  _hasHydrated: boolean
  setHasHydrated: (value: boolean) => void
  setUser: (user: AuthState['user'], token: string) => void
  logout: () => void
}
 
const useAuthStore = create<AuthState>()(
  persist(
    (set) => ({
      user: null,
      accessToken: null,
      _hasHydrated: false,
      setHasHydrated: (value) => set({ _hasHydrated: value }),
      setUser: (user, accessToken) => set({ user, accessToken }),
      logout: () => set({ user: null, accessToken: null }),
    }),
    {
      name: 'auth-store-v2',
      storage: createJSONStorage(() => webCryptoStorage),
      partialize: (state) => ({
        user: state.user,
        accessToken: state.accessToken,
      }),
      // onRehydrateStorage 내부 콜백의 state는 복원된 상태 스냅샷입니다.
      // state?.setState(...)는 동작하지 않으며, 스토어에 정의한 액션을 호출해야 합니다.
      onRehydrateStorage: () => (state) => {
        state?.setHasHydrated(true)
      },
    }
  )
)
 
// 컴포넌트에서 하이드레이션 대기
function AuthGuard({ children }: { children: React.ReactNode }) {
  const hasHydrated = useAuthStore((s) => s._hasHydrated)
  if (!hasHydrated) return <LoadingSpinner />
  return <>{children}</>
}

// encryptedStorage는 예시 1에서 정의한 CryptoJS 기반 어댑터입니다.
 
// 민감 데이터 전용 스토어 — 암호화 어댑터 사용
const useSecureStore = create<SecureState>()(
  persist(
    (set) => ({
      accessToken: null,
      refreshToken: null,
      setTokens: (access, refresh) => set({ accessToken: access, refreshToken: refresh }),
      clearTokens: () => set({ accessToken: null, refreshToken: null }),
    }),
    {
      name: 'secure-store',
      storage: createJSONStorage(() => encryptedStorage),
      partialize: (s) => ({
        accessToken: s.accessToken,
        refreshToken: s.refreshToken,
      }),
    }
  )
)
 
// 비민감 설정 스토어 — 일반 localStorage 사용
const useSettingsStore = create<SettingsState>()(
  persist(
    (set) => ({
      theme: 'light' as const,
      locale: 'ko',
      setTheme: (theme) => set({ theme }),
      setLocale: (locale) => set({ locale }),
    }),
    {
      name: 'settings-store',
      storage: createJSONStorage(() => localStorage),
      partialize: (s) => ({ theme: s.theme, locale: s.locale }),
    }
  )
)