Add FileIO.writeFile (#44)

* FileIO can write `HBRequestBody` now

* Use new stream.consumeAll(on:_)

* Add HBRequest.Context and use in FileIO

HBRequest.Context is the context which a request is being processed in. It includes eventLoop, logger, and allocate references

Sources/Hummingbird/Middleware/Middleware.swift

@@ -1,7 +1,8 @@
 import NIO
 
-/// Applied to request before it is dealt with by the router. Middleware passes the processed request onto the next responder
-/// by calling `next.apply(to: request)`. If you want to shortcut the request you can return a response immediately
+/// Applied to `HBRequest` before it is dealt with by the router. Middleware passes the processed request onto the next responder
+/// (either the next middleware or the router) by calling `next.apply(to: request)`. If you want to shortcut the request you
+/// can return a response immediately
 ///
 /// Middleware is added to the application by calling `app.middleware.add(MyMiddleware()`.
 ///

Sources/Hummingbird/Server/Request.swift

@@ -91,6 +91,21 @@ public final class HBRequest: HBExtensible {
         return self.eventLoop.makeSucceededFuture(value)
     }
 
+    /// Return context request is running in
+    public var context: Context {
+        .init(logger: self.logger, eventLoop: self.eventLoop, allocator: self.allocator)
+    }
+
+    /// Context request is running in
+    public struct Context {
+        /// Logger to use
+        public var logger: Logger
+        /// EventLoop request is running on
+        public var eventLoop: EventLoop
+        /// ByteBuffer allocator used by request
+        public var allocator: ByteBufferAllocator
+    }
+
     private static func loggerWithRequestId(_ logger: Logger, uri: String, method: String) -> Logger {
         var logger = logger
         logger[metadataKey: "hb_id"] = .string(String(describing: Self.globalRequestID.add(1)))

Sources/Hummingbird/Server/Responder.swift

@@ -15,7 +15,7 @@ public struct HBCallbackResponder: HBResponder {
     public init(callback: @escaping (HBRequest) -> EventLoopFuture<HBResponse>) {
         self.callback = callback
     }
-    
+
     /// Return EventLoopFuture that will be fulfilled with response to the request supplied
     public func respond(to request: HBRequest) -> EventLoopFuture<HBResponse> {
         return self.callback(request)

Sources/HummingbirdFoundation/Files/FileIO.swift

@@ -1,4 +1,5 @@
 import Hummingbird
+import Logging
 import NIO
 
 /// Manages File loading. Can either stream or load file.
@@ -18,13 +19,13 @@ public struct HBFileIO {
     ///   - request: request for file
     ///   - path: System file path
     /// - Returns: Response including file details
-    public func headFile(for request: HBRequest, path: String) -> EventLoopFuture<HBResponse> {
-        return self.fileIO.openFile(path: path, eventLoop: request.eventLoop).flatMap { handle, region in
-            request.logger.debug("[FileIO] HEAD", metadata: ["file": .string(path)])
+    public func headFile(path: String, context: HBRequest.Context) -> EventLoopFuture<HBResponse> {
+        return self.fileIO.openFile(path: path, eventLoop: context.eventLoop).flatMap { handle, region in
+            context.logger.debug("[FileIO] HEAD", metadata: ["file": .string(path)])
             let headers: HTTPHeaders = ["content-length": region.readableBytes.description]
             let response = HBResponse(status: .ok, headers: headers, body: .empty)
             try? handle.close()
-            return request.eventLoop.makeSucceededFuture(response)
+            return context.eventLoop.makeSucceededFuture(response)
         }.flatMapErrorThrowing { _ in
             throw HBHTTPError(.notFound)
         }
@@ -38,40 +39,81 @@ public struct HBFileIO {
     ///   - request: request for file
     ///   - path: System file path
     /// - Returns: Response include file
-    public func loadFile(for request: HBRequest, path: String) -> EventLoopFuture<HBResponse> {
-        return self.fileIO.openFile(path: path, eventLoop: request.eventLoop).flatMap { handle, region in
-            request.logger.debug("[FileIO] GET", metadata: ["file": .string(path)])
-            let futureResponse: EventLoopFuture<HBResponse>
+    public func loadFile(path: String, context: HBRequest.Context) -> EventLoopFuture<HBResponseBody> {
+        return self.fileIO.openFile(path: path, eventLoop: context.eventLoop).flatMap { handle, region in
+            context.logger.debug("[FileIO] GET", metadata: ["file": .string(path)])
+            let futureResult: EventLoopFuture<HBResponseBody>
             if region.readableBytes > self.chunkSize {
-                futureResponse = streamFile(for: request, handle: handle, region: region)
+                futureResult = streamFile(handle: handle, region: region, context: context)
             } else {
-                futureResponse = loadFile(for: request, handle: handle, region: region)
+                futureResult = loadFile(handle: handle, region: region, context: context)
+                // only close file handle for load, as streamer hasn't loaded data at this point
+                futureResult.whenComplete { _ in
+                    try? handle.close()
+                }
             }
-            return futureResponse
+            return futureResult
         }.flatMapErrorThrowing { _ in
             throw HBHTTPError(.notFound)
         }
     }
 
-    func loadFile(for request: HBRequest, handle: NIOFileHandle, region: FileRegion) -> EventLoopFuture<HBResponse> {
-        return self.fileIO.read(fileHandle: handle, byteCount: region.readableBytes, allocator: request.allocator, eventLoop: request.eventLoop).map { buffer in
-            return HBResponse(status: .ok, headers: [:], body: .byteBuffer(buffer))
+    /// Write contents of request body to file
+    ///
+    /// This can be used to save arbitrary ByteBuffers by passing in `.byteBuffer(ByteBuffer)` as contents
+    /// - Parameters:
+    ///   - contents: Request body to write.
+    ///   - path: Path to write to
+    ///   - eventLoop: EventLoop everything runs on
+    ///   - logger: Logger
+    /// - Returns: EventLoopFuture fulfilled when everything is done
+    public func writeFile(contents: HBRequestBody, path: String, context: HBRequest.Context) -> EventLoopFuture<Void> {
+        return self.fileIO.openFile(path: path, mode: .write, flags: .allowFileCreation(), eventLoop: context.eventLoop).flatMap { handle in
+            context.logger.debug("[FileIO] PUT", metadata: ["file": .string(path)])
+            let futureResult: EventLoopFuture<Void>
+            switch contents {
+            case .byteBuffer(let buffer):
+                guard let buffer = buffer else { return context.eventLoop.makeSucceededFuture(()) }
+                futureResult = writeFile(buffer: buffer, handle: handle, on: context.eventLoop)
+            case .stream(let streamer):
+                futureResult = writeFile(stream: streamer, handle: handle, on: context.eventLoop)
+            }
+            futureResult.whenComplete { _ in
+                try? handle.close()
+            }
+            return futureResult
         }
-        .always { _ in
-            try? handle.close()
+    }
+
+    /// Load file as ByteBuffer
+    func loadFile(handle: NIOFileHandle, region: FileRegion, context: HBRequest.Context) -> EventLoopFuture<HBResponseBody> {
+        return self.fileIO.read(fileHandle: handle, byteCount: region.readableBytes, allocator: context.allocator, eventLoop: context.eventLoop).map { buffer in
+            return .byteBuffer(buffer)
         }
     }
 
-    func streamFile(for request: HBRequest, handle: NIOFileHandle, region: FileRegion) -> EventLoopFuture<HBResponse> {
+    /// Return streamer that will load file
+    func streamFile(handle: NIOFileHandle, region: FileRegion, context: HBRequest.Context) -> EventLoopFuture<HBResponseBody> {
         let fileStreamer = FileStreamer(
             handle: handle,
             fileSize: region.readableBytes,
             fileIO: self.fileIO,
             chunkSize: self.chunkSize,
-            allocator: request.allocator
+            allocator: context.allocator
         )
-        let response = HBResponse(status: .ok, headers: [:], body: .stream(fileStreamer))
-        return request.eventLoop.makeSucceededFuture(response)
+        return context.eventLoop.makeSucceededFuture(.stream(fileStreamer))
+    }
+
+    /// write byte buffer to file
+    func writeFile(buffer: ByteBuffer, handle: NIOFileHandle, on eventLoop: EventLoop) -> EventLoopFuture<Void> {
+        return self.fileIO.write(fileHandle: handle, buffer: buffer, eventLoop: eventLoop)
+    }
+
+    /// write output of streamer to file
+    func writeFile(stream: HBRequestBodyStreamer, handle: NIOFileHandle, on eventLoop: EventLoop) -> EventLoopFuture<Void> {
+        return stream.consumeAll(on: eventLoop) { buffer in
+            return self.fileIO.write(fileHandle: handle, buffer: buffer, eventLoop: eventLoop)
+        }
     }
 
     /// class used to stream files
@@ -96,7 +138,13 @@ public struct HBFileIO {
                 self.bytesLeft -= bytesToRead
                 return self.fileIO.read(fileHandle: self.handle, byteCount: bytesToRead, allocator: self.allocator, eventLoop: eventLoop)
                     .map { .byteBuffer($0) }
+                    .flatMapErrorThrowing { error in
+                        // close handle on error being returned
+                        try? self.handle.close()
+                        throw error
+                    }
             } else {
+                // close handle now streamer has finished
                 try? self.handle.close()
                 return eventLoop.makeSucceededFuture(.end)
             }

Sources/HummingbirdFoundation/Files/FileMiddleware.swift

@@ -52,10 +52,11 @@ public struct HBFileMiddleware: HBMiddleware {
 
             switch request.method {
             case .GET:
-                return fileIO.loadFile(for: request, path: fullPath)
+                return fileIO.loadFile(path: fullPath, context: request.context)
+                    .map { HBResponse(status: .ok, body: $0) }
 
             case .HEAD:
-                return fileIO.headFile(for: request, path: fullPath)
+                return fileIO.headFile(path: fullPath, context: request.context)
 
             default:
                 return request.eventLoop.makeFailedFuture(error)

Tests/HummingbirdFoundationTests/FilesTests.swift

@@ -6,7 +6,13 @@ import HummingbirdXCT
 import XCTest
 
 class HummingbirdFilesTests: XCTestCase {
-    func testGet() {
+    func randomBuffer(size: Int) -> ByteBuffer {
+        var data = [UInt8](repeating: 0, count: size)
+        data = data.map { _ in UInt8.random(in: 0...255) }
+        return ByteBufferAllocator().buffer(bytes: data)
+    }
+
+    func testRead() {
         let app = HBApplication(testing: .live)
         app.middleware.add(HBFileMiddleware(".", application: app))
 
@@ -25,6 +31,25 @@ class HummingbirdFilesTests: XCTestCase {
         }
     }
 
+    func testReadLargeFile() {
+        let app = HBApplication(testing: .live)
+        app.middleware.add(HBFileMiddleware(".", application: app))
+
+        let buffer = self.randomBuffer(size: 380_000)
+        let data = Data(buffer: buffer)
+        let fileURL = URL(fileURLWithPath: "test.txt")
+        XCTAssertNoThrow(try data.write(to: fileURL))
+        defer { XCTAssertNoThrow(try FileManager.default.removeItem(at: fileURL)) }
+
+        app.XCTStart()
+        defer { app.XCTStop() }
+
+        app.XCTExecute(uri: "/test.txt", method: .GET) { response in
+            let body = try XCTUnwrap(response.body)
+            XCTAssertEqual(body, buffer)
+        }
+    }
+
     func testHead() throws {
         let app = HBApplication(testing: .live)
         app.middleware.add(HBFileMiddleware(".", application: app))
@@ -43,4 +68,50 @@ class HummingbirdFilesTests: XCTestCase {
             XCTAssertEqual(response.headers["Content-Length"].first, text.utf8.count.description)
         }
     }
+
+    func testWrite() throws {
+        let filename = "testWrite.txt"
+        let app = HBApplication(testing: .live)
+        app.router.put("store") { request -> EventLoopFuture<HTTPResponseStatus> in
+            let fileIO = HBFileIO(application: request.application)
+            return fileIO.writeFile(contents: request.body, path: filename, context: request.context)
+                .map { .ok }
+        }
+
+        app.XCTStart()
+        defer { app.XCTStop() }
+
+        let buffer = ByteBufferAllocator().buffer(string: "This is a test")
+        app.XCTExecute(uri: "/store", method: .PUT, body: buffer) { response in
+            XCTAssertEqual(response.status, .ok)
+        }
+
+        let fileURL = URL(fileURLWithPath: filename)
+        let data = try Data(contentsOf: fileURL)
+        defer { XCTAssertNoThrow(try FileManager.default.removeItem(at: fileURL)) }
+        XCTAssertEqual(String(decoding: data, as: Unicode.UTF8.self), "This is a test")
+    }
+
+    func testWriteLargeFile() throws {
+        let filename = "testWriteLargeFile.txt"
+        let app = HBApplication(testing: .live)
+        app.router.put("store") { request -> EventLoopFuture<HTTPResponseStatus> in
+            let fileIO = HBFileIO(application: request.application)
+            return fileIO.writeFile(contents: request.body, path: filename, context: request.context)
+                .map { .ok }
+        }
+
+        app.XCTStart()
+        defer { app.XCTStop() }
+
+        let buffer = self.randomBuffer(size: 400_000)
+        app.XCTExecute(uri: "/store", method: .PUT, body: buffer) { response in
+            XCTAssertEqual(response.status, .ok)
+        }
+
+        let fileURL = URL(fileURLWithPath: filename)
+        let data = try Data(contentsOf: fileURL)
+        defer { XCTAssertNoThrow(try FileManager.default.removeItem(at: fileURL)) }
+        XCTAssertEqual(Data(buffer: buffer), data)
+    }
 }