gopkg.in/hugelgupf/u-root.v2@v2.0.0-20180831055005-3f8fdb0ce09d/pkg/bb/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 bb 6 7 import ( 8 "errors" 9 "fmt" 10 "os" 11 ) 12 13 // ErrNotRegistered is returned by Run if the given command is not registered. 14 var ErrNotRegistered = errors.New("command not registered") 15 16 // Noop is a noop function. 17 var Noop = func() {} 18 19 type bbCmd struct { 20 init, main func() 21 } 22 23 var bbCmds = map[string]bbCmd{} 24 25 var defaultCmd *bbCmd 26 27 // Register registers an init and main function for name. 28 func Register(name string, init, main func()) { 29 if _, ok := bbCmds[name]; ok { 30 panic(fmt.Sprintf("cannot register two commands with name %q", name)) 31 } 32 bbCmds[name] = bbCmd{ 33 init: init, 34 main: main, 35 } 36 } 37 38 // RegisterDefault registers a default init and main function. 39 func RegisterDefault(init, main func()) { 40 defaultCmd = &bbCmd{ 41 init: init, 42 main: main, 43 } 44 } 45 46 // Run runs the command with the given name. 47 // 48 // If the command's main exits without calling os.Exit, Run will exit with exit 49 // code 0. 50 func Run(name string) error { 51 var cmd *bbCmd 52 if c, ok := bbCmds[name]; ok { 53 cmd = &c 54 } else if defaultCmd != nil { 55 cmd = defaultCmd 56 } else { 57 return ErrNotRegistered 58 } 59 cmd.init() 60 cmd.main() 61 os.Exit(0) 62 // Unreachable. 63 return nil 64 }