[test commit]

Signed-off-by: Ludvig Liljenberg <4257730+ludfjig@users.noreply.github.com>

Author: ludfjig

src/hyperlight_host/src/hypervisor/hyperlight_vm.rs

@@ -623,6 +623,10 @@ impl HyperlightVm {
             idtr: sregs.idt,
             tr: sregs.tr,
             cs: sregs.cs,
+            ds: sregs.ds,
+            es: sregs.es,
+            fs: sregs.fs,
+            gs: sregs.gs,
             ss: sregs.ss,
         })
     }
@@ -948,6 +952,10 @@ impl HyperlightVm {
             sregs.idt = _segment_state.idtr;
             sregs.tr = _segment_state.tr;
             sregs.cs = _segment_state.cs;
+            sregs.ds = _segment_state.ds;
+            sregs.es = _segment_state.es;
+            sregs.fs = _segment_state.fs;
+            sregs.gs = _segment_state.gs;
             sregs.ss = _segment_state.ss;
             self.vm.set_sregs(&sregs)?;
         }
@@ -1673,6 +1681,10 @@ mod tests {
             idtr: CommonTableRegister::default(),
             tr: CommonSegmentRegister::default(),
             cs: CommonSegmentRegister::default(),
+            ds: CommonSegmentRegister::default(),
+            es: CommonSegmentRegister::default(),
+            fs: CommonSegmentRegister::default(),
+            gs: CommonSegmentRegister::default(),
             ss: CommonSegmentRegister::default(),
         }
     }

src/hyperlight_host/src/sandbox/snapshot.rs

@@ -41,6 +41,10 @@ pub(crate) struct GuestSegmentState {
     pub(crate) idtr: CommonTableRegister,
     pub(crate) tr: CommonSegmentRegister,
     pub(crate) cs: CommonSegmentRegister,
+    pub(crate) ds: CommonSegmentRegister,
+    pub(crate) es: CommonSegmentRegister,
+    pub(crate) fs: CommonSegmentRegister,
+    pub(crate) gs: CommonSegmentRegister,
     pub(crate) ss: CommonSegmentRegister,
 }
 
@@ -601,6 +605,10 @@ mod tests {
             idtr: CommonTableRegister::default(),
             tr: CommonSegmentRegister::default(),
             cs: CommonSegmentRegister::default(),
+            ds: CommonSegmentRegister::default(),
+            es: CommonSegmentRegister::default(),
+            fs: CommonSegmentRegister::default(),
+            gs: CommonSegmentRegister::default(),
             ss: CommonSegmentRegister::default(),
         }
     }

src/hyperlight_host/tests/integration_test.rs

@@ -234,7 +234,7 @@ fn interrupt_same_thread_no_barrier() {
             // Only allow successful calls or interrupted.
             // The call can be successful in case the call is finished before kill() is called.
             Ok(_) | Err(HyperlightError::ExecutionCanceledByHost()) => {}
-            _ => panic!("Unexpected return"),
+            other => panic!("Unexpected return: {:?}", other),
         };
         if sbox2.poisoned() {
             sbox2.restore(snapshot2.clone()).unwrap();