Merge pull request #103 from kzys/patch-network

Add UpdateGuestNetworkInterface

Author: mxpv

firecracker.go

@@ -141,6 +141,27 @@ func (f *Client) PutGuestNetworkInterfaceByID(ctx context.Context, ifaceID strin
 	return f.client.Operations.PutGuestNetworkInterfaceByID(cfg)
 }
 
+// PatchGuestNetworkInterfaceByIDOpt is a functional option to be used for the
+// PatchGuestNetworkInterfaceByID API in setting any additional optional fields.
+type PatchGuestNetworkInterfaceByIDOpt func(*ops.PatchGuestNetworkInterfaceByIDParams)
+
+// PatchGuestNetworkInterfaceByID is a wrapper for the swagger generated client to make calling of the
+// API easier.
+func (f *Client) PatchGuestNetworkInterfaceByID(ctx context.Context, ifaceID string, ifaceCfg *models.PartialNetworkInterface, opts ...PatchGuestNetworkInterfaceByIDOpt) (*ops.PatchGuestNetworkInterfaceByIDNoContent, error) {
+	timeout, cancel := context.WithTimeout(ctx, firecrackerRequestTimeout)
+	defer cancel()
+
+	cfg := ops.NewPatchGuestNetworkInterfaceByIDParamsWithContext(timeout)
+	cfg.SetBody(ifaceCfg)
+	cfg.SetIfaceID(ifaceID)
+
+	for _, opt := range opts {
+		opt(cfg)
+	}
+
+	return f.client.Operations.PatchGuestNetworkInterfaceByID(cfg)
+}
+
 // PutGuestDriveByIDOpt is a functional option to be used for the
 // PutGuestDriveByID API in setting any additional optional fields.
 type PutGuestDriveByIDOpt func(*ops.PutGuestDriveByIDParams)

machine.go

@@ -192,6 +192,14 @@ type NetworkInterface struct {
 	OutRateLimiter *models.RateLimiter
 }
 
+// RateLimiterSet represents a pair of RateLimiters (inbound and outbound)
+type RateLimiterSet struct {
+	// InRateLimiter limits the incoming bytes.
+	InRateLimiter *models.RateLimiter
+	// OutRateLimiter limits the outgoing bytes.
+	OutRateLimiter *models.RateLimiter
+}
+
 // VsockDevice represents a vsock connection between the host and the guest
 // microVM.
 type VsockDevice struct {
@@ -564,6 +572,26 @@ func (m *Machine) createNetworkInterface(ctx context.Context, iface NetworkInter
 	return err
 }
 
+// UpdateGuestNetworkInterfaceRateLimit modifies the specified network interface's rate limits
+func (m *Machine) UpdateGuestNetworkInterfaceRateLimit(ctx context.Context, ifaceID string, rateLimiters RateLimiterSet, opts ...PatchGuestNetworkInterfaceByIDOpt) error {
+	iface := models.PartialNetworkInterface{
+		IfaceID: &ifaceID,
+	}
+	if rateLimiters.InRateLimiter != nil {
+		iface.RxRateLimiter = rateLimiters.InRateLimiter
+	}
+	if rateLimiters.OutRateLimiter != nil {
+		iface.TxRateLimiter = rateLimiters.InRateLimiter
+	}
+	if _, err := m.client.PatchGuestNetworkInterfaceByID(ctx, ifaceID, &iface, opts...); err != nil {
+		m.logger.Errorf("Update network interface failed: %s: %v", ifaceID, err)
+		return err
+	}
+
+	m.logger.Infof("Updated network interface: %s", ifaceID)
+	return nil
+}
+
 // attachDrive attaches a secondary block device
 func (m *Machine) attachDrive(ctx context.Context, dev models.Drive) error {
 	hostPath := StringValue(dev.PathOnHost)

machine_test.go

@@ -275,6 +275,13 @@ func TestMicroVMExecution(t *testing.T) {
 
 	vmlinuxPath := getVmlinuxPath(t)
 
+	networkIfaces := []NetworkInterface{
+		{
+			MacAddress:  "01-23-45-67-89-AB-CD-EF",
+			HostDevName: "tap0",
+		},
+	}
+
 	cfg := Config{
 		SocketPath:  socketPath,
 		LogFifo:     logFifo,
@@ -288,6 +295,7 @@ func TestMicroVMExecution(t *testing.T) {
 		},
 		Debug:             true,
 		DisableValidation: true,
+		NetworkInterfaces: networkIfaces,
 	}
 
 	ctx := context.Background()
@@ -332,6 +340,7 @@ func TestMicroVMExecution(t *testing.T) {
 	t.Run("TestAttachVsock", func(t *testing.T) { testAttachVsock(ctx, t, m) })
 	t.Run("SetMetadata", func(t *testing.T) { testSetMetadata(ctx, t, m) })
 	t.Run("TestUpdateGuestDrive", func(t *testing.T) { testUpdateGuestDrive(ctx, t, m) })
+	t.Run("TestUpdateGuestNetworkInterface", func(t *testing.T) { testUpdateGuestNetworkInterface(ctx, t, m) })
 	t.Run("TestStartInstance", func(t *testing.T) { testStartInstance(ctx, t, m) })
 
 	// Let the VMM start and stabilize...
@@ -489,6 +498,18 @@ func testUpdateGuestDrive(ctx context.Context, t *testing.T, m *Machine) {
 	}
 }
 
+func testUpdateGuestNetworkInterface(ctx context.Context, t *testing.T, m *Machine) {
+	rateLimitSet := RateLimiterSet{
+		InRateLimiter: NewRateLimiter(
+			TokenBucketBuilder{}.WithBucketSize(10).WithRefillDuration(10).Build(),
+			TokenBucketBuilder{}.WithBucketSize(10).WithRefillDuration(10).Build(),
+		),
+	}
+	if err := m.UpdateGuestNetworkInterfaceRateLimit(ctx, "1", rateLimitSet); err != nil {
+		t.Fatalf("Failed to update the network interface %v", err)
+	}
+}
+
 func testCreateNetworkInterfaceByID(ctx context.Context, t *testing.T, m *Machine) {
 	if skipTuntap {
 		t.Skip("Skipping: tuntap tests explicitly disabled")

machineiface.go

@@ -29,4 +29,5 @@ type MachineIface interface {
 	Wait(context.Context) error
 	SetMetadata(context.Context, interface{}) error
 	UpdateGuestDrive(context.Context, string, string, ...PatchGuestDriveByIDOpt) error
+	UpdateGuestNetworkInterfaceRateLimit(context.Context, string, RateLimiterSet, ...PatchGuestNetworkInterfaceByIDOpt) error
 }