Secure Memory For Swift Objects

I am writing a swift application that requires handling private keys in memory. Because of the sensitivity of such objects, the keys need to be cleared (a.k.a. written to all zeros) when the object is deallocated, and the memory cannot be paged to disk (which is typically done using mlock()).

In Objective-C, you can provide a custom CFAllocator object, which allows you to use your own functions to allocate/deallocate/reallocate the memory used by an object.

So one solution is to just implement a "SecureData" object in objective-c, which internally creates an NSMutableData object using a custom CFAllocator (also in objective-c).

However, is there any way for me to provide my own custom memory allocation functions for a pure swift object (for example, a struct or a [UInt8])? Or is there a better, "proper" way to implement secure memory like this in swift?

If you want complete control over a region of memory you allocate yourself, you can use UnsafePointer and co:

// allocate enough memory for ten Ints
var ump = UnsafeMutablePointer<Int>.alloc(10)
// memory is in an uninitialized raw state

// initialize that memory with Int objects
// (here, from a collection)
ump.initializeFrom(reverse(0..<10))

// memory property gives you access to the underlying value
ump.memory // 9

// UnsafeMutablePointer acts like an IndexType
ump.successor().memory // 8
// and it has a subscript, but it's not a CollectionType
ump[3] // = 6

// wrap it in an UnsafeMutableBufferPointer to treat it
// like a collection (or UnsafeBufferPointer if you don't
// need to be able to alter the values)
let col = UnsafeMutableBufferPointer(start: ump, count: 10)
col[3] = 99
println(",".join(map(col,toString)))
// prints 9,8,7,99,5,4,3,2,1,0

ump.destroy(10)
// now the allocated memory is back in a raw state
// you could re-allocate it...
ump.initializeFrom(0..<10)
ump.destroy(10)

// when you're done, deallocate the memory
ump.dealloc(10)

You can also have UnsafePointer point to other memory, such as memory you’re handed by some C API.

UnsafePointer can be passed into C functions that take a pointer to a contiguous block of memory. So for your purposes, you could then pass this pointer into a function like mlock:

let count = 10
let ump = UnsafeMutablePointer.allocate<Int>(count)
mlock(ump, UInt(sizeof(Int) * count))
// initialize, use, and destroy the memory
munlock(ump, UInt(sizeof(Int) * count))
ump.dealloc(count)

You can even hold your own custom types:

struct MyStruct {
    let a: Int
    let b: Int
}

var pointerToStruct = UnsafeMutablePointer<MyStruct>.alloc(1)
pointerToStruct.initialize(MyStruct(a: 1, b: 2))
pointerToStruct.memory.b  // 2
pointerToStruct.destroy()
pointerToStruct.dealloc(1)

However be aware if doing this with classes, or even arrays or strings (or a struct that contains them), that all you will be holding in your memory is pointers to other memory that these objects allocate and own. If this matters to you (i.e. you are doing something special to this memory such as securing it, in your example), this is probably not what you want.

So either you need to use fixed-size objects, or make further use of UnsafePointer to hold pointers to more memory regions. If they don't need to dynamically resize, then just a single allocation of an unsafe pointer, possibly wrapped in a UnsafeBufferPointer for a collection interface, could do it.

If you need more dynamic behavior, below is a very bare-bones implementation of a collection that can resize as necessary, that could be enhanced to cover specialty memory-handling logic:

// Note this is a class not a struct, so it does NOT have value semantics,
// changing a copy changes all copies.
public class UnsafeCollection<T> {
    private var _len: Int = 0
    private var _buflen: Int = 0
    private var _buf: UnsafeMutablePointer<T> = nil

    public func removeAll(keepCapacity: Bool = false) {
        _buf.destroy(_len)
        _len = 0
        if !keepCapacity {
            _buf.dealloc(_buflen)
            _buflen = 0
            _buf = nil
        }
    }

    public required init() { }
    deinit { self.removeAll(keepCapacity: false) }

    public var count: Int { return _len }
    public var isEmpty: Bool { return _len == 0 }
}

To cover the requirements of MutableCollectionType (i.e. CollectionType plus assignable subscript):

extension UnsafeCollection: MutableCollectionType {
    typealias Index = Int
    public var startIndex: Int { return 0 }
    public var endIndex: Int { return _len }

    public subscript(idx: Int) -> T {
        get {
            precondition(idx < _len)
            return _buf[idx]
        }
        set(newElement) {
            precondition(idx < _len)
            let ptr = _buf.advancedBy(idx)
            ptr.destroy()
            ptr.initialize(newElement)
        }
    }

    typealias Generator = IndexingGenerator<UnsafeCollection>
    public func generate() -> Generator {
        return Generator(self)
    }
}

And ExtensibleCollectionType, to allow for dynamic growth:

extension UnsafeCollection: ExtensibleCollectionType {
    public func reserveCapacity(n: Index.Distance) {
        if n > _buflen {
            let newBuf = UnsafeMutablePointer<T>.alloc(n)
            newBuf.moveInitializeBackwardFrom(_buf, count: _len)
            _buf.dealloc(_buflen)
            _buf = newBuf
            _buflen = n
        }
    }

    public func append(x: T) {
        if _len == _buflen {
            reserveCapacity(Int(Double(_len) * 1.6) + 1)
        }
        _buf.advancedBy(_len++).initialize(x)
    }

    public func extend<S: SequenceType where S.Generator.Element == T>
      (newElements: S) {
        var g = newElements.generate()
        while let x: T = g.next() {
            self.append(x)
        }
    }
}