Add support for the MMDS configuration.

Signed-off-by: Yutaka Juba <[email protected]>

Author: yjuba

firecracker.go

@@ -332,6 +332,16 @@ func (f *Client) PatchMmds(ctx context.Context, metadata interface{}, opts ...Pa
 	return f.client.Operations.PatchMmds(params)
 }
 
+// PutMmdsConfig is a wrapper for the swagger generated client to make calling of the
+// API easier.
+func (f *Client) PutMmdsConfig(ctx context.Context, config *models.MmdsConfig) (*ops.PutMmdsConfigNoContent, error) {
+	params := ops.NewPutMmdsConfigParams()
+	params.SetContext(ctx)
+	params.SetBody(config)
+
+	return f.client.Operations.PutMmdsConfig(params)
+}
+
 // GetMachineConfigurationOpt  is a functional option to be used for the
 // GetMachineConfiguration API in setting any additional optional fields.
 type GetMachineConfigurationOpt func(*ops.GetMachineConfigurationParams)

handlers.go

@@ -16,6 +16,7 @@ package firecracker
 import (
 	"context"
 	"fmt"
+	"net"
 	"os"
 )
 
@@ -30,6 +31,7 @@ const (
 	CreateNetworkInterfacesHandlerName = "fcinit.CreateNetworkInterfaces"
 	AddVsocksHandlerName               = "fcinit.AddVsocks"
 	SetMetadataHandlerName             = "fcinit.SetMetadata"
+	ConfigMmdsHandlerName              = "fcinit.ConfigMmds"
 	LinkFilesToRootFSHandlerName       = "fcinit.LinkFilesToRootFS"
 	SetupNetworkHandlerName            = "fcinit.SetupNetwork"
 	SetupKernelArgsHandlerName         = "fcinit.SetupKernelArgs"
@@ -258,6 +260,17 @@ func NewSetMetadataHandler(metadata interface{}) Handler {
 	}
 }
 
+// NewConfigMmdsHandler is a named handler that puts the MMDS config into the
+// firecracker process.
+func NewConfigMmdsHandler(address net.IP) Handler {
+	return Handler{
+		Name: ConfigMmdsHandlerName,
+		Fn: func(ctx context.Context, m *Machine) error {
+			return m.setMmdsConfig(ctx, address)
+		},
+	}
+}
+
 var defaultFcInitHandlerList = HandlerList{}.Append(
 	SetupNetworkHandler,
 	SetupKernelArgsHandler,

handlers_test.go

@@ -3,6 +3,7 @@ package firecracker
 import (
 	"context"
 	"fmt"
+	"net"
 	"os"
 	"path/filepath"
 	"reflect"
@@ -515,6 +516,8 @@ func TestHandlers(t *testing.T) {
 		"foo": "bar",
 		"baz": "qux",
 	}
+	mmdsAddress := net.IPv4(169, 254, 169, 254)
+	mmdsConfig := &models.MmdsConfig{IPV4Address: String(mmdsAddress.String())}
 
 	cases := []struct {
 		Handler Handler
@@ -619,6 +622,19 @@ func TestHandlers(t *testing.T) {
 			},
 			Config: Config{},
 		},
+		{
+			Handler: NewConfigMmdsHandler(mmdsAddress),
+			Client: fctesting.MockClient{
+				PutMmdsConfigFn: func(params *ops.PutMmdsConfigParams) (*ops.PutMmdsConfigNoContent, error) {
+					called = ConfigMmdsHandlerName
+					if !reflect.DeepEqual(mmdsConfig, params.Body) {
+						return nil, fmt.Errorf("incorrect mmds config value: %v", params.Body)
+					}
+					return &ops.PutMmdsConfigNoContent{}, nil
+				},
+			},
+			Config: Config{},
+		},
 	}
 
 	ctx := context.Background()

machine.go

@@ -18,6 +18,7 @@ import (
 	"encoding/json"
 	"fmt"
 	"io"
+	"net"
 	"os"
 	"os/exec"
 	"os/signal"
@@ -870,6 +871,19 @@ func (m *Machine) sendCtrlAltDel(ctx context.Context) error {
 	return err
 }
 
+func (m *Machine) setMmdsConfig(ctx context.Context, address net.IP) error {
+	mmdsCfg := models.MmdsConfig{
+		IPV4Address: String(address.String()),
+	}
+	if _, err := m.client.PutMmdsConfig(ctx, &mmdsCfg); err != nil {
+		m.logger.Errorf("Setting mmds configuration failed: %s: %v", address, err)
+		return err
+	}
+
+	m.logger.Debug("SetMmdsConfig successful")
+	return nil
+}
+
 // SetMetadata sets the machine's metadata for MDDS
 func (m *Machine) SetMetadata(ctx context.Context, metadata interface{}) error {
 	if _, err := m.client.PutMmds(ctx, metadata); err != nil {