netboot/netboot.go

// Copyright 2017 The Fuchsia Authors. All rights reserved.
// Use of this source code is governed by a BSD-style license that can be
// found in the LICENSE file.
package netboot

import (
	"bytes"
	"context"
	"encoding/binary"
	"fmt"
	"net"
	"strings"
	"time"
)

const (
	Magic = 0xAA774217
	ServerPort = 33330
	AdvertPort = 33331
	CmdPortStart = 33332
	CmdPortEnd = 33339
)

var (
	Cookie = uint32(0x12345678)
)

const (
	NetbootBoot = 4
	NetbootQuery = 5
)

type Netboot struct {
	Port int
	Timeout time.Duration
	Wait bool
}

type NetbootHeader struct {
	Magic uint32
	Cookie uint32
	Cmd uint32
	Arg uint32
}

type msg struct {
	hdr NetbootHeader
	data [1024]byte
}

func (n *Netboot) bind() (conn *net.UDPConn, err error) {
	for p := CmdPortStart; p <= CmdPortEnd; p++ {
		conn, err = net.ListenUDP("udp6", &net.UDPAddr{IP: net.IPv6zero, Port: p})
		if err == nil {
			break
		}
	}
	return
}

func (n *Netboot) Discover(_ context.Context, host string, fuchsia bool) (*net.UDPAddr, error) {
	conn, err := n.bind()
	if err != nil {
		return nil, fmt.Errorf("failed to bind to udp6 port: %v\n", err)
	}

	Cookie++
	req := msg{
		hdr: NetbootHeader{
			Magic: Magic,
			Cookie: Cookie,
			Cmd: NetbootQuery,
			Arg: 0,
		},
	}
	copy(req.data[:], host)

	buf := new(bytes.Buffer)
	if err := binary.Write(buf, binary.LittleEndian, req); err != nil {
		return nil, fmt.Errorf("binary write failed: %v\n", err)
	}

	_, err = conn.WriteToUDP(buf.Bytes(), &net.UDPAddr{IP: net.IPv6linklocalallnodes, Port: ServerPort, Zone: "eno1"})
	if err != nil {
		return nil, err
	}

	conn.SetReadDeadline(time.Now().Add(n.Timeout))

	b := make([]byte, 4096)
	_, addr, err := conn.ReadFromUDP(b)
	if err != nil {
		return nil, err
	}

	buf2 := bytes.NewReader(b)

	var res NetbootHeader
	if err := binary.Read(buf2, binary.LittleEndian, &res); err != nil {
		return nil, fmt.Errorf("binary read failed: %v\n", err)
	}

	data := make([]byte, 4096)
	buf2.Read(data)

	fields := strings.Split(string(data[:]), ";")
	for _, f := range fields {
		vars := strings.SplitN(f, "=", 2)
		if vars[0] == "nodename" && vars[1] != host {
			return nil, fmt.Errorf("invalid hostname `%s`", vars[1])
		}
	}

	if fuchsia {
		addr.IP[11] = 0xFF;
	}

	return addr, nil
}

func (n *Netboot) Beacon(_ context.Context) (*net.UDPAddr, error) {
	conn, err := net.ListenUDP("udp6", &net.UDPAddr{IP: net.IPv6zero, Port: AdvertPort})

	conn.SetReadDeadline(time.Now().Add(n.Timeout))

	b := make([]byte, 4096)
	_, addr, err := conn.ReadFromUDP(b)
	if err != nil {
		return nil, err
	}

	buf := bytes.NewReader(b)

	var msg NetbootHeader
	if err := binary.Read(buf, binary.LittleEndian, &msg); err != nil {
		return nil, fmt.Errorf("binary read failed: %v\n", err)
	}

	data := make([]byte, 4096)
	if _, err := buf.Read(data); err != nil {
		return nil, fmt.Errorf("read failed: %v\n", err)
	}

	fields := strings.Split(string(data[:]), ";")
	for _, f := range fields {
		vars := strings.SplitN(f, "=", 2)
		if vars[0] == "nodename" {
			return addr, nil
		}
	}

	return nil, fmt.Errorf("no valid beacon")
}

func (n *Netboot) Boot(_ context.Context, addr *net.UDPAddr) error {
	Cookie++
	msg := NetbootHeader{
		Magic: Magic,
		Cookie: Cookie,
		Cmd: NetbootBoot,
		Arg: 0,
	}

	buf := new(bytes.Buffer)
	if err := binary.Write(buf, binary.LittleEndian, msg); err != nil {
		return fmt.Errorf("binary write failed: %v\n", err)
	}

	conn, err := net.ListenUDP("udp6", &net.UDPAddr{IP: net.IPv6zero, Port: AdvertPort})
	if err != nil {
		return fmt.Errorf("failed to create a socket: %v\n", err)
	}

	_, err = conn.WriteToUDP(buf.Bytes(), &net.UDPAddr{IP: addr.IP, Port: ServerPort, Zone: addr.Zone})
	if err != nil {
		return fmt.Errorf("failed to send boot command: %v\n", err)
	}

	return nil
}