github.com/coreos/mantle@v0.13.0/system/exec/multicall.go (about)

     1  // Copyright 2015 CoreOS, Inc.
     2  //
     3  // Licensed under the Apache License, Version 2.0 (the "License");
     4  // you may not use this file except in compliance with the License.
     5  // You may obtain a copy of the License at
     6  //
     7  //     http://www.apache.org/licenses/LICENSE-2.0
     8  //
     9  // Unless required by applicable law or agreed to in writing, software
    10  // distributed under the License is distributed on an "AS IS" BASIS,
    11  // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    12  // See the License for the specific language governing permissions and
    13  // limitations under the License.
    14  
    15  // inspired by github.com/docker/docker/pkg/reexec
    16  
    17  package exec
    18  
    19  import (
    20  	"fmt"
    21  	"os"
    22  	"strings"
    23  	"syscall"
    24  )
    25  
    26  // prefix of first argument if it is defining an entrypoint to be called.
    27  const entryArgPrefix = "_MULTICALL_ENTRYPOINT_"
    28  
    29  var exePath string
    30  
    31  func init() {
    32  	// save the program path
    33  	var err error
    34  	exePath, err = os.Readlink("/proc/self/exe")
    35  	if err != nil {
    36  		panic("cannot get current executable")
    37  	}
    38  }
    39  
    40  type entrypointFn func(args []string) error
    41  
    42  var entrypoints = make(map[string]entrypointFn)
    43  
    44  // Entrypoint provides the access to a multicall command.
    45  type Entrypoint string
    46  
    47  // NewEntrypoint adds a new multicall command. name is the command name
    48  // and fn is the function that will be executed for the specified
    49  // command. It returns the related Entrypoint. Packages adding new
    50  // multicall commands should call Add in their init function.
    51  func NewEntrypoint(name string, fn entrypointFn) Entrypoint {
    52  	if _, ok := entrypoints[name]; ok {
    53  		panic(fmt.Errorf("command with name %q already exists", name))
    54  	}
    55  	entrypoints[name] = fn
    56  	return Entrypoint(name)
    57  }
    58  
    59  // MaybeExec should be called at the start of the program, if the process argv[0] is
    60  // a name registered with multicall, the related function will be executed.
    61  // If the functions returns an error, it will be printed to stderr and will
    62  // exit with an exit status of 1, otherwise it will exit with a 0 exit status.
    63  func MaybeExec() {
    64  	if len(os.Args) < 2 || !strings.HasPrefix(os.Args[1], entryArgPrefix) {
    65  		return
    66  	}
    67  	name := os.Args[1][len(entryArgPrefix):]
    68  	if err := entrypoints[name](os.Args[2:]); err != nil {
    69  		fmt.Fprintln(os.Stderr, err)
    70  		os.Exit(1)
    71  	}
    72  	os.Exit(0)
    73  }
    74  
    75  // Command will prepare the *ExecCmd for the given entrypoint, configured with
    76  // the provided args.
    77  func (e Entrypoint) Command(args ...string) *ExecCmd {
    78  	args = append([]string{entryArgPrefix + string(e)}, args...)
    79  	cmd := Command(exePath, args...)
    80  	cmd.SysProcAttr = &syscall.SysProcAttr{
    81  		Pdeathsig: syscall.SIGTERM,
    82  	}
    83  	return cmd
    84  }
    85  
    86  // Sudo will prepare the *ExecCmd for the given entrypoint to be run as root
    87  // via sudo with the provided args.
    88  func (e Entrypoint) Sudo(args ...string) *ExecCmd {
    89  	args = append([]string{"-E", "-p", "sudo password for %p: ", "--",
    90  		exePath, entryArgPrefix + string(e)}, args...)
    91  	cmd := Command("sudo", args...)
    92  	cmd.SysProcAttr = &syscall.SysProcAttr{
    93  		Pdeathsig: syscall.SIGTERM,
    94  	}
    95  	return cmd
    96  }