/**
 * Copyright 2020 German Tischler-Höhle
 *
 * Redistribution and use in source and binary forms, with or without
 * modification, are permitted provided that the following conditions
 * are met:
 *
 * 1. Redistributions of source code must retain the above copyright
 *    notice, this list of conditions and the following disclaimer.
 * 2. Redistributions in binary form must reproduce the above copyright
 *    notice, this list of conditions and the following disclaimer in the
 *    documentation and/or other materials provided with the distribution.
 *
 * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
 * "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
 * LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A
 * PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT
 * HOLDER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,
 * SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT
 * LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
 * DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY
 * THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
 * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
 * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
 **/
#if ! defined(LIBSECRECY_LOCKEDMEMORY_HPP)
#define LIBSECRECY_LOCKEDMEMORY_HPP

#include <libsecrecy/NumToHex.hpp>
#include <cstdlib>
#include <cerrno>
#include <climits>
#include <cstring>
#include <memory>
#include <sstream>
#include <stdexcept>
#include <sys/mman.h>

namespace libsecrecy
{
	/**
	 * class for locked memory region (memory which will not be swapped to
	 * secondary storage)
	 **/
	template<typename _type>
	struct LockedMemory : public NumToHex
	{
		typedef _type type;

		private:
		// number of elements
		std::size_t n;
		// memory (space for n elements)
		std::shared_ptr<type[]> A;


		void doUnlock()
		{
			while ( true )
			{
				int const r = ::munlock(A.get(),n*sizeof(type));

				if ( r == 0 )
					break;

				int const error = errno;

				switch ( error )
				{
					case EAGAIN:
					case EINTR:
						break;
					default:
					{
						std::ostringstream ostr;
						ostr << "libsecrecy::LockedMemory: munlock call failed: " << strerror(errno) << std::endl;
						throw std::runtime_error(ostr.str());
					}
				}
			}

		}

		void doLock()
		{
			while ( true )
			{
				int const r = ::mlock(A.get(),n*sizeof(type));

				if ( r == 0 )
					break;

				int const error = errno;

				switch ( error )
				{
					case EAGAIN:
					case EINTR:
						break;
					default:
					{
						std::ostringstream ostr;
						ostr << "libsecrecy::LockedMemory: mlock call failed: " << strerror(errno) << std::endl;
						throw std::runtime_error(ostr.str());
					}
				}
			}
		}

		void cleanup()
		{
			if ( A.get() )
			{
				memset(A.get(),0,n*sizeof(type));
				doUnlock();
				A = std::shared_ptr<type[]>();
				n = 0;
			}
		}

		public:
		/**
		 * allocate region of 0 elements
		 **/
		LockedMemory()
		: n(0), A()
		{

		}
		/**
		 * allocate and lock region of rn elements
		 **/
		LockedMemory(::std::size_t const rn)
		: n(rn), A(new type[n])
		{
			doLock();
		}
		/**
		 * construct locked copy of L
		 **/
		LockedMemory(LockedMemory const & L)
		: n(L.n), A(new type[n])
		{
			doLock();
			std::copy(L.begin(),L.end(),A.get());
		}
		/**
		 * assign locked copy of L
		 **/
		LockedMemory & operator=(LockedMemory const & L)
		{
			if ( this != &L )
			{
				cleanup();
				std::shared_ptr<type[]> tptr(new type[L.n]);
				A = tptr;
				n = L.n;

				try
				{
					doLock();
				}
				catch(...)
				{
					A = std::shared_ptr<type[]>();
					n = 0;
					throw;
				}

				std::copy(L.begin(),L.end(),A.get());
			}
			return *this;
		}
		/**
		 * erase, unlock and deallocate
		 **/
		~LockedMemory()
		{
			try
			{
				cleanup();
			}
			catch(std::exception const & ex)
			{
				// std::cerr << ex.what() << std::endl;
			}
		}

		std::shared_ptr< LockedMemory<type> > bump()
		{
			std::size_t const one = 1;
			std::size_t const nsize = std::max(one,2*n);

			std::shared_ptr< LockedMemory<type> > L(new LockedMemory<type>(nsize));
			std::copy(begin(),end(),L->begin());

			return L;
		}

		type * get()
		{
			return A.get();
		}

		type const * get() const
		{
			return A.get();
		}

		type * begin()
		{
			return get();
		}

		type * end()
		{
			return begin()+n;
		}

		type const * begin() const
		{
			return get();
		}

		type const * end() const
		{
			return begin()+n;
		}

		type const * cbegin() const
		{
			return get();
		}

		type const * cend() const
		{
			return begin()+n;
		}

		type & operator[](std::size_t const i)
		{
			return get()[i];
		}

		type const & operator[](std::size_t const i) const
		{
			return get()[i];
		}

		/**
		 * return size of array
		 **/
		::std::size_t size() const
		{
			return n;
		}

		/**
		 * convert contents to a sequence of ASCII hex numbers and return
		 * as character array
		 **/
		std::shared_ptr< LockedMemory<char> > toHex() const
		{
			static unsigned int const bitsPerValue = CHAR_BIT * sizeof(type);
			static unsigned int const runsPerValue = (bitsPerValue+3) / 4;

			std::shared_ptr< LockedMemory<char> > sptr(new LockedMemory<char>(runsPerValue * size()));
			char * p = sptr->get();

			for ( ::std::size_t i = 0; i < size(); ++i )
				for ( unsigned int j = 0; j < runsPerValue; ++j )
					*p++ = numToHex((A[i]>>(bitsPerValue - (j+1)*4))&0xF);

			return sptr;
		}

	};
}
#endif
