@@ -8,97 +8,101 @@ use uefi::{
88} ; 
99
1010pub  fn  test ( bt :  & BootServices )  { 
11+     // Skip the test if the `pxe` feature is not enabled. 
12+     if  cfg ! ( not( feature = "pxe" ) )  { 
13+         return ; 
14+     } 
15+ 
1116    info ! ( "Testing The PXE base code protocol" ) ; 
1217
13-     if  let  Ok ( handles)  = bt. find_handles :: < BaseCode > ( )  { 
14-         for  handle in  handles { 
15-             let  mut  base_code = bt. open_protocol_exclusive :: < BaseCode > ( handle) . unwrap ( ) ; 
16- 
17-             info ! ( "Starting PXE Base Code" ) ; 
18-             base_code
19-                 . start ( false ) 
20-                 . expect ( "failed to start PXE Base Code" ) ; 
21-             base_code
22-                 . dhcp ( false ) 
23-                 . expect ( "failed to complete a dhcpv4 handshake" ) ; 
24- 
25-             assert ! ( base_code. mode( ) . dhcp_ack_received) ; 
26-             let  dhcp_ack:  & DhcpV4Packet  = base_code. mode ( ) . dhcp_ack . as_ref ( ) ; 
27-             let  server_ip = dhcp_ack. bootp_si_addr ; 
28-             let  server_ip = IpAddress :: new_v4 ( server_ip) ; 
29- 
30-             const  EXAMPLE_FILE_NAME :  & [ u8 ]  = b"example-file.txt\0 " ; 
31-             const  EXAMPLE_FILE_CONTENT :  & [ u8 ]  = b"Hello world!" ; 
32-             let  example_file_name = CStr8 :: from_bytes_with_nul ( EXAMPLE_FILE_NAME ) . unwrap ( ) ; 
33- 
34-             info ! ( "Getting remote file size" ) ; 
35-             let  file_size = base_code
36-                 . tftp_get_file_size ( & server_ip,  example_file_name) 
37-                 . expect ( "failed to query file size" ) ; 
38-             assert_eq ! ( file_size,  EXAMPLE_FILE_CONTENT . len( )  as  u64 ) ; 
39- 
40-             info ! ( "Reading remote file" ) ; 
41-             let  mut  buffer = [ 0 ;  512 ] ; 
42-             let  len = base_code
43-                 . tftp_read_file ( & server_ip,  example_file_name,  Some ( & mut  buffer) ) 
44-                 . expect ( "failed to read file" ) ; 
45-             let  len = usize:: try_from ( len) . unwrap ( ) ; 
46-             assert_eq ! ( EXAMPLE_FILE_CONTENT ,  & buffer[ ..len] ) ; 
47- 
48-             base_code
49-                 . set_ip_filter ( & IpFilter :: new ( IpFilters :: STATION_IP ,  & [ ] ) ) 
50-                 . expect ( "failed to set IP filter" ) ; 
51- 
52-             const  EXAMPLE_SERVICE_PORT :  u16  = 21572 ; 
53- 
54-             info ! ( "Writing UDP packet to example service" ) ; 
55- 
56-             let  payload = [ 1 ,  2 ,  3 ,  4 ] ; 
57-             let  header = [ payload. len ( )  as  u8 ] ; 
58-             let  mut  write_src_port = 0 ; 
59-             base_code
60-                 . udp_write ( 
61-                     UdpOpFlags :: ANY_SRC_PORT , 
62-                     & server_ip, 
63-                     EXAMPLE_SERVICE_PORT , 
64-                     None , 
65-                     None , 
66-                     Some ( & mut  write_src_port) , 
67-                     Some ( & header) , 
68-                     & payload, 
69-                 ) 
70-                 . expect ( "failed to write UDP packet" ) ; 
71- 
72-             info ! ( "Reading UDP packet from example service" ) ; 
73- 
74-             let  mut  src_ip = server_ip; 
75-             let  mut  src_port = EXAMPLE_SERVICE_PORT ; 
76-             let  mut  dest_ip = base_code. mode ( ) . station_ip ; 
77-             let  mut  dest_port = write_src_port; 
78-             let  mut  header = [ 0 ;  1 ] ; 
79-             let  mut  received = [ 0 ;  4 ] ; 
80-             base_code
81-                 . udp_read ( 
82-                     UdpOpFlags :: USE_FILTER , 
83-                     Some ( & mut  dest_ip) , 
84-                     Some ( & mut  dest_port) , 
85-                     Some ( & mut  src_ip) , 
86-                     Some ( & mut  src_port) , 
87-                     Some ( & mut  header) , 
88-                     & mut  received, 
89-                 ) 
90-                 . unwrap ( ) ; 
91- 
92-             // Check the header. 
93-             assert_eq ! ( header[ 0 ]  as  usize ,  payload. len( ) ) ; 
94-             // Check that we receive the reversed payload. 
95-             received. reverse ( ) ; 
96-             assert_eq ! ( payload,  received) ; 
97- 
98-             info ! ( "Stopping PXE Base Code" ) ; 
99-             base_code. stop ( ) . expect ( "failed to stop PXE Base Code" ) ; 
100-         } 
101-     }  else  { 
102-         warn ! ( "PXE Base Code protocol is not supported" ) ; 
18+     let  handles = bt
19+         . find_handles :: < BaseCode > ( ) 
20+         . expect ( "failed to get PXE base code handles" ) ; 
21+     for  handle in  handles { 
22+         let  mut  base_code = bt. open_protocol_exclusive :: < BaseCode > ( handle) . unwrap ( ) ; 
23+ 
24+         info ! ( "Starting PXE Base Code" ) ; 
25+         base_code
26+             . start ( false ) 
27+             . expect ( "failed to start PXE Base Code" ) ; 
28+         base_code
29+             . dhcp ( false ) 
30+             . expect ( "failed to complete a dhcpv4 handshake" ) ; 
31+ 
32+         assert ! ( base_code. mode( ) . dhcp_ack_received) ; 
33+         let  dhcp_ack:  & DhcpV4Packet  = base_code. mode ( ) . dhcp_ack . as_ref ( ) ; 
34+         let  server_ip = dhcp_ack. bootp_si_addr ; 
35+         let  server_ip = IpAddress :: new_v4 ( server_ip) ; 
36+ 
37+         const  EXAMPLE_FILE_NAME :  & [ u8 ]  = b"example-file.txt\0 " ; 
38+         const  EXAMPLE_FILE_CONTENT :  & [ u8 ]  = b"Hello world!" ; 
39+         let  example_file_name = CStr8 :: from_bytes_with_nul ( EXAMPLE_FILE_NAME ) . unwrap ( ) ; 
40+ 
41+         info ! ( "Getting remote file size" ) ; 
42+         let  file_size = base_code
43+             . tftp_get_file_size ( & server_ip,  example_file_name) 
44+             . expect ( "failed to query file size" ) ; 
45+         assert_eq ! ( file_size,  EXAMPLE_FILE_CONTENT . len( )  as  u64 ) ; 
46+ 
47+         info ! ( "Reading remote file" ) ; 
48+         let  mut  buffer = [ 0 ;  512 ] ; 
49+         let  len = base_code
50+             . tftp_read_file ( & server_ip,  example_file_name,  Some ( & mut  buffer) ) 
51+             . expect ( "failed to read file" ) ; 
52+         let  len = usize:: try_from ( len) . unwrap ( ) ; 
53+         assert_eq ! ( EXAMPLE_FILE_CONTENT ,  & buffer[ ..len] ) ; 
54+ 
55+         base_code
56+             . set_ip_filter ( & IpFilter :: new ( IpFilters :: STATION_IP ,  & [ ] ) ) 
57+             . expect ( "failed to set IP filter" ) ; 
58+ 
59+         const  EXAMPLE_SERVICE_PORT :  u16  = 21572 ; 
60+ 
61+         info ! ( "Writing UDP packet to example service" ) ; 
62+ 
63+         let  payload = [ 1 ,  2 ,  3 ,  4 ] ; 
64+         let  header = [ payload. len ( )  as  u8 ] ; 
65+         let  mut  write_src_port = 0 ; 
66+         base_code
67+             . udp_write ( 
68+                 UdpOpFlags :: ANY_SRC_PORT , 
69+                 & server_ip, 
70+                 EXAMPLE_SERVICE_PORT , 
71+                 None , 
72+                 None , 
73+                 Some ( & mut  write_src_port) , 
74+                 Some ( & header) , 
75+                 & payload, 
76+             ) 
77+             . expect ( "failed to write UDP packet" ) ; 
78+ 
79+         info ! ( "Reading UDP packet from example service" ) ; 
80+ 
81+         let  mut  src_ip = server_ip; 
82+         let  mut  src_port = EXAMPLE_SERVICE_PORT ; 
83+         let  mut  dest_ip = base_code. mode ( ) . station_ip ; 
84+         let  mut  dest_port = write_src_port; 
85+         let  mut  header = [ 0 ;  1 ] ; 
86+         let  mut  received = [ 0 ;  4 ] ; 
87+         base_code
88+             . udp_read ( 
89+                 UdpOpFlags :: USE_FILTER , 
90+                 Some ( & mut  dest_ip) , 
91+                 Some ( & mut  dest_port) , 
92+                 Some ( & mut  src_ip) , 
93+                 Some ( & mut  src_port) , 
94+                 Some ( & mut  header) , 
95+                 & mut  received, 
96+             ) 
97+             . unwrap ( ) ; 
98+ 
99+         // Check the header. 
100+         assert_eq ! ( header[ 0 ]  as  usize ,  payload. len( ) ) ; 
101+         // Check that we receive the reversed payload. 
102+         received. reverse ( ) ; 
103+         assert_eq ! ( payload,  received) ; 
104+ 
105+         info ! ( "Stopping PXE Base Code" ) ; 
106+         base_code. stop ( ) . expect ( "failed to stop PXE Base Code" ) ; 
103107    } 
104108} 
0 commit comments