diff --git a/client.go b/client.go index 90d0ead..d9e1041 100644 --- a/client.go +++ b/client.go @@ -148,6 +148,12 @@ func newClient() (*client, error) { return c, nil } +func (c *client) getIsClosed() bool { + c.closeLock.Lock() + defer c.closeLock.Unlock() + return c.closed +} + // Close is used to cleanup the client func (c *client) Close() error { c.closeLock.Lock() @@ -326,7 +332,7 @@ func (c *client) recv(l *net.UDPConn, msgCh chan *dns.Msg) { return } buf := make([]byte, 65536) - for !c.closed { + for !c.getIsClosed() { n, err := l.Read(buf) if err != nil { log.Printf("[ERR] mdns: Failed to read packet: %v", err) diff --git a/server.go b/server.go index 55988c4..272b439 100644 --- a/server.go +++ b/server.go @@ -81,6 +81,12 @@ func NewServer(config *Config) (*Server, error) { return s, nil } +func (s *Server) getIsShutdown() bool { + s.shutdownLock.Lock() + defer s.shutdownLock.Unlock() + return s.shutdown +} + // Shutdown is used to shutdown the listener func (s *Server) Shutdown() error { s.shutdownLock.Lock() @@ -107,7 +113,7 @@ func (s *Server) recv(c *net.UDPConn) { return } buf := make([]byte, 65536) - for !s.shutdown { + for !s.getIsShutdown() { n, from, err := c.ReadFrom(buf) if err != nil { continue