michael@0: /* This Source Code Form is subject to the terms of the Mozilla Public
michael@0:  * License, v. 2.0. If a copy of the MPL was not distributed with this
michael@0:  * file, You can obtain one at http://mozilla.org/MPL/2.0/. */
michael@0: 
michael@0: package org.mozilla.gecko.background.datareporting;
michael@0: 
michael@0: import java.io.BufferedOutputStream;
michael@0: import java.io.File;
michael@0: import java.io.FileOutputStream;
michael@0: import java.io.IOException;
michael@0: import java.io.OutputStream;
michael@0: import java.io.UnsupportedEncodingException;
michael@0: import java.security.MessageDigest;
michael@0: import java.security.NoSuchAlgorithmException;
michael@0: 
michael@0: import org.json.JSONObject;
michael@0: import org.mozilla.gecko.background.common.log.Logger;
michael@0: 
michael@0: import android.util.Base64;
michael@0: 
michael@0: /**
michael@0:  * Writes telemetry ping to file.
michael@0:  *
michael@0:  * Also creates and updates a SHA-256 checksum for the payload to be included in the ping
michael@0:  * file.
michael@0:  *
michael@0:  * A saved telemetry ping file consists of JSON in the following format,
michael@0:  *   {
michael@0:  *     "slug": "<uuid-string>",
michael@0:  *     "payload": "<escaped-json-data-string>",
michael@0:  *     "checksum": "<base64-sha-256-string>"
michael@0:  *   }
michael@0:  *
michael@0:  * This class writes first to a temporary file and then, after finishing the contents of the ping,
michael@0:  * moves that to the file specified by the caller. This is to avoid uploads of partially written
michael@0:  * ping files.
michael@0:  *
michael@0:  * The API provided by this class:
michael@0:  * startPingFile() - opens stream to a tmp File in the Android cache directory and writes the slug header
michael@0:  * appendPayload(String payloadContent) - appends to the payload of the ping and updates the checksum
michael@0:  * finishPingFile() - writes the checksum to the tmp file and moves it to the File specified by the caller.
michael@0:  *
michael@0:  * In the case of errors, we try to close the stream and File.
michael@0:  */
michael@0: public class TelemetryRecorder {
michael@0:   private final String LOG_TAG = "TelemetryRecorder";
michael@0: 
michael@0:   private final File parentDir;
michael@0:   private final String filename;
michael@0: 
michael@0:   private File tmpFile;
michael@0:   private File destFile;
michael@0:   private File cacheDir;
michael@0: 
michael@0:   private OutputStream  outputStream;
michael@0:   private MessageDigest checksum;
michael@0:   private String        base64Checksum;
michael@0: 
michael@0:   /**
michael@0:    * Charset to use for writing pings; default is us-ascii.
michael@0:    *
michael@0:    * When telemetry calculates the checksum for the ping file, it lossily
michael@0:    * converts utf-16 to ascii. Therefore we have to treat characters in the
michael@0:    * traces as ascii rather than say utf-8. Otherwise we will get a "wrong"
michael@0:    * checksum.
michael@0:    */
michael@0:   private String charset = "us-ascii";
michael@0: 
michael@0:   /**
michael@0:    * Override blockSize in constructor if desired.
michael@0:    * Default block size is that of BufferedOutputStream.
michael@0:    */
michael@0:   private int blockSize = 0;
michael@0: 
michael@0:   /**
michael@0:    * Constructs a TelemetryRecorder for writing a ping file. A temporary file will be written first,
michael@0:    * and then moved to the destination file location specified by the caller.
michael@0:    *
michael@0:    * The directory for writing the temporary file is highly suggested to be the Android internal cache directory,
michael@0:    * fetched by <code>context.getCacheDir()</code>
michael@0:    *
michael@0:    * If the destination file already exists, it will be deleted and overwritten.
michael@0:    *
michael@0:    * Default charset: "us-ascii"
michael@0:    * Default block size: uses constructor default of 8192 bytes (see javadocs for
michael@0:    *                     <code>BufferedOutputStream</code>
michael@0:    * @param parentPath
michael@0:    *        path of parent directory of ping file to be written
michael@0:    * @param cacheDir
michael@0:    *        path of cache directory for writing temporary files.
michael@0:    * @param filename
michael@0:    *        name of ping file to be written
michael@0:    */
michael@0:   public TelemetryRecorder(File parentDir, File cacheDir, String filename) {
michael@0:     if (!parentDir.isDirectory()) {
michael@0:       throw new IllegalArgumentException("Expecting directory, got non-directory File instead.");
michael@0:     }
michael@0:     this.parentDir = parentDir;
michael@0:     this.filename = filename;
michael@0:     this.cacheDir = cacheDir;
michael@0:   }
michael@0: 
michael@0:   public TelemetryRecorder(File parentDir, File cacheDir, String filename, String charset) {
michael@0:     this(parentDir, cacheDir, filename);
michael@0:     this.charset = charset;
michael@0:   }
michael@0: 
michael@0:   public TelemetryRecorder(File parentDir, File cacheDir, String filename, String charset, int blockSize) {
michael@0:     this(parentDir, cacheDir, filename, charset);
michael@0:     this.blockSize = blockSize;
michael@0:   }
michael@0: 
michael@0:   /**
michael@0:    * Start the temporary ping file and write the "slug" header and payload key, of the
michael@0:    * format:
michael@0:    *
michael@0:    *     { "slug": "< filename >", "payload":
michael@0:    *
michael@0:    * @throws Exception
michael@0:    *           Checked exceptions <code>NoSuchAlgorithmException</code>,
michael@0:    *           <code>UnsupportedEncodingException</code>, or
michael@0:    *           <code>IOException</code> and unchecked exception that
michael@0:    *           are rethrown to caller
michael@0:    */
michael@0:   public void startPingFile() throws Exception {
michael@0: 
michael@0:     // Open stream to temporary file for writing.
michael@0:     try {
michael@0:       tmpFile = File.createTempFile(filename, "tmp", cacheDir);
michael@0:     } catch (IOException e) {
michael@0:       // Try to create the temporary file in the ping directory.
michael@0:       tmpFile = new File(parentDir, filename + ".tmp");
michael@0:       try {
michael@0:         tmpFile.createNewFile();
michael@0:       } catch (IOException e1) {
michael@0:         cleanUpAndRethrow("Failed to create tmp file in temp directory or ping directory.", e1);
michael@0:       }
michael@0:     }
michael@0: 
michael@0:     try {
michael@0:       if (blockSize > 0) {
michael@0:         outputStream = new BufferedOutputStream(new FileOutputStream(tmpFile), blockSize);
michael@0:       } else {
michael@0:         outputStream = new BufferedOutputStream(new FileOutputStream(tmpFile));
michael@0:       }
michael@0: 
michael@0:       // Create checksum for ping.
michael@0:       checksum = MessageDigest.getInstance("SHA-256");
michael@0: 
michael@0:       // Write ping header.
michael@0:       byte[] header = makePingHeader(filename);
michael@0:       outputStream.write(header);
michael@0:       Logger.debug(LOG_TAG, "Wrote " + header.length + " header bytes.");
michael@0: 
michael@0:     } catch (NoSuchAlgorithmException e) {
michael@0:       cleanUpAndRethrow("Error creating checksum digest", e);
michael@0:     } catch (UnsupportedEncodingException e) {
michael@0:       cleanUpAndRethrow("Error writing header", e);
michael@0:     } catch (IOException e) {
michael@0:       cleanUpAndRethrow("Error writing to stream", e);
michael@0:     }
michael@0:   }
michael@0: 
michael@0:   private byte[] makePingHeader(String slug)
michael@0:       throws UnsupportedEncodingException {
michael@0:     return ("{\"slug\":" + JSONObject.quote(slug) + "," + "\"payload\":\"")
michael@0:         .getBytes(charset);
michael@0:   }
michael@0: 
michael@0:   /**
michael@0:    * Append payloadContent to ping file and update the checksum.
michael@0:    *
michael@0:    * @param payloadContent
michael@0:    *          String content to be written
michael@0:    * @return number of bytes written, or -1 if writing failed
michael@0:    * @throws Exception
michael@0:    *           Checked exceptions <code>UnsupportedEncodingException</code> or
michael@0:    *           <code>IOException</code> and unchecked exception that
michael@0:    *           are rethrown to caller
michael@0:    */
michael@0:   public int appendPayload(String payloadContent) throws Exception {
michael@0:     if (payloadContent == null) {
michael@0:       cleanUpAndRethrow("Payload is null", new Exception());
michael@0:       return -1;
michael@0:     }
michael@0: 
michael@0:     try {
michael@0:       byte[] payloadBytes = payloadContent.getBytes(charset);
michael@0:       // If we run into an error, we'll throw and abort, so checksum won't be stale.
michael@0:       checksum.update(payloadBytes);
michael@0: 
michael@0:       byte[] quotedPayloadBytes = JSONObject.quote(payloadContent).getBytes(charset);
michael@0: 
michael@0:       // First and last bytes are quotes inserted by JSONObject.quote; discard
michael@0:       // them.
michael@0:       int numBytes = quotedPayloadBytes.length - 2;
michael@0:       outputStream.write(quotedPayloadBytes, 1, numBytes);
michael@0:       return numBytes;
michael@0: 
michael@0:     } catch (UnsupportedEncodingException e) {
michael@0:       cleanUpAndRethrow("Error encoding payload", e);
michael@0:       return -1;
michael@0:     } catch (IOException e) {
michael@0:       cleanUpAndRethrow("Error writing to stream", e);
michael@0:       return -1;
michael@0:     }
michael@0:   }
michael@0: 
michael@0:   /**
michael@0:    * Add the checksum of the payload to the ping file and close the stream.
michael@0:    *
michael@0:    * @throws Exception
michael@0:    *          Checked exceptions <code>UnsupportedEncodingException</code> or
michael@0:    *          <code>IOException</code> and unchecked exception that
michael@0:    *          are rethrown to caller
michael@0:    */
michael@0:   public void finishPingFile() throws Exception {
michael@0:     try {
michael@0:       byte[] footer = makePingFooter(checksum);
michael@0:       outputStream.write(footer);
michael@0:       // We're done writing, so force the stream to flush the buffer.
michael@0:       outputStream.flush();
michael@0:       Logger.debug(LOG_TAG, "Wrote " + footer.length + " footer bytes.");
michael@0:     } catch (UnsupportedEncodingException e) {
michael@0:       cleanUpAndRethrow("Checksum encoding exception", e);
michael@0:     } catch (IOException e) {
michael@0:       cleanUpAndRethrow("Error writing footer to stream", e);
michael@0:     } finally {
michael@0:       try {
michael@0:         outputStream.close();
michael@0:       } catch (IOException e) {
michael@0:         // Failed to close, nothing we can do except discard the reference to the stream.
michael@0:         outputStream = null;
michael@0:       }
michael@0:     }
michael@0: 
michael@0:     // Move temp file to destination specified by caller.
michael@0:     try {
michael@0:       File destFile = new File(parentDir, filename);
michael@0:       // Delete file if it exists - docs state that rename may fail if the File already exists.
michael@0:       if (destFile.exists()) {
michael@0:         destFile.delete();
michael@0:       }
michael@0:       boolean result = tmpFile.renameTo(destFile);
michael@0:       if (!result) {
michael@0:         throw new IOException("Could not move tmp file to destination.");
michael@0:       }
michael@0:     } finally {
michael@0:       cleanUp();
michael@0:     }
michael@0:   }
michael@0: 
michael@0:   private byte[] makePingFooter(MessageDigest checksum)
michael@0:       throws UnsupportedEncodingException {
michael@0:     base64Checksum = Base64.encodeToString(checksum.digest(), Base64.NO_WRAP);
michael@0:     return ("\",\"checksum\":" + JSONObject.quote(base64Checksum) + "}")
michael@0:         .getBytes(charset);
michael@0:   }
michael@0: 
michael@0:   /**
michael@0:    * Get final digested Base64 checksum.
michael@0:    *
michael@0:    * @return String checksum if it has been calculated, null otherwise.
michael@0:    */
michael@0:   protected String getFinalChecksum() {
michael@0:     return base64Checksum;
michael@0:   }
michael@0: 
michael@0:   public String getCharset() {
michael@0:     return this.charset;
michael@0:   }
michael@0: 
michael@0:   /**
michael@0:    * Clean up checksum and delete the temporary file.
michael@0:    */
michael@0:   private void cleanUp() {
michael@0:     // Discard checksum.
michael@0:     checksum.reset();
michael@0: 
michael@0:     // Clean up files.
michael@0:     if (tmpFile != null && tmpFile.exists()) {
michael@0:       tmpFile.delete();
michael@0:     }
michael@0:     tmpFile = null;
michael@0:   }
michael@0: 
michael@0:   /**
michael@0:    * Log message and error and clean up, then rethrow exception to caller.
michael@0:    *
michael@0:    * @param message
michael@0:    *          Error message
michael@0:    * @param e
michael@0:    *          Exception
michael@0:    *
michael@0:    * @throws Exception
michael@0:    *           Exception to be rethrown to caller
michael@0:    */
michael@0:   private void cleanUpAndRethrow(String message, Exception e) throws Exception {
michael@0:     Logger.error(LOG_TAG, message, e);
michael@0:     cleanUp();
michael@0: 
michael@0:     if (outputStream != null) {
michael@0:       try {
michael@0:         outputStream.close();
michael@0:       } catch (IOException exception) {
michael@0:         // Failed to close stream; nothing we can do, and we're aborting anyways.
michael@0:       }
michael@0:     }
michael@0: 
michael@0:     if (destFile != null && destFile.exists()) {
michael@0:       destFile.delete();
michael@0:     }
michael@0:     // Rethrow the exception.
michael@0:     throw e;
michael@0:   }
michael@0: }