talos

1
// This Source Code Form is subject to the terms of the Mozilla Public
2
// License, v. 2.0. If a copy of the MPL was not distributed with this
3
// file, You can obtain one at http://mozilla.org/MPL/2.0/.
4

5
// Package ipxe provides utility to deliver iPXE images and build iPXE scripts.
6
package ipxe
7

8
import (
9
	"bytes"
10
	"embed"
11
	"errors"
12
	"fmt"
13
	"io"
14
	"log"
15
	"path/filepath"
16
	"text/template"
17
)
18

19
//go:embed "data/*"
20
var ipxeFiles embed.FS
21

22
// TFTPHandler is called when client starts file download from the TFTP server.
23
//
24
// TFTP handler also patches the iPXE binary on the fly with the script
25
// which chainloads next handler.
26
func TFTPHandler(next string) func(filename string, rf io.ReaderFrom) error {
27
	return func(filename string, rf io.ReaderFrom) error {
28
		log.Printf("tftp request: %s", filename)
29

30
		file, err := ipxeFiles.Open(filepath.Join("data", filename))
31
		if err != nil {
32
			return err
33
		}
34

35
		defer file.Close() //nolint:errcheck
36

37
		contents, err := io.ReadAll(file)
38
		if err != nil {
39
			return err
40
		}
41

42
		var script bytes.Buffer
43

44
		if err = scriptTemplate.Execute(&script, struct {
45
			Next string
46
		}{
47
			Next: next,
48
		}); err != nil {
49
			return err
50
		}
51

52
		contents, err = patchScript(contents, script.Bytes())
53
		if err != nil {
54
			return fmt.Errorf("error patching %q: %w", filename, err)
55
		}
56

57
		_, err = rf.ReadFrom(bytes.NewReader(contents))
58

59
		return err
60
	}
61
}
62

63
// scriptTemplate to run DHCP and chain the boot to the .Next endpoint.
64
var scriptTemplate = template.Must(template.New("iPXE embedded").Parse(`#!ipxe
65
prompt --key 0x02 --timeout 2000 Press Ctrl-B for the iPXE command line... && shell ||
66

67
{{/* print interfaces */}}
68
ifstat
69

70
{{/* retry 10 times overall */}}
71
set attempts:int32 10
72
set x:int32 0
73

74
:retry_loop
75

76
	set idx:int32 0
77

78
	:loop
79
		{{/* try DHCP on each interface */}}
80
		isset ${net${idx}/mac} || goto exhausted
81

82
		ifclose
83
		iflinkwait --timeout 5000 net${idx} || goto next_iface
84
		dhcp net${idx} || goto next_iface
85
		goto boot
86

87
	:next_iface
88
		inc idx && goto loop
89

90
	:boot
91
		{{/* attempt boot, if fails try next iface */}}
92
		route
93

94
		chain --replace {{ .Next }} || goto next_iface
95

96
:exhausted
97
	echo
98
	echo Failed to iPXE boot successfully via all interfaces
99

100
	iseq ${x} ${attempts} && goto fail ||
101

102
	echo Retrying...
103
	echo
104

105
	inc x
106
	goto retry_loop
107

108
:fail
109
	echo
110
	echo Failed to get a valid response after ${attempts} attempts
111
	echo
112

113
	echo Rebooting in 5 seconds...
114
	sleep 5
115
	reboot
116
`))
117

118
var (
119
	placeholderStart = []byte("# *PLACEHOLDER START*")
120
	placeholderEnd   = []byte("# *PLACEHOLDER END*")
121
)
122

123
// patchScript patches the iPXE script into the iPXE binary.
124
//
125
// The iPXE binary should be built uncompressed with an embedded
126
// script stub which contains abovementioned placeholders.
127
func patchScript(contents, script []byte) ([]byte, error) {
128
	start := bytes.Index(contents, placeholderStart)
129
	if start == -1 {
130
		return nil, errors.New("placeholder start not found")
131
	}
132

133
	end := bytes.Index(contents, placeholderEnd)
134
	if end == -1 {
135
		return nil, errors.New("placeholder end not found")
136
	}
137

138
	if end < start {
139
		return nil, errors.New("placeholder end before start")
140
	}
141

142
	end += len(placeholderEnd)
143

144
	length := end - start
145

146
	if len(script) > length {
147
		return nil, fmt.Errorf("script size %d is larger than placeholder space %d", len(script), length)
148
	}
149

150
	script = append(script, bytes.Repeat([]byte{'\n'}, length-len(script))...)
151

152
	copy(contents[start:end], script)
153

154
	return contents, nil
155
}
156