Add SendFile support to Chunked, Static.

owin · Sep 17, 2012 · ab68fbc · ab68fbc
1 parent befdc10
commit ab68fbc
Show file tree

Hide file tree

Showing 3 changed files with 70 additions and 7 deletions.
diff --git a/src/Main/Gate.Middleware/Chunked.cs b/src/Main/Gate.Middleware/Chunked.cs
@@ -23,6 +23,7 @@ public static IAppBuilder UseChunked(this IAppBuilder builder)
 namespace Gate.Middleware
 {
     using AppFunc = Func<IDictionary<string, object>, Task>;
+    using SendFileFunc = Func<string, long, long, Task>;
 
     // This middleware implements chunked response body encoding as the default encoding 
     // if the application does not specify Content-Length or Transfer-Encoding.
@@ -47,11 +48,11 @@ public Task Invoke(IDictionary<string, object> env)
             TriggerStream triggerStream = new TriggerStream(orriginalStream);
             response.OutputStream = triggerStream;
             FilterStream filterStream = null;
-            bool onFirstWriteExecuted = false;
+            bool finalizeHeadersExecuted = false;
 
-            triggerStream.OnFirstWrite = () =>
+            Action finalizeHeaders = () =>
             {
-                onFirstWriteExecuted = true;
+                finalizeHeadersExecuted = true;
                 if (IsStatusWithNoNoEntityBody(response.StatusCode)
                     || response.Headers.ContainsKey("Content-Length")
                     || response.Headers.ContainsKey("Transfer-Encoding"))
@@ -74,13 +75,42 @@ public Task Invoke(IDictionary<string, object> env)
                 }
             };
 
+            // Hook first write
+            triggerStream.OnFirstWrite = finalizeHeaders;
             env[OwinConstants.ResponseBody] = triggerStream;
 
+            // Hook SendFileFunc
+            SendFileFunc sendFile = env.Get<SendFileFunc>("sendfile.Func");
+            if (sendFile != null)
+            {
+                SendFileFunc sendFileChunked = (name, offset, count) =>
+                {
+                    if (!finalizeHeadersExecuted)
+                    {
+                        finalizeHeaders();
+                    }
+
+                    if (filterStream == null)
+                    {
+                        // Due to headers we are not doing chunked, just pass through.
+                        return sendFile(name, offset, count);
+                    }
+
+                    // Insert chunking around the file body
+                    ArraySegment<byte> prefix = ChunkPrefix((uint)count);
+                    return orriginalStream.WriteAsync(prefix.Array, prefix.Offset, prefix.Count)
+                        .Then(() => orriginalStream.FlushAsync()) // Flush to ensure the data hits the wire before sendFile.
+                        .Then(() => sendFile(name, offset, count))
+                        .Then(() => orriginalStream.WriteAsync(EndOfChunk.Array, EndOfChunk.Offset, EndOfChunk.Count));
+                };
+                env["sendfile.Func"] = sendFileChunked;
+            }
+
             return nextApp(env).Then(() =>
             {
-                if (!onFirstWriteExecuted)
+                if (!finalizeHeadersExecuted)
                 {
-                    triggerStream.OnFirstWrite();
+                    finalizeHeaders();
                 }
 
                 if (filterStream != null)

diff --git a/src/Main/Gate.Middleware/StaticFiles/FileBody.cs b/src/Main/Gate.Middleware/StaticFiles/FileBody.cs
@@ -7,6 +7,8 @@
 
 namespace Gate.Middleware.StaticFiles
 {
+    using SendFileFunc = Func<string, long, long, Task>;
+
     public class FileBody
     {
         private Stream fileStream;
@@ -22,7 +24,16 @@ public FileBody(string path, Tuple<long, long> range)
         public Task Start(Stream stream)
         {
             this.OpenFileStream();
-            return this.fileStream.CopyToAsync(stream, (int)(range.Item2 - range.Item1 + 1));
+            try
+            {
+                return this.fileStream.CopyToAsync(stream, (int)(range.Item2 - range.Item1 + 1))
+                    .Finally(() => CloseFileStream());
+            }
+            catch (Exception)
+            {
+                CloseFileStream();
+                throw;
+            }
         }
 
         private void OpenFileStream()
@@ -33,5 +44,18 @@ private void OpenFileStream()
                 this.fileStream.Seek(range.Item1, SeekOrigin.Begin);
             }
         }
+
+        private void CloseFileStream()
+        {
+            if (this.fileStream != null)
+            {
+                this.fileStream.Close();
+            }
+        }
+
+        public Task Start(SendFileFunc sendFile)
+        {
+            return sendFile(path, range.Item1, (range.Item2 - range.Item1 + 1));
+        }
     }
 }
diff --git a/src/Main/Gate.Middleware/StaticFiles/FileServer.cs b/src/Main/Gate.Middleware/StaticFiles/FileServer.cs
@@ -11,6 +11,7 @@
 namespace Gate.Middleware.StaticFiles
 {
     using AppFunc = Func<IDictionary<string, object>, Task>;
+    using SendFileFunc = Func<string, long, long, Task>;
 
     // Used by the Static middleware to send static files to the client.
     public class FileServer
@@ -140,7 +141,15 @@ private Task Serve(IDictionary<string, object> env, string path)
                 return TaskHelpers.Completed();
             }
 
-            return new FileBody(path, range).Start(response.OutputStream);
+            FileBody body = new FileBody(path, range);
+
+            SendFileFunc sendFile = env.Get<SendFileFunc>("sendfile.Func");
+            if (sendFile != null)
+            {
+                return body.Start(sendFile);
+            }
+
+            return body.Start(response.OutputStream);
         }
     }
 }