1// secblock.h - written and placed in the public domain by Wei Dai
2
3#ifndef CRYPTOPP_SECBLOCK_H
4#define CRYPTOPP_SECBLOCK_H
5
6#include "config.h"
7#include "misc.h"
8#include <assert.h>
9
10#if defined(CRYPTOPP_MEMALIGN_AVAILABLE) || defined(CRYPTOPP_MM_MALLOC_AVAILABLE) || defined(QNX)
11	#include <malloc.h>
12#else
13	#include <stdlib.h>
14#endif
15
16NAMESPACE_BEGIN(CryptoPP)
17
18// ************** secure memory allocation ***************
19
20template<class T>
21class AllocatorBase
22{
23public:
24	typedef T value_type;
25	typedef size_t size_type;
26#ifdef CRYPTOPP_MSVCRT6
27	typedef ptrdiff_t difference_type;
28#else
29	typedef std::ptrdiff_t difference_type;
30#endif
31	typedef T * pointer;
32	typedef const T * const_pointer;
33	typedef T & reference;
34	typedef const T & const_reference;
35
36	pointer address(reference r) const {return (&r);}
37	const_pointer address(const_reference r) const {return (&r); }
38	void construct(pointer p, const T& val) {new (p) T(val);}
39	void destroy(pointer p) {p->~T();}
40	size_type max_size() const {return ~size_type(0)/sizeof(T);}	// switch to std::numeric_limits<T>::max later
41
42protected:
43	static void CheckSize(size_t n)
44	{
45		if (n > ~size_t(0) / sizeof(T))
46			throw InvalidArgument("AllocatorBase: requested size would cause integer overflow");
47	}
48};
49
50#define CRYPTOPP_INHERIT_ALLOCATOR_TYPES	\
51typedef typename AllocatorBase<T>::value_type value_type;\
52typedef typename AllocatorBase<T>::size_type size_type;\
53typedef typename AllocatorBase<T>::difference_type difference_type;\
54typedef typename AllocatorBase<T>::pointer pointer;\
55typedef typename AllocatorBase<T>::const_pointer const_pointer;\
56typedef typename AllocatorBase<T>::reference reference;\
57typedef typename AllocatorBase<T>::const_reference const_reference;
58
59#if defined(_MSC_VER) && (_MSC_VER < 1300)
60// this pragma causes an internal compiler error if placed immediately before std::swap(a, b)
61#pragma warning(push)
62#pragma warning(disable: 4700)	// VC60 workaround: don't know how to get rid of this warning
63#endif
64
65template <class T, class A>
66typename A::pointer StandardReallocate(A& a, T *p, typename A::size_type oldSize, typename A::size_type newSize, bool preserve)
67{
68	if (oldSize == newSize)
69		return p;
70
71	if (preserve)
72	{
73		typename A::pointer newPointer = a.allocate(newSize, NULL);
74		memcpy_s(newPointer, sizeof(T)*newSize, p, sizeof(T)*STDMIN(oldSize, newSize));
75		a.deallocate(p, oldSize);
76		return newPointer;
77	}
78	else
79	{
80		a.deallocate(p, oldSize);
81		return a.allocate(newSize, NULL);
82	}
83}
84
85#if defined(_MSC_VER) && (_MSC_VER < 1300)
86#pragma warning(pop)
87#endif
88
89template <class T, bool T_Align16 = false>
90class AllocatorWithCleanup : public AllocatorBase<T>
91{
92public:
93	CRYPTOPP_INHERIT_ALLOCATOR_TYPES
94
95	pointer allocate(size_type n, const void * = NULL)
96	{
97		CheckSize(n);
98		if (n == 0)
99			return NULL;
100
101		if (CRYPTOPP_BOOL_ALIGN16_ENABLED && T_Align16 && n*sizeof(T) >= 16)
102		{
103			byte *p;
104		#ifdef CRYPTOPP_MM_MALLOC_AVAILABLE
105			while (!(p = (byte *)_mm_malloc(sizeof(T)*n, 16)))
106		#elif defined(CRYPTOPP_MEMALIGN_AVAILABLE)
107			while (!(p = (byte *)memalign(16, sizeof(T)*n)))
108		#elif defined(CRYPTOPP_MALLOC_ALIGNMENT_IS_16)
109			while (!(p = (byte *)malloc(sizeof(T)*n)))
110		#else
111			while (!(p = (byte *)malloc(sizeof(T)*n + 16)))
112		#endif
113				CallNewHandler();
114
115		#ifdef CRYPTOPP_NO_ALIGNED_ALLOC
116			size_t adjustment = 16-((size_t)p%16);
117			p += adjustment;
118			p[-1] = (byte)adjustment;
119		#endif
120
121			assert(IsAlignedOn(p, 16));
122			return (pointer)p;
123		}
124
125		pointer p;
126		while (!(p = (pointer)malloc(sizeof(T)*n)))
127			CallNewHandler();
128		return p;
129	}
130
131	void deallocate(void *p, size_type n)
132	{
133		memset_z(p, 0, n*sizeof(T));
134
135		if (CRYPTOPP_BOOL_ALIGN16_ENABLED && T_Align16 && n*sizeof(T) >= 16)
136		{
137		#ifdef CRYPTOPP_MM_MALLOC_AVAILABLE
138			_mm_free(p);
139		#elif defined(CRYPTOPP_NO_ALIGNED_ALLOC)
140			p = (byte *)p - ((byte *)p)[-1];
141			free(p);
142		#else
143			free(p);
144		#endif
145			return;
146		}
147
148		free(p);
149	}
150
151	pointer reallocate(T *p, size_type oldSize, size_type newSize, bool preserve)
152	{
153		return StandardReallocate(*this, p, oldSize, newSize, preserve);
154	}
155
156	// VS.NET STL enforces the policy of "All STL-compliant allocators have to provide a
157	// template class member called rebind".
158    template <class U> struct rebind { typedef AllocatorWithCleanup<U, T_Align16> other; };
159#if _MSC_VER >= 1500
160	AllocatorWithCleanup() {}
161	template <class U, bool A> AllocatorWithCleanup(const AllocatorWithCleanup<U, A> &) {}
162#endif
163};
164
165CRYPTOPP_DLL_TEMPLATE_CLASS AllocatorWithCleanup<byte>;
166CRYPTOPP_DLL_TEMPLATE_CLASS AllocatorWithCleanup<word16>;
167CRYPTOPP_DLL_TEMPLATE_CLASS AllocatorWithCleanup<word32>;
168CRYPTOPP_DLL_TEMPLATE_CLASS AllocatorWithCleanup<word64>;
169#if CRYPTOPP_BOOL_X86
170CRYPTOPP_DLL_TEMPLATE_CLASS AllocatorWithCleanup<word, true>;	// for Integer
171#endif
172
173template <class T>
174class NullAllocator : public AllocatorBase<T>
175{
176public:
177	CRYPTOPP_INHERIT_ALLOCATOR_TYPES
178
179	pointer allocate(size_type n, const void * = NULL)
180	{
181		assert(false);
182		return NULL;
183	}
184
185	void deallocate(void *p, size_type n)
186	{
187		assert(false);
188	}
189
190	size_type max_size() const {return 0;}
191};
192
193// This allocator can't be used with standard collections because
194// they require that all objects of the same allocator type are equivalent.
195// So this is for use with SecBlock only.
196template <class T, size_t S, class A = NullAllocator<T>, bool T_Align16 = false>
197class FixedSizeAllocatorWithCleanup : public AllocatorBase<T>
198{
199public:
200	CRYPTOPP_INHERIT_ALLOCATOR_TYPES
201
202	FixedSizeAllocatorWithCleanup() : m_allocated(false) {}
203
204	pointer allocate(size_type n)
205	{
206		assert(IsAlignedOn(m_array, 8));
207
208		if (n <= S && !m_allocated)
209		{
210			m_allocated = true;
211			return GetAlignedArray();
212		}
213		else
214			return m_fallbackAllocator.allocate(n);
215	}
216
217	pointer allocate(size_type n, const void *hint)
218	{
219		if (n <= S && !m_allocated)
220		{
221			m_allocated = true;
222			return GetAlignedArray();
223		}
224		else
225			return m_fallbackAllocator.allocate(n, hint);
226	}
227
228	void deallocate(void *p, size_type n)
229	{
230		if (p == GetAlignedArray())
231		{
232			assert(n <= S);
233			assert(m_allocated);
234			m_allocated = false;
235			memset(p, 0, n*sizeof(T));
236		}
237		else
238			m_fallbackAllocator.deallocate(p, n);
239	}
240
241	pointer reallocate(pointer p, size_type oldSize, size_type newSize, bool preserve)
242	{
243		if (p == GetAlignedArray() && newSize <= S)
244		{
245			assert(oldSize <= S);
246			if (oldSize > newSize)
247				memset(p + newSize, 0, (oldSize-newSize)*sizeof(T));
248			return p;
249		}
250
251		pointer newPointer = allocate(newSize, NULL);
252		if (preserve)
253			memcpy(newPointer, p, sizeof(T)*STDMIN(oldSize, newSize));
254		deallocate(p, oldSize);
255		return newPointer;
256	}
257
258	size_type max_size() const {return STDMAX(m_fallbackAllocator.max_size(), S);}
259
260private:
261#ifdef __BORLANDC__
262	T* GetAlignedArray() {return m_array;}
263	T m_array[S];
264#else
265	T* GetAlignedArray() {return (CRYPTOPP_BOOL_ALIGN16_ENABLED && T_Align16) ? (T*)(((byte *)m_array) + (0-(size_t)m_array)%16) : m_array;}
266	CRYPTOPP_ALIGN_DATA(8) T m_array[(CRYPTOPP_BOOL_ALIGN16_ENABLED && T_Align16) ? S+8/sizeof(T) : S];
267#endif
268	A m_fallbackAllocator;
269	bool m_allocated;
270};
271
272//! a block of memory allocated using A
273template <class T, class A = AllocatorWithCleanup<T> >
274class SecBlock
275{
276public:
277	typedef typename A::value_type value_type;
278	typedef typename A::pointer iterator;
279	typedef typename A::const_pointer const_iterator;
280	typedef typename A::size_type size_type;
281
282	explicit SecBlock(size_type size=0)
283		: m_size(size) {m_ptr = m_alloc.allocate(size, NULL);}
284	SecBlock(const SecBlock<T, A> &t)
285		: m_size(t.m_size) {m_ptr = m_alloc.allocate(m_size, NULL); memcpy_s(m_ptr, m_size*sizeof(T), t.m_ptr, m_size*sizeof(T));}
286	SecBlock(const T *t, size_type len)
287		: m_size(len)
288	{
289		m_ptr = m_alloc.allocate(len, NULL);
290		if (t == NULL)
291			memset_z(m_ptr, 0, len*sizeof(T));
292		else
293			memcpy(m_ptr, t, len*sizeof(T));
294	}
295
296	~SecBlock()
297		{m_alloc.deallocate(m_ptr, m_size);}
298
299#ifdef __BORLANDC__
300	operator T *() const
301		{return (T*)m_ptr;}
302#else
303	operator const void *() const
304		{return m_ptr;}
305	operator void *()
306		{return m_ptr;}
307
308	operator const T *() const
309		{return m_ptr;}
310	operator T *()
311		{return m_ptr;}
312#endif
313
314//	T *operator +(size_type offset)
315//		{return m_ptr+offset;}
316
317//	const T *operator +(size_type offset) const
318//		{return m_ptr+offset;}
319
320//	T& operator[](size_type index)
321//		{assert(index >= 0 && index < m_size); return m_ptr[index];}
322
323//	const T& operator[](size_type index) const
324//		{assert(index >= 0 && index < m_size); return m_ptr[index];}
325
326	iterator begin()
327		{return m_ptr;}
328	const_iterator begin() const
329		{return m_ptr;}
330	iterator end()
331		{return m_ptr+m_size;}
332	const_iterator end() const
333		{return m_ptr+m_size;}
334
335	typename A::pointer data() {return m_ptr;}
336	typename A::const_pointer data() const {return m_ptr;}
337
338	size_type size() const {return m_size;}
339	bool empty() const {return m_size == 0;}
340
341	byte * BytePtr() {return (byte *)m_ptr;}
342	const byte * BytePtr() const {return (const byte *)m_ptr;}
343	size_type SizeInBytes() const {return m_size*sizeof(T);}
344
345	//! set contents and size
346	void Assign(const T *t, size_type len)
347	{
348		New(len);
349		memcpy_s(m_ptr, m_size*sizeof(T), t, len*sizeof(T));
350	}
351
352	//! copy contents and size from another SecBlock
353	void Assign(const SecBlock<T, A> &t)
354	{
355		New(t.m_size);
356		memcpy_s(m_ptr, m_size*sizeof(T), t.m_ptr, m_size*sizeof(T));
357	}
358
359	SecBlock<T, A>& operator=(const SecBlock<T, A> &t)
360	{
361		Assign(t);
362		return *this;
363	}
364
365	// append to this object
366	SecBlock<T, A>& operator+=(const SecBlock<T, A> &t)
367	{
368		size_type oldSize = m_size;
369		Grow(m_size+t.m_size);
370		memcpy_s(m_ptr+oldSize, m_size*sizeof(T), t.m_ptr, t.m_size*sizeof(T));
371		return *this;
372	}
373
374	// append operator
375	SecBlock<T, A> operator+(const SecBlock<T, A> &t)
376	{
377		SecBlock<T, A> result(m_size+t.m_size);
378		memcpy_s(result.m_ptr, result.m_size*sizeof(T), m_ptr, m_size*sizeof(T));
379		memcpy_s(result.m_ptr+m_size, t.m_size*sizeof(T), t.m_ptr, t.m_size*sizeof(T));
380		return result;
381	}
382
383	bool operator==(const SecBlock<T, A> &t) const
384	{
385		return m_size == t.m_size && VerifyBufsEqual(m_ptr, t.m_ptr, m_size*sizeof(T));
386	}
387
388	bool operator!=(const SecBlock<T, A> &t) const
389	{
390		return !operator==(t);
391	}
392
393	//! change size, without preserving contents
394	void New(size_type newSize)
395	{
396		m_ptr = m_alloc.reallocate(m_ptr, m_size, newSize, false);
397		m_size = newSize;
398	}
399
400	//! change size and set contents to 0
401	void CleanNew(size_type newSize)
402	{
403		New(newSize);
404		memset_z(m_ptr, 0, m_size*sizeof(T));
405	}
406
407	//! change size only if newSize > current size. contents are preserved
408	void Grow(size_type newSize)
409	{
410		if (newSize > m_size)
411		{
412			m_ptr = m_alloc.reallocate(m_ptr, m_size, newSize, true);
413			m_size = newSize;
414		}
415	}
416
417	//! change size only if newSize > current size. contents are preserved and additional area is set to 0
418	void CleanGrow(size_type newSize)
419	{
420		if (newSize > m_size)
421		{
422			m_ptr = m_alloc.reallocate(m_ptr, m_size, newSize, true);
423			memset(m_ptr+m_size, 0, (newSize-m_size)*sizeof(T));
424			m_size = newSize;
425		}
426	}
427
428	//! change size and preserve contents
429	void resize(size_type newSize)
430	{
431		m_ptr = m_alloc.reallocate(m_ptr, m_size, newSize, true);
432		m_size = newSize;
433	}
434
435	//! swap contents and size with another SecBlock
436	void swap(SecBlock<T, A> &b)
437	{
438		std::swap(m_alloc, b.m_alloc);
439		std::swap(m_size, b.m_size);
440		std::swap(m_ptr, b.m_ptr);
441	}
442
443//private:
444	A m_alloc;
445	size_type m_size;
446	T *m_ptr;
447};
448
449typedef SecBlock<byte> SecByteBlock;
450typedef SecBlock<byte, AllocatorWithCleanup<byte, true> > AlignedSecByteBlock;
451typedef SecBlock<word> SecWordBlock;
452
453//! a SecBlock with fixed size, allocated statically
454template <class T, unsigned int S, class A = FixedSizeAllocatorWithCleanup<T, S> >
455class FixedSizeSecBlock : public SecBlock<T, A>
456{
457public:
458	explicit FixedSizeSecBlock() : SecBlock<T, A>(S) {}
459};
460
461template <class T, unsigned int S, bool T_Align16 = true>
462class FixedSizeAlignedSecBlock : public FixedSizeSecBlock<T, S, FixedSizeAllocatorWithCleanup<T, S, NullAllocator<T>, T_Align16> >
463{
464};
465
466//! a SecBlock that preallocates size S statically, and uses the heap when this size is exceeded
467template <class T, unsigned int S, class A = FixedSizeAllocatorWithCleanup<T, S, AllocatorWithCleanup<T> > >
468class SecBlockWithHint : public SecBlock<T, A>
469{
470public:
471	explicit SecBlockWithHint(size_t size) : SecBlock<T, A>(size) {}
472};
473
474template<class T, bool A, class U, bool B>
475inline bool operator==(const CryptoPP::AllocatorWithCleanup<T, A>&, const CryptoPP::AllocatorWithCleanup<U, B>&) {return (true);}
476template<class T, bool A, class U, bool B>
477inline bool operator!=(const CryptoPP::AllocatorWithCleanup<T, A>&, const CryptoPP::AllocatorWithCleanup<U, B>&) {return (false);}
478
479NAMESPACE_END
480
481NAMESPACE_BEGIN(std)
482template <class T, class A>
483inline void swap(CryptoPP::SecBlock<T, A> &a, CryptoPP::SecBlock<T, A> &b)
484{
485	a.swap(b);
486}
487
488#if defined(_STLP_DONT_SUPPORT_REBIND_MEMBER_TEMPLATE) || (defined(_STLPORT_VERSION) && !defined(_STLP_MEMBER_TEMPLATE_CLASSES))
489// working for STLport 5.1.3 and MSVC 6 SP5
490template <class _Tp1, class _Tp2>
491inline CryptoPP::AllocatorWithCleanup<_Tp2>&
492__stl_alloc_rebind(CryptoPP::AllocatorWithCleanup<_Tp1>& __a, const _Tp2*)
493{
494	return (CryptoPP::AllocatorWithCleanup<_Tp2>&)(__a);
495}
496#endif
497
498NAMESPACE_END
499
500#endif
501