pdfbox-trunk/pdfbox/src/main/java/org/apache/pdfbox/pdfparser/XrefTrailerResolver.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.pdfparser;

import java.util.ArrayList;
import java.util.Collections;
import java.util.HashMap;
import java.util.HashSet;
import java.util.List;
import java.util.Map;
import java.util.Set;
import java.util.SortedSet;
import java.util.TreeSet;

import org.apache.commons.logging.Log;
import org.apache.commons.logging.LogFactory;
import org.apache.pdfbox.cos.COSDictionary;
import org.apache.pdfbox.cos.COSName;
import org.apache.pdfbox.cos.COSObjectKey;

/**
 * This class will collect all XRef/trailer objects and creates correct
 * xref/trailer information after all objects are read using startxref
 * and 'Prev' information (unused XRef/trailer objects are discarded).
 *
 * In case of missing startxref or wrong startxref pointer all
 * XRef/trailer objects are used to create xref table / trailer dictionary
 * in order they occur.
 *
 * For each new xref object/XRef stream method {@link #nextXrefObj(long, XRefType)}
 * must be called with start byte position. All following calls to
 * {@link #setXRef(COSObjectKey, long)} or {@link #setTrailer(COSDictionary)}
 * will add the data for this byte position.
 *
 * After all objects are parsed the startxref position must be provided
 * using {@link #setStartxref(long)}. This is used to build the chain of
 * active xref/trailer objects used for creating document trailer and xref table.
 *
 * @author Timo Böhme
 */
public class XrefTrailerResolver
{

    /**
     * A class which represents a xref/trailer object.
     */
    private static class XrefTrailerObj
    {
        protected COSDictionary trailer = null;

        private XRefType xrefType;

        private final Map<COSObjectKey, Long> xrefTable = new HashMap<>();
        
        /**
         *  Default constructor.
         */
        private XrefTrailerObj()
        {
            xrefType = XRefType.TABLE;
        }

        public void reset()
        {
            xrefTable.clear();
        }
    }

    /** 
     * The XRefType of a trailer.
     */
    public enum XRefType
    {
        /**
         * XRef table type.
         */
        TABLE, 
        /**
         * XRef stream type.
         */
        STREAM
    }
    
    private final Map<Long, XrefTrailerObj> bytePosToXrefMap = new HashMap<>();
    private XrefTrailerObj curXrefTrailerObj   = null;
    private XrefTrailerObj resolvedXrefTrailer = null;

    /** Log instance. */
    private static final Log LOG = LogFactory.getLog( XrefTrailerResolver.class );

    /**
     * Returns the first trailer if at least one exists.
     * 
     * @return the first trailer or null
     */
    public final COSDictionary getFirstTrailer() 
    {
        if (bytePosToXrefMap.isEmpty())
        {
            return null;
        }
        Set<Long> offsets = bytePosToXrefMap.keySet();
        SortedSet<Long> sortedOffset = new TreeSet<>(offsets);
        return bytePosToXrefMap.get(sortedOffset.first()).trailer;
    }
    
    /**
     * Returns the last trailer if at least one exists.
     * 
     * @return the last trailer ir null
     */
    public final COSDictionary getLastTrailer() 
    {
        if (bytePosToXrefMap.isEmpty()) 
        {
            return null;
        }
        Set<Long> offsets = bytePosToXrefMap.keySet();
        SortedSet<Long> sortedOffset = new TreeSet<>(offsets);
        return bytePosToXrefMap.get(sortedOffset.last()).trailer;
    }

    /**
     * Returns the count of trailers.
     *
     * @return the count of trailers.
     */
    public final int getTrailerCount()
    {
        return bytePosToXrefMap.size();
    }

    /**
     * Signals that a new XRef object (table or stream) starts.
     * @param startBytePos the offset to start at
     * @param type the type of the Xref object
     */
    public void nextXrefObj( final long startBytePos, XRefType type )
    {
        curXrefTrailerObj = new XrefTrailerObj();
        bytePosToXrefMap.put(startBytePos, curXrefTrailerObj);
        curXrefTrailerObj.xrefType = type;
    }

    /**
     * Returns the XRefTxpe of the resolved trailer.
     * 
     * @return the XRefType or null.
     */
    public XRefType getXrefType()
    { 
        return ( resolvedXrefTrailer == null ) ? null : resolvedXrefTrailer.xrefType; 
    } 
    
    /**
     * Populate XRef HashMap of current XRef object.
     * Will add an Xreftable entry that maps ObjectKeys to byte offsets in the file.
     * @param objKey The objkey, with id and gen numbers
     * @param offset The byte offset in this file
     */
    public void setXRef( COSObjectKey objKey, long offset )
    {
        if ( curXrefTrailerObj == null )
        {
            // should not happen...
            LOG.warn( "Cannot add XRef entry for '" + objKey.getNumber() + "' because XRef start was not signalled." );
            return;
        }
        // PDFBOX-3506 check before adding to the map, to avoid entries from the table being 
        // overwritten by obsolete entries in hybrid files (/XRefStm entry)
        if (!curXrefTrailerObj.xrefTable.containsKey(objKey) )
        {
            curXrefTrailerObj.xrefTable.put(objKey, offset);
        }
    }

    /**
     * Adds trailer information for current XRef object.
     *
     * @param trailer the current document trailer dictionary
     */
    public void setTrailer( COSDictionary trailer )
    {
        if ( curXrefTrailerObj == null )
        {
            // should not happen...
            LOG.warn( "Cannot add trailer because XRef start was not signalled." );
            return;
        }
        curXrefTrailerObj.trailer = trailer;
    }

    /**
     * Returns the trailer last set by {@link #setTrailer(COSDictionary)}.
     * 
     * @return the current trailer.
     * 
     */
    public COSDictionary getCurrentTrailer() 
    {
        return curXrefTrailerObj.trailer;
    }

    /**
     * Sets the byte position of the first XRef
     * (has to be called after very last startxref was read).
     * This is used to resolve chain of active XRef/trailer.
     *
     * In case startxref position is not found we output a
     * warning and use all XRef/trailer objects combined
     * in byte position order.
     * Thus for incomplete PDF documents with missing
     * startxref one could call this method with parameter value -1.
     * 
     * @param startxrefBytePosValue starting position of the first XRef
     * 
     */
    public void setStartxref( long startxrefBytePosValue )
    {
        if ( resolvedXrefTrailer != null )
        {
            LOG.warn( "Method must be called only ones with last startxref value." );
            return;
        }

        resolvedXrefTrailer = new XrefTrailerObj();
        resolvedXrefTrailer.trailer = new COSDictionary();

        XrefTrailerObj curObj = bytePosToXrefMap.get( startxrefBytePosValue );
        List<Long>  xrefSeqBytePos = new ArrayList<>();

        if ( curObj == null )
        {
            // no XRef at given position
            LOG.warn( "Did not found XRef object at specified startxref position " + startxrefBytePosValue );

            // use all objects in byte position order (last entries overwrite previous ones)
            xrefSeqBytePos.addAll( bytePosToXrefMap.keySet() );
            Collections.sort( xrefSeqBytePos );
        }
        else
        {
            // copy xref type
            resolvedXrefTrailer.xrefType = curObj.xrefType;
            // found starting Xref object
            // add this and follow chain defined by 'Prev' keys
            xrefSeqBytePos.add( startxrefBytePosValue );
            while ( curObj.trailer != null )
            {
                long prevBytePos = curObj.trailer.getLong( COSName.PREV, -1L );
                if ( prevBytePos == -1 )
                {
                    break;
                }

                curObj = bytePosToXrefMap.get( prevBytePos );
                if ( curObj == null )
                {
                    LOG.warn( "Did not found XRef object pointed to by 'Prev' key at position " + prevBytePos );
                    break;
                }
                xrefSeqBytePos.add( prevBytePos );

                // prevent infinite loops
                if ( xrefSeqBytePos.size() >= bytePosToXrefMap.size() )
                {
                    break;
                }
            }
            // have to reverse order so that later XRefs will overwrite previous ones
            Collections.reverse( xrefSeqBytePos );
        }

        // merge used and sorted XRef/trailer
        for ( Long bPos : xrefSeqBytePos )
        {
            curObj = bytePosToXrefMap.get( bPos );
            if ( curObj.trailer != null )
            {
                resolvedXrefTrailer.trailer.addAll( curObj.trailer );
            }
            resolvedXrefTrailer.xrefTable.putAll( curObj.xrefTable );
        }

    }

    /**
     * Gets the resolved trailer. Might return <code>null</code> in case
     * {@link #setStartxref(long)} was not called before.
     *
     * @return the trailer if available
     */
    public COSDictionary getTrailer()
    {
        return ( resolvedXrefTrailer == null ) ? null : resolvedXrefTrailer.trailer;
    }

    /**
     * Gets the resolved xref table. Might return <code>null</code> in case
     *  {@link #setStartxref(long)} was not called before.
     *
     * @return the xrefTable if available
     */
    public Map<COSObjectKey, Long> getXrefTable()
    {
        return ( resolvedXrefTrailer == null ) ? null : resolvedXrefTrailer.xrefTable;
    }
    
    /** Returns object numbers which are referenced as contained
     *  in object stream with specified object number.
     *  
     *  This will scan resolved xref table for all entries having negated
     *  stream object number as value.
     *
     *  @param objstmObjNr  object number of object stream for which contained object numbers
     *                      should be returned
     *                       
     *  @return set of object numbers referenced for given object stream
     *          or <code>null</code> if {@link #setStartxref(long)} was not
     *          called before so that no resolved xref table exists
     */
    public Set<Long> getContainedObjectNumbers( final int objstmObjNr ) 
    {
        if ( resolvedXrefTrailer == null )
        {
            return null;
        }
        final Set<Long> refObjNrs = new HashSet<>();
        final long cmpVal = - objstmObjNr;
        
        resolvedXrefTrailer.xrefTable.forEach((key, value) ->
        {
            if (value == cmpVal)
            {
                refObjNrs.add(key.getNumber());
            }
        });
        return refObjNrs;
    }

    /**
     * Reset all data so that it can be used to rebuild the trailer.
     * 
     */
    protected void reset()
    {
        for (XrefTrailerObj trailerObj : bytePosToXrefMap.values())
        {
            trailerObj.reset();
        }
        curXrefTrailerObj = null;
        resolvedXrefTrailer = null;
    }
}
