NetworkPkg: Handling timeout case in httpboot driver

[mirror_edk2.git] / NetworkPkg / HttpBootDxe / HttpBootSupport.c

diff --git a/NetworkPkg/HttpBootDxe/HttpBootSupport.c b/NetworkPkg/HttpBootDxe/HttpBootSupport.c
index 66eca783af712e13457bc79ab9884b35e5a62b23..9410bf9e15209cf2c21eaf0121fef297985247ad 100644 (file)
--- a/NetworkPkg/HttpBootDxe/HttpBootSupport.c
+++ b/NetworkPkg/HttpBootDxe/HttpBootSupport.c
@@ -754,6 +754,21 @@ HttpIoCreateIo (
   HttpIo->RspToken.Event = Event;\r
   HttpIo->RspToken.Message = &HttpIo->RspMessage;\r
 \r
+  //\r
+  // Create TimeoutEvent for response\r
+  //\r
+  Status = gBS->CreateEvent (\r
+                  EVT_TIMER,\r
+                  TPL_CALLBACK,\r
+                  NULL,\r
+                  NULL,\r
+                  &Event\r
+                  );\r
+  if (EFI_ERROR (Status)) {\r
+    goto ON_ERROR;\r
+  }\r
+  HttpIo->TimeoutEvent = Event;\r
+\r
   return EFI_SUCCESS;\r
   \r
 ON_ERROR:\r
@@ -789,6 +804,11 @@ HttpIoDestroyIo (
   if (Event != NULL) {\r
     gBS->CloseEvent (Event);\r
   }\r
+\r
+  Event = HttpIo->TimeoutEvent;\r
+  if (Event != NULL) {\r
+    gBS->CloseEvent (Event);\r
+  }\r
   \r
   Http = HttpIo->Http;\r
   if (Http != NULL) {\r
@@ -902,6 +922,14 @@ HttpIoRecvResponse (
     return EFI_INVALID_PARAMETER;\r
   }\r
 \r
+  //\r
+  // Start the timer, and wait Timeout seconds to receive the header packet.\r
+  //\r
+  Status = gBS->SetTimer (HttpIo->TimeoutEvent, TimerRelative, HTTP_BOOT_RESPONSE_TIMEOUT * TICKS_PER_MS);\r
+  if (EFI_ERROR (Status)) {\r
+    return Status;\r
+  }\r
+\r
   //\r
   // Queue the response token to HTTP instances.\r
   //\r
@@ -924,16 +952,32 @@ HttpIoRecvResponse (
                    );\r
   \r
   if (EFI_ERROR (Status)) {\r
+    gBS->SetTimer (HttpIo->TimeoutEvent, TimerCancel, 0);\r
     return Status;\r
   }\r
 \r
   //\r
   // Poll the network until receive finish.\r
   //\r
-  while (!HttpIo->IsRxDone) {\r
+  while (!HttpIo->IsRxDone && ((HttpIo->TimeoutEvent == NULL) || EFI_ERROR (gBS->CheckEvent (HttpIo->TimeoutEvent)))) {\r
     Http->Poll (Http);\r
   }\r
 \r
+  gBS->SetTimer (HttpIo->TimeoutEvent, TimerCancel, 0);\r
+\r
+  if (!HttpIo->IsRxDone) {\r
+    //\r
+    // Timeout occurs, cancel the response token.\r
+    //\r
+    Http->Cancel (Http, &HttpIo->RspToken);\r
+   \r
+    Status = EFI_TIMEOUT;\r
+    \r
+    return Status;\r
+  } else {\r
+    HttpIo->IsRxDone = FALSE;\r
+  }\r
+\r
   //\r
   // Store the received data into the wrapper.\r
   //\r