Generic base64 encode/decode

generic-base64.swift

// This is a version of swift-corelibs-foundation's base64 encoder/decoder,
// adapted to encode/decode any Collection of UInt8s.
//
// This means you can directly encode/decode:
// Array<UInt8>, String.UTF8View, Unsafe{Mutable}BufferPointer<UInt8>, Unsafe{Raw}BufferPointer,
// SwiftNIO's ByteBuffer, Data, NSData and DispatchData (of course), and more.
//
// e.g.
// let encoded = "Hello, this is an encoded String! 🏆".utf8.base64EncodedString()
// print(encoded)
// -> SGVsbG8sIHRoaXMgaXMgYW4gZW5jb2RlZCBTdHJpbmchIPCfj4Y=

import Foundation // for Data.Base64DecodingOptions, Data.Base64EncodingOptions

/// The ranges of ASCII characters that are used to encode data in Base64.
private let base64ByteMappings: [Range<UInt8>] = [
    65 ..< 91,      // A-Z
    97 ..< 123,     // a-z
    48 ..< 58,      // 0-9
    43 ..< 44,      // +
    47 ..< 48,      // /
]
/**
 Padding character used when the number of bytes to encode is not divisible by 3
 */
private let base64Padding : UInt8 = 61 // =

/**
 This method takes a byte with a character from Base64-encoded string
 and gets the binary value that the character corresponds to.
 
 - parameter byte:       The byte with the Base64 character.
 - returns:              Base64DecodedByte value containing the result (Valid , Invalid, Padding)
 */
private enum Base64DecodedByte {
    case valid(UInt8)
    case invalid
    case padding
}

private func base64DecodeByte(_ byte: UInt8) -> Base64DecodedByte {
    guard byte != base64Padding else {return .padding}
    var decodedStart: UInt8 = 0
    for range in base64ByteMappings {
        if range.contains(byte) {
            let result = decodedStart + (byte - range.lowerBound)
            return .valid(result)
        }
        decodedStart += range.upperBound - range.lowerBound
    }
    return .invalid
}

private func base64EncodeByte(_ byte: UInt8) -> UInt8 {
    assert(byte < 64)
    var decodedStart: UInt8 = 0
    for range in base64ByteMappings {
        let decodedRange = decodedStart ..< decodedStart + (range.upperBound - range.lowerBound)
        if decodedRange.contains(byte) {
            return range.lowerBound + (byte - decodedStart)
        }
        decodedStart += range.upperBound - range.lowerBound
    }
    return 0
}

/**
 This method decodes Base64-encoded data.
 
 If the input contains any bytes that are not valid Base64 characters,
 this will return nil.
 
 - parameter bytes:      The Base64 bytes
 - parameter options:    Options for handling invalid input
 - returns:              The decoded bytes.
 */
private func base64DecodeBytes<Input, Output>(_ bytes: Input, options: Data.Base64DecodingOptions = [], into output: inout Output) -> Bool
  where Input: Collection, Input.Element == UInt8, Output: RangeReplaceableCollection, Output.Element == UInt8 {
    let endIdx = output.endIndex
    output.reserveCapacity(output.count + (bytes.count/3)*2)
    
    var currentByte : UInt8 = 0
    var validCharacterCount = 0
    var paddingCount = 0
    var index = 0
    
    func removePartiallyDecodedBytes() -> Bool {
      output.removeSubrange(endIdx..<output.endIndex)
      return false
    }
    
    for base64Char in bytes {
        
        let value : UInt8
        
        switch base64DecodeByte(base64Char) {
        case .valid(let v):
            value = v
            validCharacterCount += 1
        case .invalid:
            if options.contains(.ignoreUnknownCharacters) {
                continue
            } else {
              return removePartiallyDecodedBytes()
            }
        case .padding:
            paddingCount += 1
            continue
        }
        
        //padding found in the middle of the sequence is invalid
        if paddingCount > 0 {
            return removePartiallyDecodedBytes()
        }
        
        switch index%4 {
        case 0:
            currentByte = (value << 2)
        case 1:
            currentByte |= (value >> 4)
            output.append(currentByte)
            currentByte = (value << 4)
        case 2:
            currentByte |= (value >> 2)
            output.append(currentByte)
            currentByte = (value << 6)
        case 3:
            currentByte |= value
            output.append(currentByte)
        default:
            fatalError()
        }
        
        index += 1
    }
    
    guard (validCharacterCount + paddingCount)%4 == 0 else {
        //invalid character count
        return removePartiallyDecodedBytes()
    }
    return true
}

/**
This method encodes data in Base64.

- parameter bytes:      The bytes you want to encode
- parameter options:    Options for formatting the result
- returns:              The Base64-encoding for those bytes.
*/
private func base64EncodeBytes<Input, Output>(_ bytes: Input, options: Data.Base64EncodingOptions = [], into output: inout Output)
  where Input: Collection, Input.Element == UInt8, Output: RangeReplaceableCollection, Output.Element == UInt8 {
    output.reserveCapacity(output.count + (bytes.count/3)*4)
    
    let lineOptions : (lineLength : Int, separator : [UInt8])? = {
        let lineLength: Int
        
        if options.contains(.lineLength64Characters) { lineLength = 64 }
        else if options.contains(.lineLength76Characters) { lineLength = 76 }
        else {
            return nil
        }
        
        var separator = [UInt8]()
        if options.contains(.endLineWithCarriageReturn) { separator.append(13) }
        if options.contains(.endLineWithLineFeed) { separator.append(10) }
        
        //if the kind of line ending to insert is not specified, the default line ending is Carriage Return + Line Feed.
        if separator.isEmpty { separator = [13,10] }
        
        return (lineLength,separator)
    }()
    
    var currentLineCount = 0
    func appendByteToResult(_ byte: UInt8) -> Void {
        output.append(byte)
        currentLineCount += 1
        if let options = lineOptions, currentLineCount == options.lineLength {
            output.append(contentsOf: options.separator)
            currentLineCount = 0
        }
    }
    
    var currentByte : UInt8 = 0
    
    for (index,value) in bytes.enumerated() {
        switch index%3 {
        case 0:
            currentByte = (value >> 2)
            appendByteToResult(base64EncodeByte(currentByte))
            currentByte = ((value << 6) >> 2)
        case 1:
            currentByte |= (value >> 4)
            appendByteToResult(base64EncodeByte(currentByte))
            currentByte = ((value << 4) >> 2)
        case 2:
            currentByte |= (value >> 6)
            appendByteToResult(base64EncodeByte(currentByte))
            currentByte = ((value << 2) >> 2)
            appendByteToResult(base64EncodeByte(currentByte))
        default:
            fatalError()
        }
    }
    //add padding
    switch bytes.count%3 {
    case 0: break //no padding needed
    case 1:
        appendByteToResult(base64EncodeByte(currentByte))
        appendByteToResult(base64Padding)
        appendByteToResult(base64Padding)
    case 2:
        appendByteToResult(base64EncodeByte(currentByte))
        appendByteToResult(base64Padding)
    default:
        fatalError()
    }
}

// - Encoding to base64.

extension RangeReplaceableCollection where Element == UInt8 {
  
  /// Encodes the given collection as base64 and appends the encoded contents to this collection.
  public mutating func appendBase64EncodedBytes<C>(_ bytes: C, options: Data.Base64EncodingOptions = [])
    where C: Collection, C.Element == UInt8 {
      bytes.withContiguousStorageIfAvailable { base64EncodeBytes($0, options: options, into: &self) }
        ?? base64EncodeBytes(bytes, options: options, into: &self)
  }
  
  /// Initialises this colleciton with the contents of the given collection, encoded as base64.
  public init<C>(encodingToBase64 bytes: C, options: Data.Base64EncodingOptions = [])
    where C: Collection, C.Element == UInt8 {
      self.init()
      self.appendBase64EncodedBytes(bytes, options: options)
  }
}

// Encoding to common collections (existing Foundation API).

extension Collection where Element == UInt8 {
  
  public func base64EncodedData(options: Data.Base64EncodingOptions = []) -> Data {
    return Data(encodingToBase64: self, options: options)
  }
  
  public func base64EncodedString(options: Data.Base64EncodingOptions = []) -> String {
    // String's UTF8View is not a RangeReplaceableCollection, so we need an intermediate Array.
    let array = Array(encodingToBase64: self, options: options)
    return String(bytes: array, encoding: .utf8)!
  }
}

// - Decoding from base64.

extension RangeReplaceableCollection where Element == UInt8 {
  
  /// Decodes the given collection from base64 and appends the decoded contents to this collection.
  public mutating func appendBase64DecodedBytes<C>(_ bytes: C, options: Data.Base64DecodingOptions = []) -> Bool
    where C: Collection, C.Element == UInt8 {
      return bytes.withContiguousStorageIfAvailable { base64DecodeBytes($0, options: options, into: &self) } ??
               base64DecodeBytes(bytes, options: options, into: &self)
  }
  
  /// Initialises this collection with the contents of the given collection, once it has been decoded from base64.
  public init?<C>(decodingFromBase64 bytes: C, options: Data.Base64DecodingOptions = [])
    where C: Collection, C.Element == UInt8 {
      self.init()
      guard self.appendBase64DecodedBytes(bytes, options: options) else { return nil }
  }
}

// Decoding common base64-encoded collections (existing Foundation API).

extension RangeReplaceableCollection where Element == UInt8 {
  
  /// Initializes a data object with the given Base64 encoded string.
  public init?(base64Encoded base64String: String, options: Data.Base64DecodingOptions = []) {
    self.init(decodingFromBase64: base64String.utf8, options: options)
  }
  
  /// Initializes a data object with the given Base64 encoded data.
  public init?(base64Encoded base64Data: Data, options: Data.Base64DecodingOptions = []) {
    self.init(decodingFromBase64: base64Data, options: options)
  }
}

// Decoding as a String.

extension String {
  
  public init?<C>(decodingFromBase64 bytes: C, options: Data.Base64DecodingOptions = [])
    where C: Collection, C.Element == UInt8 {
      // String's UTF8View is not a RangeReplaceableCollection, so we need an intermediate Array.
      guard let array = Array(decodingFromBase64: bytes, options: options) else { return nil }
      self.init(bytes: array, encoding: .utf8)
  }
  
  public init?<C>(decodingFromBase64 string: C, options: Data.Base64DecodingOptions = [])
    where C: StringProtocol {
      self.init(decodingFromBase64: string.utf8, options: options)
  }
}


func test() {
  let str = "Hello, this is an encoded String! 🏆"

  // Round-trip String.
  let enc_str = str.utf8.base64EncodedString()
  print(enc_str)
  let dec_str = String(decodingFromBase64: enc_str)
  print(dec_str)

  print("")
  print(repeatElement("=", count: 20).joined())
  print("")

  // Invalid bytes should fail without leaving partially-decoded content around.
  var fail_str = enc_str
  fail_str.append("🤷‍♂️")
  var dec_fail = [UInt8]()
  dec_fail.appendBase64DecodedBytes(fail_str.utf8)
  print(dec_fail)

  print("")
  print(repeatElement("=", count: 20).joined())
  print("")

  // Direct encoding to-/decoding from- an Array.
  let enc_arr = Array(encodingToBase64: str.utf8)
  print(enc_arr)
  let dec_arr = Array(decodingFromBase64: enc_arr)!
  print(dec_arr)
  let dec_arr_str = String(bytes: dec_arr, encoding: .utf8)!
  print(dec_arr_str)
}
test()

And yes, I'm sure the names could be better/clearer.