1#!/usr/bin/env python 

2# This file is part of Xpra. 

3# Copyright (C) 2011-2020 Antoine Martin <antoine@xpra.org> 

4# Copyright (C) 2008, 2009, 2010 Nathaniel Smith <njs@pobox.com> 

5# Xpra is released under the terms of the GNU GPL v2, or, at your option, any 

6# later version. See the file COPYING for details. 

7 

8from collections import namedtuple 

9 

10from xpra.util import envbool 

11from xpra.net.header import LZ4_FLAG, ZLIB_FLAG, LZO_FLAG, BROTLI_FLAG 

12 

13 

14MAX_SIZE = 256*1024*1024 

15 

16#all the compressors we know about, in best compatibility order: 

17ALL_COMPRESSORS = ("zlib", "lz4", "lzo", "brotli", "none") 

18#order for performance: 

19PERFORMANCE_ORDER = ("none", "lz4", "lzo", "zlib", "brotli") 

20 

21 

22Compression = namedtuple("Compression", ["name", "version", "python_version", "compress", "decompress"]) 

23 

24COMPRESSION = {} 

25 

26 

27def init_lz4(): 

28 from lz4 import VERSION 

29 from lz4.version import version 

30 from lz4.block import compress, decompress 

31 import struct 

32 LZ4_HEADER = struct.Struct(b'<L') 

33 def lz4_compress(packet, level): 

34 flag = min(15, level) | LZ4_FLAG 

35 if level>=7: 

36 return flag, compress(packet, mode="high_compression", compression=level) 

37 if level<=3: 

38 return flag, compress(packet, mode="fast", acceleration=8-level*2) 

39 return flag, compress(packet) 

40 def lz4_decompress(data): 

41 size = LZ4_HEADER.unpack_from(data[:4])[0] 

42 #it would be better to use the max_size we have in protocol, 

43 #but this hardcoded value will do for now 

44 if size>MAX_SIZE: 

45 sizemb = size//1024//1024 

46 maxmb = MAX_SIZE//1024//1024 

47 raise Exception("uncompressed data is too large: %iMB, limit is %iMB" % (sizemb, maxmb)) 

48 return decompress(data) 

49 return Compression("lz4", version, VERSION.encode("latin1"), lz4_compress, lz4_decompress) 

50 

51def init_lzo(): 

52 import lzo #@UnresolvedImport 

53 def lzo_compress(packet, level): 

54 if isinstance(packet, memoryview): 

55 packet = packet.tobytes() 

56 return level | LZO_FLAG, lzo.compress(packet) 

57 return Compression("lzo", lzo.LZO_VERSION_STRING, lzo.__version__, lzo_compress, lzo.decompress) 

58 

59def init_brotli(): 

60 from brotli import compress, decompress, __version__ 

61 def brotli_compress(packet, level): 

62 if len(packet)>1024*1024: 

63 level = min(9, level) 

64 else: 

65 level = min(11, level) 

66 if not isinstance(packet, bytes): 

67 packet = bytes(str(packet), 'UTF-8') 

68 return level | BROTLI_FLAG, compress(packet, quality=level) 

69 return Compression("brotli", None, __version__, brotli_compress, decompress) 

70 

71def init_zlib(): 

72 from zlib import compress, decompress, __version__ 

73 def zlib_compress(packet, level): 

74 level = max(1, level//2) 

75 if isinstance(packet, memoryview): 

76 packet = packet.tobytes() 

77 elif not isinstance(packet, bytes): 

78 packet = bytes(str(packet), 'UTF-8') 

79 return level + ZLIB_FLAG, compress(packet, level) 

80 def zlib_decompress(data): 

81 if isinstance(data, memoryview): 

82 data = data.tobytes() 

83 return decompress(data) 

84 return Compression("zlib", None, __version__, zlib_compress, zlib_decompress) 

85 

86def init_none(): 

87 def nocompress(packet, _level): 

88 if not isinstance(packet, bytes): 

89 packet = bytes(str(packet), 'UTF-8') 

90 return 0, packet 

91 def decompress(v): 

92 return v 

93 return Compression("none", None, None, nocompress, decompress) 

94 

95 

96def init_all(): 

97 for x in list(ALL_COMPRESSORS)+["none"]: 

98 if not envbool("XPRA_%s" % (x.upper()), True): 

99 continue 

100 fn = globals().get("init_%s" % x) 

101 try: 

102 c = fn() 

103 assert c 

104 COMPRESSION[x] = c 

105 except (ImportError, AttributeError): 

106 from xpra.log import Logger 

107 logger = Logger("network", "protocol") 

108 logger.debug("no %s", x, exc_info=True) 

109init_all() 

110 

111 

112def use(compressor) -> bool: 

113 return compressor in COMPRESSION 

114 

115 

116def get_compression_caps() -> dict: 

117 caps = {} 

118 for x in ALL_COMPRESSORS: 

119 c = COMPRESSION.get(x) 

120 if c is None: 

121 continue 

122 ccaps = caps.setdefault(x, {}) 

123 if c.version: 

124 ccaps["version"] = c.version 

125 if c.python_version: 

126 pcaps = ccaps.setdefault("python-%s" % x, {}) 

127 pcaps[""] = True 

128 if c.python_version is not None: 

129 pcaps["version"] = c.python_version 

130 #legacy format - only used for zlib: 

131 if x=="zlib": 

132 ccaps[""] = True 

133 return caps 

134 

135def get_enabled_compressors(order=ALL_COMPRESSORS): 

136 return tuple(x for x in order if x in COMPRESSION) 

137 

138def get_compressor(name): 

139 c = COMPRESSION.get(name) 

140 assert c is not None, "'%s' compression is not supported" % name 

141 return c.compress 

142 

143 

144def sanity_checks(): 

145 if not use("lzo") and not use("lz4"): 

146 from xpra.log import Logger 

147 logger = Logger("network", "protocol") 

148 if not use("zlib"): 

149 logger.warn("Warning: all the compressors are unavailable or disabled,") 

150 logger.warn(" performance may suffer in some cases") 

151 else: 

152 logger.warn("Warning: zlib is the only compressor enabled") 

153 logger.warn(" install and enable lz4 support for better performance") 

154 

155 

156class Compressed: 

157 def __init__(self, datatype, data, can_inline=False): 

158 assert data is not None, "compressed data cannot be set to None" 

159 self.datatype = datatype 

160 self.data = data 

161 self.can_inline = can_inline 

162 def __len__(self): 

163 return len(self.data) 

164 def __repr__(self): 

165 return "Compressed(%s: %i bytes)" % (self.datatype, len(self.data)) 

166 

167 

168class LevelCompressed(Compressed): 

169 def __init__(self, datatype, data, level, algo, can_inline): 

170 super().__init__(datatype, data, can_inline) 

171 self.level = level 

172 self.algorithm = algo 

173 def __repr__(self): 

174 return "LevelCompressed(%s: %i bytes as %s/%i)" % (self.datatype, len(self.data), self.algorithm, self.level) 

175 

176 

177class LargeStructure: 

178 def __init__(self, datatype, data): 

179 self.datatype = datatype 

180 self.data = data 

181 def __len__(self): 

182 return len(self.data) 

183 def __repr__(self): 

184 return "LargeStructure(%s: %i bytes)" % (self.datatype, len(self.data)) 

185 

186class Compressible(LargeStructure): 

187 #wrapper for data that should be compressed at some point, 

188 #to use this class, you must override compress() 

189 def __repr__(self): 

190 return "Compressible(%s: %i bytes)" % (self.datatype, len(self.data)) 

191 def compress(self): 

192 raise Exception("compress() not defined on %s" % self) 

193 

194 

195def compressed_wrapper(datatype, data, level=5, zlib=False, lz4=False, lzo=False, brotli=False, none=False, can_inline=True): 

196 size = len(data) 

197 if size>MAX_SIZE: 

198 sizemb = size//1024//1024 

199 maxmb = MAX_SIZE//1024//1024 

200 raise Exception("uncompressed data is too large: %iMB, limit is %iMB" % (sizemb, maxmb)) 

201 if lz4 and use("lz4"): 

202 algo = "lz4" 

203 elif lzo and use("lzo"): 

204 algo = "lzo" 

205 elif brotli and use("brotli"): 

206 algo = "brotli" 

207 elif zlib and use("zlib"): 

208 algo = "zlib" 

209 elif none and use("none"): 

210 algo = "none" 

211 else: 

212 raise InvalidCompressionException("no compressors available") 

213 c = COMPRESSION[algo] 

214 cl, cdata = c.compress(data, level) 

215 return LevelCompressed(datatype, cdata, cl, algo, can_inline=can_inline) 

216 

217 

218class InvalidCompressionException(Exception): 

219 pass 

220 

221 

222def get_compression_type(level) -> str: 

223 if level & LZ4_FLAG: 

224 return "lz4" 

225 if level & LZO_FLAG: 

226 return "lzo" 

227 if level & BROTLI_FLAG: 

228 return "brotli" 

229 return "zlib" 

230 

231 

232def decompress(data, level): 

233 #log.info("decompress(%s bytes, %s) type=%s", len(data), get_compression_type(level)) 

234 if level & LZ4_FLAG: 

235 algo = "lz4" 

236 elif level & LZO_FLAG: 

237 algo = "lzo" 

238 elif level & BROTLI_FLAG: 

239 algo = "brotli" 

240 else: 

241 algo = "zlib" 

242 return decompress_by_name(data, algo) 

243 

244def decompress_by_name(data, algo): 

245 c = COMPRESSION.get(algo) 

246 if c is None: 

247 raise InvalidCompressionException("%s is not available" % algo) 

248 return c.decompress(data) 

249 

250 

251def main(): # pragma: no cover 

252 from xpra.util import print_nested_dict 

253 from xpra.platform import program_context 

254 with program_context("Compression", "Compression Info"): 

255 print_nested_dict(get_compression_caps()) 

256 

257 

258if __name__ == "__main__": # pragma: no cover 

259 main()