sow-tools/internal/erf/erf.go

package erf

import (
	"bytes"
	"crypto/sha256"
	"encoding/binary"
	"encoding/hex"
	"fmt"
	"io"
	"os"
	"regexp"
	"sort"
	"strings"
)

var expectedSHA256Pattern = regexp.MustCompile(`^[0-9a-f]{64}$`)

const (
	headerSize = 160
	versionV10 = "V1.0"
	strRefNone = 0xFFFFFFFF
	typeMOD    = "MOD "
	typeHAK    = "HAK "
)

type Archive struct {
	FileType  string
	Version   string
	Resources []Resource
}

type Resource struct {
	Name       string
	Type       uint16
	Data       []byte
	SourcePath string
	Size       int64
	// ExpectedSHA256, when set, is the lowercase hex SHA-256 the streamed
	// SourcePath payload must hash to. A mismatch fails the write.
	ExpectedSHA256 string
}

type header struct {
	FileType              [4]byte
	Version               [4]byte
	LanguageCount         uint32
	LocalizedStringSize   uint32
	EntryCount            uint32
	LocalizedStringOffset uint32
	KeyListOffset         uint32
	ResourceListOffset    uint32
	BuildYear             uint32
	BuildDay              uint32
	DescriptionStrRef     uint32
	Reserved              [116]byte
}

type keyEntry struct {
	ResRef       [16]byte
	ResourceID   uint32
	ResourceType uint16
	Unused       uint16
}

type resourceEntry struct {
	Offset uint32
	Size   uint32
}

// extensionTypes and typeExtensions must stay exact inverses of each other;
// init() panics if they drift apart.
//
// Numbers below 0x0BB8 follow upstream neverwinter.nim
// (neverwinter/restype.nim). The 0x0BB8 and up entries are lyt/vis/mdx: real
// Aurora archive types that NWN1 ships in its own data/*.bif but upstream
// happens not to register. xoreos corroborates those three numbers
// (src/aurora/types.h).
//
// This table is NWN:EE only. Do not add NWN2 formats (mdb, gr2, wlk, xml):
// NWN:EE either gives that number to something else (2070 is xbc here and mdb
// in NWN2) or has no number for it at all, so packing one into a HAK writes a
// resource the game misreads. See the reserved list in erf_test.go.
var extensionTypes = map[string]uint16{
	"res": 0x0000,
	"bmp": 0x0001,
	"mve": 0x0002,
	"tga": 0x0003,
	"wav": 0x0004,
	"plt": 0x0006,
	"ini": 0x0007,
	"bmu": 0x0008,
	"txt": 0x000A,
	"mdl": 0x07D2,
	"nss": 0x07D9,
	"ncs": 0x07DA,
	"are": 0x07DC,
	"set": 0x07DD,
	"ifo": 0x07DE,
	"bic": 0x07DF,
	"wok": 0x07E0,
	"2da": 0x07E1,
	"tlk": 0x07E2,
	"txi": 0x07E6,
	"git": 0x07E7,
	"uti": 0x07E9,
	"utc": 0x07EB,
	"dlg": 0x07ED,
	"itp": 0x07EE,
	"utt": 0x07F0,
	"dds": 0x07F1,
	"uts": 0x07F3,
	"ltr": 0x07F4,
	"gff": 0x07F5,
	"fac": 0x07F6,
	"ute": 0x07F8,
	"utd": 0x07FA,
	"utp": 0x07FC,
	"dft": 0x07FD,
	"gic": 0x07FE,
	"gui": 0x07FF,
	"utm": 0x0803,
	"dwk": 0x0804,
	"pwk": 0x0805,
	"utg": 0x0807,
	"jrl": 0x0808,
	"utw": 0x080A,
	"ssf": 0x080C,
	"hak": 0x080D,
	"nwm": 0x080E,
	"bik": 0x080F,
	"ndb": 0x0810,
	"ptm": 0x0811,
	"ptt": 0x0812,
	"shd": 0x0815,
	"mtr": 0x0818,
	"lod": 0x081E,
	"gif": 0x081F,
	"png": 0x0820,
	"jpg": 0x0821,
	"lyt": 0x0BB8,
	"vis": 0x0BB9,
	"mdx": 0x0BC0,
}

var typeExtensions = map[uint16]string{
	0x0000: "res",
	0x0001: "bmp",
	0x0002: "mve",
	0x0003: "tga",
	0x0004: "wav",
	0x0006: "plt",
	0x0007: "ini",
	0x0008: "bmu",
	0x000A: "txt",
	0x07D2: "mdl",
	0x07D9: "nss",
	0x07DA: "ncs",
	0x07DC: "are",
	0x07DD: "set",
	0x07DE: "ifo",
	0x07DF: "bic",
	0x07E0: "wok",
	0x07E1: "2da",
	0x07E2: "tlk",
	0x07E6: "txi",
	0x07E7: "git",
	0x07E9: "uti",
	0x07EB: "utc",
	0x07ED: "dlg",
	0x07EE: "itp",
	0x07F0: "utt",
	0x07F1: "dds",
	0x07F3: "uts",
	0x07F4: "ltr",
	0x07F5: "gff",
	0x07F6: "fac",
	0x07F8: "ute",
	0x07FA: "utd",
	0x07FC: "utp",
	0x07FD: "dft",
	0x07FE: "gic",
	0x07FF: "gui",
	0x0803: "utm",
	0x0804: "dwk",
	0x0805: "pwk",
	0x0807: "utg",
	0x0808: "jrl",
	0x080A: "utw",
	0x080C: "ssf",
	0x080D: "hak",
	0x080E: "nwm",
	0x080F: "bik",
	0x0810: "ndb",
	0x0811: "ptm",
	0x0812: "ptt",
	0x0815: "shd",
	0x0818: "mtr",
	0x081E: "lod",
	0x081F: "gif",
	0x0820: "png",
	0x0821: "jpg",
	0x0BB8: "lyt",
	0x0BB9: "vis",
	0x0BC0: "mdx",
}

func init() {
	for ext, resourceType := range extensionTypes {
		canonicalExt, ok := typeExtensions[resourceType]
		if !ok {
			panic(fmt.Sprintf("missing canonical extension for resource type 0x%04X", resourceType))
		}
		if canonicalExt != ext {
			panic(fmt.Sprintf("resource type 0x%04X is %q in extensionTypes but %q in typeExtensions", resourceType, ext, canonicalExt))
		}
	}
	for resourceType, ext := range typeExtensions {
		if registered, ok := extensionTypes[ext]; !ok || registered != resourceType {
			panic(fmt.Sprintf("extension %q is 0x%04X in typeExtensions but 0x%04X in extensionTypes (registered=%v)", ext, resourceType, registered, ok))
		}
	}
}

func New(fileType string, resources []Resource) Archive {
	normalized := make([]Resource, len(resources))
	copy(normalized, resources)
	sort.Slice(normalized, func(i, j int) bool {
		if normalized[i].Name == normalized[j].Name {
			return normalized[i].Type < normalized[j].Type
		}
		return normalized[i].Name < normalized[j].Name
	})

	return Archive{
		FileType:  padFour(fileType),
		Version:   versionV10,
		Resources: normalized,
	}
}

func ArchiveSize(resources []Resource) int64 {
	return int64(headerSize+len(resources)*24+len(resources)*8) + totalResourceBytes(resources)
}

func Write(w io.Writer, archive Archive) error {
	if archive.Version == "" {
		archive.Version = versionV10
	}
	if archive.FileType == "" {
		archive.FileType = typeMOD
	}

	keys := make([]keyEntry, 0, len(archive.Resources))
	entries := make([]resourceEntry, 0, len(archive.Resources))

	dataOffset := uint32(headerSize + len(archive.Resources)*24 + len(archive.Resources)*8)
	for index, resource := range archive.Resources {
		if len(resource.Name) > 16 {
			return fmt.Errorf("resource %q exceeds 16-byte resref limit", resource.Name)
		}
		size, err := payloadSize(resource)
		if err != nil {
			return err
		}
		if size > int64(^uint32(0)) {
			return fmt.Errorf("resource %q exceeds 4 GiB resource size limit", resource.Name)
		}

		entry := resourceEntry{
			Offset: dataOffset,
			Size:   uint32(size),
		}
		entries = append(entries, entry)
		dataOffset += uint32(size)

		var key keyEntry
		copy(key.ResRef[:], []byte(resource.Name))
		key.ResourceID = uint32(index)
		key.ResourceType = resource.Type
		keys = append(keys, key)
	}

	hdr := header{
		LanguageCount:         0,
		LocalizedStringSize:   0,
		EntryCount:            uint32(len(archive.Resources)),
		LocalizedStringOffset: headerSize,
		KeyListOffset:         headerSize,
		ResourceListOffset:    uint32(headerSize + len(keys)*24),
		BuildYear:             0,
		BuildDay:              0,
		DescriptionStrRef:     strRefNone,
	}
	copy(hdr.FileType[:], []byte(padFour(archive.FileType)))
	copy(hdr.Version[:], []byte(padFour(archive.Version)))

	if err := binary.Write(w, binary.LittleEndian, hdr); err != nil {
		return err
	}
	if err := binary.Write(w, binary.LittleEndian, keys); err != nil {
		return err
	}
	if err := binary.Write(w, binary.LittleEndian, entries); err != nil {
		return err
	}
	for _, resource := range archive.Resources {
		if err := writeResourceData(w, resource); err != nil {
			return err
		}
	}
	return nil
}

// IndexEntry locates one resource inside an archive without holding its
// payload. Streaming callers read one payload at a time from these, so peak
// memory tracks the largest resource instead of the whole archive.
type IndexEntry struct {
	Name   string
	Type   uint16
	Offset int64
	Size   int64
}

// Index is the header plus the resource table of an ERF: everything except the
// payloads.
type Index struct {
	FileType string
	Version  string
	Entries  []IndexEntry
}

// ReadIndex parses the tables of an ERF of the given size, reading only the
// header, the key list and the resource list.
func ReadIndex(r io.ReaderAt, size int64) (Index, error) {
	if size < headerSize {
		return Index{}, fmt.Errorf("erf file too small: %d bytes", size)
	}

	var hdr header
	if err := binary.Read(io.NewSectionReader(r, 0, headerSize), binary.LittleEndian, &hdr); err != nil {
		return Index{}, fmt.Errorf("decode erf header: %w", err)
	}

	if int64(hdr.KeyListOffset)+int64(hdr.EntryCount)*24 > size {
		return Index{}, fmt.Errorf("erf key list exceeds file bounds")
	}
	keys := make([]keyEntry, hdr.EntryCount)
	keyReader := io.NewSectionReader(r, int64(hdr.KeyListOffset), int64(hdr.EntryCount)*24)
	if err := binary.Read(keyReader, binary.LittleEndian, &keys); err != nil {
		return Index{}, fmt.Errorf("decode key list: %w", err)
	}

	if int64(hdr.ResourceListOffset)+int64(hdr.EntryCount)*8 > size {
		return Index{}, fmt.Errorf("erf resource list exceeds file bounds")
	}
	entries := make([]resourceEntry, hdr.EntryCount)
	entryReader := io.NewSectionReader(r, int64(hdr.ResourceListOffset), int64(hdr.EntryCount)*8)
	if err := binary.Read(entryReader, binary.LittleEndian, &entries); err != nil {
		return Index{}, fmt.Errorf("decode resource list: %w", err)
	}

	index := Index{
		FileType: string(hdr.FileType[:]),
		Version:  string(hdr.Version[:]),
		Entries:  make([]IndexEntry, 0, hdr.EntryCount),
	}
	for position, key := range keys {
		entry := entries[position]
		if int64(entry.Offset)+int64(entry.Size) > size {
			return Index{}, fmt.Errorf("resource %d exceeds file bounds", position)
		}
		index.Entries = append(index.Entries, IndexEntry{
			Name:   string(bytes.TrimRight(key.ResRef[:], "\x00")),
			Type:   key.ResourceType,
			Offset: int64(entry.Offset),
			Size:   int64(entry.Size),
		})
	}
	return index, nil
}

// ReadPayload returns one resource's bytes.
func ReadPayload(r io.ReaderAt, entry IndexEntry) ([]byte, error) {
	payload := make([]byte, entry.Size)
	if _, err := r.ReadAt(payload, entry.Offset); err != nil {
		return nil, fmt.Errorf("read resource %q: %w", entry.Name, err)
	}
	return payload, nil
}

// Read materialises a whole archive. Payloads are subslices of the buffer the
// archive was read into, so nothing is copied twice; callers that only need one
// resource at a time should use ReadIndex instead.
func Read(r io.Reader) (Archive, error) {
	data, err := io.ReadAll(r)
	if err != nil {
		return Archive{}, fmt.Errorf("read erf: %w", err)
	}
	index, err := ReadIndex(bytes.NewReader(data), int64(len(data)))
	if err != nil {
		return Archive{}, err
	}

	resources := make([]Resource, 0, len(index.Entries))
	for _, entry := range index.Entries {
		resources = append(resources, Resource{
			Name: entry.Name,
			Type: entry.Type,
			Data: data[entry.Offset : entry.Offset+entry.Size],
			Size: entry.Size,
		})
	}

	return Archive{
		FileType:  index.FileType,
		Version:   index.Version,
		Resources: resources,
	}, nil
}

func ResourceTypeForExtension(extension string) (uint16, bool) {
	resourceType, ok := extensionTypes[strings.TrimPrefix(strings.ToLower(extension), ".")]
	return resourceType, ok
}

func HAKResourceTypeForExtension(extension string) (uint16, bool) {
	return ResourceTypeForExtension(extension)
}

func ExtensionForResourceType(resourceType uint16) (string, bool) {
	extension, ok := typeExtensions[resourceType]
	return extension, ok
}

func padFour(value string) string {
	value = strings.ToUpper(value)
	if len(value) >= 4 {
		return value[:4]
	}
	return value + strings.Repeat(" ", 4-len(value))
}

func totalResourceBytes(resources []Resource) int64 {
	var total int64
	for _, resource := range resources {
		switch {
		case resource.Size > 0:
			total += resource.Size
		default:
			total += int64(len(resource.Data))
		}
	}
	return total
}

func payloadSize(resource Resource) (int64, error) {
	if resource.Size > 0 {
		return resource.Size, nil
	}
	if resource.SourcePath != "" && len(resource.Data) == 0 {
		info, err := os.Stat(resource.SourcePath)
		if err != nil {
			return 0, fmt.Errorf("stat resource %q: %w", resource.SourcePath, err)
		}
		return info.Size(), nil
	}
	return int64(len(resource.Data)), nil
}

func writeResourceData(w io.Writer, resource Resource) error {
	if len(resource.Data) > 0 || resource.SourcePath == "" {
		_, err := w.Write(resource.Data)
		return err
	}

	if resource.ExpectedSHA256 != "" && !expectedSHA256Pattern.MatchString(resource.ExpectedSHA256) {
		return fmt.Errorf("resource %q has invalid expected sha256", resource.SourcePath)
	}

	file, err := os.Open(resource.SourcePath)
	if err != nil {
		return fmt.Errorf("open resource %q: %w", resource.SourcePath, err)
	}
	defer file.Close()

	// The hash is computed while streaming into w, so size/SHA mismatches are
	// only detected AFTER the (bad) bytes have already been written. A non-nil
	// return therefore means w holds partially-written, unverified output; the
	// caller must discard it. writeHAKArchive does this by writing to a temp file
	// and removing it on any Write error rather than renaming it into place.
	hash := sha256.New()
	written, err := io.Copy(io.MultiWriter(w, hash), file)
	if err != nil {
		return fmt.Errorf("copy resource %q: %w", resource.SourcePath, err)
	}
	if resource.Size > 0 && written != resource.Size {
		return fmt.Errorf("copy resource %q: expected %d bytes, wrote %d", resource.SourcePath, resource.Size, written)
	}
	if resource.ExpectedSHA256 != "" {
		got := hex.EncodeToString(hash.Sum(nil))
		if got != resource.ExpectedSHA256 {
			return fmt.Errorf("resource %q sha256 mismatch: expected %s, got %s", resource.SourcePath, resource.ExpectedSHA256, got)
		}
	}
	return nil
}