Add CloseAllConns() and SetRefuse()

proxy.go

@@ -13,6 +13,7 @@ type Proxy struct {
 	listener *net.TCPListener
 
 	conns  map[*net.TCPConn]bool
+	refuse bool
 	mu     sync.Mutex
 }
 
@@ -41,10 +42,8 @@ func (p *Proxy) Addr() *net.TCPAddr {
 	return p.listener.Addr().(*net.TCPAddr)
 }
 
-func (p *Proxy) Close() {
-	p.t.Logf("* -> %s -> [proxy] -> * -> %s closing...", p.listener.Addr(), p.backend)
-
-	p.listener.Close()
+func (p *Proxy) CloseAllConns() {
+	p.t.Logf("* -> %s -> [proxy] -> * -> %s closing all connections...", p.listener.Addr(), p.backend)
 
 	p.mu.Lock()
 	defer p.mu.Unlock()
@@ -54,6 +53,20 @@ func (p *Proxy) Close() {
 	}
 
 	p.conns = map[*net.TCPConn]bool{}
+}
+
+func (p *Proxy) SetRefuse(refuse bool) {
+	p.mu.Lock()
+	defer p.mu.Unlock()
+
+	p.refuse = true
+}
+
+func (p *Proxy) Close() {
+	p.t.Logf("* -> %s -> [proxy] -> * -> %s closing...", p.listener.Addr(), p.backend)
+
+	p.listener.Close()
+	p.CloseAllConns()
 
 	p.t.Logf("* -> %s -> [proxy] -> * -> %s closed", p.listener.Addr(), p.backend)
 }
@@ -65,8 +78,17 @@ func (p *Proxy) accept() {
 			return
 		}
 
+		p.mu.Lock()
+
+		if p.refuse {
+			p.t.Logf("%s -> %s -> [proxy] -> * -> %s refusing...", frontConn.RemoteAddr(), frontConn.LocalAddr(), p.backend)
+			frontConn.Close()
+		} else {
 			go p.dial(frontConn)
 		}
+
+		p.mu.Unlock()
+	}
 }
 
 func (p *Proxy) dial(frontConn *net.TCPConn) {