From 82294d7f01f64c7e49cdb368437e5abb6db6c7e1 Mon Sep 17 00:00:00 2001
From: Murilo Santana <mvrilo@gmail.com>
Date: Thu, 12 Dec 2013 10:31:26 -0200
Subject: [PATCH] locks

---
 cache.go   | 13 ++++++++++++-
 handler.go |  8 ++++++--
 2 files changed, 18 insertions(+), 3 deletions(-)

diff --git a/cache.go b/cache.go
index 9c2f823..5bee9ec 100644
--- a/cache.go
+++ b/cache.go
@@ -6,6 +6,7 @@ import (
 	"fmt"
 	"github.com/hoisie/redis"
 	"github.com/miekg/dns"
+	"sync"
 	"time"
 )
 
@@ -56,11 +57,13 @@ type MemoryCache struct {
 	Backend  map[string]Mesg
 	Expire   time.Duration
 	Maxcount int
+	mu       *sync.RWMutex
 }
 
 func (c *MemoryCache) Get(key string) (*dns.Msg, error) {
-
+	c.mu.RLock()
 	mesg, ok := c.Backend[key]
+	c.mu.RUnlock()
 	if !ok {
 		return nil, KeyNotFound{key}
 	}
@@ -81,20 +84,28 @@ func (c *MemoryCache) Set(key string, msg *dns.Msg) error {
 
 	expire := time.Now().Add(c.Expire)
 	mesg := Mesg{msg, expire}
+	c.mu.Lock()
 	c.Backend[key] = mesg
+	c.mu.Unlock()
 	return nil
 }
 
 func (c *MemoryCache) Remove(key string) {
+	c.mu.RLock()
 	delete(c.Backend, key)
+	c.mu.RUnlock()
 }
 
 func (c *MemoryCache) Exists(key string) bool {
+	c.mu.RLock()
 	_, ok := c.Backend[key]
+	c.mu.RUnlock()
 	return ok
 }
 
 func (c *MemoryCache) Length() int {
+	c.mu.RLock()
+	defer c.mu.RUnlock()
 	return len(c.Backend)
 }
 
diff --git a/handler.go b/handler.go
index a4258da..41dc174 100644
--- a/handler.go
+++ b/handler.go
@@ -3,6 +3,7 @@ package main
 import (
 	"github.com/miekg/dns"
 	"net"
+	"sync"
 	"time"
 )
 
@@ -20,6 +21,7 @@ type GODNSHandler struct {
 	resolver *Resolver
 	cache    Cache
 	hosts    Hosts
+	mu       *sync.Mutex
 }
 
 func NewHandler() *GODNSHandler {
@@ -48,6 +50,7 @@ func NewHandler() *GODNSHandler {
 			Backend:  make(map[string]Mesg),
 			Expire:   time.Duration(cacheConfig.Expire) * time.Second,
 			Maxcount: cacheConfig.Maxcount,
+			mu:       new(sync.RWMutex),
 		}
 	case "redis":
 		// cache = &MemoryCache{
@@ -64,11 +67,10 @@ func NewHandler() *GODNSHandler {
 
 	hosts := NewHosts(settings.Hosts, settings.Redis)
 
-	return &GODNSHandler{resolver, cache, hosts}
+	return &GODNSHandler{resolver, cache, hosts, new(sync.Mutex)}
 }
 
 func (h *GODNSHandler) do(Net string, w dns.ResponseWriter, req *dns.Msg) {
-
 	q := req.Question[0]
 	Q := Question{UnFqdn(q.Name), dns.TypeToString[q.Qtype], dns.ClassToString[q.Qclass]}
 
@@ -97,8 +99,10 @@ func (h *GODNSHandler) do(Net string, w dns.ResponseWriter, req *dns.Msg) {
 			Debug("%s didn't hit cache: %s", Q.String(), err)
 		} else {
 			Debug("%s hit cache", Q.String())
+			h.mu.Lock()
 			mesg.Id = req.Id
 			w.WriteMsg(mesg)
+			h.mu.Unlock()
 			return
 		}