Coverage for wasmtime/_wasi.py: 92%

1import ctypes 

2import errno 

3from ctypes import POINTER, c_char, c_char_p, cast, CFUNCTYPE, c_void_p 

4from enum import Enum 

5from os import PathLike 

6from typing import Iterable, List, Union, Callable 

7 

8from wasmtime import Managed, WasmtimeError 

9 

10from . import _ffi as ffi 

11from ._config import setter_property 

12from ._slab import Slab 

13 

14 

15def _encode_path(path: Union[str, bytes, PathLike]) -> bytes: 

16 if isinstance(path, (bytes, str)): 

17 path2 = path 

18 else: 

19 path2 = path.__fspath__() 

20 if isinstance(path2, bytes): 

21 return path2 

22 return path2.encode('utf8') 

23 

24class DirPerms(Enum): 

25 READ_ONLY = ffi.wasi_dir_perms_flags.WASMTIME_WASI_DIR_PERMS_READ.value 

26 WRITE_ONLY = ffi.wasi_dir_perms_flags.WASMTIME_WASI_DIR_PERMS_WRITE.value 

27 READ_WRITE = ffi.wasi_dir_perms_flags.WASMTIME_WASI_DIR_PERMS_READ.value | ffi.wasi_dir_perms_flags.WASMTIME_WASI_DIR_PERMS_WRITE.value 

28 

29class FilePerms(Enum): 

30 READ_ONLY = ffi.wasi_file_perms_flags.WASMTIME_WASI_FILE_PERMS_READ.value 

31 WRITE_ONLY = ffi.wasi_file_perms_flags.WASMTIME_WASI_FILE_PERMS_WRITE.value 

32 READ_WRITE = ffi.wasi_file_perms_flags.WASMTIME_WASI_FILE_PERMS_READ.value | ffi.wasi_file_perms_flags.WASMTIME_WASI_FILE_PERMS_WRITE.value 

33 

34 

35CustomOutput = Callable[[bytes], Union[int, None]] 

36CUSTOM_OUTPUTS: Slab[CustomOutput] = Slab() 

37 

38 

39class WasiConfig(Managed["ctypes._Pointer[ffi.wasi_config_t]"]): 

40 

41 def __init__(self) -> None: 

42 self._set_ptr(ffi.wasi_config_new()) 

43 

44 def _delete(self, ptr: "ctypes._Pointer[ffi.wasi_config_t]") -> None: 

45 ffi.wasi_config_delete(ptr) 

46 

47 @setter_property 

48 def argv(self, argv: List[str]) -> None: 

49 """ 

50 Explicitly configure the `argv` for this WASI configuration 

51 """ 

52 ptrs = to_char_array(argv) 

53 if not ffi.wasi_config_set_argv(self.ptr(), len(argv), ptrs): 

54 raise WasmtimeError("failed to configure argv") 

55 

56 def inherit_argv(self) -> None: 

57 ffi.wasi_config_inherit_argv(self.ptr()) 

58 

59 @setter_property 

60 def env(self, pairs: Iterable[Iterable]) -> None: 

61 """ 

62 Configure environment variables to be returned for this WASI 

63 configuration. 

64 

65 The `pairs` provided must be an iterable list of key/value pairs of 

66 environment variables. 

67 """ 

68 names = [] 

69 values = [] 

70 for name, value in pairs: 

71 names.append(name) 

72 values.append(value) 

73 name_ptrs = to_char_array(names) 

74 value_ptrs = to_char_array(values) 

75 if not ffi.wasi_config_set_env(self.ptr(), len(names), name_ptrs, value_ptrs): 

76 raise WasmtimeError("failed to configure environment") 

77 

78 def inherit_env(self) -> None: 

79 """ 

80 Configures the environment variables available within WASI to be those 

81 in this own process's environment. All environment variables are 

82 inherited. 

83 """ 

84 ffi.wasi_config_inherit_env(self.ptr()) 

85 

86 @setter_property 

87 def stdin_file(self, path: Union[str, bytes, PathLike]) -> None: 

88 """ 

89 Configures a file to be used as the stdin stream of this WASI 

90 configuration. 

91 

92 Reads of the stdin stream will read the path specified. 

93 

94 The file must already exist on the filesystem. If it cannot be 

95 opened then `WasmtimeError` is raised. 

96 """ 

97 

98 res = ffi.wasi_config_set_stdin_file( 

99 self.ptr(), c_char_p(_encode_path(path))) 

100 if not res: 

101 raise WasmtimeError("failed to set stdin file") 

102 

103 def inherit_stdin(self) -> None: 

104 """ 

105 Configures this own process's stdin to be used as the WASI program's 

106 stdin. 

107 

108 Reads of the stdin stream will read this process's stdin. 

109 """ 

110 ffi.wasi_config_inherit_stdin(self.ptr()) 

111 

112 @setter_property 

113 def stdout_file(self, path: str) -> None: 

114 """ 

115 Configures a file to be used as the stdout stream of this WASI 

116 configuration. 

117 

118 Writes to stdout will be written to the file specified. 

119 

120 The file specified will be created if it doesn't exist, or truncated if 

121 it already exists. It must be available to open for writing. If it 

122 cannot be opened for writing then `WasmtimeError` is raised. 

123 """ 

124 res = ffi.wasi_config_set_stdout_file( 

125 self.ptr(), c_char_p(_encode_path(path))) 

126 if not res: 

127 raise WasmtimeError("failed to set stdout file") 

128 

129 @setter_property 

130 def stdout_custom(self, callback: CustomOutput) -> None: 

131 """ 

132 Sets a custom `callback` that is invoked whenever stdout is written to. 

133 """ 

134 ffi.wasi_config_set_stdout_custom( 

135 self.ptr(), custom_call, 

136 CUSTOM_OUTPUTS.allocate(callback), custom_finalize) 

137 

138 def inherit_stdout(self) -> None: 

139 """ 

140 Configures this own process's stdout to be used as the WASI program's 

141 stdout. 

142 

143 Writes to stdout stream will write to this process's stdout. 

144 """ 

145 ffi.wasi_config_inherit_stdout(self.ptr()) 

146 

147 @setter_property 

148 def stderr_file(self, path: str) -> None: 

149 """ 

150 Configures a file to be used as the stderr stream of this WASI 

151 configuration. 

152 

153 Writes to stderr will be written to the file specified. 

154 

155 The file specified will be created if it doesn't exist, or truncated if 

156 it already exists. It must be available to open for writing. If it 

157 cannot be opened for writing then `WasmtimeError` is raised. 

158 """ 

159 res = ffi.wasi_config_set_stderr_file( 

160 self.ptr(), c_char_p(_encode_path(path))) 

161 if not res: 

162 raise WasmtimeError("failed to set stderr file") 

163 

164 @setter_property 

165 def stderr_custom(self, callback: CustomOutput) -> None: 

166 """ 

167 Sets a custom `callback` that is invoked whenever stderr is written to. 

168 """ 

169 ffi.wasi_config_set_stderr_custom( 

170 self.ptr(), custom_call, 

171 CUSTOM_OUTPUTS.allocate(callback), custom_finalize) 

172 

173 def inherit_stderr(self) -> None: 

174 """ 

175 Configures this own process's stderr to be used as the WASI program's 

176 stderr. 

177 

178 Writes to stderr stream will write to this process's stderr. 

179 """ 

180 ffi.wasi_config_inherit_stderr(self.ptr()) 

181 

182 def preopen_dir(self, path: str, guest_path: str, dir_perms: DirPerms = DirPerms.READ_WRITE, file_perms: FilePerms = FilePerms.READ_WRITE) -> None: 

183 """ 

184 Allows the WASI program to access the directory at `path` using the 

185 path `guest_path` within the WASI program. 

186 

187 `dir_perms` specifies the permissions that wasm will have to operate on 

188 `guest_path`. This can be used, for example, to provide readonly access to a 

189 directory. 

190 

191 `file_perms` specifies the maximum set of permissions that can be used for 

192 any file in this directory. 

193 """ 

194 path_ptr = c_char_p(path.encode('utf-8')) 

195 guest_path_ptr = c_char_p(guest_path.encode('utf-8')) 

196 if not ffi.wasi_config_preopen_dir(self.ptr(), path_ptr, guest_path_ptr, dir_perms.value, file_perms.value): 

197 raise WasmtimeError('failed to add preopen dir') 

198 

199 

200def to_char_array(strings: List[str]) -> "ctypes._Pointer[ctypes._Pointer[c_char]]": 

201 ptrs = (c_char_p * len(strings))() 

202 for i, s in enumerate(strings): 

203 ptrs[i] = c_char_p(s.encode('utf-8')) 

204 return cast(ptrs, POINTER(POINTER(c_char))) 

205 

206 

207@CFUNCTYPE(ctypes.c_ssize_t, c_void_p, POINTER(ctypes.c_ubyte), ctypes.c_size_t) 

208def custom_call(idx, ptr, size): # type: ignore 

209 try: 

210 ty = ctypes.c_uint8 * size 

211 arg = bytes(ty.from_address(ctypes.addressof(ptr.contents))) 

212 ret = CUSTOM_OUTPUTS.get(idx or 0)(arg) 

213 if ret is None: 

214 return size 

215 return ret 

216 except Exception as e: 

217 print('failed custom output, required to catch exception:', e) 

218 return -errno.EIO 

219 

220 

221@CFUNCTYPE(None, c_void_p) 

222def custom_finalize(idx): # type: ignore 

223 if CUSTOM_OUTPUTS: 

224 CUSTOM_OUTPUTS.deallocate(idx or 0) 

225 return None