// See https://dev.to/fastly/filter-pngs-for-acropalypse-using-computeedge-1c58

addEventListener("fetch", event => event.respondWith(handleRequest(event)));

async function handleRequest(event) {
  const response = await fetch(event.request, {backend: 'origin_0'});
  return acropalypseFilter(response);
}

function acropalypseFilter(response) {
  // Define the byte sequences for the PNG magic bytes and the IEND marker
  // that identifies the end of the image
  const pngMarker = new Uint8Array([0x89,0x50,0x4e,0x47]);
  const pngIEND = new Uint8Array([0x00,0x00,0x00,0x00,0x49,0x45,0x4e,0x44]);

  // Define an async function so we can use await when processing the stream
  async function processChunks(reader, writer) {
    let isPNG = false;
    let isIEND = false;
    while (true) {

      // Fetch a chunk from the input stream
      const { done, value: chunk } = await reader.read();
      if (done) break;

      // If we have not yet found a PNG marker, see if there's one
      // in this chunk.  If there is, we have a PNG that is potentially
      // vulnerable
      if (!isPNG && seqFind(chunk, pngMarker) !== -1) {
        console.log("It's a PNG");
        isPNG = true;
      }

      // If we know we're past the end of the PNG, any remaining data
      // in the file is the hidden data we want to remove.  Since we already
      // sent the Content-Length header, we'll pad the rest of the response
      // with zeroes.
      if (isIEND) { 
        writer.write(new Uint8Array(chunk.length));
        continue;
      
      // If it's a PNG but we're yet to get to the end of it...
      } else if (isPNG) {

        // See if this chunk contains the IEND marker
        // If so, output the data up to the marker and replace the rest of the
        // chunk with zeroes.
        const idx = seqFind(chunk, pngIEND);
        if (idx > 0) {
          console.log(`Found IEND at ${idx}`);
          isIEND = true;
          writer.write(chunk.slice(0, idx));
          writer.write(new Uint8Array(chunk.length-idx));
          continue;
        }
      }

      // Either we're not dealing with a PNG, or we're in a PNG but have not
      // reached the IEND marker yet. Either way, we can simply copy the
      // chunk directly to the output.
      writer.write(chunk);
    }

    // After the input stream ends, we should do cleanup.
    writer.close();
    reader.releaseLock();
  }

  if (response.body) {
    const {readable, writable} = new TransformStream();
    const writer = writable.getWriter();
    const reader = response.body.getReader();
    processChunks(reader, writer);
    return new Response(readable, response);
  }
  return response;
}

// Find a sequence of elements in an array
function seqFind(input, target, fromIndex = 0) {
  const index = input.indexOf(target[0], fromIndex);
  if (target.length === 1 || index === -1) return index;
  let i, j;
  for (i = index, j = 0; j < target.length && i < input.length; i++, j++) {
    if (input[i] !== target[j]) {
      return seqFind(input, target, index + 1);
    }
  }
  return (i === index + target.length) ? index : -1;
}