Merge pull request swiftlang#33698 from atrick/fix-smallstring

Fix undefined behavior in SmallString.withUTF8

Author: atrick

lib/SILOptimizer/Utils/SILInliner.cpp

@@ -688,6 +688,7 @@ InlineCost swift::instructionInlineCost(SILInstruction &I) {
   case SILInstructionKind::BaseAddrForOffsetInst:
   case SILInstructionKind::EndLifetimeInst:
   case SILInstructionKind::UncheckedOwnershipConversionInst:
+  case SILInstructionKind::BindMemoryInst:
     return InlineCost::Free;
 
   // Typed GEPs are free.
@@ -798,7 +799,6 @@ InlineCost swift::instructionInlineCost(SILInstruction &I) {
   case SILInstructionKind::AllocRefDynamicInst:
   case SILInstructionKind::AllocStackInst:
   case SILInstructionKind::AllocValueBufferInst:
-  case SILInstructionKind::BindMemoryInst:
   case SILInstructionKind::BeginApplyInst:
   case SILInstructionKind::ValueMetatypeInst:
   case SILInstructionKind::WitnessMethodInst:

stdlib/public/core/SmallString.swift

@@ -197,26 +197,30 @@ extension _SmallString {
   internal func withUTF8<Result>(
     _ f: (UnsafeBufferPointer<UInt8>) throws -> Result
   ) rethrows -> Result {
+    let count = self.count
     var raw = self.zeroTerminatedRawCodeUnits
-    return try Swift.withUnsafeBytes(of: &raw) { rawBufPtr in
-      let ptr = rawBufPtr.baseAddress._unsafelyUnwrappedUnchecked
-        .assumingMemoryBound(to: UInt8.self)
-      return try f(UnsafeBufferPointer(start: ptr, count: self.count))
+    return try Swift.withUnsafeBytes(of: &raw) {
+      let rawPtr = $0.baseAddress._unsafelyUnwrappedUnchecked
+      // Rebind the underlying (UInt64, UInt64) tuple to UInt8 for the
+      // duration of the closure. Accessing self after this rebind is undefined.
+      let ptr = rawPtr.bindMemory(to: UInt8.self, capacity: count)
+      defer {
+        // Restore the memory type of self._storage
+        _ = rawPtr.bindMemory(to: RawBitPattern.self, capacity: 1)
+      }
+      return try f(UnsafeBufferPointer(start: ptr, count: count))
     }
   }
 
   // Overwrite stored code units, including uninitialized. `f` should return the
   // new count.
   @inline(__always)
   internal mutating func withMutableCapacity(
-    _ f: (UnsafeMutableBufferPointer<UInt8>) throws -> Int
+    _ f: (UnsafeMutableRawBufferPointer) throws -> Int
   ) rethrows {
     let len = try withUnsafeMutableBytes(of: &self._storage) {
       (rawBufPtr: UnsafeMutableRawBufferPointer) -> Int in
-      let ptr = rawBufPtr.baseAddress._unsafelyUnwrappedUnchecked
-        .assumingMemoryBound(to: UInt8.self)
-      return try f(UnsafeMutableBufferPointer(
-        start: ptr, count: _SmallString.capacity))
+      return try f(rawBufPtr)
     }
     if len == 0 {
       self = _SmallString()
@@ -273,7 +277,17 @@ extension _SmallString {
   ) rethrows {
     self.init()
     try self.withMutableCapacity {
-      return try initializer($0)
+      let capacity = $0.count
+      let rawPtr = $0.baseAddress._unsafelyUnwrappedUnchecked
+      // Rebind the underlying (UInt64, UInt64) tuple to UInt8 for the
+      // duration of the closure. Accessing self after this rebind is undefined.
+      let ptr = rawPtr.bindMemory(to: UInt8.self, capacity: capacity)
+      defer {
+        // Restore the memory type of self._storage
+        _ = rawPtr.bindMemory(to: RawBitPattern.self, capacity: 1)
+      }
+      return try initializer(
+        UnsafeMutableBufferPointer<UInt8>(start: ptr, count: capacity))
     }
     self._invariantCheck()
   }

stdlib/public/core/StringBridge.swift

@@ -169,7 +169,7 @@ internal func _cocoaStringSubscript(
 @_effects(releasenone)
 private func _NSStringCopyUTF8(
   _ o: _StringSelectorHolder,
-  into bufPtr: UnsafeMutableBufferPointer<UInt8>
+  into bufPtr: UnsafeMutableRawBufferPointer
 ) -> Int? {
   let ptr = bufPtr.baseAddress._unsafelyUnwrappedUnchecked
   let len = o.length
@@ -193,7 +193,7 @@ private func _NSStringCopyUTF8(
 @_effects(releasenone)
 internal func _cocoaStringCopyUTF8(
   _ target: _CocoaString,
-  into bufPtr: UnsafeMutableBufferPointer<UInt8>
+  into bufPtr: UnsafeMutableRawBufferPointer
 ) -> Int? {
   return _NSStringCopyUTF8(_objc(target), into: bufPtr)
 }
@@ -206,7 +206,7 @@ private func _NSStringUTF8Count(
   var remainingRange = _SwiftNSRange(location: 0, length: 0)
   var usedLen = 0
   let success = 0 != o.getBytes(
-    UnsafeMutablePointer<UInt8>(Builtin.inttoptr_Word(0._builtinWordValue)),
+    UnsafeMutableRawPointer(Builtin.inttoptr_Word(0._builtinWordValue)),
     maxLength: 0,
     usedLength: &usedLen,
     encoding: _cocoaUTF8Encoding,
@@ -340,7 +340,7 @@ internal enum _KnownCocoaString {
 @_effects(releasenone) // @opaque
 internal func _bridgeTagged(
   _ cocoa: _CocoaString,
-  intoUTF8 bufPtr: UnsafeMutableBufferPointer<UInt8>
+  intoUTF8 bufPtr: UnsafeMutableRawBufferPointer
 ) -> Int? {
   _internalInvariant(_isObjCTaggedPointer(cocoa))
   return _cocoaStringCopyUTF8(cocoa, into: bufPtr)

stdlib/public/core/StringGuts.swift

@@ -250,7 +250,9 @@ extension _StringGuts {
   ) -> Int? {
     #if _runtime(_ObjC)
     // Currently, foreign  means NSString
-    if let res = _cocoaStringCopyUTF8(_object.cocoaObject, into: mbp) {
+    if let res = _cocoaStringCopyUTF8(_object.cocoaObject,
+      into: UnsafeMutableRawBufferPointer(start: mbp.baseAddress,
+                                          count: mbp.count)) {
       return res
     }