1 // Copyright (c) 2016, the R8 project authors. Please see the AUTHORS file 2 // for details. All rights reserved. Use of this source code is governed by a 3 // BSD-style license that can be found in the LICENSE file. 4 package com.android.tools.r8; 5 6 import com.android.tools.r8.dex.DexFileReader; 7 import com.android.tools.r8.dex.Segment; 8 import java.awt.image.BufferedImage; 9 import java.io.ByteArrayInputStream; 10 import java.io.IOException; 11 import java.io.InputStream; 12 import java.io.OutputStream; 13 import java.nio.ByteBuffer; 14 import java.nio.file.Files; 15 import java.nio.file.Path; 16 import java.nio.file.Paths; 17 import javax.imageio.ImageIO; 18 import org.apache.commons.compress.compressors.CompressorException; 19 import org.apache.commons.compress.compressors.bzip2.BZip2CompressorInputStream; 20 21 public class BSPatch { 22 23 private final static char[] BSDIFF_MAGIC = "BSDIFF40".toCharArray(); 24 private final static int BSDIFF_HEADER_LENGTH = 32; 25 26 private final ByteBuffer patchInput; 27 private final ByteBuffer oldInput; 28 private final Path output; 29 private final Path dexPath; 30 31 private InputStream controlStream; 32 private InputStream diffStream; 33 private InputStream extraStream; 34 35 private long controlBytesRead; 36 private long diffBytesRead; 37 private long extraBytesRead; 38 private int controlBlockLen; 39 private int diffBlockLen; 40 private int extraBlockLen; 41 main(String[] args)42 public static void main(String[] args) { 43 boolean imageMode = args.length > 1 && args[0].equals("-i"); 44 int argOffset = imageMode ? 1 : 0; 45 if (args.length < argOffset + 3) { 46 System.out.println("Usage: [-i] <patch file> <original dex> <output file> [target dex]"); 47 System.exit(1); 48 } 49 try { 50 new BSPatch(Paths.get(args[argOffset]), Paths.get(args[argOffset + 1]), 51 Paths.get(args[argOffset + 2]), 52 args.length != argOffset + 4 ? null : Paths.get(args[argOffset + 3])) 53 .apply(imageMode); 54 } catch (IOException e) { 55 System.err.println("File I/O error: " + e.toString()); 56 } catch (CompressorException e) { 57 System.err.println("BZIP error: " + e.toString()); 58 } 59 } 60 BSPatch(Path patchInput, Path oldInput, Path output, Path dexPath)61 private BSPatch(Path patchInput, Path oldInput, Path output, Path dexPath) throws IOException { 62 this.patchInput = ByteBuffer.wrap(Files.readAllBytes(patchInput)); 63 this.oldInput = ByteBuffer.wrap(Files.readAllBytes(oldInput)); 64 this.output = output; 65 this.dexPath = dexPath; 66 } 67 apply(boolean imageMode)68 public void apply(boolean imageMode) throws CompressorException, IOException { 69 PatchExecutor executor = imageMode ? new ImageExecutor(dexPath) : new FileExecutor(); 70 checkHeader(); 71 setupSegmentsAndOutput(executor); 72 processControl(executor); 73 executor.writeResult(); 74 printStats(); 75 } 76 percentOf(long a, long b)77 private int percentOf(long a, long b) { 78 return (int) ((((double) a) / ((double) b)) * 100); 79 } 80 printStats()81 private void printStats() { 82 System.out.println("Size of control block (compressed bytes): " + controlBlockLen); 83 System.out.println("Size of control block (read bytes): " + controlBytesRead); 84 System.out 85 .println("Compression of control block: " + percentOf(controlBlockLen, controlBytesRead)); 86 System.out.println("Size of diff data block (compressed bytes): " + diffBlockLen); 87 System.out.println("Size of diff data block (read bytes): " + diffBytesRead); 88 System.out 89 .println("Compression of diff data block: " + percentOf(diffBlockLen, diffBytesRead)); 90 System.out.println("Size of extra data block (compressed bytes): " + extraBlockLen); 91 System.out.println("Size of extra data block (read bytes): " + extraBytesRead); 92 System.out 93 .println("Compression of extra data block: " + percentOf(extraBlockLen, extraBytesRead)); 94 } 95 processControl(PatchExecutor executor)96 private void processControl(PatchExecutor executor) throws IOException { 97 int blockSize; 98 while ((blockSize = readNextControlEntry()) != Integer.MIN_VALUE) { 99 int extraSize = readNextControlEntry(); 100 int advanceOld = readNextControlEntry(); 101 executor.copyDiff(blockSize); 102 executor.copyOld(blockSize); 103 executor.submitBlock(blockSize); 104 executor.copyExtra(extraSize); 105 executor.skipOld(advanceOld); 106 } 107 } 108 checkHeader()109 private void checkHeader() { 110 for (int i = 0; i < BSDIFF_MAGIC.length; i++) { 111 if (patchInput.get() != BSDIFF_MAGIC[i]) { 112 throw new RuntimeException("Illegal patch, wrong magic!"); 113 } 114 } 115 } 116 setupSegmentsAndOutput(PatchExecutor executor)117 private void setupSegmentsAndOutput(PatchExecutor executor) 118 throws CompressorException, IOException { 119 controlBlockLen = readOffset(); 120 diffBlockLen = readOffset(); 121 int newFileSize = readOffset(); 122 123 extraBlockLen = 124 patchInput.array().length - (BSDIFF_HEADER_LENGTH + controlBlockLen + diffBlockLen); 125 126 executor.createOutput(newFileSize); 127 controlStream = new BZip2CompressorInputStream( 128 new ByteArrayInputStream(patchInput.array(), BSDIFF_HEADER_LENGTH, controlBlockLen)); 129 diffStream = new BZip2CompressorInputStream( 130 new ByteArrayInputStream(patchInput.array(), BSDIFF_HEADER_LENGTH + controlBlockLen, 131 diffBlockLen)); 132 extraStream = new BZip2CompressorInputStream( 133 new ByteArrayInputStream(patchInput.array(), 134 BSDIFF_HEADER_LENGTH + controlBlockLen + diffBlockLen, 135 extraBlockLen)); 136 } 137 readOffset()138 private int readOffset() { 139 byte[] buffer = new byte[8]; 140 patchInput.get(buffer); 141 return decodeOffset(buffer); 142 } 143 readNextControlEntry()144 private int readNextControlEntry() throws IOException { 145 byte[] buffer = new byte[8]; 146 int read = controlStream.read(buffer); 147 if (read == -1) { 148 return Integer.MIN_VALUE; 149 } 150 controlBytesRead += read; 151 assert read == buffer.length; 152 return decodeOffset(buffer); 153 } 154 decodeOffset(byte[] buffer)155 private static int decodeOffset(byte[] buffer) { 156 long offset = buffer[7] & 0x7F; 157 for (int i = 6; i >= 0; i--) { 158 offset = (offset << 8) | (((int) buffer[i]) & 0xff); 159 } 160 if ((buffer[7] & 0x80) != 0) { 161 offset = -offset; 162 } 163 assert offset < Integer.MAX_VALUE && offset > Integer.MIN_VALUE; 164 return (int) offset; 165 } 166 167 private abstract class PatchExecutor { 168 createOutput(int newFileSize)169 public abstract void createOutput(int newFileSize); 170 copyDiff(int blockSize)171 public abstract void copyDiff(int blockSize) throws IOException; 172 copyOld(int blockSize)173 public abstract void copyOld(int blockSize) throws IOException; 174 submitBlock(int blockSize)175 public abstract void submitBlock(int blockSize); 176 copyExtra(int extraSize)177 public abstract void copyExtra(int extraSize) throws IOException; 178 skipOld(int advanceOld)179 public abstract void skipOld(int advanceOld) throws IOException; 180 writeResult()181 public abstract void writeResult() throws IOException; 182 } 183 184 private class FileExecutor extends PatchExecutor { 185 186 private ByteBuffer resultBuffer; 187 private byte[] mergeBuffer = null; 188 189 @Override createOutput(int newFileSize)190 public void createOutput(int newFileSize) { 191 resultBuffer = ByteBuffer.allocate(newFileSize); 192 } 193 194 @Override copyDiff(int blockSize)195 public void copyDiff(int blockSize) throws IOException { 196 assert mergeBuffer == null; 197 mergeBuffer = new byte[blockSize]; 198 int read = diffStream.read(mergeBuffer); 199 diffBytesRead += read; 200 assert read == blockSize; 201 } 202 203 @Override copyOld(int blockSize)204 public void copyOld(int blockSize) throws IOException { 205 assert mergeBuffer != null; 206 assert mergeBuffer.length == blockSize; 207 byte[] data = new byte[blockSize]; 208 oldInput.get(data); 209 for (int i = 0; i < mergeBuffer.length; i++) { 210 mergeBuffer[i] = (byte) ((((int) mergeBuffer[i]) & 0xff) + (((int) data[i]) & 0xff)); 211 } 212 } 213 214 @Override submitBlock(int blockSize)215 public void submitBlock(int blockSize) { 216 assert mergeBuffer != null; 217 assert mergeBuffer.length == blockSize; 218 resultBuffer.put(mergeBuffer); 219 mergeBuffer = null; 220 } 221 222 @Override copyExtra(int extraSize)223 public void copyExtra(int extraSize) throws IOException { 224 byte[] data = new byte[extraSize]; 225 int read = extraStream.read(data); 226 assert read == extraSize; 227 extraBytesRead += read; 228 resultBuffer.put(data); 229 } 230 231 @Override skipOld(int delta)232 public void skipOld(int delta) throws IOException { 233 oldInput.position(oldInput.position() + delta); 234 } 235 236 @Override writeResult()237 public void writeResult() throws IOException { 238 OutputStream outputStream = Files.newOutputStream(output); 239 outputStream.write(resultBuffer.array()); 240 outputStream.close(); 241 } 242 } 243 244 private class ImageExecutor extends PatchExecutor { 245 246 private final Path dexPath; 247 248 BufferedImage image; 249 int position = 0; 250 int width; 251 int height; 252 ImageExecutor(Path dexPath)253 private ImageExecutor(Path dexPath) { 254 this.dexPath = dexPath; 255 } 256 257 @Override createOutput(int newFileSize)258 public void createOutput(int newFileSize) { 259 int root = (int) Math.sqrt(newFileSize); 260 width = newFileSize / root; 261 height = newFileSize / width + 1; 262 image = new BufferedImage(width, height, BufferedImage.TYPE_3BYTE_BGR); 263 } 264 265 @Override copyDiff(int blockSize)266 public void copyDiff(int blockSize) throws IOException { 267 byte[] buffer = new byte[blockSize]; 268 int read = diffStream.read(buffer); 269 assert read == blockSize; 270 diffBytesRead += read; 271 for (int i = 0; i < buffer.length; i++) { 272 if (buffer[i] != 0) { 273 int y = (position + i) / width; 274 int x = (position + i) % width; 275 int rgb = image.getRGB(x, y); 276 rgb = rgb | 0xFF0000; 277 image.setRGB(x, y, rgb); 278 } 279 } 280 } 281 282 @Override copyOld(int blockSize)283 public void copyOld(int blockSize) throws IOException { 284 for (int i = 0; i < blockSize; i++) { 285 int y = (position + i) / width; 286 int x = (position + i) % width; 287 int rgb = image.getRGB(x, y); 288 if ((rgb & 0xFF0000) == 0) { 289 rgb = rgb | 0xFF00; 290 } 291 image.setRGB(x, y, rgb); 292 } 293 } 294 295 @Override submitBlock(int blockSize)296 public void submitBlock(int blockSize) { 297 position += blockSize; 298 } 299 300 @Override copyExtra(int extraSize)301 public void copyExtra(int extraSize) throws IOException { 302 long skipped = extraStream.skip(extraSize); 303 assert skipped == extraSize; 304 extraBytesRead += skipped; 305 for (int i = 0; i < extraSize; i++) { 306 int y = (position + i) / width; 307 int x = (position + i) % width; 308 int rgb = image.getRGB(x, y); 309 rgb = rgb | 0xFF; 310 image.setRGB(x, y, rgb); 311 } 312 position += extraSize; 313 } 314 315 @Override skipOld(int advanceOld)316 public void skipOld(int advanceOld) throws IOException { 317 } 318 319 @Override writeResult()320 public void writeResult() throws IOException { 321 if (dexPath != null) { 322 Segment[] segments = DexFileReader.parseMapFrom(dexPath); 323 for (Segment segment : segments) { 324 int y = segment.offset / width; 325 for (int x = 0; x < width; x++) { 326 int val = (x / 10) % 2 == 0 ? 0 : 0xffffff; 327 image.setRGB(x, y, val); 328 } 329 System.out.println(segment); 330 } 331 } 332 ImageIO.write(image, "png", output.toFile()); 333 } 334 } 335 } 336