atlsrecordstage.java

来自「ATLS原代码库,实现的很不错 ATLS原代码库,实现的很不错」· Java 代码 · 共 517 行 · 第 1/2 页

		if (DEBUG) System.err.println ("aTLSRecordStage: DONE Calling recordReader()");

		if (records.isEmpty()) {
		    if (DEBUG) System.err.println ("aTLSRecordStage: recordReader returned null, so not enough info to generate a Record yet.");
		    // not enough data to generate an entire record yet.
		    return;
		}
	    
		while (!records.isEmpty()) {
		    // multiple records could have been sent in one packet, so need to handle all of them individually
		    aTLSRecord record = (aTLSRecord) records.removeFirst();
		 		
		    if (record instanceof aTLSCipherSpecRecord) {
			changeCipherSpecMessage(atlsconn, record);
		    }
		    else if (record instanceof aTLSAlertRecord) {
			alertMessage(atlsconn, record);
		    }
		    else if (record instanceof aTLSHandshakeRecord) {
			handshakeMessage(atlsconn, record);
		    }
		    else if (record instanceof aTLSAppDataRecord) {
			dataMessage(atlsconn, record);
		    }
		    else {
			System.err.println ("aTLSRecordStage: Received Bad Packet. Internal Error, contact [e-mail]");
		    }
		}
	    }
	    
	}
    }
   
    public void handleEvents(QueueElementIF[] qelarr) {
	for (int i = 0; i < qelarr.length; i++) {
	    if (qelarr[i] == null) {
		System.err.println ("aTLSRecordStage: Received a null in handleEvents(). Internal Error, contact [e-mail]");
	    }
	    else {
		handleEvent(qelarr[i]);
	    }
	}
    }

    /**
     * Helper function to deal with alert messages. Calls PureTLS to decode the alert message and to handle is
     * appropriately in SSLRecordReader.readRecord().
     */
    private void alertMessage(aTLSConnection atlsconn, aTLSRecord record) {
	if (DEBUG) System.err.println ("aTLSRecordStage: received an alert message.");

	ByteArrayInputStream bais = new ByteArrayInputStream (record.data);
	atlsconn.conn.sock_in = new PushbackInputStream (bais);
	
	try {
	    atlsconn.rr.readRecord();
	}
	catch (Exception e) {
	    System.err.println ("aTLSRecordStage: Exception trying to readrecord " + e);
	    try {
		atlsconn.getConnection().close(mySink);
	    }
	    catch (SinkClosedException sce) {
		System.err.println ("aTLSRecordStage: Exception trying to close the connection " + sce);
	    }
	}
    }

    /**
     * Helper function that posts handshake messages to the handshake stage.
     * If the server is attemptiing to resume a session, then handleResume is called.
     */
    private void handshakeMessage(aTLSConnection atlsconn, aTLSRecord record) {
	if (DEBUG) System.err.println ("aTLSREcordStage: received a handshake message.");

	if ((atlsconn.isServer) && ((SSLHandshakeServer)atlsconn.conn.hs).resume) {
	    handleResume (atlsconn, record);
	}
	else {
	    if (PROFILE) {
		atlsconn.prof.addMeasurements(System.currentTimeMillis(), "aTLSRecordStage sending handshake message to handshake stage");
	    }

	    handshakeSink.enqueue_lossy(new aTLSHandshakePacket (atlsconn, (aTLSHandshakeRecord)record));
	}
    }  

    /**
     * Helper function that does the same things has handshake message, but needs to differentiate
     * between a cipher spec message and a regular handshake message because must be handled differently.
     */
    private void changeCipherSpecMessage(aTLSConnection atlsconn, aTLSRecord record) {
	if (DEBUG) System.err.println ("aTLSRecordStage: dealing with a changeCipherSpec Message");
	
	if ((atlsconn.isServer) && ((SSLHandshakeServer)atlsconn.conn.hs).resume) {
	    handleResume (atlsconn, record);
	}
	else {
	    handshakeSink.enqueue_lossy(new aTLSCipherSpecPacket (atlsconn, (aTLSCipherSpecRecord)record));
	}
    }

    /**
     * This function will decrypt the application data record by calling PureTLS, and will then
     * pass the decrypted record back to the user. But if startReader for the connection has not been called,
     * then must enqueue onto the dataQueue rather than post to the user sink. The dataSinkLock needed to 
     * synchronize actions with aTLSConnection.
     */
    private void dataMessage(aTLSConnection atlsconn, aTLSRecord record) {
	ByteArrayInputStream bais = new ByteArrayInputStream (record.data);
	atlsconn.conn.sock_in = new PushbackInputStream (bais);

	try {
	    atlsconn.rr.readRecord();
	}
	catch (Exception e) {
	    System.err.println ("aTLSRecordStage: Exception trying to readrecord " + e + " ");
	    byte[] temp = atlsconn.conn.hs.session_id;
	    for (int i = 0; i < temp.length; i++) {
		System.err.print (temp[i] + " ");
	    }

	    e.printStackTrace();
	    try {
		atlsconn.getConnection().close(mySink);
	    }
	    catch (SinkClosedException sce) {
		System.err.println ("aTLSRecordStage: Exception trying to close the connection " + sce);
	    }
	}
	// after read record is called, all the data resides on the sock_in_data stream of PureTLS.
	// so need to read all the data off of that stream and pass it to the user.
	int length = ((aTLSAppDataRecord)record).length;
	byte[] data = new byte[length];

	try {
	    atlsconn.conn.sock_in_data.read (data, 0, length);
	}
	catch (Exception e) {
	    System.err.println ("aTLSRecordStage: Exception trying to read data from sock_in_data");
	    try {
		atlsconn.getConnection().close(mySink);
	    }
	    catch (SinkClosedException sce) {
		System.err.println ("aTLSRecordStage: Exception trying to close the connection " + sce);
	    }
	}

	record.data = data;

	// if startReader hasn't been called, then no dataSink has been assigned, so just enqueue onto
	// the temporary queue until user has called startReader().
	synchronized (atlsconn.dataSinkLock) {
	    if (atlsconn.getDataSink() != null) {
		atlsconn.getDataSink().enqueue_lossy (new ATcpInPacket (atlsconn, new BufferElement(record.data)));
	    }
	    else {
		atlsconn.dataQueue.add(record);
	    }
	}
    }

    /**
     * PureTLS will decide if a connection can be resumed once the clientHello has been received. If
     * the sessionID is valid, then will set the boolean resume to true in SSLHandshakeServer. So
     * a resumed session does not need to go through the same handshake process, so no need to
     * enqueue anything to handshake stage. Best way is to handle it right here because only a few steps, 
     * and decryption can occur here too, rather than in the handshake stage.
     * For a resume, only need a change cipher spec record and a finished record.
     */
    private void handleResume(aTLSConnection atlsconn, aTLSRecord record) {
	if (DEBUG) System.err.println ("aTLSRecordStage: Server is going to resume session with client");
		
	ByteArrayInputStream bais = new ByteArrayInputStream (record.data);
	atlsconn.conn.sock_in = new PushbackInputStream (bais);
	
	try {
	    atlsconn.rr.readRecord();
	}
	catch (Exception e) {
	    System.err.println ("aTLSRecordStage: Exception trying to readrecord " + e);
	    try {
		atlsconn.getConnection().close(mySink);
	    }
	    catch (SinkClosedException sce) {
		System.err.println ("aTLSRecordStage: Exception trying to close the connection " + sce);
	    }
	}
	
	if (record instanceof aTLSCipherSpecRecord) {
	    return;
	    // because readRecord already handled it for us.
	}
	else if (record instanceof aTLSHandshakeRecord) {
	    // then it better have been a finished message!
	    try {
		if (DEBUG) System.err.println ("aTLSRecordStage: Calling handshakecontinue()");
		atlsconn.conn.hs.handshakeContinue();
	    }
	    catch (Exception e) {
		System.err.println ("aTLSRecordStage: Exception in handshakeContinue " + e);
		try {
		    atlsconn.getConnection().close(mySink);
		}
		catch (SinkClosedException sce) {
		    System.err.println ("aTLSRecordStage: Exception trying to close the connection " + sce);
		}
	    }
	    
	    if (atlsconn.conn.hs.state == SSL_HANDSHAKE_FINISHED) {
		atlsconn.conn.sock_out_external=new SSLOutputStream(atlsconn.conn);
		// creating the output stream for the data
		
		// pass the atlsconn back up to the user who created the serversocket
		atlsconn.getNewConnSink().enqueue_lossy(atlsconn);
	    } 
	    else {
		System.err.println ("aTLSRecordStage: Error, trying to resume in bad state with bad packet." + 
				    " Internal error, contact [e-mail]");
	    }
	}
	else {
	    System.err.println ("aTLSRecordStage: Internal Error, contact [e-mail]");
	}
    }
    
    /**
     * This function cleans up the state for a closed connection. Necessary to remove from
     * hash tables to avoid memory leaks.
     *
     * NOTE: if the server receives a sink closed event, should not remove itself from the
     * hash tables. The sinkclosedevent is simply used to notify server of a closed connection,
     * no need to clean up state because needs to remain intact.
     */
    void cleanupConnection(aTLSConnection atlsconn) {
	if (DEBUG) System.err.println ("aTLSRecordStage: Closing and cleaning up connection.");
	// have to clean up two hashtables
	ATcpConnection atcpconn = atlsconn.getConnection();
	connTable.remove(atcpconn);

	if (!(atlsconn.isServer)) 
	    clientSocketTable.remove(atcpconn.getClientSocket());
    }
    
    /** 
     * Return mySink
     */
    SinkIF getSink() {
	return mySink;
    }

    /** 
     * The Sandstorm stage destroy method.
     */
    public void destroy() {
    }
}