feat: implemented disconnect command

Author: chetan

.gitignore

@@ -30,3 +30,4 @@ _testmain.go
 dist/
 .vproxy.conf
 t/
+*.pprof

bin/vproxy/flags.go

@@ -98,7 +98,7 @@ vproxy connect hello.local:8888 -- vproxy hello
 				Name:    "connect",
 				Aliases: []string{"client", "c", "add"},
 				Usage:   "Add a new vhost",
-				Action:  startClient,
+				Action:  connectVhost,
 				Before:  loadClientConfig,
 				Flags: []cli.Flag{
 					&cli.StringFlag{
@@ -121,6 +121,30 @@ vproxy connect hello.local:8888 -- vproxy hello
 					},
 				},
 			},
+			{
+				Name:      "disconnect",
+				Aliases:   []string{"remove"},
+				Usage:     "Remove vhost",
+				Action:    disconnectVhost,
+				Before:    loadClientConfig,
+				UsageText: `vproxy disconnect [command options] <hostname>`,
+				Flags: []cli.Flag{
+					&cli.StringFlag{
+						Name:  "host",
+						Value: "127.0.0.1",
+						Usage: "Daemon host IP",
+					},
+					&cli.IntFlag{
+						Name:  "http",
+						Value: 80,
+						Usage: "Daemon HTTP port",
+					},
+					&cli.BoolFlag{
+						Name:  "all",
+						Usage: "Remove all vhosts",
+					},
+				},
+			},
 			{
 				Name:      "tail",
 				Aliases:   []string{"stream", "attach"},

bin/vproxy/main.go

@@ -37,7 +37,7 @@ func printVersion(c *cli.Context) error {
 	return nil
 }
 
-func startClient(c *cli.Context) error {
+func connectVhost(c *cli.Context) error {
 
 	// collect and validate binds
 	args := c.Args().Slice()
@@ -74,6 +74,13 @@ func startClient(c *cli.Context) error {
 	return nil
 }
 
+func disconnectVhost(c *cli.Context) error {
+	hostname := c.Args().First()
+	client := createClient(c)
+	client.RemoveVhost(hostname, c.Bool("all"))
+	return nil
+}
+
 func createClient(c *cli.Context) *vproxy.Client {
 	host := c.String("host")
 	httpPort := c.Int("http")

client.go

@@ -10,6 +10,7 @@ import (
 	"os"
 	"os/exec"
 	"os/signal"
+	"strconv"
 	"strings"
 	"sync"
 )
@@ -71,7 +72,7 @@ func (c *Client) addBinding(bind string, detach bool) {
 
 	s := strings.Split(bind, ":")
 	if len(s) >= 2 {
-		fmt.Printf("[*] registering vhost: %s -> https://%s\n", bind, s[0])
+		fmt.Printf("[*] registering vhost: https://%s -> %s\n", s[0], bind)
 	} else {
 		fmt.Println("[*] registering vhost:", bind)
 	}
@@ -94,8 +95,7 @@ func (c *Client) Tail(hostname string, follow bool) {
 	data.Add("host", hostname)
 	res, err := http.DefaultClient.PostForm(c.uri("/clients/stream"), data)
 	if err != nil {
-		stopCommand(c.cmd)
-		log.Fatalf("error registering client: %s\n", err)
+		log.Fatalf("error: %s\n", err)
 	}
 	fmt.Printf("[*] streaming logs for %s\n", hostname)
 	streamLogs(res, follow)
@@ -126,6 +126,20 @@ func streamLogs(res *http.Response, follow bool) {
 	}
 }
 
+func (c *Client) RemoveVhost(hostname string, all bool) {
+	data := url.Values{}
+	data.Add("host", hostname)
+	data.Add("all", strconv.FormatBool(all))
+	res, err := http.DefaultClient.PostForm(c.uri("/clients/remove"), data)
+	if err != nil {
+		log.Fatalf("error: %s\n", err)
+	}
+	defer res.Body.Close()
+	r := bufio.NewReader(res.Body)
+	b, _ := ioutil.ReadAll(r)
+	fmt.Println(string(b))
+}
+
 // IsDaemonRunning tries to check if a vproxy daemon is already running on the given addr
 func (c *Client) IsDaemonRunning() bool {
 	res, err := http.DefaultClient.Get(c.uri("/hello"))

daemon.go

@@ -110,6 +110,7 @@ func (d *Daemon) Run() {
 	d.loggedHandler.HandleFunc("/_vproxy/clients", d.listClients)
 	d.loggedHandler.HandleFunc("/_vproxy/clients/add", d.registerVhost)
 	d.loggedHandler.HandleFunc("/_vproxy/clients/stream", d.streamLogs)
+	d.loggedHandler.HandleFunc("/_vproxy/clients/remove", d.removeVhost)
 	d.wg.Add(1) // ensure we don't exit immediately
 
 	if d.enableHTTP() {
@@ -237,6 +238,20 @@ 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)
+		return
+	}
+
+	fmt.Printf("[*] removing vhost: %s -> %d\n", vhost.Host, vhost.Port)
+	fmt.Fprintf(w, "removing vhost: %s -> %d", vhost.Host, vhost.Port)
+	vhost.Close()
+	d.loggedHandler.RemoveVhost(vhost.Host)
+}
+
 // addVhost for the given binding to the LoggedHandler
 func (d *Daemon) addVhost(binding string, w http.ResponseWriter) *Vhost {
 	vhost, err := CreateVhost(binding, d.enableTLS())