dprox 0.1.2.1 → 0.2.0
raw patch · 5 files changed
+159/−19 lines, 5 filesdep +hashabledep +psqueuesdep +time
Dependencies added: hashable, psqueues, time
Files
- README.md +2/−1
- dprox.cabal +11/−3
- src/Config.hs +27/−5
- src/LRU.hs +78/−0
- src/Main.hs +41/−10
README.md view
@@ -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
dprox.cabal view
@@ -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 <https://github.com/bjin/dprox#readme> 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)
src/Config.hs view
@@ -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 "<ipaddr>" <> help "Listen on the given IP address") + cacheOption = option auto+ ( long "cache-size"+ <> metavar "<integer>"+ <> 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 "<seconds>"+ <> 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 <domain> and <ip> 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
+ src/LRU.hs view
@@ -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)
src/Main.hs view
@@ -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