github.com/ratrocket/u-root@v0.0.0-20180201221235-1cf9f48ee2cf/cmds/pxeboot/pxeboot.go (about)

     1  package main
     2  
     3  import (
     4  	"flag"
     5  	"fmt"
     6  	"io"
     7  	"io/ioutil"
     8  	"log"
     9  	"net/url"
    10  	"os"
    11  	"path"
    12  	"time"
    13  
    14  	"github.com/u-root/u-root/pkg/dhclient"
    15  	"github.com/u-root/u-root/pkg/kexec"
    16  	"github.com/u-root/u-root/pkg/pxe"
    17  	"github.com/vishvananda/netlink"
    18  )
    19  
    20  var (
    21  	verbose = flag.Bool("v", true, "print all kinds of things out, more than Chris wants")
    22  	dryRun  = flag.Bool("dry-run", false, "download kernel, but don't kexec it")
    23  	debug   = func(string, ...interface{}) {}
    24  )
    25  
    26  func copyToFile(r io.Reader) (*os.File, error) {
    27  	f, err := ioutil.TempFile("", "nerf-netboot")
    28  	if err != nil {
    29  		return nil, err
    30  	}
    31  	if _, err := io.Copy(f, r); err != nil {
    32  		f.Close()
    33  		return nil, err
    34  	}
    35  	if err := f.Sync(); err != nil {
    36  		f.Close()
    37  		return nil, err
    38  	}
    39  	return f, nil
    40  }
    41  
    42  func attemptDHCPLease(iface netlink.Link, timeout time.Duration, retry int) dhclient.Packet {
    43  	if _, err := dhclient.IfUp(iface.Attrs().Name); err != nil {
    44  		return nil
    45  	}
    46  
    47  	client, err := dhclient.NewV4(iface, timeout, retry)
    48  	if err != nil {
    49  		return nil
    50  	}
    51  
    52  	p, err := client.Solicit()
    53  	if err != nil {
    54  		return nil
    55  	}
    56  	return p
    57  }
    58  
    59  func Netboot() error {
    60  	ifs, err := netlink.LinkList()
    61  	if err != nil {
    62  		return err
    63  	}
    64  
    65  	for _, iface := range ifs {
    66  		// TODO: Do 'em all in parallel.
    67  		if iface.Attrs().Name != "eth0" {
    68  			continue
    69  		}
    70  
    71  		log.Printf("Attempting to get DHCP lease on %s", iface.Attrs().Name)
    72  		packet := attemptDHCPLease(iface, 30*time.Second, 5)
    73  		if packet == nil {
    74  			log.Printf("No lease on %s", iface.Attrs().Name)
    75  			continue
    76  		}
    77  		log.Printf("Got lease on %s", iface.Attrs().Name)
    78  		if err := dhclient.HandlePacket(iface, packet); err != nil {
    79  			log.Printf("shit: %v", err)
    80  			continue
    81  		}
    82  
    83  		// We may have to make this DHCPv6 and DHCPv4-specific anyway.
    84  		// Only tested with v4 right now; and assuming the uri points
    85  		// to a pxelinux.0.
    86  		//
    87  		// Or rather, we need to make this option-specific. DHCPv6 has
    88  		// options for passing a kernel and cmdline directly. v4
    89  		// usually just passes a pxelinux.0. But what about an initrd?
    90  		uri, _, err := packet.Boot()
    91  		if err != nil {
    92  			log.Printf("Got DHCP lease, but no valid PXE information.")
    93  			continue
    94  		}
    95  
    96  		wd := &url.URL{
    97  			Scheme: uri.Scheme,
    98  			Host:   uri.Host,
    99  			Path:   path.Dir(uri.Path),
   100  		}
   101  		pc := pxe.NewConfig(wd)
   102  		if err := pc.FindConfigFile(iface.Attrs().HardwareAddr, packet.IPs()[0]); err != nil {
   103  			return fmt.Errorf("failed to parse pxelinux config: %v", err)
   104  		}
   105  
   106  		label := pc.Entries[pc.DefaultEntry]
   107  		log.Printf("Got configuration: %v", label)
   108  
   109  		k, err := copyToFile(label.Kernel)
   110  		if err != nil {
   111  			return err
   112  		}
   113  		defer k.Close()
   114  		var i *os.File
   115  		if label.Initrd != nil {
   116  			i, err = copyToFile(label.Initrd)
   117  			if err != nil {
   118  				return err
   119  			}
   120  			defer i.Close()
   121  		}
   122  
   123  		if *dryRun {
   124  			log.Printf("Kernel: %s", k.Name())
   125  			if i != nil {
   126  				log.Printf("Initrd: %s", i.Name())
   127  			}
   128  			log.Printf("Command line: %s", label.Cmdline)
   129  		} else {
   130  			if err := kexec.FileLoad(k, i, label.Cmdline); err != nil {
   131  				return err
   132  			}
   133  			return kexec.Reboot()
   134  		}
   135  	}
   136  	return nil
   137  }
   138  
   139  func main() {
   140  	flag.Parse()
   141  	if *verbose {
   142  		debug = log.Printf
   143  	}
   144  
   145  	if err := Netboot(); err != nil {
   146  		log.Fatal(err)
   147  	}
   148  }