...

Source file src/github.com/amenzhinsky/go-memexec/cmd/memexec-gen/main.go

Documentation: github.com/amenzhinsky/go-memexec/cmd/memexec-gen

     1  //go:build linux
     2  
     3  package main
     4  
     5  import (
     6  	"debug/elf"
     7  	"flag"
     8  	"fmt"
     9  	"io"
    10  	"os"
    11  	"path/filepath"
    12  	"runtime"
    13  	"strings"
    14  )
    15  
    16  var (
    17  	destFlag    string
    18  	packageFlag string
    19  )
    20  
    21  func main() {
    22  	flag.Usage = func() {
    23  		fmt.Fprint(os.Stderr, `Usage: memexec-gen PATH
    24  
    25  Environment variables:
    26    LD_LIBRARY_PATH
    27    	list of directories where to search libraries, see 'man 7 ld.so'
    28  
    29  Options:
    30  `)
    31  		flag.PrintDefaults()
    32  	}
    33  	flag.StringVar(&destFlag, "dest", "",
    34  		"`path` where binaries and go file are written to, defaults to PATH basename",
    35  	)
    36  	flag.StringVar(&packageFlag, "package", "",
    37  		"generated package `name`, defaults to PATH basename",
    38  	)
    39  	flag.Parse()
    40  	if flag.NArg() != 1 {
    41  		flag.Usage()
    42  		os.Exit(2)
    43  	}
    44  	if err := run(flag.Arg(0)); err != nil {
    45  		fmt.Fprintf(os.Stderr, "error: %s\n", err)
    46  		os.Exit(127)
    47  	}
    48  }
    49  
    50  func run(bin string) error {
    51  	base := filepath.Base(bin)
    52  	if destFlag == "" {
    53  		destFlag = base
    54  	}
    55  	if packageFlag == "" {
    56  		packageFlag = base
    57  	}
    58  
    59  	err := os.MkdirAll(destFlag, 0o755)
    60  	if err != nil && !os.IsExist(err) {
    61  		return err
    62  	}
    63  	output, err = os.OpenFile(
    64  		filepath.Join(destFlag, "memexec.go"),
    65  		os.O_CREATE|os.O_WRONLY|os.O_TRUNC,
    66  		0o644,
    67  	)
    68  	if err != nil {
    69  		return err
    70  	}
    71  
    72  	libs := map[string]string{}
    73  	if err := ldd(libs, bin, func(name string) (string, error) {
    74  		return name, nil
    75  	}); err != nil {
    76  		return err
    77  	}
    78  	for name, path := range libs {
    79  		src, err := os.Open(path)
    80  		if err != nil {
    81  			return err
    82  		}
    83  		if name == bin {
    84  			name = base
    85  		}
    86  		dst, err := os.OpenFile(
    87  			filepath.Join(destFlag, name),
    88  			os.O_CREATE|os.O_WRONLY|os.O_TRUNC,
    89  			0o755,
    90  		)
    91  		if err != nil {
    92  			return err
    93  		}
    94  		if _, err := io.Copy(dst, src); err != nil {
    95  			return err
    96  		}
    97  		if err := dst.Close(); err != nil {
    98  			return err
    99  		}
   100  		if err := src.Close(); err != nil {
   101  			return err
   102  		}
   103  	}
   104  
   105  	out("// Code generated by memexec-gen. DO NOT EDIT.")
   106  	out("package ", packageFlag)
   107  	out()
   108  	out("import (")
   109  	out(`	_ "embed"`)
   110  	out(`	"os"`)
   111  	out(`	"os/exec"`)
   112  	out(`	"path/filepath"`)
   113  	out()
   114  	out(`	"github.com/amenzhinsky/go-memexec"`)
   115  	out(")")
   116  	out()
   117  
   118  	for name := range libs {
   119  		if name == bin {
   120  			name = base
   121  		}
   122  		out("//go:embed ", name)
   123  		out("var ", goName(name), " []byte")
   124  		out()
   125  	}
   126  
   127  	out("func New(opts ...memexec.Option) (*memexec.Exec, error) {")
   128  	out(`	temp, err := os.MkdirTemp("", "")`)
   129  	out("	if err != nil {")
   130  	out("		return nil, err")
   131  	out("	}")
   132  	for name := range libs {
   133  		if name == bin {
   134  			continue
   135  		}
   136  		out(`	if err := os.WriteFile(filepath.Join(temp, "`, name, `"), `, goName(name), ", 0o755); err != nil {")
   137  		out("		return nil, err")
   138  		out("	}")
   139  	}
   140  	out("	return memexec.New(", goName(base), ",  append([]memexec.Option{")
   141  	out("		memexec.WithPrepare(func(exe *exec.Cmd) {")
   142  	out(`			exe.Env = append(exe.Env, "LD_LIBRARY_PATH="+temp)`)
   143  	out("		}),")
   144  	out("		memexec.WithCleanup(func() error {")
   145  	out("			return os.RemoveAll(temp)")
   146  	out("		}),")
   147  	out("	}, opts...)...)")
   148  	out("}")
   149  
   150  	return nil
   151  }
   152  
   153  var output io.Writer
   154  
   155  func out(strs ...string) {
   156  	for _, str := range strs {
   157  		if _, err := fmt.Fprint(output, str); err != nil {
   158  			panic(err)
   159  		}
   160  	}
   161  	if _, err := fmt.Fprintln(output); err != nil {
   162  		panic(err)
   163  	}
   164  }
   165  
   166  func goName(name string) string {
   167  	name = strings.ReplaceAll(name, ".", "_")
   168  	name = strings.ReplaceAll(name, "-", "_")
   169  	return "memexec_" + name
   170  }
   171  
   172  func ldd(
   173  	deps map[string]string, name string, look func(name string) (string, error),
   174  ) error {
   175  	path, err := look(name)
   176  	if err != nil {
   177  		return err
   178  	}
   179  	deps[name] = path
   180  	e, err := elf.Open(path)
   181  	if err != nil {
   182  		return err
   183  	}
   184  	l, err := e.ImportedLibraries()
   185  	if err != nil {
   186  		return err
   187  	}
   188  	rpaths, err := e.DynString(elf.DT_RPATH)
   189  	if err != nil {
   190  		return err
   191  	}
   192  	runpaths, err := e.DynString(elf.DT_RUNPATH)
   193  	if err != nil {
   194  		return err
   195  	}
   196  	if err := e.Close(); err != nil {
   197  		return err
   198  	}
   199  	for _, lib := range l {
   200  		if _, ok := deps[lib]; ok {
   201  			continue
   202  		}
   203  		if err := ldd(deps, lib, func(name string) (string, error) {
   204  			return lookLibrary(name, rpaths, runpaths)
   205  		}); err != nil {
   206  			return err
   207  		}
   208  	}
   209  	return nil
   210  }
   211  
   212  var ldPaths = []string{"/lib", "/usr/local/lib", "/usr/lib"}
   213  
   214  func init() {
   215  	// TODO: os.ReadFile(fmt.Sprintf("/etc/ld-musl-$(ARCH).path"))
   216  	switch runtime.GOARCH {
   217  	case "amd64":
   218  		ldPaths = append(ldPaths,
   219  			"/lib64",
   220  			"/usr/lib64",
   221  			"/usr/local/lib/x86_64-linux-gnu",
   222  			"/lib/x86_64-linux-gnu",
   223  			"/usr/lib/x86_64-linux-gnu",
   224  		)
   225  	case "arm64":
   226  		ldPaths = append(ldPaths,
   227  			"/usr/local/lib/aarch64-linux-gnu",
   228  			"/lib/aarch64-linux-gnu",
   229  			"/usr/lib/aarch64-linux-gnu",
   230  		)
   231  	}
   232  	if s := os.Getenv("LD_LIBRARY_PATH"); s != "" {
   233  		ldPaths = append(ldPaths, filepath.SplitList(s)...)
   234  	}
   235  }
   236  
   237  func lookLibrary(name string, rpaths, runpaths []string) (string, error) {
   238  	for _, dir := range append(rpaths, append(ldPaths, runpaths...)...) {
   239  		path := filepath.Join(dir, name)
   240  		stat, err := os.Stat(path)
   241  		if err != nil {
   242  			if os.IsNotExist(err) {
   243  				continue
   244  			}
   245  			return "", err
   246  		}
   247  		if !stat.IsDir() {
   248  			return path, nil
   249  		}
   250  	}
   251  	return "", fmt.Errorf("dynamic library %s not found", name)
   252  }
   253  

View as plain text