Merge pull request #125 from Tieske/sanity-checks

chore(*) add some sanity checks and auto-coercion
EvandroLG · Feb 27, 2023 · d44a9a1 · d44a9a1
2 parents e9a741f + 545b80e
commit d44a9a1
Show file tree

Hide file tree

Showing 4 changed files with 79 additions and 16 deletions.
diff --git a/spec/unit/request_spec.lua b/spec/unit/request_spec.lua
@@ -79,6 +79,18 @@ describe('require', function()
     )
   end)
 
+  test('header lookup is case-insensitive', function()
+    local request = getInstance(
+      { 'GET / HTTP/1.1', 'hello: A', 'WORLD: B', 'Hello-World: X', '' }
+    )
+
+    local headers = request:headers()
+    assert.equal("A", headers["hello"])
+    assert.equal("A", headers["hELLo"])
+    assert.equal("B", headers["world"])
+    assert.equal("X", headers["hello-world"])
+  end)
+
   test('find value with = signal', function()
     local headers = { 'GET /Makefile?a=b= HTTP/1.1', 'a: A=', '' }
     local request = getInstance(headers)

diff --git a/spec/unit/response_spec.lua b/spec/unit/response_spec.lua
@@ -104,6 +104,16 @@ describe('response', function()
                     .. "My-Header: value 2" .. "\r\n"
       assert.equal(expected, response:_getHeaders())
     end)
+
+    it("fails if headers were already sent", function()
+      local response = Response:new({
+        send = function() end,
+      })
+      response:sendHeaders()
+      assert.has.error(function()
+        response:addHeader("Hi", "there")
+      end)
+    end)
   end)
 
   describe('status code', function()
@@ -124,9 +134,31 @@ describe('response', function()
       verifyStatus(200, nil, 'OK')
     end)
 
+    it('should add status code if passed as string', function()
+      verifyStatus("200", nil, 'OK')
+    end)
+
     it('should add status and message passed as parameters', function()
       verifyStatus(200, 'Perfect!', 'Perfect!')
     end)
+
+    it('fails for an unknown status code', function()
+      local response = Response:new({})
+      assert.has.error(function()
+        response:statusCode(9999)
+      end, "http status code '9999' is unknown")
+    end)
+
+    it("fails if headers were already sent", function()
+      local response = Response:new({
+        send = function() end,
+      })
+      response:sendHeaders()
+      assert.has.error(function()
+        response:statusCode(200)
+      end)
+    end)
+
   end)
 
   describe('set default headers', function()

diff --git a/src/pegasus/request.lua b/src/pegasus/request.lua
@@ -137,11 +137,20 @@ function Request:headers()
 
   local data = self.client:receive()
 
-  local headers = {}
+  local headers = setmetatable({},{ -- add metatable to do case-insensitive lookup
+    __index = function(self, key)
+      if type(key) == "string" then
+        key = key:lower()
+        return rawget(self, key)
+      end
+    end
+  })
+
   while (data ~= nil) and (data:len() > 0) do
     local key, value = string.match(data, Request.PATTERN_HEADER)
 
     if key and value then
+      key = key:lower()
       local v = headers[key]
       if not v then
         headers[key] = value
@@ -156,7 +165,7 @@ function Request:headers()
   end
 
   self._headerParsed = true
-  self._contentLength = tonumber(headers["Content-Length"] or 0)
+  self._contentLength = tonumber(headers["content-length"] or 0)
   self._headers = headers
 
   return headers

diff --git a/src/pegasus/response.lua b/src/pegasus/response.lua
@@ -10,7 +10,7 @@ local function toHex(dec)
   return table.concat(tmp)
 end
 
-local STATUS_TEXT = {
+local STATUS_TEXT = setmetatable({
   [100] = 'Continue',
   [101] = 'Switching Protocols',
   [200] = 'OK',
@@ -51,7 +51,18 @@ local STATUS_TEXT = {
   [503] = 'Service Unavailable',
   [504] = 'Gateway Time-out',
   [505] = 'HTTP Version not supported',
-}
+}, {
+  __index = function(self, statusCode)
+    -- if the lookup failed, try coerce to a number and try again
+    if type(statusCode) == "string" then
+      local result = rawget(self, tonumber(statusCode) or -1)
+      if result then
+        return result
+      end
+    end
+    error("http status code '"..tostring(statusCode).."' is unknown", 2)
+  end,
+})
 
 local DEFAULT_ERROR_MESSAGE = [[
   <!DOCTYPE HTML PUBLIC '-//W3C//DTD HTML 4.01//EN'
@@ -87,26 +98,27 @@ function Response:new(client, writeHandler)
 end
 
 function Response:addHeader(key, value)
+  assert(not self._headersSended, "can't add header, they were already sent")
   self._headers[key] = value
   return self
 end
 
 function Response:addHeaders(params)
   for key, value in pairs(params) do
-    self._headers[key] = value
+    self:addHeader(key, value)
   end
 
   return self
 end
 
 function Response:contentType(value)
-  self._headers['Content-Type'] = value
-  return self
+  return self:addHeader('Content-Type', value)
 end
 
 function Response:statusCode(statusCode, statusText)
+  assert(not self._headersSended, "can't set status code, it was already sent")
   self.status = statusCode
-  self._headFirstLine = string.gsub(self._templateFirstLine, '{{ STATUS_CODE }}', statusCode)
+  self._headFirstLine = string.gsub(self._templateFirstLine, '{{ STATUS_CODE }}', tostring(statusCode))
   self._headFirstLine = string.gsub(self._headFirstLine, '{{ STATUS_TEXT }}', statusText or STATUS_TEXT[statusCode])
 
   return self
@@ -140,13 +152,11 @@ function Response:close()
   local body = self._writeHandler:processBodyData(nil, true, self)
 
   if body and #body > 0 then
-    self._client:send(
-      toHex(#body) .. '\r\n' .. body .. '\r\n'
-    )
+    self._client:send(toHex(#body) .. '\r\n' .. body .. '\r\n')
   end
 
   self._client:send('0\r\n\r\n')
-  self.close = true
+  self.close = true  -- TODO: this seems unused??
 
   return self
 end
@@ -175,9 +185,9 @@ function Response:sendHeaders(stayOpen, body)
     self:addHeader('Content-Type', 'text/html')
   end
 
-  self._client:send(self._headFirstLine .. self:_getHeaders())
-  self._client:send('\r\n')
   self._headersSended = true
+  self._client:send(self._headFirstLine .. self:_getHeaders() .. '\r\n')
+  self._chunked = stayOpen
 
   return self
 end
@@ -186,7 +196,7 @@ function Response:write(body, stayOpen)
   body = self._writeHandler:processBodyData(body or '', stayOpen, self)
   self:sendHeaders(stayOpen, body)
 
-  self._isClosed = not(stayOpen or false)
+  self._isClosed = not stayOpen
 
   if self._isClosed then
     self._client:send(body)
@@ -197,7 +207,7 @@ function Response:write(body, stayOpen)
   end
 
   if self._isClosed then
-    self._client:close()
+    self._client:close() -- TODO: remove this, a non-chunked body can also be sent in multiple pieces
   end
 
   return self