Skip to content

Commit

Permalink
Status save. Resume after commute
Browse files Browse the repository at this point in the history
  • Loading branch information
Marco Dondero committed Oct 7, 2018
1 parent d3ecc1f commit 23967df
Show file tree
Hide file tree
Showing 5 changed files with 42 additions and 50 deletions.
34 changes: 18 additions & 16 deletions api.go
View file
Expand Up @@ -9,7 +9,7 @@ import (
)

// StartAPIServer launches the API server
func StartAPIServer(config *Config) error {
func StartAPIServer(config *Config, blockCache *MemoryBlockCache, questionCache *MemoryQuestionCache) error {
if config.LogLevel == 0 {
gin.SetMode(gin.ReleaseMode)
}
Expand All @@ -18,60 +18,60 @@ func StartAPIServer(config *Config) error {
router.Use(cors.Default())

router.GET("/blockcache", func(c *gin.Context) {
c.IndentedJSON(http.StatusOK, gin.H{"length": BlockCache.Length(), "items": BlockCache.Backend})
c.IndentedJSON(http.StatusOK, gin.H{"length": blockCache.Length(), "items": blockCache.Backend})
})

router.GET("/blockcache/exists/:key", func(c *gin.Context) {
c.IndentedJSON(http.StatusOK, gin.H{"exists": BlockCache.Exists(c.Param("key"))})
c.IndentedJSON(http.StatusOK, gin.H{"exists": blockCache.Exists(c.Param("key"))})
})

router.GET("/blockcache/get/:key", func(c *gin.Context) {
if ok, _ := BlockCache.Get(c.Param("key")); !ok {
if ok, _ := blockCache.Get(c.Param("key")); !ok {
c.IndentedJSON(http.StatusOK, gin.H{"error": c.Param("key") + " not found"})
} else {
c.IndentedJSON(http.StatusOK, gin.H{"success": ok})
}
})

router.GET("/blockcache/length", func(c *gin.Context) {
c.IndentedJSON(http.StatusOK, gin.H{"length": BlockCache.Length()})
c.IndentedJSON(http.StatusOK, gin.H{"length": blockCache.Length()})
})

router.GET("/blockcache/remove/:key", func(c *gin.Context) {
// Removes from BlockCache only. If the domain has already been queried and placed into MemoryCache, will need to wait until item is expired.
BlockCache.Remove(c.Param("key"))
blockCache.Remove(c.Param("key"))
c.IndentedJSON(http.StatusOK, gin.H{"success": true})
})

router.GET("/blockcache/set/:key", func(c *gin.Context) {
// MemoryBlockCache Set() always returns nil, so ignoring response.
_ = BlockCache.Set(c.Param("key"), true)
_ = blockCache.Set(c.Param("key"), true)
c.IndentedJSON(http.StatusOK, gin.H{"success": true})
})

router.GET("/questioncache", func(c *gin.Context) {
c.IndentedJSON(http.StatusOK, gin.H{"length": QuestionCache.Length(), "items": QuestionCache.Backend})
c.IndentedJSON(http.StatusOK, gin.H{"length": questionCache.Length(), "items": questionCache.Backend})
})

router.GET("/questioncache/length", func(c *gin.Context) {
c.IndentedJSON(http.StatusOK, gin.H{"length": QuestionCache.Length()})
c.IndentedJSON(http.StatusOK, gin.H{"length": questionCache.Length()})
})

router.GET("/questioncache/clear", func(c *gin.Context) {
QuestionCache.Clear()
questionCache.Clear()
c.IndentedJSON(http.StatusOK, gin.H{"success": true})
})

router.GET("/questioncache/client/:client", func(c *gin.Context) {
var filteredCache []QuestionCacheEntry

QuestionCache.mu.RLock()
for _, entry := range QuestionCache.Backend {
questionCache.mu.RLock()
for _, entry := range questionCache.Backend {
if entry.Remote == c.Param("client") {
filteredCache = append(filteredCache, entry)
}
}
QuestionCache.mu.RUnlock()
questionCache.mu.RUnlock()

c.IndentedJSON(http.StatusOK, filteredCache)
})
Expand Down Expand Up @@ -103,8 +103,8 @@ func StartAPIServer(config *Config) error {
grimdActivation.set(false)
c.IndentedJSON(http.StatusOK, gin.H{"active": grimdActive})
case "Snooze":
timeout_string := c.DefaultQuery("timeout", "300")
timeout, err := strconv.ParseUint(timeout_string, 0, 0)
timeoutString := c.DefaultQuery("timeout", "300")
timeout, err := strconv.ParseUint(timeoutString, 0, 0)
if err != nil {
c.JSON(http.StatusBadRequest, gin.H{"error": "Illegal value for 'timeout'"})
} else {
Expand All @@ -121,7 +121,9 @@ func StartAPIServer(config *Config) error {
})

router.POST("/blocklist/update", func(c *gin.Context) {
PerformUpdate(true, config)
here we need to regenerate everything by firing some signal to a channel
otherwise the pointers will be old
blockcache = PerformUpdate(true, config)
c.AbortWithStatus(http.StatusOK)
})

Expand Down
15 changes: 6 additions & 9 deletions handler.go
View file
Expand Up @@ -50,7 +50,7 @@ type DNSOperationData struct {
}

// NewHandler returns a new DNSHandler
func NewHandler(config *Config) *DNSHandler {
func NewHandler(config *Config, blockCache *MemoryBlockCache, questionCache *MemoryQuestionCache) *DNSHandler {
var (
clientConfig *dns.ClientConfig
resolver *Resolver
Expand All @@ -76,15 +76,14 @@ func NewHandler(config *Config) *DNSHandler {
negCache: negCache,
}

go handler.do(config)
go handler.do(config, blockCache, questionCache)

return handler
}

func (h *DNSHandler) do(config *Config) {
func (h *DNSHandler) do(config *Config, blockCache *MemoryBlockCache, questionCache *MemoryQuestionCache) {
for {
data := <-h.requestChannel
logger.Debugf("New request: %s", data.req)
func(Net string, w dns.ResponseWriter, req *dns.Msg) {
defer w.Close()
q := req.Question[0]
Expand Down Expand Up @@ -143,7 +142,7 @@ func (h *DNSHandler) do(config *Config) {
var blacklisted = false

if IPQuery > 0 {
blacklisted = BlockCache.Exists(Q.Qname)
blacklisted = blockCache.Exists(Q.Qname)

if grimdActive && blacklisted {
m := new(dns.Msg)
Expand Down Expand Up @@ -179,7 +178,7 @@ func (h *DNSHandler) do(config *Config) {

// log query
NewEntry := QuestionCacheEntry{Date: time.Now().Unix(), Remote: remote.String(), Query: Q, Blocked: true}
go QuestionCache.Add(NewEntry)
go questionCache.Add(NewEntry)

// cache the block; we don't know the true TTL for blocked entries: we just enforce our config
err := h.cache.Set(key, m, true)
Expand All @@ -194,7 +193,7 @@ func (h *DNSHandler) do(config *Config) {

// log query
NewEntry := QuestionCacheEntry{Date: time.Now().Unix(), Remote: remote.String(), Query: Q, Blocked: false}
go QuestionCache.Add(NewEntry)
go questionCache.Add(NewEntry)

mesg, err := h.resolver.Lookup(Net, req, config.Timeout, config.Interval, config.Nameservers)

Expand Down Expand Up @@ -256,13 +255,11 @@ func (h *DNSHandler) do(config *Config) {

// DoTCP begins a tcp query
func (h *DNSHandler) DoTCP(w dns.ResponseWriter, req *dns.Msg) {
logger.Debugf("Sending tcp request %s", req)
h.requestChannel <- DNSOperationData{"tcp", w, req}
}

// DoUDP begins a udp query
func (h *DNSHandler) DoUDP(w dns.ResponseWriter, req *dns.Msg) {
logger.Debugf("Sending udp request %s", req)
h.requestChannel <- DNSOperationData{"udp", w, req}
}

Expand Down
28 changes: 13 additions & 15 deletions main.go
View file
Expand Up @@ -13,15 +13,13 @@ var (
forceUpdate bool
grimdActive bool
grimdActivation ActivationHandler

// BlockCache contains all blocked domains
BlockCache = &MemoryBlockCache{Backend: make(map[string]bool)}

// QuestionCache contains all queries to the dns server
QuestionCache = &MemoryQuestionCache{Backend: make([]QuestionCacheEntry, 0), Maxcount: 1000}
)

func main() {
// BlockCache contains all blocked domains
blockCache := &MemoryBlockCache{Backend: make(map[string]bool)}
// QuestionCache contains all queries to the dns server
questionCache := &MemoryQuestionCache{Backend: make([]QuestionCacheEntry, 0), Maxcount: 1000}

flag.Parse()

Expand All @@ -30,7 +28,7 @@ func main() {
logger.Fatal(err)
}

QuestionCache.Maxcount = config.QuestionCacheCap
questionCache.Maxcount = config.QuestionCacheCap

logFile, err := LoggerInit(config.LogLevel, config.Log)
if err != nil {
Expand All @@ -39,36 +37,36 @@ func main() {
defer logFile.Close()

// delay updating the blocklists, cache until the server starts and can serve requests as the local resolver
start_update := make(chan bool, 1)
startUpdate := make(chan bool, 1)

//abort if the server does not come up in 10 seconds
timer := time.NewTimer(time.Second * 10)
go func() {
<-timer.C
start_update <- false
startUpdate <- false
}()

go func() {
run := <-start_update
run := <-startUpdate
if !run {
panic("The DNS server did not start in 10 seconds")
}
PerformUpdate(forceUpdate, config)
}()

grimdActive = true
quit_activation := make(chan bool)
go grimdActivation.loop(quit_activation, config.ReactivationDelay)
quitActivation := make(chan bool)
go grimdActivation.loop(quitActivation, config.ReactivationDelay)

server := &Server{
host: config.Bind,
rTimeout: 5 * time.Second,
wTimeout: 5 * time.Second,
}

server.Run(start_update, config)
server.Run(startUpdate, config, blockCache, questionCache)

if err := StartAPIServer(config); err != nil {
if err := StartAPIServer(config, blockCache, questionCache); err != nil {
logger.Fatal(err)
}

Expand All @@ -80,7 +78,7 @@ forever:
select {
case <-sig:
logger.Error("signal received, stopping\n")
quit_activation <- true
quitActivation <- true
break forever
}
}
Expand Down
6 changes: 3 additions & 3 deletions server.go
View file
Expand Up @@ -14,8 +14,8 @@ type Server struct {
}

// Run starts the server
func (s *Server) Run(update_lists chan bool, config *Config) {
Handler := NewHandler(config)
func (s *Server) Run(updateLists chan bool, config *Config, blockCache *MemoryBlockCache, questionCache *MemoryQuestionCache) {
Handler := NewHandler(config, blockCache, questionCache)

tcpHandler := dns.NewServeMux()
tcpHandler.HandleFunc(".", Handler.DoTCP)
Expand All @@ -39,7 +39,7 @@ func (s *Server) Run(update_lists chan bool, config *Config) {
go s.start(udpServer)
go s.start(tcpServer)

update_lists <- true
updateLists <- true
}

func (s *Server) start(ds *dns.Server) {
Expand Down
9 changes: 2 additions & 7 deletions updater.go
View file
Expand Up @@ -154,7 +154,7 @@ func parseHostFile(fileName string, blockCache *MemoryBlockCache) error {

// PerformUpdate updates the block cache by building a new one and swapping
// it for the old cache.
func PerformUpdate(forceUpdate bool, config *Config) {
func PerformUpdate(forceUpdate bool, config *Config, blockCache *MemoryBlockCache) *MemoryBlockCache {
newBlockCache := &MemoryBlockCache{Backend: make(map[string]bool)}
if _, err := os.Stat("lists"); os.IsNotExist(err) || forceUpdate {
if err := update(newBlockCache, config.Whitelist, config.Blocklist, config.Sources); err != nil {
Expand All @@ -165,10 +165,5 @@ func PerformUpdate(forceUpdate bool, config *Config) {
logger.Fatal(err)
}

oldBlockCache := BlockCache
oldBlockCache.mu.Lock()
newBlockCache.mu.Lock()
BlockCache = newBlockCache
newBlockCache.mu.Unlock()
oldBlockCache.mu.Unlock()
return newBlockCache
}

0 comments on commit 23967df

Please sign in to comment.