lychee-slicer/asar/asar.go

package asar

import (
	"bytes"
	"encoding/json"
	"fmt"
	"io"
	"math"
	"path/filepath"
	"strconv"
	"strings"

	"vaclive.party/software/lychee-slicer/asar/pickle"
)

type EntryMetadata struct {
	Unpacked *bool `json:"unpacked,omitempty"`
}

type FilesystemDirectoryEntry struct {
	EntryMetadata
	Files    map[string]FilesystemEntry
	nezuMeta *map[string]string
}

func (f *FilesystemDirectoryEntry) UnmarshalJSON(data []byte) error {
	var raw struct {
		Unpacked *bool                       `json:"unpacked"`
		Files    *map[string]json.RawMessage `json:"files"`
		NezuMeta *map[string]string          `json:"nezuMeta,omitempty"`
	}
	if err := json.Unmarshal(data, &raw); err != nil {
		return err
	}
	if raw.Files == nil {
		return fmt.Errorf("missing files")
	}
	f.Unpacked = raw.Unpacked
	f.Files = make(map[string]FilesystemEntry)
	f.nezuMeta = raw.NezuMeta
	for name, entry := range *raw.Files {
		var dir FilesystemDirectoryEntry
		if err := json.Unmarshal(entry, &dir); err == nil {
			f.Files[name] = &dir
			continue
		}

		var file FilesystemFileEntry
		if err := json.Unmarshal(entry, &file); err == nil && file.Size != nil {
			f.Files[name] = &file
			continue
		}

		var link FilesystemLinkEntry
		if err := json.Unmarshal(entry, &link); err == nil && link.Link != nil {
			f.Files[name] = &link
			continue
		}

		return fmt.Errorf("unexpected entry type: %s", name)
	}
	return nil
}

func (f *FilesystemDirectoryEntry) MarshalJSON() ([]byte, error) {
	raw := struct {
		Unpacked *bool                      `json:"unpacked,omitempty"`
		Files    map[string]json.RawMessage `json:"files"`
		NezuMeta *map[string]string         `json:"nezuMeta,omitempty"`
	}{
		Unpacked: f.Unpacked,
		Files:    make(map[string]json.RawMessage),
		NezuMeta: f.nezuMeta,
	}
	for name, entry := range f.Files {
		data, err := json.Marshal(entry)
		if err != nil {
			return nil, err
		}
		raw.Files[name] = data
	}
	return json.Marshal(raw)
}

type FilesystemFileEntry struct {
	EntryMetadata
	Executable *bool          `json:"executable,omitempty"`
	Offset     *string        `json:"offset,omitempty"`
	Size       *int64         `json:"size"`
	Integrity  *FileIntegrity `json:"integrity,omitempty"`
}

type FilesystemLinkEntry struct {
	EntryMetadata
	Link *string `json:"link"`
}

type FilesystemEntry interface{}

type Filesystem struct {
	storage FilesystemStorage
	header  *FilesystemDirectoryEntry
}

func newFileSystemDirectoryEntry() *FilesystemDirectoryEntry {
	return &FilesystemDirectoryEntry{Files: make(map[string]FilesystemEntry)}
}

func newFileSystemFileEntry() *FilesystemFileEntry {
	return &FilesystemFileEntry{}
}

func newFileSystemLinkEntry() *FilesystemLinkEntry {
	return &FilesystemLinkEntry{}
}

func NewFilesystem() *Filesystem {
	return &Filesystem{
		storage: *NewFilesystemStorage(bytes.NewReader([]byte{}), NewBytesReadAtWriter([]byte{})),
		header:  newFileSystemDirectoryEntry(),
	}
}

func OpenFilesystem(reader io.ReaderAt, readerSize int64) (*Filesystem, error) {
	sizeBytes := make([]byte, 8)
	if _, err := reader.ReadAt(sizeBytes, 0); err != nil {
		return nil, err
	}

	sizePickle := pickle.CreateFromBuffer(sizeBytes)
	headerSize := sizePickle.CreateIterator().ReadUInt32()

	headerBytes := make([]byte, headerSize)
	if _, err := reader.ReadAt(headerBytes, 8); err != nil {
		return nil, err
	}

	headerPickle := pickle.CreateFromBuffer(headerBytes)
	headerJson := headerPickle.CreateIterator().ReadString()

	var header FilesystemDirectoryEntry
	if err := json.Unmarshal([]byte(headerJson), &header); err != nil {
		return nil, err
	}

	sectionReader := io.NewSectionReader(reader, 8+int64(headerSize), readerSize-8-int64(headerSize))

	return &Filesystem{
		storage: *NewFilesystemStorage(sectionReader, NewBytesReadAtWriter([]byte{})),
		header:  &header,
	}, nil
}

func (fs *Filesystem) searchNodeFromDirectory(p string) (*FilesystemDirectoryEntry, error) {
	entry := fs.header
	// windows moment
	dirs := strings.FieldsFunc(p, func(r rune) bool { return r == '/' || r == '\\' })
	for _, dir := range dirs {
		if dir != "." {
			if entry.Files != nil {
				if _, ok := entry.Files[dir]; !ok {
					entry.Files[dir] = newFileSystemDirectoryEntry()
				}
				entry = entry.Files[dir].(*FilesystemDirectoryEntry)
			} else {
				return nil, fmt.Errorf("unexpected directory state while traversing: %s", p)
			}
		}
	}
	return entry, nil
}

func (fs *Filesystem) searchNodeFromPath(p string) (FilesystemEntry, error) {
	if p == "." {
		return fs.header, nil
	}
	println(filepath.Base(p), filepath.Dir(p))
	name := filepath.Base(p)
	dirNode, err := fs.searchNodeFromDirectory(filepath.Dir(p))
	if err != nil {
		return nil, err
	}
	if dirNode.Files == nil {
		dirNode.Files = make(map[string]FilesystemEntry)
	}
	if _, ok := dirNode.Files[name]; !ok {
		dirNode.Files[name] = newFileSystemFileEntry()
	}
	return dirNode.Files[name], nil
}

func (fs *Filesystem) GetMetaFlag(flag string) string {
	if fs.header.nezuMeta != nil {
		if value, ok := (*fs.header.nezuMeta)[flag]; ok {
			return value
		}
	}
	return ""
}

func (fs *Filesystem) SetMetaFlag(flag string, value string) error {
	if fs.header.nezuMeta == nil {
		fs.header.nezuMeta = &map[string]string{}
	}
	(*fs.header.nezuMeta)[flag] = value
	return nil
}

func (fs *Filesystem) InsertDirectory(p string) (map[string]FilesystemEntry, error) {
	node, err := fs.searchNodeFromPath(p)
	if err != nil {
		return nil, err
	}
	dirNode := node.(*FilesystemDirectoryEntry)
	return dirNode.Files, nil
}

func (fs *Filesystem) InsertFile(p string, data []byte, executable bool) error {
	size := int64(len(data))
	if size > math.MaxUint32 {
		return fmt.Errorf("%s: file size exceeds 4GB", p)
	}

	dirNode, err := fs.searchNodeFromDirectory(filepath.Dir(p))
	if err != nil {
		return err
	}
	if dirNode.Unpacked != nil && *dirNode.Unpacked {
		return fmt.Errorf("unpacked directories are not supported: %s", p)
	}

	node, err := fs.searchNodeFromPath(p)
	if err != nil {
		return err
	}
	fileNode := node.(*FilesystemFileEntry)

	fileNode.Size = &size
	offset := fmt.Sprintf("%d", fs.storage.Size())
	fileNode.Offset = &offset
	fileNode.Integrity = getFileIntegrity(data)
	if executable {
		fileNode.Executable = &executable
	}
	fs.storage.Write(data)
	return nil
}

func (fs *Filesystem) ReadFile(p string, followLinks bool) ([]byte, error) {
	file, err := fs.GetFile(p, followLinks)
	if err != nil {
		return nil, err
	}
	if file == nil {
		return nil, fmt.Errorf("file not found: %s", p)
	}
	if linkEntry, ok := file.(*FilesystemLinkEntry); ok {
		return []byte(*linkEntry.Link), nil
	}

	fileEntry := file.(*FilesystemFileEntry)
	if fileEntry.Unpacked != nil && *fileEntry.Unpacked {
		return nil, fmt.Errorf("unpacked files are not supported: %s", p)
	}

	offset, err := strconv.ParseInt(*fileEntry.Offset, 10, 64)
	if err != nil {
		return nil, err
	}

	ret := make([]byte, *fileEntry.Size)
	if _, err := fs.storage.ReadAt(ret, offset); err != nil {
		return nil, err
	}

	// FIXME: integrity check

	return ret, nil
}

func (fs *Filesystem) InsertLink(p string, target string) error {
	dirNode, err := fs.searchNodeFromDirectory(filepath.Dir(p))
	if err != nil {
		return err
	}
	if dirNode.Unpacked != nil && *dirNode.Unpacked {
		return fmt.Errorf("unexpected directories not set to unpack: %s", p)
	}

	name := filepath.Base(p)
	if _, ok := dirNode.Files[name]; !ok {
		dirNode.Files[name] = newFileSystemLinkEntry()
	}

	linkNode := dirNode.Files[name].(*FilesystemLinkEntry)
	linkNode.Link = &target
	return nil
}

func (fs *Filesystem) ListFiles() []string {
	files := []string{}

	// recursive function to traverse the filesystem
	var traverse func(string, FilesystemEntry)
	traverse = func(p string, entry FilesystemEntry) {
		switch e := entry.(type) {
		case *FilesystemDirectoryEntry:
			for name, child := range e.Files {
				path := name
				if p != "" {
					path = p + "/" + name
				}
				traverse(path, child)
			}
		case *FilesystemFileEntry:
			files = append(files, p)
		case *FilesystemLinkEntry:
			files = append(files, p)
		}
	}

	traverse("", fs.header)
	return files
}

func (fs *Filesystem) GetNode(p string) (FilesystemEntry, error) {
	node, err := fs.searchNodeFromDirectory(filepath.Dir(p))
	if err != nil {
		return nil, err
	}
	name := filepath.Base(p)
	if name != "" {
		fmt.Println(p, filepath.Base(p), filepath.Dir(p), name, len(node.Files))
		return node.Files[name], nil
	}
	return node, nil
}

func (fs *Filesystem) GetFile(p string, followLinks bool) (FilesystemEntry, error) {
	info, err := fs.GetNode(p)
	if err != nil {
		return nil, err
	}
	if linkEntry, ok := info.(*FilesystemLinkEntry); ok && followLinks {
		return fs.GetFile(*linkEntry.Link, followLinks)
	}
	if _, ok := info.(*FilesystemDirectoryEntry); ok {
		return nil, fmt.Errorf("not a file: %s", p)
	}
	return info, nil
}

func (fs *Filesystem) Save(writer io.Writer) error {
	headerJson, err := json.Marshal(fs.header)
	if err != nil {
		return err
	}

	headerPickle := pickle.CreateEmpty()
	headerPickle.WriteString(string(headerJson))
	headerBytes := headerPickle.ToBuffer()

	sizePickle := pickle.CreateEmpty()
	sizePickle.WriteUInt32(uint32(len(headerBytes)))
	sizeBytes := sizePickle.ToBuffer()

	if _, err := writer.Write(sizeBytes); err != nil {
		return err
	}

	if _, err := writer.Write(headerBytes); err != nil {
		return err
	}

	if _, err := fs.storage.WriteTo(writer); err != nil {
		return err
	}

	return nil
}