github.com/oweisse/u-root@v0.0.0-20181109060735-d005ad25fef1/integration/testcmd/pxeserver/main.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 main
     6  
     7  import (
     8  	"flag"
     9  	"fmt"
    10  	"io"
    11  	"log"
    12  	"net"
    13  	"os"
    14  	"path/filepath"
    15  	"sync"
    16  
    17  	"github.com/u-root/dhcp4/dhcp4server"
    18  	"pack.ag/tftp"
    19  )
    20  
    21  var (
    22  	selfIP    = flag.String("ip", "192.168.0.1", "IP of self")
    23  	subnet    = flag.String("subnet", "192.168.1.0/24", "CIDR of network to assign to clients")
    24  	directory = flag.String("dir", "", "Directory to serve")
    25  )
    26  
    27  func serve(w tftp.ReadRequest) {
    28  	path := filepath.Join(*directory, filepath.Clean(w.Name()))
    29  
    30  	file, err := os.Open(path)
    31  	if err != nil {
    32  		w.WriteError(tftp.ErrCodeFileNotFound, fmt.Sprintf("File %q does not exist", w.Name()))
    33  		return
    34  	}
    35  	defer file.Close()
    36  
    37  	finfo, _ := file.Stat()
    38  	w.WriteSize(finfo.Size())
    39  	if _, err = io.Copy(w, file); err != nil {
    40  		log.Println(err)
    41  	}
    42  }
    43  
    44  func main() {
    45  	flag.Parse()
    46  
    47  	var wg sync.WaitGroup
    48  	if len(*directory) != 0 {
    49  		wg.Add(1)
    50  		go func() {
    51  			defer wg.Done()
    52  			server, err := tftp.NewServer(":69")
    53  			if err != nil {
    54  				log.Fatalf("Could not start TFTP server: %v", err)
    55  			}
    56  
    57  			log.Println("starting file server")
    58  			server.ReadHandler(tftp.ReadHandlerFunc(serve))
    59  			log.Fatal(server.ListenAndServe())
    60  		}()
    61  	}
    62  
    63  	wg.Add(1)
    64  	go func() {
    65  		defer wg.Done()
    66  		_, sn, err := net.ParseCIDR(*subnet)
    67  		if err != nil {
    68  			log.Fatal(err)
    69  		}
    70  
    71  		l, err := net.ListenPacket("udp4", ":67")
    72  		if err != nil {
    73  			log.Fatal(err)
    74  		}
    75  		defer l.Close()
    76  
    77  		self := net.ParseIP(*selfIP)
    78  		log.Printf("Self IP: %v", self)
    79  		s := dhcp4server.New(self, sn, "", "pxelinux.0")
    80  
    81  		log.Fatal(s.Serve(log.New(os.Stdout, "", log.LstdFlags), l))
    82  	}()
    83  
    84  	wg.Wait()
    85  }