add settings modal, annotations, fix memory leak, improve UX, improve decompiler

Author: 0xInception

app.go

@@ -2,21 +2,81 @@ package main
 
 import (
 	"context"
+	"encoding/json"
 	"fmt"
+	"os"
 
 	"github.com/0xInception/wasmspy/pkg/decompile"
 	"github.com/0xInception/wasmspy/pkg/wasm"
 	"github.com/wailsapp/wails/v2/pkg/runtime"
 )
 
+type FunctionAnnotation struct {
+	Name    string `json:"name,omitempty"`
+	Comment string `json:"comment,omitempty"`
+}
+
+type Annotations struct {
+	Version           int                            `json:"version"`
+	Functions         map[string]*FunctionAnnotation `json:"functions,omitempty"`
+	Comments          map[string]string              `json:"comments,omitempty"`
+	DecompileComments map[string]string              `json:"decompileComments,omitempty"`
+	Bookmarks         []uint32                       `json:"bookmarks,omitempty"`
+}
+
+func NewAnnotations() *Annotations {
+	return &Annotations{
+		Version:           1,
+		Functions:         make(map[string]*FunctionAnnotation),
+		Comments:          make(map[string]string),
+		DecompileComments: make(map[string]string),
+		Bookmarks:         []uint32{},
+	}
+}
+
+func annotationsPath(wasmPath string) string {
+	return wasmPath + ".wasmspy"
+}
+
+func loadAnnotationsFromFile(wasmPath string) *Annotations {
+	data, err := os.ReadFile(annotationsPath(wasmPath))
+	if err != nil {
+		return NewAnnotations()
+	}
+	var ann Annotations
+	if err := json.Unmarshal(data, &ann); err != nil {
+		return NewAnnotations()
+	}
+	if ann.Functions == nil {
+		ann.Functions = make(map[string]*FunctionAnnotation)
+	}
+	if ann.Comments == nil {
+		ann.Comments = make(map[string]string)
+	}
+	if ann.DecompileComments == nil {
+		ann.DecompileComments = make(map[string]string)
+	}
+	return &ann
+}
+
+func saveAnnotationsToFile(wasmPath string, ann *Annotations) error {
+	data, err := json.MarshalIndent(ann, "", "  ")
+	if err != nil {
+		return err
+	}
+	return os.WriteFile(annotationsPath(wasmPath), data, 0644)
+}
+
 type App struct {
-	ctx     context.Context
-	modules map[string]*wasm.ResolvedModule
+	ctx         context.Context
+	modules     map[string]*wasm.ResolvedModule
+	annotations map[string]*Annotations
 }
 
 func NewApp() *App {
 	return &App{
-		modules: make(map[string]*wasm.ResolvedModule),
+		modules:     make(map[string]*wasm.ResolvedModule),
+		annotations: make(map[string]*Annotations),
 	}
 }
 
@@ -35,6 +95,11 @@ func (a *App) onDomReady(ctx context.Context) {
 	})
 }
 
+func (a *App) SetWindowSize(width, height int) {
+	runtime.WindowSetSize(a.ctx, width, height)
+	runtime.WindowCenter(a.ctx)
+}
+
 func (a *App) OpenFileDialog() (string, error) {
 	path, err := runtime.OpenFileDialog(a.ctx, runtime.OpenDialogOptions{
 		Title: "Open WebAssembly Module",
@@ -101,6 +166,7 @@ func (a *App) loadModule(path string, mod *wasm.Module) (*ModuleInfo, error) {
 		return nil, err
 	}
 	a.modules[path] = resolved
+	a.annotations[path] = loadAnnotationsFromFile(path)
 
 	info := &ModuleInfo{}
 
@@ -130,10 +196,17 @@ func (a *App) loadModule(path string, mod *wasm.Module) (*ModuleInfo, error) {
 		})
 	}
 
+	ann := a.annotations[path]
 	for _, fn := range resolved.Functions {
+		name := fn.Name
+		if ann != nil {
+			if fa := ann.Functions[fmt.Sprintf("%d", fn.Index)]; fa != nil && fa.Name != "" {
+				name = fa.Name
+			}
+		}
 		fi := FunctionInfo{
 			Index:    fn.Index,
-			Name:     fn.Name,
+			Name:     name,
 			Imported: fn.Imported,
 		}
 		if fn.Type != nil {
@@ -174,7 +247,8 @@ func (a *App) DisassembleFunction(path string, index uint32, indent bool) (strin
 	if fn == nil {
 		return "", fmt.Errorf("function %d not found", index)
 	}
-	return formatFunctionDisasm(fn, module, indent), nil
+	ann := a.annotations[path]
+	return formatFunctionDisasm(fn, module, indent, ann), nil
 }
 
 func (a *App) GetFunctionWAT(path string, index uint32) (string, error) {
@@ -438,13 +512,29 @@ func getDataSegOffset(instrs []wasm.Instruction) uint32 {
 	return 0
 }
 
-func formatFunctionDisasm(fn *wasm.ResolvedFunction, rm *wasm.ResolvedModule, indented bool) string {
+func formatFunctionDisasm(fn *wasm.ResolvedFunction, rm *wasm.ResolvedModule, indented bool, ann *Annotations) string {
 	if fn.Imported {
 		return fmt.Sprintf("; imported: %s", fn.Name)
 	}
 
+	name := fn.Name
+	funcComment := ""
+	if ann != nil {
+		if fa := ann.Functions[fmt.Sprintf("%d", fn.Index)]; fa != nil {
+			if fa.Name != "" {
+				name = fa.Name
+			}
+			if fa.Comment != "" {
+				funcComment = fa.Comment
+			}
+		}
+	}
+
 	var result string
-	result = fmt.Sprintf("; Function %d: %s\n", fn.Index, fn.Name)
+	result = fmt.Sprintf("; Function %d: %s\n", fn.Index, name)
+	if funcComment != "" {
+		result += fmt.Sprintf("; %s\n", funcComment)
+	}
 
 	if fn.Type != nil {
 		result += fmt.Sprintf("; Params: %d, Results: %d\n", len(fn.Type.Params), len(fn.Type.Results))
@@ -465,6 +555,13 @@ func formatFunctionDisasm(fn *wasm.ResolvedFunction, rm *wasm.ResolvedModule, in
 
 		indent := 0
 		for _, instr := range fn.Body.Instructions {
+			comment := ""
+			if ann != nil {
+				if c := ann.Comments[fmt.Sprintf("0x%x", instr.Offset)]; c != "" {
+					comment = " ; " + c
+				}
+			}
+
 			if indented {
 				if instr.Opcode == wasm.OpEnd || instr.Opcode == wasm.OpElse {
 					if indent > 0 {
@@ -476,14 +573,14 @@ func formatFunctionDisasm(fn *wasm.ResolvedFunction, rm *wasm.ResolvedModule, in
 				for j := 0; j < indent; j++ {
 					prefix += "  "
 				}
-				result += fmt.Sprintf("%08x: %s%s\n", instr.Offset, prefix, formatInstrWithImm(&instr))
+				result += fmt.Sprintf("%08x: %s%s%s\n", instr.Offset, prefix, formatInstrWithImm(&instr), comment)
 
 				if instr.Opcode == wasm.OpBlock || instr.Opcode == wasm.OpLoop ||
 					instr.Opcode == wasm.OpIf || instr.Opcode == wasm.OpElse {
 					indent++
 				}
 			} else {
-				result += fmt.Sprintf("%08x: %s\n", instr.Offset, formatInstrWithImm(&instr))
+				result += fmt.Sprintf("%08x: %s%s\n", instr.Offset, formatInstrWithImm(&instr), comment)
 			}
 		}
 	}
@@ -574,3 +671,168 @@ func formatInstr(instr *wasm.Instruction) string {
 	}
 	return result
 }
+
+func (a *App) GetAnnotations(path string) *Annotations {
+	if ann := a.annotations[path]; ann != nil {
+		return ann
+	}
+	return NewAnnotations()
+}
+
+func (a *App) SetFunctionName(path string, index uint32, name string) error {
+	ann := a.annotations[path]
+	if ann == nil {
+		ann = NewAnnotations()
+		a.annotations[path] = ann
+	}
+	key := fmt.Sprintf("%d", index)
+	if ann.Functions[key] == nil {
+		ann.Functions[key] = &FunctionAnnotation{}
+	}
+	ann.Functions[key].Name = name
+	return nil
+}
+
+func (a *App) SetFunctionComment(path string, index uint32, comment string) error {
+	ann := a.annotations[path]
+	if ann == nil {
+		ann = NewAnnotations()
+		a.annotations[path] = ann
+	}
+	key := fmt.Sprintf("%d", index)
+	if ann.Functions[key] == nil {
+		ann.Functions[key] = &FunctionAnnotation{}
+	}
+	ann.Functions[key].Comment = comment
+	return nil
+}
+
+func (a *App) SetOffsetComment(path string, offset uint64, comment string, isDecompile bool) error {
+	ann := a.annotations[path]
+	if ann == nil {
+		ann = NewAnnotations()
+		a.annotations[path] = ann
+	}
+	key := fmt.Sprintf("0x%x", offset)
+	target := ann.Comments
+	if isDecompile {
+		target = ann.DecompileComments
+	}
+	if comment == "" {
+		delete(target, key)
+	} else {
+		target[key] = comment
+	}
+	return nil
+}
+
+func (a *App) SetBookmarks(path string, bookmarks []uint32) error {
+	ann := a.annotations[path]
+	if ann == nil {
+		ann = NewAnnotations()
+		a.annotations[path] = ann
+	}
+	ann.Bookmarks = bookmarks
+	return nil
+}
+
+func (a *App) SaveAnnotations(path string) error {
+	ann := a.annotations[path]
+	if ann == nil {
+		return nil
+	}
+	return saveAnnotationsToFile(path, ann)
+}
+
+func (a *App) ClearAnnotations(path string) *Annotations {
+	a.annotations[path] = NewAnnotations()
+	return a.annotations[path]
+}
+
+func (a *App) ExportAnnotations(path string) (string, error) {
+	ann := a.annotations[path]
+	if ann == nil {
+		ann = NewAnnotations()
+	}
+	data, err := json.MarshalIndent(ann, "", "  ")
+	if err != nil {
+		return "", err
+	}
+	return string(data), nil
+}
+
+func (a *App) ImportAnnotations(path string, data string) error {
+	var ann Annotations
+	if err := json.Unmarshal([]byte(data), &ann); err != nil {
+		return err
+	}
+	if ann.Functions == nil {
+		ann.Functions = make(map[string]*FunctionAnnotation)
+	}
+	if ann.Comments == nil {
+		ann.Comments = make(map[string]string)
+	}
+	if ann.DecompileComments == nil {
+		ann.DecompileComments = make(map[string]string)
+	}
+	a.annotations[path] = &ann
+	return nil
+}
+
+func (a *App) ExportAnnotationsToFile(wasmPath string) (string, error) {
+	savePath, err := runtime.SaveFileDialog(a.ctx, runtime.SaveDialogOptions{
+		Title:           "Export Annotations",
+		DefaultFilename: "annotations.wasmspy",
+		Filters: []runtime.FileFilter{
+			{DisplayName: "WasmSpy Annotations", Pattern: "*.wasmspy"},
+			{DisplayName: "JSON Files", Pattern: "*.json"},
+		},
+	})
+	if err != nil || savePath == "" {
+		return "", err
+	}
+	ann := a.annotations[wasmPath]
+	if ann == nil {
+		ann = NewAnnotations()
+	}
+	data, err := json.MarshalIndent(ann, "", "  ")
+	if err != nil {
+		return "", err
+	}
+	if err := os.WriteFile(savePath, data, 0644); err != nil {
+		return "", err
+	}
+	return savePath, nil
+}
+
+func (a *App) ImportAnnotationsFromFile(wasmPath string) (string, error) {
+	openPath, err := runtime.OpenFileDialog(a.ctx, runtime.OpenDialogOptions{
+		Title: "Import Annotations",
+		Filters: []runtime.FileFilter{
+			{DisplayName: "WasmSpy Annotations", Pattern: "*.wasmspy"},
+			{DisplayName: "JSON Files", Pattern: "*.json"},
+		},
+	})
+	if err != nil || openPath == "" {
+		return "", err
+	}
+	data, err := os.ReadFile(openPath)
+	if err != nil {
+		return "", err
+	}
+	var ann Annotations
+	if err := json.Unmarshal(data, &ann); err != nil {
+		return "", err
+	}
+	if ann.Functions == nil {
+		ann.Functions = make(map[string]*FunctionAnnotation)
+	}
+	if ann.Comments == nil {
+		ann.Comments = make(map[string]string)
+	}
+	if ann.DecompileComments == nil {
+		ann.DecompileComments = make(map[string]string)
+	}
+	a.annotations[wasmPath] = &ann
+	return openPath, nil
+}