From e32c5b9a536f7ae51230d7c6882d605c9c8b3e5c Mon Sep 17 00:00:00 2001
From: Michael Snoyman <michael@snoyman.com>
Date: Wed, 30 Jun 2010 21:33:32 +0300
Subject: [PATCH] Static file serving: extensible mime-type dictionary

---
 Yesod.hs                |  2 ++
 Yesod/Content.hs        | 35 ++++++++++++++++++++++-------------
 Yesod/Helpers/Static.hs | 21 +++++++++++++++------
 3 files changed, 39 insertions(+), 19 deletions(-)

diff --git a/Yesod.hs b/Yesod.hs
index 073a625a..a36f348a 100644
--- a/Yesod.hs
+++ b/Yesod.hs
@@ -13,6 +13,7 @@ module Yesod
     , module Yesod.Formable
     , Application
     , liftIO
+    , mempty
     ) where
 
 #if TEST
@@ -33,3 +34,4 @@ import Yesod.Handler hiding (runHandler)
 import Network.Wai (Application)
 import Yesod.Hamlet
 import "transformers" Control.Monad.IO.Class (liftIO)
+import Data.Monoid (mempty)
diff --git a/Yesod/Content.hs b/Yesod/Content.hs
index bbd6ac5f..0dea1d78 100644
--- a/Yesod/Content.hs
+++ b/Yesod/Content.hs
@@ -85,9 +85,14 @@ data Content = ContentFile FilePath
                           -> a
                           -> IO (Either a a))
 
+-- | Zero-length enumerator.
 emptyContent :: Content
 emptyContent = ContentEnum $ \_ -> return . Right
 
+-- | Anything which can be converted into 'Content'. Most of the time, you will
+-- want to use the 'ContentEnum' constructor. An easier approach will be to use
+-- a pre-defined 'toContent' function, such as converting your data into a lazy
+-- bytestring and then calling 'toContent' on that.
 class ToContent a where
     toContent :: a -> Content
 
@@ -140,6 +145,9 @@ instance HasReps ChooseRep where
 instance HasReps () where
     chooseRep = defChooseRep [(typePlain, const $ return $ toContent "")]
 
+instance HasReps (ContentType, Content) where
+    chooseRep = const . return
+
 instance HasReps [(ContentType, Content)] where
     chooseRep a cts = return $
         case filter (\(ct, _) -> go ct `elem` map go cts) a of
@@ -218,19 +226,20 @@ typeOctet = "application/octet-stream"
 simpleContentType :: String -> String
 simpleContentType = fst . span (/= ';')
 
--- | Determine a mime-type based on the file extension.
-typeByExt :: String -> ContentType
-typeByExt "jpg" = typeJpeg
-typeByExt "jpeg" = typeJpeg
-typeByExt "js" = typeJavascript
-typeByExt "css" = typeCss
-typeByExt "html" = typeHtml
-typeByExt "png" = typePng
-typeByExt "gif" = typeGif
-typeByExt "txt" = typePlain
-typeByExt "flv" = typeFlv
-typeByExt "ogv" = typeOgv
-typeByExt _ = typeOctet
+-- | A default extension to mime-type dictionary.
+typeByExt :: [(String, ContentType)]
+typeByExt =
+    [ ("jpg", typeJpeg)
+    , ("jpeg", typeJpeg)
+    , ("js", typeJavascript)
+    , ("css", typeCss)
+    , ("html", typeHtml)
+    , ("png", typePng)
+    , ("gif", typeGif)
+    , ("txt", typePlain)
+    , ("flv", typeFlv)
+    , ("ogv", typeOgv)
+    ]
 
 -- | Get a file extension (everything after last period).
 ext :: String -> String
diff --git a/Yesod/Helpers/Static.hs b/Yesod/Helpers/Static.hs
index 8a496bdb..da0f3bc4 100644
--- a/Yesod/Helpers/Static.hs
+++ b/Yesod/Helpers/Static.hs
@@ -38,6 +38,7 @@ module Yesod.Helpers.Static
 
 import System.Directory
 import Control.Monad
+import Data.Maybe (fromMaybe)
 
 import Yesod
 import Data.List (intercalate)
@@ -51,7 +52,11 @@ import Test.HUnit hiding (Test)
 
 -- | A function for looking up file contents. For serving from the file system,
 -- see 'fileLookupDir'.
-data Static = Static (FilePath -> IO (Maybe (Either FilePath Content)))
+data Static = Static
+    { staticLookup :: FilePath -> IO (Maybe (Either FilePath Content))
+    -- | Mapping from file extension to content type. See 'typeByExt'.
+    , staticTypes :: [(String, ContentType)]
+    }
 
 mkYesodSub "Static" [] [$parseRoutes|
 *Strings StaticRoute GET
@@ -63,7 +68,7 @@ mkYesodSub "Static" [] [$parseRoutes|
 -- probably are), the handler itself checks that no unsafe paths are being
 -- requested. In particular, no path segments may begin with a single period,
 -- so hidden files and parent directories are safe.
-fileLookupDir :: FilePath -> Static
+fileLookupDir :: FilePath -> [(String, ContentType)] -> Static
 fileLookupDir dir = Static $ \fp -> do
     let fp' = dir ++ '/' : fp
     exists <- doesFileExist fp'
@@ -72,16 +77,20 @@ fileLookupDir dir = Static $ \fp -> do
         else return Nothing
 
 getStaticRoute :: [String]
-               -> GHandler Static master [(ContentType, Content)]
+               -> GHandler Static master (ContentType, Content)
 getStaticRoute fp' = do
-    Static fl <- getYesodSub
+    Static fl ctypes <- getYesodSub
     when (any isUnsafe fp') notFound
     let fp = intercalate "/" fp'
     content <- liftIO $ fl fp
     case content of
         Nothing -> notFound
-        Just (Left fp'') -> sendFile (typeByExt $ ext fp'') fp''
-        Just (Right bs) -> return [(typeByExt $ ext fp, bs)]
+        Just (Left fp'') -> do
+            let ctype = fromMaybe typeOctet $ lookup (ext fp'') ctypes
+            sendFile ctype fp''
+        Just (Right bs) -> do
+            let ctype = fromMaybe typeOctet $ lookup (ext fp) ctypes
+            return (ctype, bs)
   where
     isUnsafe [] = True
     isUnsafe ('.':_) = True