src/main/java/org/apache/pdfbox/jbig2/JBIG2Document.java

/**
 * Licensed to the Apache Software Foundation (ASF) under one or more
 * contributor license agreements.  See the NOTICE file distributed with
 * this work for additional information regarding copyright ownership.
 * The ASF licenses this file to You under the Apache License, Version 2.0
 * (the "License"); you may not use this file except in compliance with
 * the License.  You may obtain a copy of the License at
 *
 *      http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

package org.apache.pdfbox.jbig2;

import java.io.EOFException;
import java.io.IOException;
import java.util.LinkedList;
import java.util.List;
import java.util.Map;
import java.util.TreeMap;

import javax.imageio.stream.ImageInputStream;

import org.apache.pdfbox.jbig2.io.SubInputStream;
import org.apache.pdfbox.jbig2.util.log.Logger;
import org.apache.pdfbox.jbig2.util.log.LoggerFactory;

/**
 * This class represents the document structure with its pages and global segments.
 */
class JBIG2Document {
  private static final Logger log = LoggerFactory.getLogger(JBIG2Document.class);

  /** ID string in file header, see ISO/IEC 14492:2001, D.4.1 */
  private int[] FILE_HEADER_ID = {
      0x97, 0x4A, 0x42, 0x32, 0x0D, 0x0A, 0x1A, 0x0A
  };

  /**
   * This map contains all pages of this document. The key is the number of the page.
   */
  private final Map<Integer, JBIG2Page> pages = new TreeMap<Integer, JBIG2Page>();


  /** BASIC INFORMATION ABOUT THE CURRENT JBIG2 DOCUMENT */

  /** The length of the file header if exists */
  private short fileHeaderLength = 9;

  /**
   * According to D.4.2 - File header bit 0
   * <p>
   * This flag contains information about the file organization:
   * <ul>
   * <li>{@code 1} for sequential</li>
   * <li>{@code 0} for random-access</li>
   * </ul>
   * You can use the constants {@link #RANDOM} and {@link JBIG2Document#SEQUENTIAL}.
   */
  private short organisationType = SEQUENTIAL;

  public static final int RANDOM = 0;
  public static final int SEQUENTIAL = 1;

  /**
   * According to D.4.2 - Bit 1<br>
   * <ul>
   * <li>{@code true} if amount of pages is unknown, amount of pages field is not present</li>
   * <li>{@code false} if there is a field in the file header where the amount of pages can be read</li>
   * </ul>
   */
  private boolean amountOfPagesUnknown = true;

  /**
   * According to D.4.3 - Number of pages field (4 bytes). Only present if
   * {@link #amountOfPagesUnknown} is {@code false}.
   */
  private int amountOfPages;

  /** Defines whether extended Template is used. */
  private boolean gbUseExtTemplate;

  /**
   * This is the source data stream wrapped into a {@link SubInputStream}.
   */
  private final SubInputStream subInputStream;

  /**
   * Flag:
   * <ul>
   * <li>{@code true} if stream is embedded in another file format and the file header is missing</li>
   * <li>{@code false} if stream is created of a native jbig2 file and the file header should be
   * present</li>
   * </ul>
   */

  /**
   * Holds a load of segments, that aren't associated with a page.
   */
  private JBIG2Globals globalSegments;

  protected JBIG2Document(ImageInputStream input) throws IOException {
    this(input, null);
  }

  protected JBIG2Document(ImageInputStream input, JBIG2Globals globals) throws IOException {
    if (input == null)
      throw new IllegalArgumentException("imageInputStream must not be null");

    this.subInputStream = new SubInputStream(input, 0, Long.MAX_VALUE);
    this.globalSegments = globals;

    mapStream();
  }

  /**
   * Retrieves the segment with the given segment number considering only segments that aren't
   * associated with a page.
   * 
   * @param segmentNr - The number of the wanted segment.
   * @return The requested {@link SegmentHeader}.
   */
  SegmentHeader getGlobalSegment(int segmentNr) {
    if (null != globalSegments) {
      return globalSegments.getSegment(segmentNr);
    }

    if (log.isErrorEnabled()) {
      log.error("Segment not found. Returning null.");
    }

    return null;
  }

  /**
   * Retrieves a {@link JBIG2Page} specified by the given page number.
   * 
   * @param pageNumber - The page number of the wanted {@link JBIG2Page}.
   * 
   * @return The requested {@link JBIG2Page}.
   */
  protected JBIG2Page getPage(int pageNumber) {
    return pages.get(pageNumber);
  }

  /**
   * Retrieves the amount of pages in this JBIG2 document. If the pages are striped, the document
   * will be completely parsed and the amount of pages will be gathered.
   * 
   * @return The amount of pages in this JBIG2 document.
   * @throws IOException
   */
  protected int getAmountOfPages() throws IOException {
    if (amountOfPagesUnknown || amountOfPages == 0) {
      if (pages.size() == 0) {
        mapStream();
      }

      return pages.size();
    } else {
      return amountOfPages;
    }
  }

  /**
   * This method maps the stream and stores all segments.
   */
  private void mapStream() throws IOException {
    final List<SegmentHeader> segments = new LinkedList<SegmentHeader>();

    long offset = 0;
    int segmentType = 0;

    /*
     * Parse the file header if there is one.
     */
    if (isFileHeaderPresent()) {
      parseFileHeader();
      offset += fileHeaderLength;
    }

    if (globalSegments == null) {
      globalSegments = new JBIG2Globals();
    }

    JBIG2Page page = null;

    /*
     * If organisation type is random-access: walk through the segment headers until EOF segment
     * appears (specified with segment number 51)
     */
    while (segmentType != 51 && !reachedEndOfStream(offset)) {
      SegmentHeader segment = new SegmentHeader(this, subInputStream, offset, organisationType);

      final int associatedPage = segment.getPageAssociation();
      segmentType = segment.getSegmentType();

      if (associatedPage != 0) {
        page = getPage(associatedPage);
        if (page == null) {
          page = new JBIG2Page(this, associatedPage);
          pages.put(associatedPage, page);
        }
        page.add(segment);
      } else {
        globalSegments.addSegment(segment.getSegmentNr(), segment);
      }
      segments.add(segment);

      if (JBIG2ImageReader.DEBUG) {
        if (log.isDebugEnabled()) {
          log.debug(segment.toString());
        }
      }

      offset = subInputStream.getStreamPosition();

      // Sequential organization skips data part and sets the offset
      if (organisationType == SEQUENTIAL) {
        offset += segment.getSegmentDataLength();
      }
    }

    /*
     * Random organization: segment headers are finished. Data part starts and the offset can be
     * set.
     */
    determineRandomDataOffsets(segments, offset);
  }

  private boolean isFileHeaderPresent() throws IOException {
    final SubInputStream input = subInputStream;
    input.mark();

    for (int magicByte : FILE_HEADER_ID) {
      if (magicByte != input.read()) {
        input.reset();
        return false;
      }
    }

    input.reset();
    return true;
  }

  /**
   * Determines the start of the data parts and sets the offset.
   * 
   * @param segments
   * @param offset
   */
  private void determineRandomDataOffsets(List<SegmentHeader> segments, long offset) {
    if (organisationType == RANDOM) {
      for (SegmentHeader s : segments) {
        s.setSegmentDataStartOffset(offset);
        offset += s.getSegmentDataLength();
      }
    }
  }

  /**
   * This method reads the stream and sets variables for information about organization type and
   * length etc.
   * 
   * @return
   * @throws IOException
   */
  private void parseFileHeader() throws IOException {
    subInputStream.seek(0);

    /* D.4.1 - ID string, read will be skipped */
    subInputStream.skipBytes(8);

    /*
     * D.4.2 Header flag (1 byte)
     */

    // Bit 3-7 are reserved and must be 0
    subInputStream.readBits(5);

    // Bit 2 - Indicates if extended templates are used
    if (subInputStream.readBit() == 1) {
      gbUseExtTemplate = true;
    }

    // Bit 1 - Indicates if amount of pages are unknown
    if (subInputStream.readBit() != 1) {
      amountOfPagesUnknown = false;
    }

    // Bit 0 - Indicates file organisation type
    organisationType = (short) subInputStream.readBit();

    // fileHeaderLength = 9;

    /*
     * D.4.3 Number of pages (field is only present if amount of pages are 'NOT unknown')
     */
    if (!amountOfPagesUnknown) {
      amountOfPages = (int) subInputStream.readUnsignedInt();
      fileHeaderLength = 13;
    }

  }

  /**
   * This method checks, if the stream is at its end to avoid {@link EOFException}s and reads 32
   * bits.
   * 
   * @param offset
   * @return <li>{@code true} if end of stream reached <li>{@code false} if there are more bytes to
   *         read
   * @throws IOException
   */
  private boolean reachedEndOfStream(long offset) throws IOException {
    try {
      subInputStream.seek(offset);
      subInputStream.readBits(32);
      return false;
    } catch (EOFException e) {
      return true;
    } catch (IndexOutOfBoundsException e) {
      return true;
    }
  }

  protected JBIG2Globals getGlobalSegments() {
    return globalSegments;
  }

  protected boolean isAmountOfPagesUnknown() {
    return amountOfPagesUnknown;
  }

  boolean isGbUseExtTemplate() {
    return gbUseExtTemplate;
  }

}