Fun with String Interpolation — For more information read my article at https://oleb.net/blog/2017/01/fun-with-string-interpolation/. — Dependencies: Foundation

fun-with-string-interpolation.swift

/// An unescaped string from a potentially unsafe
/// source (such as user input)
struct UnsafeString {
    var value: String
}

/// A string that either comes from a safe source
/// (e.g. a string literal in the source code)
/// or has been escaped.
struct SanitizedHTML {
    fileprivate(set) var value: String
    // Required for string interpolation processing
    fileprivate var interpolationSegment: Any? = nil

    init(unsafe input: UnsafeString) {
        value = SanitizedHTML.escaping(unsafe: input.value)
    }
}

import Foundation // required for String.replacingOccurrences(of:with:)

extension SanitizedHTML {
    /// Escapes a string.
    fileprivate static func escaping(unsafe input: String) -> String {
        return input
            .replacingOccurrences(of: "<", with: "&lt;")
            .replacingOccurrences(of: ">", with: "&gt;")
    }
}

extension SanitizedHTML {
    mutating func append(_ other: SanitizedHTML) {
        value.append(other.value)
    }

    mutating func append(_ other: UnsafeString) {
        let sanitized = SanitizedHTML(unsafe: other)
        append(sanitized)
    }
}

// Initialization with a string literal should not escape the input.
extension SanitizedHTML: ExpressibleByStringLiteral {
    init(stringLiteral value: String) {
        self.value = value
    }
    init(unicodeScalarLiteral value: String) {
        self.init(stringLiteral: value)
    }
    init(extendedGraphemeClusterLiteral value: String) {
        self.init(stringLiteral: value)
    }
}

extension UnsafeString: ExpressibleByStringLiteral {
    init(stringLiteral value: String) {
        self.value = value
    }
    init(unicodeScalarLiteral value: String) {
        self.init(stringLiteral: value)
    }
    init(extendedGraphemeClusterLiteral value: String) {
        self.init(stringLiteral: value)
    }
}

let userInput: UnsafeString = "<script>alert('P0wn3d');</script>"
var sanitized: SanitizedHTML = "<strong>Name:</strong> "
sanitized.append(userInput)
sanitized.value

//error: cannot convert value of type 'String' to specified type 'SanitizedHTML'
//let sanitized2: SanitizedHTML = "<strong>Name:</strong> \(userInput)"

// Initialization via string interpolation: variables in the interpolation string should be escaped, but the string literal parts shouldn't be. The current string interpolation API is very weird (see https://twitter.com/jckarter/status/815054320106283008 ). We need a helper property `stringInterpolationSegment` to manage it (see https://twitter.com/jtbandes/status/815096781872607233).
extension SanitizedHTML: ExpressibleByStringInterpolation {
    // Step 1
    public init<T>(stringInterpolationSegment expr: T) {
        // Store the segment
        interpolationSegment = expr
        // Dummy initialization, this is never used
        value = ""
    }

    // Step 2
    public init(stringInterpolation segments: SanitizedHTML...) {
        let stringSegments = segments.enumerated()
            .map { index, segment -> String in
                guard let segment = segment.interpolationSegment else {
                    fatalError("Invalid interpolation sequence")
                }
                if index % 2 == 0 {
                    // Even indices are literal segments
                    // and thus already safe.
                    if let string = segment as? String {
                        return string
                    } else {
                        return String(describing: segment)
                    }
                } else {
                    // Odd indices are variable expressions
                    switch segment {
                    case let safe as SanitizedHTML:
                        // Already safe
                        return safe.value
                    case let unsafe as UnsafeString:
                        return SanitizedHTML.escaping(unsafe: unsafe.value)
                    default:
                        // All other types are treated as unsafe too.
                        let unsafe = UnsafeString(value: String(describing: segment))
                        return SanitizedHTML(unsafe: unsafe).value
                    }
                }
        }
        value = stringSegments.joined()
    }
}

let sanitized2: SanitizedHTML = "<strong>Name:</strong> \(userInput)"
sanitized2.value


// MARK: - Convenience conformances

extension UnsafeString: CustomStringConvertible {
    var description: String { return String(describing: value) }
}

extension UnsafeString: CustomDebugStringConvertible {
    var debugDescription: String { return "UnsafeString \(String(reflecting: value))" }
}

extension SanitizedHTML: CustomStringConvertible {
    var description: String {
        if let segment = interpolationSegment {
            return "Segment: \"\(segment)\""
        } else {
            return String(describing: value)
        }
    }
}
extension SanitizedHTML: CustomDebugStringConvertible {
    var debugDescription: String {
        if let segment = interpolationSegment {
            return "<SanitizedHTML> segment: \"\(segment)\""
        } else {
            return "<SanitizedHTML> \(String(reflecting: value))"
        }
    }
}