From 164637bad61b89ef3bbbd25d6642dc9ea4860daa Mon Sep 17 00:00:00 2001
From: Chetan Sarva <chetan@pixelcop.net>
Date: Thu, 4 Nov 2021 14:50:48 -0400
Subject: [PATCH] feat: implemented disconect --all flag

---
 bin/vproxy/main.go |  8 +++++++-
 client.go          |  5 ++++-
 daemon.go          | 26 ++++++++++++++++++++++----
 3 files changed, 33 insertions(+), 6 deletions(-)

diff --git a/bin/vproxy/main.go b/bin/vproxy/main.go
index 96d09f7..6a479ee 100644
--- a/bin/vproxy/main.go
+++ b/bin/vproxy/main.go
@@ -76,8 +76,14 @@ func connectVhost(c *cli.Context) error {
 
 func disconnectVhost(c *cli.Context) error {
 	hostname := c.Args().First()
+	all := c.Bool("all")
+
+	if !all && hostname == "" {
+		return fmt.Errorf("missing hostname or --all")
+	}
+
 	client := createClient(c)
-	client.RemoveVhost(hostname, c.Bool("all"))
+	client.RemoveVhost(hostname, all)
 	return nil
 }
 
diff --git a/client.go b/client.go
index 5c7276d..93b74b1 100644
--- a/client.go
+++ b/client.go
@@ -137,7 +137,10 @@ func (c *Client) RemoveVhost(hostname string, all bool) {
 	defer res.Body.Close()
 	r := bufio.NewReader(res.Body)
 	b, _ := ioutil.ReadAll(r)
-	fmt.Println(string(b))
+	s := strings.TrimSpace(string(b))
+	if s != "" {
+		fmt.Println(s)
+	}
 }
 
 // IsDaemonRunning tries to check if a vproxy daemon is already running on the given addr
diff --git a/daemon.go b/daemon.go
index 7de7639..4fb8c46 100644
--- a/daemon.go
+++ b/daemon.go
@@ -7,6 +7,7 @@ import (
 	"net"
 	"net/http"
 	"os"
+	"strconv"
 	"strings"
 	"sync"
 	"syscall"
@@ -240,14 +241,31 @@ func (d *Daemon) relayLogsUntilClose(vhost *Vhost, w http.ResponseWriter, reqCtx
 
 func (d *Daemon) removeVhost(w http.ResponseWriter, r *http.Request) {
 	hostname := r.PostFormValue("host")
-	vhost := d.loggedHandler.GetVhost(hostname)
-	if vhost == nil {
-		fmt.Fprintf(w, "error: host '%s' not found", hostname)
+	all, _ := strconv.ParseBool(r.PostFormValue("all"))
+
+	if all {
+		for _, vhost := range d.loggedHandler.vhostMux.Servers {
+			d.doRemoveVhost(vhost, w)
+		}
+
+	} else if hostname != "" {
+		vhost := d.loggedHandler.GetVhost(hostname)
+		if vhost == nil {
+			fmt.Fprintf(w, "error: host '%s' not found", hostname)
+			return
+		}
+		d.doRemoveVhost(vhost, w)
+
+	} else {
+		fmt.Fprint(w, "error: missing hostname")
 		return
 	}
 
+}
+
+func (d *Daemon) doRemoveVhost(vhost *Vhost, w http.ResponseWriter) {
 	fmt.Printf("[*] removing vhost: %s -> %d\n", vhost.Host, vhost.Port)
-	fmt.Fprintf(w, "removing vhost: %s -> %d", vhost.Host, vhost.Port)
+	fmt.Fprintf(w, "removing vhost: %s -> %d\n", vhost.Host, vhost.Port)
 	vhost.Close()
 	d.loggedHandler.RemoveVhost(vhost.Host)
 }