From 88b035214935fd01e3c4bb594a82abdbc9bd9c48 Mon Sep 17 00:00:00 2001
From: Dan Norris <protochron@users.noreply.github.com>
Date: Thu, 21 Sep 2023 16:25:48 -0400
Subject: [PATCH] Add entropy device method

---
 firecracker.go  | 14 ++++++++++++++
 machine.go      |  9 +++++++++
 machine_test.go |  6 ++++++
 3 files changed, 29 insertions(+)

diff --git a/firecracker.go b/firecracker.go
index 0df8b55b..f2b3c2d8 100644
--- a/firecracker.go
+++ b/firecracker.go
@@ -519,3 +519,17 @@ func (f *Client) GetExportVMConfig(opts ...GetExportVMConfigOpt) (*ops.GetExport
 
 	return f.client.Operations.GetExportVMConfig(p)
 }
+
+type PutEntropyDeviceOpt func(*ops.PutEntropyDeviceParams)
+
+// PutEntropyDevice is a wrapper for the swagger generated client to make calling of the
+// API easier.
+func (f *Client) PutEntropyDevice(ctx context.Context, opts ...PutEntropyDeviceOpt) (*ops.PutEntropyDeviceNoContent, error) {
+	p := ops.NewPutEntropyDeviceParams()
+	p.SetTimeout(time.Duration(f.firecrackerRequestTimeout) * time.Millisecond)
+	for _, opt := range opts {
+		opt(p)
+	}
+
+	return f.client.Operations.PutEntropyDevice(p)
+}
diff --git a/machine.go b/machine.go
index 67313f42..0d6be9c8 100644
--- a/machine.go
+++ b/machine.go
@@ -1281,3 +1281,12 @@ func (m *Machine) UpdateBalloonStats(ctx context.Context, statsPollingIntervals
 	m.logger.Debug("UpdateBalloonStats successful")
 	return nil
 }
+
+func (m *Machine) CreateEntropyDevice(ctx context.Context, opts ...PutEntropyDeviceOpt) error {
+	if _, err := m.client.PutEntropyDevice(ctx, opts...); err != nil {
+		m.logger.Errorf("CreateEntropyDevice failed: %v", err)
+		return err
+	}
+	m.logger.Debug("UpdateEntropyDevice successful")
+	return nil
+}
diff --git a/machine_test.go b/machine_test.go
index b8d7b502..8eb75ac3 100644
--- a/machine_test.go
+++ b/machine_test.go
@@ -2422,3 +2422,9 @@ func testUpdateBalloonStats(ctx context.Context, t *testing.T, m *Machine) {
 		t.Errorf("Updating balloon staistics failed from testUpdateBalloonStats: %s", err)
 	}
 }
+
+func testPutEntropyDevice(ctx context.Context, t *testing.T, m *Machine) {
+	if _, err := m.client.PutEntropyDevice(ctx); err != nil {
+		t.Errorf("Put entropy device failed from testPutEntropyDevice: %s", err)
+	}
+}