| 1 | package worktree |
| 2 | |
| 3 | import ( |
| 4 | "context" |
| 5 | "crypto/sha256" |
| 6 | "encoding/hex" |
| 7 | "encoding/json" |
| 8 | "errors" |
| 9 | "fmt" |
| 10 | "io" |
| 11 | "os" |
| 12 | "path/filepath" |
| 13 | "runtime" |
| 14 | "sort" |
| 15 | "strings" |
| 16 | |
| 17 | "reasonix/internal/fileutil" |
| 18 | ) |
| 19 | |
| 20 | const ( |
| 21 | cleanupStateVersion = 2 |
| 22 | legacyCleanupStateVersion = 1 |
| 23 | cleanupStateName = "cleanup-state.json" |
| 24 | |
| 25 | cleanupStagePlanned = "planned" |
| 26 | cleanupStageRetained = "retained" |
| 27 | |
| 28 | legacyCleanupStagePrepared = "prepared" |
| 29 | legacyCleanupStageUnregistered = "unregistered" |
| 30 | ) |
| 31 | |
| 32 | type cleanupManifestEntry struct { |
| 33 | Path string `json:"path"` |
| 34 | Mode uint32 `json:"mode"` |
| 35 | Digest string `json:"digest,omitempty"` |
| 36 | } |
| 37 | |
| 38 | // cleanupState v2 records a retained, registered recovery checkout. It never |
| 39 | // describes deletion work: old and unknown writers can only preserve it. |
| 40 | type cleanupState struct { |
| 41 | Version int `json:"version"` |
| 42 | OriginalRoot string `json:"originalRoot"` |
| 43 | RecoveryRoot string `json:"recoveryRoot"` |
| 44 | WorktreeBranch string `json:"worktreeBranch"` |
| 45 | WorktreeHead string `json:"worktreeHead"` |
| 46 | Stage string `json:"stage"` |
| 47 | } |
| 48 | |
| 49 | // legacyCleanupState is read-only compatibility for v1 journals that may |
| 50 | // have stopped between checkout detachment and physical deletion. |
| 51 | type legacyCleanupState struct { |
| 52 | Version int `json:"version"` |
| 53 | OriginalRoot string `json:"originalRoot"` |
| 54 | RegisteredRoot string `json:"registeredRoot"` |
| 55 | DetachedRoot string `json:"detachedRoot"` |
| 56 | WorktreeBranch string `json:"worktreeBranch"` |
| 57 | WorktreeHead string `json:"worktreeHead"` |
| 58 | Stage string `json:"stage"` |
| 59 | Manifest []cleanupManifestEntry `json:"manifest"` |
| 60 | } |
| 61 | |
| 62 | type cleanupJournal struct { |
| 63 | Current *cleanupState |
| 64 | Legacy *legacyCleanupState |
| 65 | } |
| 66 | |
| 67 | func cleanupJournalPath(metadata mergeMetadata) string { |
| 68 | return filepath.Join(filepath.Dir(metadata.WorktreeRoot), cleanupStateName) |
| 69 | } |
| 70 | |
| 71 | func writeCleanupState(metadata mergeMetadata, state cleanupState) error { |
| 72 | body, err := encodeCleanupState(state) |
| 73 | if err != nil { |
| 74 | return err |
| 75 | } |
| 76 | if err := fileutil.AtomicWriteFileStrict(cleanupJournalPath(metadata), body, 0o600); err != nil { |
| 77 | return fmt.Errorf("publish cleanup state: %w", err) |
| 78 | } |
| 79 | return nil |
| 80 | } |
| 81 | |
| 82 | func createCleanupState(metadata mergeMetadata, state cleanupState) error { |
| 83 | body, err := encodeCleanupState(state) |
| 84 | if err != nil { |
| 85 | return err |
| 86 | } |
| 87 | if err := fileutil.AtomicCreateFile(cleanupJournalPath(metadata), body, 0o600); err != nil { |
| 88 | return fmt.Errorf("publish initial cleanup state: %w", err) |
| 89 | } |
| 90 | return nil |
| 91 | } |
| 92 | |
| 93 | func encodeCleanupState(state cleanupState) ([]byte, error) { |
| 94 | body, err := json.MarshalIndent(state, "", " ") |
| 95 | if err != nil { |
| 96 | return nil, fmt.Errorf("encode cleanup state: %w", err) |
| 97 | } |
| 98 | return append(body, '\n'), nil |
| 99 | } |
| 100 | |
| 101 | func readCleanupState(metadata mergeMetadata, expectedHead string) (cleanupJournal, bool, error) { |
| 102 | path := cleanupJournalPath(metadata) |
| 103 | info, err := os.Lstat(path) |
| 104 | if errors.Is(err, os.ErrNotExist) { |
| 105 | return cleanupJournal{}, false, nil |
| 106 | } |
| 107 | if err != nil { |
| 108 | return cleanupJournal{}, false, fmt.Errorf("inspect cleanup state: %w", err) |
| 109 | } |
| 110 | if !info.Mode().IsRegular() || info.Mode()&os.ModeSymlink != 0 { |
| 111 | return cleanupJournal{}, false, errors.New("cleanup state is not a regular file") |
| 112 | } |
| 113 | if runtime.GOOS != "windows" && info.Mode().Perm()&0o077 != 0 { |
| 114 | return cleanupJournal{}, false, errors.New("cleanup state permissions are too broad") |
| 115 | } |
| 116 | body, err := os.ReadFile(path) |
| 117 | if err != nil { |
| 118 | return cleanupJournal{}, false, fmt.Errorf("read cleanup state: %w", err) |
| 119 | } |
| 120 | var envelope struct { |
| 121 | Version int `json:"version"` |
| 122 | } |
| 123 | if err := json.Unmarshal(body, &envelope); err != nil { |
| 124 | return cleanupJournal{}, false, fmt.Errorf("decode cleanup state version: %w", err) |
| 125 | } |
| 126 | switch envelope.Version { |
| 127 | case cleanupStateVersion: |
| 128 | var state cleanupState |
| 129 | if err := decodeCleanupJSON(body, &state); err != nil { |
| 130 | return cleanupJournal{}, false, err |
| 131 | } |
| 132 | if err := validateCleanupState(metadata, expectedHead, state); err != nil { |
| 133 | return cleanupJournal{}, false, err |
| 134 | } |
| 135 | return cleanupJournal{Current: &state}, true, nil |
| 136 | case legacyCleanupStateVersion: |
| 137 | var state legacyCleanupState |
| 138 | if err := decodeCleanupJSON(body, &state); err != nil { |
| 139 | return cleanupJournal{}, false, err |
| 140 | } |
| 141 | if err := validateLegacyCleanupState(metadata, expectedHead, state); err != nil { |
| 142 | return cleanupJournal{}, false, err |
| 143 | } |
| 144 | return cleanupJournal{Legacy: &state}, true, nil |
| 145 | default: |
| 146 | return cleanupJournal{}, false, fmt.Errorf("unsupported cleanup state version %d", envelope.Version) |
| 147 | } |
| 148 | } |
| 149 | |
| 150 | func decodeCleanupJSON(body []byte, destination any) error { |
| 151 | decoder := json.NewDecoder(strings.NewReader(string(body))) |
| 152 | decoder.DisallowUnknownFields() |
| 153 | if err := decoder.Decode(destination); err != nil { |
| 154 | return fmt.Errorf("decode cleanup state: %w", err) |
| 155 | } |
| 156 | if err := decoder.Decode(&struct{}{}); !errors.Is(err, io.EOF) { |
| 157 | return errors.New("decode cleanup state: trailing JSON content") |
| 158 | } |
| 159 | return nil |
| 160 | } |
| 161 | |
| 162 | func validateCleanupState(metadata mergeMetadata, expectedHead string, state cleanupState) error { |
| 163 | if state.Version != cleanupStateVersion { |
| 164 | return fmt.Errorf("unsupported cleanup state version %d", state.Version) |
| 165 | } |
| 166 | if !sameCleanupPath(state.OriginalRoot, metadata.WorktreeRoot) || |
| 167 | state.WorktreeBranch != metadata.WorktreeBranch || state.WorktreeHead != expectedHead { |
| 168 | return errors.New("cleanup state identity does not match the merge receipt") |
| 169 | } |
| 170 | if state.Stage != cleanupStagePlanned && state.Stage != cleanupStageRetained { |
| 171 | return fmt.Errorf("unsupported cleanup stage %q", state.Stage) |
| 172 | } |
| 173 | return validateCleanupRecoveryPath(metadata, state.RecoveryRoot) |
| 174 | } |
| 175 | |
| 176 | func validateLegacyCleanupState(metadata mergeMetadata, expectedHead string, state legacyCleanupState) error { |
| 177 | if state.Version != legacyCleanupStateVersion { |
| 178 | return fmt.Errorf("unsupported cleanup state version %d", state.Version) |
| 179 | } |
| 180 | if !sameCleanupPath(state.OriginalRoot, metadata.WorktreeRoot) || |
| 181 | state.WorktreeBranch != metadata.WorktreeBranch || state.WorktreeHead != expectedHead { |
| 182 | return errors.New("cleanup state identity does not match the merge receipt") |
| 183 | } |
| 184 | if state.Stage != legacyCleanupStagePrepared && state.Stage != legacyCleanupStageUnregistered { |
| 185 | return fmt.Errorf("unsupported cleanup stage %q", state.Stage) |
| 186 | } |
| 187 | if err := validateCleanupRecoveryPath(metadata, state.RegisteredRoot); err != nil { |
| 188 | return err |
| 189 | } |
| 190 | if err := validateCleanupRecoveryPath(metadata, state.DetachedRoot); err != nil { |
| 191 | return err |
| 192 | } |
| 193 | if sameCleanupPath(state.RegisteredRoot, state.DetachedRoot) { |
| 194 | return errors.New("cleanup state paths are not distinct") |
| 195 | } |
| 196 | if state.Manifest == nil { |
| 197 | return errors.New("cleanup state manifest is missing") |
| 198 | } |
| 199 | seen := map[string]struct{}{} |
| 200 | for _, entry := range state.Manifest { |
| 201 | if err := validateStatePath(entry.Path); err != nil { |
| 202 | return fmt.Errorf("invalid cleanup manifest path: %w", err) |
| 203 | } |
| 204 | if _, ok := seen[entry.Path]; ok { |
| 205 | return fmt.Errorf("duplicate cleanup manifest path %q", entry.Path) |
| 206 | } |
| 207 | seen[entry.Path] = struct{}{} |
| 208 | } |
| 209 | return nil |
| 210 | } |
| 211 | |
| 212 | func validateCleanupRecoveryPath(metadata mergeMetadata, path string) error { |
| 213 | cleanupDir := filepath.Join(filepath.Dir(metadata.WorktreeRoot), ".reasonix-cleanup") |
| 214 | cleanupInfo, err := os.Lstat(cleanupDir) |
| 215 | if err != nil || !cleanupInfo.IsDir() || cleanupInfo.Mode()&os.ModeSymlink != 0 { |
| 216 | return errors.New("cleanup recovery directory is not a real directory") |
| 217 | } |
| 218 | realCleanupDir, err := filepath.EvalSymlinks(cleanupDir) |
| 219 | if err != nil { |
| 220 | return errors.New("cleanup recovery directory cannot be resolved") |
| 221 | } |
| 222 | realPath, err := resolveMissingCleanupPath(path) |
| 223 | if err != nil { |
| 224 | return errors.New("cleanup recovery path cannot be resolved") |
| 225 | } |
| 226 | rel, err := filepath.Rel(filepath.Clean(realCleanupDir), filepath.Clean(realPath)) |
| 227 | if err != nil || rel == "." || rel == ".." || strings.HasPrefix(rel, ".."+string(filepath.Separator)) || strings.Contains(rel, string(filepath.Separator)) { |
| 228 | return errors.New("cleanup state path escapes the allocation recovery directory") |
| 229 | } |
| 230 | return nil |
| 231 | } |
| 232 | |
| 233 | func captureCleanupManifest(ctx context.Context, root string) ([]cleanupManifestEntry, error) { |
| 234 | rootInfo, err := os.Lstat(root) |
| 235 | if err != nil { |
| 236 | return nil, fmt.Errorf("inspect cleanup root: %w", err) |
| 237 | } |
| 238 | if !rootInfo.IsDir() || rootInfo.Mode()&os.ModeSymlink != 0 { |
| 239 | return nil, errors.New("cleanup root is not a real directory") |
| 240 | } |
| 241 | manifest := []cleanupManifestEntry{} |
| 242 | err = filepath.WalkDir(root, func(path string, entry os.DirEntry, walkErr error) error { |
| 243 | if walkErr != nil { |
| 244 | return walkErr |
| 245 | } |
| 246 | if err := ctx.Err(); err != nil { |
| 247 | return err |
| 248 | } |
| 249 | if path == root { |
| 250 | return nil |
| 251 | } |
| 252 | relative, err := filepath.Rel(root, path) |
| 253 | if err != nil { |
| 254 | return err |
| 255 | } |
| 256 | relative = filepath.ToSlash(relative) |
| 257 | if err := validateStatePath(relative); err != nil { |
| 258 | return err |
| 259 | } |
| 260 | info, err := os.Lstat(path) |
| 261 | if err != nil { |
| 262 | return err |
| 263 | } |
| 264 | item := cleanupManifestEntry{Path: relative, Mode: uint32(info.Mode())} |
| 265 | switch { |
| 266 | case info.IsDir(): |
| 267 | case info.Mode().IsRegular(): |
| 268 | item.Digest, err = digestCleanupFile(ctx, path) |
| 269 | case info.Mode()&os.ModeSymlink != 0: |
| 270 | var target string |
| 271 | target, err = os.Readlink(path) |
| 272 | if err == nil { |
| 273 | digest := sha256.Sum256([]byte(target)) |
| 274 | item.Digest = hex.EncodeToString(digest[:]) |
| 275 | } |
| 276 | default: |
| 277 | err = fmt.Errorf("cleanup path %q has unsupported type %s", relative, info.Mode().Type()) |
| 278 | } |
| 279 | if err != nil { |
| 280 | return err |
| 281 | } |
| 282 | manifest = append(manifest, item) |
| 283 | return nil |
| 284 | }) |
| 285 | if err != nil { |
| 286 | return nil, fmt.Errorf("snapshot cleanup checkout: %w", err) |
| 287 | } |
| 288 | sort.Slice(manifest, func(left, right int) bool { return manifest[left].Path < manifest[right].Path }) |
| 289 | return manifest, nil |
| 290 | } |
| 291 | |
| 292 | func digestCleanupFile(ctx context.Context, path string) (string, error) { |
| 293 | file, err := os.Open(path) |
| 294 | if err != nil { |
| 295 | return "", err |
| 296 | } |
| 297 | defer file.Close() |
| 298 | hash := sha256.New() |
| 299 | buffer := make([]byte, 128*1024) |
| 300 | for { |
| 301 | if err := ctx.Err(); err != nil { |
| 302 | return "", err |
| 303 | } |
| 304 | count, readErr := file.Read(buffer) |
| 305 | if count > 0 { |
| 306 | if _, err := hash.Write(buffer[:count]); err != nil { |
| 307 | return "", err |
| 308 | } |
| 309 | } |
| 310 | if errors.Is(readErr, io.EOF) { |
| 311 | return hex.EncodeToString(hash.Sum(nil)), nil |
| 312 | } |
| 313 | if readErr != nil { |
| 314 | return "", readErr |
| 315 | } |
| 316 | } |
| 317 | } |
| 318 | |
| 319 | func manifestsEqual(expected, actual []cleanupManifestEntry) bool { |
| 320 | if len(expected) != len(actual) { |
| 321 | return false |
| 322 | } |
| 323 | for index := range expected { |
| 324 | if expected[index] != actual[index] { |
| 325 | return false |
| 326 | } |
| 327 | } |
| 328 | return true |
| 329 | } |
| 330 |