/*

 * Copyright (c) 2015, Oracle and/or its affiliates. All rights reserved.

 * DO NOT ALTER OR REMOVE COPYRIGHT NOTICES OR THIS FILE HEADER.

 *

 * This code is free software; you can redistribute it and/or modify it

 * under the terms of the GNU General Public License version 2 only, as

 * published by the Free Software Foundation.  Oracle designates this

 * particular file as subject to the "Classpath" exception as provided

 * by Oracle in the LICENSE file that accompanied this code.

 *

 * This code is distributed in the hope that it will be useful, but WITHOUT

 * ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or

 * FITNESS FOR A PARTICULAR PURPOSE.  See the GNU General Public License

 * version 2 for more details (a copy is included in the LICENSE file that

 * accompanied this code).

 *

 * You should have received a copy of the GNU General Public License version

 * 2 along with this work; if not, write to the Free Software Foundation,

 * Inc., 51 Franklin St, Fifth Floor, Boston, MA 02110-1301 USA.

 *

 * Please contact Oracle, 500 Oracle Parkway, Redwood Shores, CA 94065 USA

 * or visit www.oracle.com if you need additional information or have any

 * questions.

 */

package sun.security.ssl;

import java.io.*;

import java.nio.*;

import java.util.*;

import javax.crypto.BadPaddingException;

import javax.net.ssl.*;

import static sun.security.ssl.HandshakeMessage.*;

/**

 * DTLS {@code InputRecord} implementation for {@code SSLEngine}.

 */

final class DTLSInputRecord extends InputRecord implements DTLSRecord {

    private DTLSReassembler reassembler = null;

    // Cache the session identifier for the detection of session-resuming

    // handshake.

    byte[]              prevSessionID = new byte[0];

    int                 readEpoch;

    int                 prevReadEpoch;

    Authenticator       prevReadAuthenticator;

    CipherBox           prevReadCipher;

    DTLSInputRecord() {

        this.readEpoch = 0;

        this.readAuthenticator = new MAC(true);

        this.prevReadEpoch = 0;

        this.prevReadCipher = CipherBox.NULL;

        this.prevReadAuthenticator = new MAC(true);

    }

    @Override

    void changeReadCiphers(Authenticator readAuthenticator,

            CipherBox readCipher) {

        prevReadCipher.dispose();

        this.prevReadAuthenticator = this.readAuthenticator;

        this.prevReadCipher = this.readCipher;

        this.prevReadEpoch = this.readEpoch;

        this.readAuthenticator = readAuthenticator;

        this.readCipher = readCipher;

        this.readEpoch++;

    }

    @Override

    public synchronized void close() throws IOException {

        if (!isClosed) {

            prevReadCipher.dispose();

            super.close();

        }

    }

    @Override

    boolean isEmpty() {

        return ((reassembler == null) || reassembler.isEmpty());

    }

    @Override

    int estimateFragmentSize(int packetSize) {

        int macLen = 0;

        if (readAuthenticator instanceof MAC) {

            macLen = ((MAC)readAuthenticator).MAClen();

        }

        if (packetSize > 0) {

            return readCipher.estimateFragmentSize(

                    packetSize, macLen, headerSize);

        } else {

            return Record.maxDataSize;

        }

    }

    @Override

    void expectingFinishFlight() {

        if (reassembler != null) {

            reassembler.expectingFinishFlight();

        }

    }

    @Override

    Plaintext acquirePlaintext() {

        if (reassembler != null) {

            Plaintext plaintext = reassembler.acquirePlaintext();

            if (reassembler.finished()) {

                // discard all buffered unused message.

                reassembler = null;

            }

            return plaintext;

        }

        return null;

    }

    @Override

    Plaintext decode(ByteBuffer packet) {

        if (isClosed) {

            return null;

        }

        if (debug != null && Debug.isOn("packet")) {

             Debug.printHex(

                    "[Raw read]: length = " + packet.remaining(), packet);

        }

        // The caller should have validated the record.

        int srcPos = packet.position();

        int srcLim = packet.limit();

        byte contentType = packet.get();                   // pos: 0

        byte majorVersion = packet.get();                  // pos: 1

        byte minorVersion = packet.get();                  // pos: 2

        byte[] recordEnS = new byte[8];                    // epoch + seqence

        packet.get(recordEnS);

        int recordEpoch = ((recordEnS[0] & 0xFF) << 8) |

                           (recordEnS[1] & 0xFF);          // pos: 3, 4

        long recordSeq  = Authenticator.toLong(recordEnS);

        int contentLen = ((packet.get() & 0xFF) << 8) |

                          (packet.get() & 0xFF);            // pos: 11, 12

        if (debug != null && Debug.isOn("record")) {

             System.out.println(Thread.currentThread().getName() +

                    ", READ: " +

                    ProtocolVersion.valueOf(majorVersion, minorVersion) +

                    " " + Record.contentName(contentType) + ", length = " +

                    contentLen);

        }

        int recLim = srcPos + DTLSRecord.headerSize + contentLen;

        if (this.readEpoch > recordEpoch) {

            // Discard old records delivered before this epoch.

            // Reset the position of the packet buffer.

            packet.position(recLim);

            return null;

        }

        if (this.readEpoch < recordEpoch) {

            if (contentType != Record.ct_handshake) {

                // just discard it if not a handshake message

                packet.position(recLim);

                return null;

            }

            // Not ready to decrypt this record, may be encrypted Finished

            // message, need to buffer it.

            if (reassembler == null) {

               reassembler = new DTLSReassembler();

            }

            byte[] fragment = new byte[contentLen];

            packet.get(fragment);              // copy the fragment

            RecordFragment buffered = new RecordFragment(fragment, contentType,

                    majorVersion, minorVersion,

                    recordEnS, recordEpoch, recordSeq, true);

            reassembler.queueUpFragment(buffered);

            // consume the full record in the packet buffer.

            packet.position(recLim);

            Plaintext plaintext = reassembler.acquirePlaintext();

            if (reassembler.finished()) {

                // discard all buffered unused message.

                reassembler = null;

            }

            return plaintext;

        }

        if (this.readEpoch == recordEpoch) {

            // decrypt the fragment

            packet.limit(recLim);

            packet.position(srcPos + DTLSRecord.headerSize);

            ByteBuffer plaintextFragment;

            try {

                plaintextFragment = decrypt(readAuthenticator,

                        readCipher, contentType, packet, recordEnS);

            } catch (BadPaddingException bpe) {

                if (debug != null && Debug.isOn("ssl")) {

                    System.out.println(Thread.currentThread().getName() +

                            " discard invalid record: " + bpe);

                }

                // invalid, discard this record [section 4.1.2.7, RFC 6347]

                return null;

            } finally {

                // comsume a complete record

                packet.limit(srcLim);

                packet.position(recLim);

            }

            if (contentType != Record.ct_change_cipher_spec &&

                contentType != Record.ct_handshake) {   // app data or alert

                                                        // no retransmission

               return new Plaintext(contentType, majorVersion, minorVersion,

                        recordEpoch, recordSeq, plaintextFragment);

            }

            if (contentType == Record.ct_change_cipher_spec) {

                if (reassembler == null) {

                    // handshake has not started, should be an

                    // old handshake message, discard it.

                    return null;

                }

                reassembler.queueUpFragment(

                        new RecordFragment(plaintextFragment, contentType,

                                majorVersion, minorVersion,

                                recordEnS, recordEpoch, recordSeq, false));

            } else {    // handshake record

                // One record may contain 1+ more handshake messages.

                while (plaintextFragment.remaining() > 0) {

                    HandshakeFragment hsFrag = parseHandshakeMessage(

                        contentType, majorVersion, minorVersion,

                        recordEnS, recordEpoch, recordSeq, plaintextFragment);

                    if (hsFrag == null) {

                        // invalid, discard this record

                        return null;

                    }

                    if ((reassembler == null) &&

                            isKickstart(hsFrag.handshakeType)) {

                       reassembler = new DTLSReassembler();

                    }

                    if (reassembler != null) {

                        reassembler.queueUpHandshake(hsFrag);

                    }   // else, just ignore the message.

                }

            }

            // Completed the read of the full record. Acquire the reassembled

            // messages.

            if (reassembler != null) {

                Plaintext plaintext = reassembler.acquirePlaintext();

                if (reassembler.finished()) {

                    // discard all buffered unused message.

                    reassembler = null;

                }

                return plaintext;

            }

        }

        return null;    // make the complier happy

    }

    @Override

    int bytesInCompletePacket(ByteBuffer packet) throws SSLException {

        // DTLS length field is in bytes 11/12

        if (packet.remaining() < headerSize) {

            return -1;

        }

        // Last sanity check that it's not a wild record

        int pos = packet.position();

        // Check the content type of the record.

        byte contentType = packet.get(pos);

        if (!Record.isValidContentType(contentType)) {

            throw new SSLException(

                    "Unrecognized SSL message, plaintext connection?");

        }

        // Check the protocol version of the record.

        ProtocolVersion recordVersion =

            ProtocolVersion.valueOf(packet.get(pos + 1), packet.get(pos + 2));

        checkRecordVersion(recordVersion, false);

        // Get the fragment length of the record.

        int fragLen = ((packet.get(pos + 11) & 0xFF) << 8) +

                       (packet.get(pos + 12) & 0xFF) + headerSize;

        if (fragLen > Record.maxFragmentSize) {

            throw new SSLException(

                    "Record overflow, fragment length (" + fragLen +

                    ") MUST not exceed " + Record.maxFragmentSize);

        }

        return fragLen;

    }

    @Override

    void checkRecordVersion(ProtocolVersion recordVersion,

            boolean allowSSL20Hello) throws SSLException {

        if (!recordVersion.maybeDTLSProtocol()) {

            throw new SSLException(

                    "Unrecognized record version " + recordVersion +

                    " , plaintext connection?");

        }

    }

    private static boolean isKickstart(byte handshakeType) {

        return (handshakeType == HandshakeMessage.ht_client_hello) ||

               (handshakeType == HandshakeMessage.ht_hello_request) ||

               (handshakeType == HandshakeMessage.ht_hello_verify_request);

    }

    private static HandshakeFragment parseHandshakeMessage(

            byte contentType, byte majorVersion, byte minorVersion,

            byte[] recordEnS, int recordEpoch, long recordSeq,

            ByteBuffer plaintextFragment) {

        int remaining = plaintextFragment.remaining();

        if (remaining < handshakeHeaderSize) {

            if (debug != null && Debug.isOn("ssl")) {

                System.out.println(

                        Thread.currentThread().getName() +

                        " discard invalid record: " +

                        "too small record to hold a handshake fragment");

            }

            // invalid, discard this record [section 4.1.2.7, RFC 6347]

            return null;

        }

        byte handshakeType = plaintextFragment.get();       // pos: 0

        int messageLength =

                ((plaintextFragment.get() & 0xFF) << 16) |

                ((plaintextFragment.get() & 0xFF) << 8) |

                 (plaintextFragment.get() & 0xFF);          // pos: 1-3

        int messageSeq =

                ((plaintextFragment.get() & 0xFF) << 8) |

                 (plaintextFragment.get() & 0xFF);          // pos: 4/5

        int fragmentOffset =

                ((plaintextFragment.get() & 0xFF) << 16) |

                ((plaintextFragment.get() & 0xFF) << 8) |

                 (plaintextFragment.get() & 0xFF);          // pos: 6-8

        int fragmentLength =

                ((plaintextFragment.get() & 0xFF) << 16) |

                ((plaintextFragment.get() & 0xFF) << 8) |

                 (plaintextFragment.get() & 0xFF);          // pos: 9-11

        if ((remaining - handshakeHeaderSize) < fragmentLength) {

            if (debug != null && Debug.isOn("ssl")) {

                System.out.println(

                        Thread.currentThread().getName() +

                        " discard invalid record: " +

                        "not a complete handshake fragment in the record");

            }

            // invalid, discard this record [section 4.1.2.7, RFC 6347]

            return null;

        }

        byte[] fragment = new byte[fragmentLength];

        plaintextFragment.get(fragment);

        return new HandshakeFragment(fragment, contentType,

                majorVersion, minorVersion,

                recordEnS, recordEpoch, recordSeq,

                handshakeType, messageLength,

                messageSeq, fragmentOffset, fragmentLength);

    }

    // buffered record fragment

    private static class RecordFragment implements Comparable<RecordFragment> {

        boolean         isCiphertext;

        byte            contentType;

        byte            majorVersion;

        byte            minorVersion;

        int             recordEpoch;

        long            recordSeq;

        byte[]          recordEnS;

        byte[]          fragment;

        RecordFragment(ByteBuffer fragBuf, byte contentType,

                byte majorVersion, byte minorVersion, byte[] recordEnS,

                int recordEpoch, long recordSeq, boolean isCiphertext) {

            this((byte[])null, contentType, majorVersion, minorVersion,

                    recordEnS, recordEpoch, recordSeq, isCiphertext);

            this.fragment = new byte[fragBuf.remaining()];

            fragBuf.get(this.fragment);

        }

        RecordFragment(byte[] fragment, byte contentType,

                byte majorVersion, byte minorVersion, byte[] recordEnS,

                int recordEpoch, long recordSeq, boolean isCiphertext) {

            this.isCiphertext = isCiphertext;

            this.contentType = contentType;

            this.majorVersion = majorVersion;

            this.minorVersion = minorVersion;

            this.recordEpoch = recordEpoch;

            this.recordSeq = recordSeq;

            this.recordEnS = recordEnS;

            this.fragment = fragment;       // The caller should have cloned

                                            // the buffer if necessary.

        }

        @Override

        public int compareTo(RecordFragment o) {

            return Long.compareUnsigned(this.recordSeq, o.recordSeq);

        }

    }

    // buffered handshake message

    private static final class HandshakeFragment extends RecordFragment {

        byte            handshakeType;     // handshake msg_type

        int             messageSeq;        // message_seq

        int             messageLength;     // Handshake body length

        int             fragmentOffset;    // fragment_offset

        int             fragmentLength;    // fragment_length

        HandshakeFragment(byte[] fragment, byte contentType,

                byte majorVersion, byte minorVersion, byte[] recordEnS,

                int recordEpoch, long recordSeq,

                byte handshakeType, int messageLength,

                int messageSeq, int fragmentOffset, int fragmentLength) {

            super(fragment, contentType, majorVersion, minorVersion,

                    recordEnS, recordEpoch , recordSeq, false);

            this.handshakeType = handshakeType;

            this.messageSeq = messageSeq;

            this.messageLength = messageLength;

            this.fragmentOffset = fragmentOffset;

            this.fragmentLength = fragmentLength;

        }

        @Override

        public int compareTo(RecordFragment o) {

            if (o instanceof HandshakeFragment) {

                HandshakeFragment other = (HandshakeFragment)o;

                if (this.messageSeq != other.messageSeq) {

                    // keep the insertion order for the same message

                    return this.messageSeq - other.messageSeq;

                }

            }

            return Long.compareUnsigned(this.recordSeq, o.recordSeq);

        }

    }

    private static final class HoleDescriptor {

        int offset;             // fragment_offset

        int limit;              // fragment_offset + fragment_length

        HoleDescriptor(int offset, int limit) {

            this.offset = offset;

            this.limit = limit;

        }

    }

    final class DTLSReassembler {

        TreeSet<RecordFragment> bufferedFragments = new TreeSet<>();

        HashMap<Byte, List<HoleDescriptor>> holesMap = new HashMap<>(5);

        // Epoch, sequence number and handshake message sequence of the

        // beginning message of a flight.

        byte        flightType = (byte)0xFF;

        int         flightTopEpoch = 0;

        long        flightTopRecordSeq = -1;

        int         flightTopMessageSeq = 0;

        // Epoch, sequence number and handshake message sequence of the

        // next message acquisition of a flight.

        int         nextRecordEpoch = 0;    // next record epoch

        long        nextRecordSeq = 0;      // next record sequence number

        int         nextMessageSeq = 0;     // next handshake message number

        // Expect ChangeCipherSpec and Finished messages for the final flight.

        boolean     expectCCSFlight = false;

        // Ready to process this flight if received all messages of the flight.

        boolean     flightIsReady = false;

        boolean     needToCheckFlight = false;

        // Is it a session-resuming abbreviated handshake.?

        boolean     isAbbreviatedHandshake = false;

        // The handshke fragment with the biggest record sequence number

        // in a flight, not counting the Finished message.

        HandshakeFragment lastHandshakeFragment = null;

        // Is handshake (intput) finished?

        boolean handshakeFinished = false;

        DTLSReassembler() {

            // blank

        }

        boolean finished() {

            return handshakeFinished;