From 325e9113732076466c0fabf1dda81c965cfcaeb9 Mon Sep 17 00:00:00 2001
From: Thijs Schreijer <thijs@thijsschreijer.nl>
Date: Sun, 19 Feb 2023 10:28:48 +0100
Subject: [PATCH] feat(response) add sendFile method to download files

also rewrites writeFile to take a filename instead of a file
descriptor.
---
 src/pegasus/response.lua | 52 +++++++++++++++++++++++++++++++++++-----
 1 file changed, 46 insertions(+), 6 deletions(-)

diff --git a/src/pegasus/response.lua b/src/pegasus/response.lua
index d5204c7..3b6737a 100644
--- a/src/pegasus/response.lua
+++ b/src/pegasus/response.lua
@@ -1,3 +1,5 @@
+local mimetypes = require 'mimetypes'
+
 local function toHex(dec)
   local charset = { '0','1','2','3','4','5','6','7','8','9','a','b','c','d','e','f' }
   local tmp = {}
@@ -202,12 +204,50 @@ function Response:write(body, stayOpen)
   return self
 end
 
-function Response:writeFile(file, contentType)
-  self:contentType(contentType)
-  self:statusCode(200)
-  local value = file:read('*a')
+local function readfile(filename)
+  local file, err = io.open(filename, 'rb')
+  if not file then
+    return nil, err
+  end
+
+  local value, err = file:read('*a')
   file:close()
-  self:write(value)
+  return value, err
+end
+
+-- return nil+err if not ok
+function Response:writeFile(filename, contentType)
+  if type(filename) ~= "string" then
+    -- deprecated backward compatibility; file is a file-descriptor
+    self:contentType(contentType)
+    self:statusCode(200)
+    local value = filename:read('*a')
+    filename:close()
+    self:write(value)
+    return self
+  end
+
+  local contents, err = readfile(filename)
+  if not contents then
+    return nil, err
+  end
+
+  self:statusCode(200)
+  self:contentType(contentType or mimetypes.guess(filename))
+  self:write(contents)
+  return self
+end
+
+-- download by browser, return nil+err if not ok
+function Response:sendFile(path)
+  local filename = path:match("[^/]*$") -- only filename, no path
+  self:addHeader('Content-Disposition', 'attachment; filename="' .. filename .. '"')
+
+  local ok, err = self:writeFile(path, 'application/octet-stream')
+  if not ok then
+    self:addHeader('Content-Disposition', nil)
+    return nil, err
+  end
 
   return self
 end
@@ -215,7 +255,7 @@ end
 function Response:redirect(location, temporary)
   self:statusCode(temporary and 302 or 301)
   self:addHeader('Location', location)
-  self.client:sendOnlyHeaders()
+  self:sendOnlyHeaders()
   return self
 end