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

…ddleware
hummingbird-project · Jan 15, 2025 · c69b929 · c69b929
1 parent f377115
commit c69b929
Show file tree

Hide file tree

Showing 2 changed files with 75 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 additionalMediaTypeExtensions: [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,
+            additionalMediaTypeExtensions: [:]
         )
     }
 
@@ -90,11 +94,40 @@ where Provider.FileAttributes: FileMiddlewareFileAttributes {
         urlBasePath: String? = nil,
         cacheControl: CacheControl = .init([]),
         searchForIndexHtml: Bool = false
+    ) {
+        self.init(
+            fileProvider: fileProvider,
+            urlBasePath: urlBasePath,
+            cacheControl: cacheControl,
+            searchForIndexHtml: searchForIndexHtml,
+            additionalMediaTypeExtensions: [:]
+        )
+    }
+
+    private init(
+        fileProvider: Provider,
+        urlBasePath: String? = nil,
+        cacheControl: CacheControl = .init([]),
+        searchForIndexHtml: Bool = false,
+        additionalMediaTypeExtensions: [String: MediaType]
     ) {
         self.cacheControl = cacheControl
         self.searchForIndexHtml = searchForIndexHtml
         self.urlBasePath = urlBasePath.map { String($0.dropSuffix("/")) }
         self.fileProvider = fileProvider
+        self.additionalMediaTypeExtensions = additionalMediaTypeExtensions
+    }
+
+    public func withAdditionalMediaType(_ mediaType: MediaType, forFileExtension fileExtension: String) -> FileMiddleware {
+        var extensions = additionalMediaTypeExtensions
+        extensions[fileExtension] = mediaType
+        return FileMiddleware(
+            fileProvider: fileProvider,
+            urlBasePath: urlBasePath,
+            cacheControl: cacheControl,
+            searchForIndexHtml: searchForIndexHtml,
+            additionalMediaTypeExtensions: extensions
+        )
     }
 
     /// Handle request
@@ -220,7 +253,7 @@ extension FileMiddleware {
 
         // content-type
         if let ext = self.fileExtension(for: path) {
-            if let contentType = MediaType.getMediaType(forExtension: ext) {
+            if let contentType = additionalMediaTypeExtensions[ext] ?? MediaType.getMediaType(forExtension: ext) {
                 headers[.contentType] = contentType.description
             }
         }

diff --git a/Tests/HummingbirdTests/FileMiddlewareTests.swift b/Tests/HummingbirdTests/FileMiddlewareTests.swift
@@ -512,4 +512,38 @@ 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, forFileExtension: "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))
+            }
+        }
+    }
 }