diff --git a/README.md b/README.md --- a/README.md +++ b/README.md @@ -25,7 +25,8 @@ stack install ``` -For Arch Linux users, an [AUR package](https://aur.archlinux.org/packages/dprox/) is also provided. +For Arch Linux users, an [AUR package](https://aur.archlinux.org/packages/dprox/) is provided. +Alternatively, you also can use the statically linked binary for the [latest release](https://github.com/bjin/dprox/releases). ### Usage diff --git a/dprox.cabal b/dprox.cabal --- a/dprox.cabal +++ b/dprox.cabal @@ -1,13 +1,13 @@ cabal-version: 1.12 --- This file has been generated from package.yaml by hpack version 0.31.1. +-- This file has been generated from package.yaml by hpack version 0.33.0. -- -- see: https://github.com/sol/hpack -- --- hash: fa94cf155fa0dc8474980910c4a902990e22de07b06863ade39fb8177bba7bf4 +-- hash: 5ad793b14e38bbe4cb7643ee582e2ba817f2912ee33cd2f21a68e5c3d20dadcf name: dprox -version: 0.1.2.1 +version: 0.2.0 synopsis: a lightweight DNS proxy server description: Please see the README on GitHub at category: DNS @@ -37,6 +37,8 @@ other-modules: Config DomainRoute + LRU + Paths_dprox hs-source-dirs: src ghc-options: -Wall -O2 -threaded -rtsopts -with-rtsopts=-N @@ -46,10 +48,13 @@ , bytestring >=0.10 , containers >=0.6 , dns >=3.0.4 + , hashable >=1.2 , iproute >=1.7 , network >=2.8 , optparse-applicative >=0.14 + , psqueues >=0.2 , streaming-commons >=0.2 + , time >=1.8 , unix >=2.7 , unordered-containers >=0.2 if flag(static) @@ -71,11 +76,14 @@ , bytestring >=0.10 , containers >=0.6 , dns >=3.0.4 + , hashable >=1.2 , hspec , iproute >=1.7 , network >=2.8 , optparse-applicative >=0.14 + , psqueues >=0.2 , streaming-commons >=0.2 + , time >=1.8 , unix >=2.7 , unordered-containers >=0.2 if flag(static) diff --git a/src/Config.hs b/src/Config.hs --- a/src/Config.hs +++ b/src/Config.hs @@ -23,15 +23,20 @@ isNothing) import Data.Streaming.Network (HostPreference) import Data.String (fromString) +import Data.Version (showVersion) import qualified Network.DNS as DNS import Network.Socket (PortNumber) import Options.Applicative import Text.Read (readMaybe) +import Paths_dprox (version) + data GlobalConfig = GlobalConfig { setUser :: Maybe String , localPort :: Maybe PortNumber , listenAddress :: Maybe HostPreference + , cacheSize :: Int + , cacheTTL :: DNS.TTL } deriving (Eq, Show) data Config = Server (Maybe DNS.Domain) IP (Maybe PortNumber) @@ -48,9 +53,12 @@ return (globalConfigs, confs1 ++ confs2 ++ confs) where opts = info ((,,,) <$> globalOption <*> configFileOption - <*> hostsFilesOption <*> plainOption <**> helper) - ( fullDesc <> progDesc "a lightweight DNS proxy server, supports a small subset of dnsmasq options") + <*> hostsFilesOption <*> plainOption <**> ver <**> helper) + ( fullDesc <> progDesc desc ) + desc = "a lightweight DNS proxy server, supports a small subset of dnsmasq options" + ver = infoOption (showVersion version) (long "version" <> help "show version") + readConfigFromFile file = handle handler (parseConfigFile <$> BS.readFile file) readHostsFromFile file = handle handler (parseHostsFile <$> BS.readFile file) @@ -128,6 +136,8 @@ globalOption = GlobalConfig <$> userOption <*> portOption <*> listenOption + <*> cacheOption + <*> ttlOption where userOption = optional $ strOption ( long "user" @@ -147,6 +157,18 @@ <> metavar "" <> help "Listen on the given IP address") + cacheOption = option auto + ( long "cache-size" + <> metavar "" + <> value 4096 + <> help "Size of the cache in entries (default value: 4096), setting this to zero disables cache") + + ttlOption = option auto + ( long "cache-ttl" + <> metavar "" + <> value 233 + <> help "Cache TTL in seconds (default value: 233)") + configFileOption :: Parser [FilePath] configFileOption = many $ strOption ( long "conf-file" @@ -205,10 +227,10 @@ serverValue :: P.Parser Config serverValue = do - parsedDomain <- P.option Nothing (Just <$> (P8.char '/' *> domain <* P8.char '/')) - parsedIP <- P.option Nothing (Just <$> ip) + parsedDomain <- optional (P8.char '/' *> domain <* P8.char '/') + parsedIP <- optional ip when (isNothing parsedDomain && isNothing parsedIP) $ fail "at least one of and must be specified" - parsedPort <- P.option Nothing (Just <$> (P8.char '#' *> port)) + parsedPort <- optional (P8.char '#' *> port) return (Server parsedDomain (fromMaybe invalidIPAddress parsedIP) parsedPort) addressValue :: P.Parser Config diff --git a/src/LRU.hs b/src/LRU.hs new file mode 100644 --- /dev/null +++ b/src/LRU.hs @@ -0,0 +1,78 @@ +-- SPDX-License-Identifier: BSD-3-Clause +-- +-- Copyright (C) 2021 Bin Jin. All Rights Reserved. +{-# LANGUAGE RecordWildCards #-} +{-# LANGUAGE TupleSections #-} +module LRU +( LRUCache +, newCache +, lookupCache +, updateCache +, purgeCache +) where + +import qualified Data.HashPSQ as PQ +import Data.Hashable (Hashable) +import Data.IORef (IORef, atomicModifyIORef', newIORef, + readIORef) +import Data.Int (Int64) +import Data.Time.Clock.System (SystemTime (..), getSystemTime) +import Network.DNS (TTL) + +type Time = Int64 + +getTime :: IO Time +getTime = systemSeconds <$> getSystemTime + +data LRUQueue k v = LRUQueue + { queueSize :: !Int + , queue :: !(PQ.HashPSQ k Time v) + } deriving (Show) + +newQueue :: LRUQueue k v +newQueue = LRUQueue 0 PQ.empty + +lookupQueue :: (Ord k, Hashable k) => k -> LRUQueue k v -> Maybe (Time, v) +lookupQueue k = PQ.lookup k . queue + +updateQueue :: (Ord k, Hashable k) => Int -> Time -> k -> v -> LRUQueue k v -> LRUQueue k v +updateQueue cacheSize expireTime k v LRUQueue{..} = case PQ.insertView k expireTime v queue of + (Just _, newQ) -> LRUQueue queueSize newQ + (Nothing, newQ) -> if queueSize >= cacheSize + then LRUQueue queueSize (PQ.deleteMin newQ) + else LRUQueue (queueSize+1) newQ + +purgeQueue :: (Ord k, Hashable k) => Time -> LRUQueue k v -> LRUQueue k v +purgeQueue _ q@(LRUQueue 0 _) = q +purgeQueue now LRUQueue{..} = LRUQueue (queueSize - length expired) newQ + where + (expired, newQ) = PQ.atMostView now queue + +data LRUCache k v = LRUCache + { cacheSize :: !Int + , timeToLive :: !TTL + , cacheRef :: IORef (LRUQueue k v) + } + +newCache :: Int -> TTL -> IO (LRUCache k v) +newCache sz ttl = LRUCache sz ttl <$> newIORef newQueue + +lookupCache :: (Ord k, Hashable k) => k -> LRUCache k v -> IO (Maybe (TTL, v)) +lookupCache k LRUCache{..} = do + now <- getTime + res <- lookupQueue k <$> readIORef cacheRef + return $ case res of + Nothing -> Nothing + Just (expireTime, v) -> if expireTime <= now + then Nothing + else Just (fromIntegral (expireTime - now), v) + +updateCache :: (Ord k, Hashable k) => k -> v -> LRUCache k v -> IO () +updateCache k v LRUCache{..} = do + now <- getTime + atomicModifyIORef' cacheRef ((,()) . updateQueue cacheSize (now+fromIntegral timeToLive) k v) + +purgeCache :: (Ord k, Hashable k) => LRUCache k v -> IO () +purgeCache LRUCache{..} = do + now <- getTime + atomicModifyIORef' cacheRef ((,()) . purgeQueue now) diff --git a/src/Main.hs b/src/Main.hs --- a/src/Main.hs +++ b/src/Main.hs @@ -3,13 +3,15 @@ -- Copyright (C) 2019 Bin Jin. All Rights Reserved. {-# LANGUAGE OverloadedStrings #-} {-# LANGUAGE RecordWildCards #-} +{-# LANGUAGE TupleSections #-} module Main where -import Control.Concurrent (forkIO) +import Control.Concurrent (forkIO, threadDelay) import Control.Exception (SomeException, handle) import Control.Monad (forM, forever, join) import Data.ByteString (ByteString) import qualified Data.Foldable as F +import Data.Hashable (Hashable (..)) import qualified Data.Map as M import Data.Maybe (fromMaybe) import qualified Data.Set as S @@ -21,21 +23,26 @@ import Config import DomainRoute +import LRU +instance Hashable DNS.TYPE where + hashWithSalt s = hashWithSalt s . DNS.fromTYPE + type Resolver = DNS.Domain -> DNS.TYPE -> IO (Either DNS.DNSError [DNS.RData]) +type CachedResolver = DNS.Domain -> DNS.TYPE -> IO (Either DNS.DNSError (DNS.TTL, [DNS.RData])) -processQuery :: Resolver -> DNS.Question -> IO [DNS.ResourceRecord] +processQuery :: CachedResolver -> DNS.Question -> IO [DNS.ResourceRecord] processQuery resolver (DNS.Question qd qt) = handle handler $ do res <- resolver qd qt case res of - Left _ -> return [] - Right r -> return (map wrapper r) + Left _ -> return [] + Right (ttl, r) -> return (map (wrapper ttl) r) where handler :: SomeException -> IO [DNS.ResourceRecord] handler _ = return [] - wrapper :: DNS.RData -> DNS.ResourceRecord - wrapper rdata = DNS.ResourceRecord qd (getType rdata) DNS.classIN 233 rdata + wrapper :: DNS.TTL -> DNS.RData -> DNS.ResourceRecord + wrapper ttl rdata = DNS.ResourceRecord qd (getType rdata) DNS.classIN ttl rdata getType :: DNS.RData -> DNS.TYPE getType DNS.RD_A{} = DNS.A @@ -49,7 +56,7 @@ getType DNS.RD_AAAA{} = DNS.AAAA getType _ = qt -processDNS :: Resolver -> ByteString -> IO (Either DNS.DNSError ByteString) +processDNS :: CachedResolver -> ByteString -> IO (Either DNS.DNSError ByteString) processDNS resolver bs | Left e <- parse = return (Left e) | Right q <- parse = do @@ -97,6 +104,27 @@ isBlacklisted (DNS.RD_AAAA ipv6) = IPv6 ipv6 `S.member` blacklist isBlacklisted _ = False +makeResolverCache :: Int -> DNS.TTL -> IO (Resolver -> CachedResolver) +makeResolverCache sz ttl | sz <= 0 = return $ \r qd qt -> fmap (ttl,) <$> r qd qt +makeResolverCache sz ttl = do + cache <- newCache sz ttl + _ <- forkIO $ forever $ do + threadDelay (fromInteger (fromIntegral ttl * 1000000 `div` 3)) + purgeCache cache + let process resolver qd qt = do + let k = (qd, qt) + res <- lookupCache k cache + case res of + Just (ttl', v) -> return (Right (max ttl' 3, v)) + Nothing -> do + resolved <- resolver qd qt + case resolved of + Left e -> return (Left e) + Right v -> do + updateCache k v cache + return (Right (ttl, v)) + return process + setuid :: String -> IO () setuid user = getUserEntryForName user >>= setUserID . userID @@ -113,7 +141,7 @@ hosts = [(domain, [ip]) | Hosts domain ip <- conf] bogusnx = [ip | BogusNX ip <- conf] - serverRoute = newDomainRoute (flip const) server + serverRoute = newDomainRoute (const id) server serverAddressSet = S.fromList $ F.toList serverRoute addressRoute = newDomainRoute (++) address @@ -128,7 +156,7 @@ then DNS.RCHostName (show host) else DNS.RCHostPort (show host) port , let rc = DNS.defaultResolvConf { - DNS.resolvCache = Just DNS.defaultCacheConf, + DNS.resolvCache = Nothing, DNS.resolvInfo = rsinfo } ] @@ -140,6 +168,8 @@ F.mapM_ setuid setUser + resolverCache <- makeResolverCache cacheSize cacheTTL + let processWithResolver resolver = forever $ do (bs, addr) <- recvFrom sock (fromIntegral DNS.maxUdpSize) forkIO $ do @@ -149,7 +179,8 @@ createResolvers ((k,v):xs) m = DNS.withResolver v $ \rs -> createResolvers xs (M.insert k (DNS.lookup rs) m) createResolvers [] m = let serverRoute' = fmap (`M.lookup`m) serverRoute - resolver = handleBogusNX bogusnxSet $ + resolver = resolverCache $ + handleBogusNX bogusnxSet $ handleAddressAndHosts addressRoute hostsRoute $ handleServer serverRoute' in processWithResolver resolver