add support for rate limiting

Adds support for rate limiting for inbound and outbound data streams of
firecracker. This also introduces a TokenBuilder which can be passed
into the NewRateLimiter factory function to return a new rate limiter.

example_test.go

@@ -4,6 +4,7 @@ import (
 	"context"
 	"fmt"
 	"os"
+	"time"
 
 	"github.com/firecracker-microvm/firecracker-go-sdk"
 	models "github.com/firecracker-microvm/firecracker-go-sdk/client/models"
@@ -151,3 +152,63 @@ func ExampleDrivesBuilder_DriveOpt() {
 		panic(err)
 	}
 }
+
+func ExampleNetworkInterfaces_RateLimiting() {
+	// construct the limitations of the bandwidth for firecracker
+	bandwidthBuilder := firecracker.TokenBucketBuilder{}.
+		WithInitialSize(1024 * 1024).        // Initial token amount
+		WithBucketSize(1024 * 1024).         // Max number of tokens
+		WithRefillDuration(30 * time.Second) // Refill rate
+
+	// construct the limitations of the number of operations per duration for firecracker
+	opsBuilder := firecracker.TokenBucketBuilder{}.
+		WithInitialSize(5).
+		WithBucketSize(5).
+		WithRefillDuration(5 * time.Second)
+
+	// create the inbound rate limiter
+	inbound := firecracker.NewRateLimiter(bandwidthBuilder.Build(), opsBuilder.Build())
+
+	bandwidthBuilder = bandwidthBuilder.WithBucketSize(1024 * 1024 * 10)
+	opsBuilder = opsBuilder.
+		WithBucketSize(100).
+		WithInitialSize(100)
+	// create the outbound rate limiter
+	outbound := firecracker.NewRateLimiter(bandwidthBuilder.Build(), opsBuilder.Build())
+
+	networkIfaces := []firecracker.NetworkInterface{
+		{
+			MacAddress:     "01-23-45-67-89-AB-CD-EF",
+			HostDevName:    "tap-name",
+			InRateLimiter:  inbound,
+			OutRateLimiter: outbound,
+		},
+	}
+
+	cfg := firecracker.Config{
+		SocketPath:      "/path/to/socket",
+		KernelImagePath: "/path/to/kernel",
+		Drives:          firecracker.NewDrivesBuilder("/path/to/rootfs").Build(),
+		MachineCfg: models.MachineConfiguration{
+			VcpuCount: 1,
+		},
+		NetworkInterfaces: networkIfaces,
+	}
+
+	ctx := context.Background()
+	m, err := firecracker.NewMachine(ctx, cfg)
+	if err != nil {
+		panic(fmt.Errorf("failed to create new machine: %v", err))
+	}
+
+	defer os.Remove(cfg.SocketPath)
+
+	if err := m.Start(ctx); err != nil {
+		panic(fmt.Errorf("failed to initialize machine: %v", err))
+	}
+
+	// wait for VMM to execute
+	if err := m.Wait(ctx); err != nil {
+		panic(err)
+	}
+}

machine.go

@@ -161,6 +161,11 @@ type NetworkInterface struct {
 	HostDevName string
 	// AllowMMDS makes the Firecracker MMDS available on this network interface.
 	AllowMDDS bool
+
+	// 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
@@ -475,6 +480,14 @@ func (m *Machine) createNetworkInterface(ctx context.Context, iface NetworkInter
 		AllowMmdsRequests: iface.AllowMDDS,
 	}
 
+	if iface.InRateLimiter != nil {
+		ifaceCfg.RxRateLimiter = iface.InRateLimiter
+	}
+
+	if iface.OutRateLimiter != nil {
+		ifaceCfg.TxRateLimiter = iface.OutRateLimiter
+	}
+
 	resp, err := m.client.PutGuestNetworkInterfaceByID(ctx, ifaceID, &ifaceCfg)
 	if err == nil {
 		m.logger.Printf("PutGuestNetworkInterfaceByID: %s", resp.Error())

rate_limiter.go

@@ -0,0 +1,58 @@
+package firecracker
+
+import (
+	"time"
+
+	models "github.com/firecracker-microvm/firecracker-go-sdk/client/models"
+)
+
+// RateLimiterOpt represents a functional option for rate limiting construction
+type RateLimiterOpt func(*models.RateLimiter)
+
+// NewRateLimiter will construct a new RateLimiter with given parameters.
+func NewRateLimiter(bandwidth, ops models.TokenBucket, opts ...RateLimiterOpt) *models.RateLimiter {
+	limiter := &models.RateLimiter{
+		Bandwidth: &bandwidth,
+		Ops:       &ops,
+	}
+
+	for _, opt := range opts {
+		opt(limiter)
+	}
+
+	return limiter
+}
+
+// TokenBucketBuilder is a builder that allows of building components of the
+// models.RateLimiter structure.
+type TokenBucketBuilder struct {
+	bucket models.TokenBucket
+}
+
+// WithBucketSize will set the bucket size for the given token bucket
+func (b TokenBucketBuilder) WithBucketSize(size int64) TokenBucketBuilder {
+	b.bucket.Size = &size
+
+	return b
+}
+
+// WithRefillDuration will set the given refill duration of the bucket fill rate.
+func (b TokenBucketBuilder) WithRefillDuration(dur time.Duration) TokenBucketBuilder {
+	ms := dur.Nanoseconds()
+	ms /= int64(time.Millisecond)
+	b.bucket.RefillTime = &ms
+
+	return b
+}
+
+// WithInitialSize will set the initial token bucket size
+func (b TokenBucketBuilder) WithInitialSize(size int64) TokenBucketBuilder {
+	b.bucket.OneTimeBurst = &size
+
+	return b
+}
+
+// Build will return a new token bucket
+func (b TokenBucketBuilder) Build() models.TokenBucket {
+	return b.bucket
+}

rate_limiter_test.go

@@ -0,0 +1,54 @@
+package firecracker_test
+
+import (
+	"testing"
+	"time"
+
+	"github.com/firecracker-microvm/firecracker-go-sdk"
+	models "github.com/firecracker-microvm/firecracker-go-sdk/client/models"
+	"github.com/stretchr/testify/assert"
+)
+
+func TestRateLimiter(t *testing.T) {
+	bucket := firecracker.TokenBucketBuilder{}.
+		WithRefillDuration(1 * time.Hour).
+		WithBucketSize(100).
+		WithInitialSize(100).
+		Build()
+
+	expectedBucket := models.TokenBucket{
+		OneTimeBurst: firecracker.Int64(100),
+		RefillTime:   firecracker.Int64(3600000),
+		Size:         firecracker.Int64(100),
+	}
+
+	assert.Equal(t, expectedBucket, bucket)
+}
+
+func TestRateLimiter_RefillTime(t *testing.T) {
+	cases := []struct {
+		Name                 string
+		Dur                  time.Duration
+		ExpectedMilliseconds int64
+	}{
+		{
+			Name:                 "one hour",
+			Dur:                  1 * time.Hour,
+			ExpectedMilliseconds: 3600000,
+		},
+		{
+			Name:                 "zero",
+			ExpectedMilliseconds: 0,
+		},
+	}
+
+	for _, c := range cases {
+		t.Run(c.Name, func(t *testing.T) {
+			bucket := firecracker.TokenBucketBuilder{}.
+				WithRefillDuration(c.Dur).
+				Build()
+
+			assert.Equal(t, &c.ExpectedMilliseconds, bucket.RefillTime)
+		})
+	}
+}