github.com/system-transparency/u-root@v6.0.1-0.20190919065413-ed07a650de4c+incompatible/pkg/bb/bbmain/register.go (about)

     1  // Copyright 2018 the u-root Authors. All rights reserved
     2  // Use of this source code is governed by a BSD-style
     3  // license that can be found in the LICENSE file.
     4  
     5  // Package bbmain is the package imported by all rewritten busybox
     6  // command-packages to register themselves.
     7  package bbmain
     8  
     9  import (
    10  	"errors"
    11  	"fmt"
    12  	"os"
    13  )
    14  
    15  // ErrNotRegistered is returned by Run if the given command is not registered.
    16  var ErrNotRegistered = errors.New("command not registered")
    17  
    18  // Noop is a noop function.
    19  var Noop = func() {}
    20  
    21  type bbCmd struct {
    22  	init, main func()
    23  }
    24  
    25  var bbCmds = map[string]bbCmd{}
    26  
    27  var defaultCmd *bbCmd
    28  
    29  // Register registers an init and main function for name.
    30  func Register(name string, init, main func()) {
    31  	if _, ok := bbCmds[name]; ok {
    32  		panic(fmt.Sprintf("cannot register two commands with name %q", name))
    33  	}
    34  	bbCmds[name] = bbCmd{
    35  		init: init,
    36  		main: main,
    37  	}
    38  }
    39  
    40  // RegisterDefault registers a default init and main function.
    41  func RegisterDefault(init, main func()) {
    42  	defaultCmd = &bbCmd{
    43  		init: init,
    44  		main: main,
    45  	}
    46  }
    47  
    48  // Run runs the command with the given name.
    49  //
    50  // If the command's main exits without calling os.Exit, Run will exit with exit
    51  // code 0.
    52  func Run(name string) error {
    53  	var cmd *bbCmd
    54  	if c, ok := bbCmds[name]; ok {
    55  		cmd = &c
    56  	} else if defaultCmd != nil {
    57  		cmd = defaultCmd
    58  	} else {
    59  		return ErrNotRegistered
    60  	}
    61  	cmd.init()
    62  	cmd.main()
    63  	os.Exit(0)
    64  	// Unreachable.
    65  	return nil
    66  }