Add ability for developers to support additional MIME types in FileMi…

…ddleware (#657) * Add ability for developers to support additional MIME types in FileMiddleware * rename mediaTypeFileExtensionMap for clarity * make mediaTypeFileExtensionMap public * unmake mediaTypeFileExtensionMap public * can add multiple mime types to filemiddleware in one call * swift-formatted * Fix formatting --------- Co-authored-by: Joannis Orlandos <joannis@orlandos.nl> Co-authored-by: Adam Fowler <adamfowler71@gmail.com>
hummingbird-project · Jan 24, 2025 · 5d7a747 · 5d7a747
1 parent adee390
commit 5d7a747
Show file tree

Hide file tree

Showing 2 changed files with 123 additions and 8 deletions.
diff --git a/Sources/Hummingbird/Middleware/FileMiddleware.swift b/Sources/Hummingbird/Middleware/FileMiddleware.swift
@@ -52,6 +52,7 @@ where Provider.FileAttributes: FileMiddlewareFileAttributes {
     let searchForIndexHtml: Bool
     let urlBasePath: String?
     let fileProvider: Provider
+    let mediaTypeFileExtensionMap: [String: MediaType]
 
     /// Create FileMiddleware
     /// - Parameters:
@@ -69,13 +70,16 @@ where Provider.FileAttributes: FileMiddlewareFileAttributes {
         threadPool: NIOThreadPool = NIOThreadPool.singleton,
         logger: Logger = Logger(label: "FileMiddleware")
     ) where Provider == LocalFileSystem {
-        self.cacheControl = cacheControl
-        self.searchForIndexHtml = searchForIndexHtml
-        self.urlBasePath = urlBasePath.map { String($0.dropSuffix("/")) }
-        self.fileProvider = LocalFileSystem(
-            rootFolder: rootFolder,
-            threadPool: threadPool,
-            logger: logger
+        self.init(
+            fileProvider: LocalFileSystem(
+                rootFolder: rootFolder,
+                threadPool: threadPool,
+                logger: logger
+            ),
+            urlBasePath: urlBasePath,
+            cacheControl: cacheControl,
+            searchForIndexHtml: searchForIndexHtml,
+            mediaTypeFileExtensionMap: [:]
         )
     }
 
@@ -90,11 +94,48 @@ where Provider.FileAttributes: FileMiddlewareFileAttributes {
         urlBasePath: String? = nil,
         cacheControl: CacheControl = .init([]),
         searchForIndexHtml: Bool = false
+    ) {
+        self.init(
+            fileProvider: fileProvider,
+            urlBasePath: urlBasePath,
+            cacheControl: cacheControl,
+            searchForIndexHtml: searchForIndexHtml,
+            mediaTypeFileExtensionMap: [:]
+        )
+    }
+
+    private init(
+        fileProvider: Provider,
+        urlBasePath: String? = nil,
+        cacheControl: CacheControl = .init([]),
+        searchForIndexHtml: Bool = false,
+        mediaTypeFileExtensionMap: [String: MediaType]
     ) {
         self.cacheControl = cacheControl
         self.searchForIndexHtml = searchForIndexHtml
         self.urlBasePath = urlBasePath.map { String($0.dropSuffix("/")) }
         self.fileProvider = fileProvider
+        self.mediaTypeFileExtensionMap = mediaTypeFileExtensionMap
+    }
+
+    public func withAdditionalMediaType(_ mediaType: MediaType, mappedToFileExtension fileExtension: String) -> FileMiddleware {
+        withAdditionalMediaTypes(forFileExtensions: [fileExtension: mediaType])
+    }
+
+    public func withAdditionalMediaTypes(forFileExtensions extensionToMediaTypeMap: [String: MediaType]) -> FileMiddleware {
+        let extensions = extensionToMediaTypeMap.reduce(
+            into: mediaTypeFileExtensionMap
+        ) {
+            $0[$1.key.lowercased()] = $1.value
+        }
+
+        return FileMiddleware(
+            fileProvider: fileProvider,
+            urlBasePath: urlBasePath,
+            cacheControl: cacheControl,
+            searchForIndexHtml: searchForIndexHtml,
+            mediaTypeFileExtensionMap: extensions
+        )
     }
 
     /// Handle request
@@ -220,7 +261,7 @@ extension FileMiddleware {
 
         // content-type
         if let ext = self.fileExtension(for: path) {
-            if let contentType = MediaType.getMediaType(forExtension: ext) {
+            if let contentType = mediaTypeFileExtensionMap[ext] ?? MediaType.getMediaType(forExtension: ext) {
                 headers[.contentType] = contentType.description
             }
         }

diff --git a/Tests/HummingbirdTests/FileMiddlewareTests.swift b/Tests/HummingbirdTests/FileMiddlewareTests.swift
@@ -512,4 +512,78 @@ final class FileMiddlewareTests: XCTestCase {
             }
         }
     }
+
+    func testCustomMIMEType() async throws {
+        let hlsStream = try XCTUnwrap(MediaType(from: "application/x-mpegURL"))
+        let router = Router()
+        router.middlewares.add(FileMiddleware(".").withAdditionalMediaType(hlsStream, mappedToFileExtension: "m3u8"))
+        let app = Application(responder: router.buildResponder())
+
+        let filename = "\(#function).m3u8"
+        let content = """
+            #EXTM3U
+            #EXT-X-VERSION:7
+            #EXT-X-ALLOW-CACHE:YES
+            #EXT-X-TARGETDURATION:0
+            #EXT-X-MEDIA-SEQUENCE:10
+            #EXT-X-PLAYLIST-TYPE:EVENT
+            #EXT-X-MAP:URI="init.mp4"
+            #EXT-X-DISCONTINUITY
+            #EXTINF:0.000000,
+            live000010.m4s
+            """
+        let data = Data(content.utf8)
+        let fileURL = URL(fileURLWithPath: filename)
+        XCTAssertNoThrow(try data.write(to: fileURL))
+        defer { XCTAssertNoThrow(try FileManager.default.removeItem(at: fileURL)) }
+
+        try await app.test(.router) { client in
+            try await client.execute(uri: filename, method: .get) { response in
+                XCTAssertEqual(String(buffer: response.body), content)
+                let contentType = try XCTUnwrap(response.headers[.contentType])
+                let validTypes = Set(["application/vnd.apple.mpegurl", "application/x-mpegurl"])
+                XCTAssert(validTypes.contains(contentType))
+            }
+        }
+    }
+
+    func testCustomMIMETypes() async throws {
+        let hlsStream = try XCTUnwrap(MediaType(from: "application/x-mpegURL"))
+        let router = Router()
+        let fileMiddleware = FileMiddleware<BasicRequestContext, LocalFileSystem>(".")
+            .withAdditionalMediaType(hlsStream, mappedToFileExtension: "m3u8")
+            .withAdditionalMediaTypes(forFileExtensions: [
+                "foo": MediaType(type: .any, subType: "x-foo"),
+                "M3U8": MediaType(type: .application, subType: "vnd.apple.mpegURL"),
+            ])
+        router.middlewares.add(fileMiddleware)
+        let app = Application(responder: router.buildResponder())
+
+        let filename = "\(#function).m3u8"
+        let content = """
+            #EXTM3U
+            #EXT-X-VERSION:7
+            #EXT-X-ALLOW-CACHE:YES
+            #EXT-X-TARGETDURATION:0
+            #EXT-X-MEDIA-SEQUENCE:10
+            #EXT-X-PLAYLIST-TYPE:EVENT
+            #EXT-X-MAP:URI="init.mp4"
+            #EXT-X-DISCONTINUITY
+            #EXTINF:0.000000,
+            live000010.m4s
+            """
+        let data = Data(content.utf8)
+        let fileURL = URL(fileURLWithPath: filename)
+        XCTAssertNoThrow(try data.write(to: fileURL))
+        defer { XCTAssertNoThrow(try FileManager.default.removeItem(at: fileURL)) }
+
+        try await app.test(.router) { client in
+            try await client.execute(uri: filename, method: .get) { response in
+                XCTAssertEqual(String(buffer: response.body), content)
+                let contentType = try XCTUnwrap(response.headers[.contentType])
+                let validTypes = Set(["application/vnd.apple.mpegurl", "application/vnd.apple.mpegURL"])
+                XCTAssert(validTypes.contains(contentType))
+            }
+        }
+    }
 }