//==============================================================================
// NSS thread-safety test program
//==============================================================================

// compile with:
// g++ -o nssTest nssTest.cpp -lpthread

#include <vector>
#include <iostream>
#include <cerrno>

#include <dirent.h>
#include <pthread.h>
#include <unistd.h>
#include <sys/types.h>
#include <sys/stat.h>
#include <pwd.h>

namespace {

//==============================================================================
// LOCAL DECLARATIONS
//==============================================================================

/*----- variables -----*/

int pwSize = 0;

typedef std::vector< std::string > StrVec;
typedef std::vector< uid_t > UidVec;
UidVec userIds;
StrVec userNames;
bool threadUnsafe;
pthread_mutex_t mutex = PTHREAD_MUTEX_INITIALIZER;

//==============================================================================
// LOCAL FUNCTIONS
//==============================================================================

//------------------------------------------------------------------------------
//
std::string getUserName( const uid_t userId )
{
   char* pwBuf = new char[ pwSize ];
   struct passwd pwd;
   struct passwd* pwdr = 0;
   int ret;

   do {
      if ( threadUnsafe ) {
         pthread_mutex_lock( &mutex );
      }
      ret = getpwuid_r( userId , &pwd, pwBuf, pwSize, &pwdr );
      if ( threadUnsafe ) {
         pthread_mutex_unlock( &mutex );
      }
   } while ( ret == -1 && errno == EINTR );
   if ( ret == -1 ) {
      return std::string( "--->Error while retrieving<---" );
   }
   if ( pwdr == 0 ) {
      return std::string( "--->Error in returned value<---" );
   }
   const std::string name( pwd.pw_name );
   delete [] pwBuf;
   return name;
}

//------------------------------------------------------------------------------
//
void* accessOwners( void* param )
{
   // Perform a few accesses to make sure there is some concurrent access to
   // getpwuid_r.
   for ( int i = 0; i < 50; ++i ) {
      const int idx = rand() % userIds.size();
      getUserName( userIds[ i ] );
   }

   const int idx = reinterpret_cast< int >( param );
   userNames[ idx ] = getUserName( userIds[ idx ] );
   return 0;
}

}

//==============================================================================
// GLOBAL FUNCTIONS
//==============================================================================

//------------------------------------------------------------------------------
//
int main( int argc, char* argv[] )
{
	if ( argc < 2 ) {
		std::cerr << "Missing path." << std::endl;
	}

    std::string dirStr( argv[ 1 ] );
    if ( argc > 2 ) {
       const std::string serializeAccessArg( "-s" );
       if ( argv[ 1 ] == serializeAccessArg ) {
          threadUnsafe = true;
          dirStr = argv[ 2 ];
          std::cerr << "Serializing calls to getpwuid_r." << std::endl;
       }
    }
    
	pwSize = sysconf( _SC_GETPW_R_SIZE_MAX );
    if ( pwSize < 0 ) {
       std::cerr << "_SC_GETPW_R_SIZE_MAX returns " << pwSize;
       pwSize = 4096;
       std::cerr << ".  Adjusted to " << pwSize << "." << std::endl;
    }

    DIR* dirStream = opendir( dirStr.c_str() );
    if ( dirStream == NULL ) {
       std::cerr << dirStr << " is not a directory." << std::endl;
       return -1;
    }

    if ( dirStr[ dirStr.size() - 1 ] != '/' ) {
       dirStr += "/";
    }

    StrVec fileNames;
    while ( true ) {
       struct dirent* dirEnt = readdir( dirStream );
       if ( dirEnt == NULL ) {
          break;
       }
       const std::string fileName( dirStr + std::string( dirEnt->d_name ) );
       fileNames.push_back( fileName );

       struct stat statStr;
       stat( fileName.c_str(), &statStr );
       userIds.push_back( statStr.st_uid );

       if ( userIds.size() > PTHREAD_THREADS_MAX ) {
          // Avoid doing work for too many files: we start one thread per file
          // and we don't want to go beyond the prheads limit.
          break;
       }
    }
    
    std::vector< pthread_t > threadIds( fileNames.size() );
    userNames.resize( fileNames.size() );
    for ( int i = 0; i < fileNames.size(); ++i ) {
       ::pthread_create( &threadIds[ i ], NULL, accessOwners,
                         reinterpret_cast< void* >( i ) );
    }

    for ( int i = 0; i < fileNames.size(); ++i ) {
       ::pthread_join( threadIds[ i ], NULL );
    }

    for ( int i = 0; i < fileNames.size(); ++i ) {
       std::cout << fileNames[ i ] << ": " << userNames[ i ] << std::endl;
    }

	return 0;
}