/*
    *  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.directory.api.ldap.codec.api;
  
  
  import javax.security.sasl.Sasl;
  import javax.security.sasl.SaslClient;
  import javax.security.sasl.SaslException;
  import javax.security.sasl.SaslServer;
  
  import org.apache.directory.api.ldap.model.constants.SaslQoP;
  import org.apache.mina.core.buffer.IoBuffer;
  import org.apache.mina.core.filterchain.IoFilterAdapter;
  import org.apache.mina.core.session.IoSession;
  import org.apache.mina.core.write.DefaultWriteRequest;
  import org.apache.mina.core.write.WriteRequest;
  import org.slf4j.Logger;
  import org.slf4j.LoggerFactory;
  
  
  /**
   * An {@link IoFilterAdapter} that handles integrity and confidentiality protection
   * for a SASL bound session.  The SaslFilter must be constructed with a SASL
   * context that has completed SASL negotiation.  Some SASL mechanisms, such as
   * CRAM-MD5, only support authentication and thus do not need this filter.  DIGEST-MD5
   * and GSSAPI do support message integrity and confidentiality and, therefore,
   * do need this filter.
   * 
   * @author <a href="mailto:dev@directory.apache.org">Apache Directory Project</a>
   */
  public class SaslFilter extends IoFilterAdapter
  {
      private static final Logger LOG = LoggerFactory.getLogger( SaslFilter.class );
  
      /**
       * A session attribute key that makes next one write request bypass
       * this filter (not adding a security layer).  This is a marker attribute,
       * which means that you can put whatever as its value. ({@link Boolean#TRUE}
       * is preferred.)  The attribute is automatically removed from the session
       * attribute map as soon as {@link IoSession#write(Object)} is invoked,
       * and therefore should be put again if you want to make more messages
       * bypass this filter.
       */
      public static final String DISABLE_SECURITY_LAYER_ONCE = SaslFilter.class.getName() + ".DisableSecurityLayerOnce";
  
      /**
       * A session attribute key that holds the received bytes of partially received
       * SASL message.
       */
      public static final String BYTES = SaslFilter.class.getName() + ".Buffer";
  
      /**
       * A session attribute key that holds the offset of partially received
       * SASL message.
       */
      public static final String OFFSET = SaslFilter.class.getName() + ".Offset";
  
      /** The SASL client, only set if the filter is used at the client side. */
      private final SaslClient saslClient;
  
      /** The SASL server, only set if the filter is used at the server side. */
      private final SaslServer saslServer;
  
      /** True if a security layer has been negotiated */
      private boolean hasSecurityLayer;
  
      /** The negotiated max buffer size */
      private int maxBufferSize;
  
      /**
       * Creates a new instance of SaslFilter.  The SaslFilter must be constructed
       * with a SASL client that has completed SASL negotiation.  The SASL client
       * will be used to provide message integrity and, optionally, message
       * confidentiality.
       *
       * @param saslClient The initialized SASL client.
       */
      public SaslFilter( SaslClient saslClient )
      {
          if ( saslClient == null )
          {
              throw new IllegalArgumentException();
         }
 
         this.saslServer = null;
         this.saslClient = saslClient;
         initHasSecurityLayer( ( String ) saslClient.getNegotiatedProperty( Sasl.QOP ) );
         initMaxBuffer( ( String ) saslClient.getNegotiatedProperty( Sasl.MAX_BUFFER ) );
     }
 
 
     /**
      * Creates a new instance of SaslFilter.  The SaslFilter must be constructed
      * with a SASL server that has completed SASL negotiation.  The SASL server
      * will be used to provide message integrity and, optionally, message
      * confidentiality.
      *
      * @param saslClient The initialized SASL server.
      */
     public SaslFilter( SaslServer saslServer )
     {
         if ( saslServer == null )
         {
             throw new IllegalArgumentException();
         }
 
         this.saslClient = null;
         this.saslServer = saslServer;
         initHasSecurityLayer( ( String ) saslServer.getNegotiatedProperty( Sasl.QOP ) );
         initMaxBuffer( ( String ) saslServer.getNegotiatedProperty( Sasl.MAX_BUFFER ) );
     }
 
 
     private void initHasSecurityLayer( String qop )
     {
         this.hasSecurityLayer = ( qop != null && ( qop.equals( SaslQoP.AUTH_INT.getValue() ) || qop
             .equals( SaslQoP.AUTH_CONF.getValue() ) ) );
     }
 
 
     private void initMaxBuffer( String maxBuffer )
     {
         this.maxBufferSize = maxBuffer != null ? Integer.parseInt( maxBuffer ) : 65536;
     }
 
 
     @Override
     public synchronized void messageReceived( NextFilter nextFilter, IoSession session, Object message )
         throws SaslException
     {
         LOG.debug( "Message received:  {}", message );
 
         if ( !hasSecurityLayer )
         {
             LOG.debug( "Will not use SASL on received message." );
             nextFilter.messageReceived( session, message );
             return;
         }
 
         /*
          * Unwrap the data for mechanisms that support QoP (DIGEST-MD5, GSSAPI).
          */
         IoBuffer buf = ( IoBuffer ) message;
         while ( buf.hasRemaining() )
         {
             /*
              * Check for a previously received partial SASL message which is stored in the session.
              * Otherwise read the first 4 bytes which is the length and allocate the bytes.
              * Ensure the buffer size doesn't exceed the negotiated max buffer size.
              */
             byte[] bytes = ( byte[] ) session.getAttribute( BYTES, null );
             int offset = ( int ) session.getAttribute( OFFSET, -1 );
             if ( bytes == null )
             {
                 int bufferSize = buf.getInt();
                 if ( bufferSize > maxBufferSize )
                 {
                     throw new IllegalStateException(
                         bufferSize + " exceeds the negotiated receive buffer size limit: " + maxBufferSize );
                 }
                 bytes = new byte[bufferSize];
                 offset = 0;
             }
 
             /*
              * Get the buffer as bytes. Handle the case that only a part of the SASL message was received.
              */
             int length = Math.min( bytes.length - offset, buf.remaining() );
             buf.get( bytes, offset, length );
 
             /*
              * Check if the full SASL message was received. If not store the partially received data in
              * the session so it can be resumed when the next message is received.
              */
             offset += length;
             if ( offset < bytes.length )
             {
                 LOG.debug( "Partial SASL message received:  {}/{}", offset, bytes.length );
                 session.setAttribute( BYTES, bytes );
                 session.setAttribute( OFFSET, offset );
                 break;
             }
 
             /*
              * Unwrap the SASL message and forward it to the next filter.
              */
             LOG.debug( "Will use SASL to unwrap received message of length:  {}", bytes.length );
             byte[] token = unwrap( bytes, 0, bytes.length );
             nextFilter.messageReceived( session, IoBuffer.wrap( token ) );
 
             /*
              * Finally clear the session attributes.
              */
             session.removeAttribute( BYTES );
             session.removeAttribute( OFFSET );
         }
     }
 
 
     @Override
     public synchronized void filterWrite( NextFilter nextFilter, IoSession session, WriteRequest writeRequest )
         throws SaslException
     {
         LOG.debug( "Filtering write request:  {}", writeRequest );
 
         /*
          * Check if security layer processing should be disabled once.
          */
         if ( session.containsAttribute( DISABLE_SECURITY_LAYER_ONCE ) )
         {
             // Remove the marker attribute because it is temporary.
             LOG.debug( "Disabling SaslFilter once; will not use SASL on write request." );
             session.removeAttribute( DISABLE_SECURITY_LAYER_ONCE );
             nextFilter.filterWrite( session, writeRequest );
             return;
         }
 
         if ( !hasSecurityLayer )
         {
             LOG.debug( "Will not use SASL on write request." );
             nextFilter.filterWrite( session, writeRequest );
             return;
         }
 
         /*
          * Wrap the data for mechanisms that support QoP (DIGEST-MD5, GSSAPI).
          */
 
         /*
          * Get the buffer as bytes.
          */
         IoBuffer buf = ( IoBuffer ) writeRequest.getMessage();
         int bufferLength = buf.remaining();
         byte[] bufferBytes = new byte[bufferLength];
         buf.get( bufferBytes );
 
         LOG.info( "Will use SASL to wrap message of length:  {}", bufferLength );
 
         /*
          * Ensure to not send larger SASL message than negotiated.
          */
         int max = maxBufferSize - 200;
         for ( int offset = 0; offset < bufferLength; offset += max )
         {
             int length = Math.min( bufferLength - offset, max );
             byte[] saslLayer = wrap( bufferBytes, offset, length );
 
             /*
              * Prepend 4 byte length.
              */
             IoBuffer saslLayerBuffer = IoBuffer.allocate( 4 + saslLayer.length );
             saslLayerBuffer.putInt( saslLayer.length );
             saslLayerBuffer.put( saslLayer );
             saslLayerBuffer.position( 0 );
             saslLayerBuffer.limit( 4 + saslLayer.length );
 
             LOG.debug( "Sending encrypted token of length {}.", saslLayerBuffer.limit() );
             nextFilter.filterWrite( session, new DefaultWriteRequest( saslLayerBuffer, writeRequest.getFuture() ) );
         }
     }
 
 
     private byte[] wrap( byte[] buffer, int offset, int length ) throws SaslException
     {
         if ( saslClient != null )
         {
             return saslClient.wrap( buffer, offset, length );
         }
         else
         {
             return saslServer.wrap( buffer, offset, length );
         }
     }
 
 
     private byte[] unwrap( byte[] buffer, int offset, int length ) throws SaslException
     {
         if ( saslClient != null )
         {
             return saslClient.unwrap( buffer, offset, length );
         }
         else
         {
             return saslServer.unwrap( buffer, offset, length );
         }
     }
 
 }