github.com/system-transparency/u-root@v6.0.1-0.20190919065413-ed07a650de4c+incompatible/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  	"log"
    10  	"net"
    11  	"sync"
    12  
    13  	"github.com/insomniacslk/dhcp/dhcpv4"
    14  	"github.com/insomniacslk/dhcp/dhcpv4/server4"
    15  	"pack.ag/tftp"
    16  )
    17  
    18  var (
    19  	selfIP    = flag.String("ip", "192.168.0.1", "IP of self")
    20  	yourIP    = flag.String("your-ip", "192.168.0.2", "The one and only IP to give to all clients")
    21  	directory = flag.String("dir", "", "Directory to serve")
    22  )
    23  
    24  func dhcpHandler(conn net.PacketConn, peer net.Addr, m *dhcpv4.DHCPv4) {
    25  	self := net.ParseIP(*selfIP)
    26  	you := net.ParseIP(*yourIP)
    27  	log.Printf("Handling request %v", m)
    28  
    29  	var replyType dhcpv4.MessageType
    30  	switch mt := m.MessageType(); mt {
    31  	case dhcpv4.MessageTypeDiscover:
    32  		replyType = dhcpv4.MessageTypeOffer
    33  	case dhcpv4.MessageTypeRequest:
    34  		replyType = dhcpv4.MessageTypeAck
    35  	default:
    36  		log.Printf("Can't handle type %v", mt)
    37  		return
    38  	}
    39  	reply, err := dhcpv4.NewReplyFromRequest(m,
    40  		dhcpv4.WithMessageType(replyType),
    41  		dhcpv4.WithServerIP(self),
    42  		dhcpv4.WithRouter(self),
    43  		dhcpv4.WithNetmask(net.CIDRMask(24, 32)),
    44  		dhcpv4.WithYourIP(you),
    45  	)
    46  	reply.BootFileName = "pxelinux.0"
    47  	if err != nil {
    48  		log.Printf("Could not create reply for %v: %v", m, err)
    49  		return
    50  	}
    51  	if _, err := conn.WriteTo(reply.ToBytes(), &net.UDPAddr{IP: net.IP{255, 255, 255, 255}, Port: 68}); err != nil {
    52  		log.Printf("Could not write %v: %v", reply, err)
    53  	}
    54  }
    55  
    56  func main() {
    57  	flag.Parse()
    58  
    59  	var wg sync.WaitGroup
    60  	if len(*directory) != 0 {
    61  		wg.Add(1)
    62  		go func() {
    63  			defer wg.Done()
    64  			server, err := tftp.NewServer(":69")
    65  			if err != nil {
    66  				log.Fatalf("Could not start TFTP server: %v", err)
    67  			}
    68  
    69  			log.Println("starting file server")
    70  			server.ReadHandler(tftp.FileServer(*directory))
    71  			log.Fatal(server.ListenAndServe())
    72  		}()
    73  	}
    74  
    75  	wg.Add(1)
    76  
    77  	go func() {
    78  		defer wg.Done()
    79  
    80  		laddr := &net.UDPAddr{Port: 67}
    81  		server, err := server4.NewServer("eth0", laddr, dhcpHandler)
    82  		if err != nil {
    83  			log.Fatal(err)
    84  		}
    85  		server.Serve()
    86  	}()
    87  
    88  	wg.Wait()
    89  }