self-modified.go

package main

import (
	"bytes"
	"debug/elf"
	"fmt"
	"io"
	"log"
	"os"
)

var CNT = 1

type entry struct {
	value, off uint64
}

type Entry interface {
	GetValue() uint64
	GetOff() uint64
}

func NewEntry(f *elf.File, name string) (Entry, error) {
	syms, err := f.Symbols()
	if err != nil {
		return nil, err
	}

	for _, s := range syms {
		if s.Name == name {
			sect := f.Sections[s.Section]
			bs, _ := sect.Data()
			varOff := s.Value - sect.Addr
			return &entry{
				f.ByteOrder.Uint64(bs[varOff : varOff+s.Size]),
				sect.Offset + varOff,
			}, nil
		}
	}

	return nil, fmt.Errorf("can't find symbol '%s'", name)
}

func (self *entry) GetValue() uint64 {
	return self.value
}

func (self *entry) GetOff() uint64 {
	return self.off
}

func main() {
	fmt.Println(CNT)

	exeName := os.Args[0]
	tmpName := exeName + ".tmp"

	f, err := os.Open(exeName)
	if err != nil {
		log.Fatalf("can't open file '%s': %e", exeName, err)
	}
	defer f.Close()

	data, _ := io.ReadAll(f)
	elfFile, err := elf.NewFile(bytes.NewReader(data))
	if err != nil {
		log.Fatalf("can't read ELF file: %e", err)
	}

	entry, err := NewEntry(elfFile, "main.CNT")
	if err != nil {
		log.Fatalf("can't find counter object in ELF file: %e", err)
	}

	elfFile.ByteOrder.PutUint64(data[entry.GetOff():], entry.GetValue()+1)

	fi, err := f.Stat()
	if err != nil {
		log.Fatalf("can't get file mode for '%s': %e", os.Args[0], err)
	}
	if err := os.WriteFile(tmpName, data, fi.Mode()); err != nil {
		log.Fatalf("can't write file '%s': %e", tmpName, err)
	}

	if err := os.Rename(tmpName, os.Args[0]); err != nil {
		log.Fatalln("can't rename temporary file", err)
	}
}