feat: save vhosts to disk and load on startup

Author: chetan

daemon.go

@@ -2,11 +2,13 @@ package vproxy
 
 import (
 	"context"
+	"encoding/json"
 	"fmt"
 	"log"
 	"net"
 	"net/http"
 	"os"
+	"path"
 	"strconv"
 	"strings"
 	"sync"
@@ -43,7 +45,9 @@ type Daemon struct {
 
 // NewDaemon
 func NewDaemon(lh *LoggedHandler, listen string, httpPort int, httpsPort int) *Daemon {
-	return &Daemon{loggedHandler: lh, listenHost: listen, httpPort: httpPort, httpsPort: httpsPort}
+	d := &Daemon{loggedHandler: lh, listenHost: listen, httpPort: httpPort, httpsPort: httpsPort}
+	d.loadVhosts()
+	return d
 }
 
 func rerunWithSudo(addr string) {
@@ -275,6 +279,52 @@ func (d *Daemon) doRemoveVhost(vhost *Vhost, w http.ResponseWriter) {
 	fmt.Fprintf(w, "removing vhost: %s -> %d\n", vhost.Host, vhost.ServicePort)
 	vhost.Close()
 	d.loggedHandler.RemoveVhost(vhost.Host)
+	d.saveVhosts()
+}
+
+// load saved vhosts from disk
+func (d *Daemon) loadVhosts() {
+	c := path.Join(CertPath(), "vhosts.json")
+	j, err := os.ReadFile(c)
+	if err != nil {
+		fmt.Println("[*] warning: failed to load vhosts from disk: ", err)
+		return
+	}
+	servers := map[string]*Vhost{}
+	err = json.Unmarshal(j, &servers)
+	if err != nil {
+		fmt.Println("[*] warning: failed to load vhosts from disk: ", err)
+		return
+	}
+	if len(servers) == 0 {
+		return
+	}
+	for _, vhost := range servers {
+		vhost.Init()
+		d.loggedHandler.AddVhost(vhost)
+		err = addToHosts(vhost.Host)
+		if err != nil {
+			msg := fmt.Sprintf("[*] warning: failed to add %s to system hosts file: %s\n", vhost.Host, err)
+			fmt.Println(msg)
+		}
+		// fmt.Printf("[*] loaded vhost: %s\n", vhost.String())
+	}
+	fmt.Printf("[*] loaded %d vhost(s)\n", len(servers))
+}
+
+// save vhosts to disk
+func (d *Daemon) saveVhosts() {
+	c := path.Join(CertPath(), "vhosts.json")
+	j, err := json.Marshal(d.loggedHandler.vhostMux.Servers)
+	if err != nil {
+		fmt.Println("[*] warning: failed to save vhosts to disk: ", err)
+		return
+	}
+	err = os.WriteFile(c, j, 0644)
+	if err != nil {
+		fmt.Println("[*] warning: failed to save vhosts to disk: ", err)
+		return
+	}
 }
 
 // addVhost for the given binding to the LoggedHandler
@@ -303,6 +353,8 @@ func (d *Daemon) addVhost(binding string, w http.ResponseWriter) *Vhost {
 	w.Header().Set("Access-Control-Allow-Origin", "*")
 
 	d.loggedHandler.AddVhost(vhost)
+	d.saveVhosts()
+
 	if d.enableTLS() {
 		d.restartTLS()
 	}

vhost.go

@@ -15,18 +15,18 @@ import (
 
 // Vhost represents a single backend service
 type Vhost struct {
-	Host string // virtual host name
+	Host string `json:"host"` // virtual host name
 
 	ServiceHost string // service host or IP
 	ServicePort int    // service port
 
-	Handler http.Handler
-	Cert    string // TLS Certificate
-	Key     string // TLS Private Key
+	Handler http.Handler `json:"-"`
+	Cert    string       // TLS Certificate
+	Key     string       // TLS Private Key
 
-	logRing   *deque.Deque[string]
-	logChan   LogListener
-	listeners []LogListener
+	logRing   *deque.Deque[string] `json:"-"`
+	logChan   LogListener          `json:"-"`
+	listeners []LogListener        `json:"-"`
 }
 
 type LogListener chan string
@@ -107,22 +107,13 @@ func CreateVhost(input string, useTLS bool) (*Vhost, error) {
 		return nil, fmt.Errorf("failed to parse target port: %s", err)
 	}
 	targetHost := "127.0.0.1"
-	targetURL := url.URL{Scheme: "http", Host: fmt.Sprintf("%s:%d", targetHost, targetPort)}
-
-	proxy := CreateProxy(targetURL, hostname)
 
 	vhost := &Vhost{
-		Host: hostname, ServiceHost: targetHost, ServicePort: targetPort, Handler: proxy,
-		logRing: &deque.Deque[string]{},
-		logChan: make(LogListener, 10),
+		Host:        hostname,
+		ServiceHost: targetHost,
+		ServicePort: targetPort,
 	}
 
-	// set fixed capacity at 16
-	vhost.logRing.Grow(16)
-	vhost.logRing.SetBaseCap(16)
-
-	go vhost.populateLogBuffer()
-
 	if useTLS {
 		vhost.Cert, vhost.Key, err = MakeCert(hostname)
 		if err != nil {
@@ -133,6 +124,17 @@ func CreateVhost(input string, useTLS bool) (*Vhost, error) {
 	return vhost, nil
 }
 
+func (v *Vhost) Init() {
+	targetURL := url.URL{Scheme: "http", Host: fmt.Sprintf("%s:%d", v.ServiceHost, v.ServicePort)}
+	v.Handler = CreateProxy(targetURL, v.Host)
+	v.logChan = make(LogListener, 10)
+	// set fixed capacity at 16
+	v.logRing = &deque.Deque[string]{}
+	v.logRing.Grow(16)
+	v.logRing.SetBaseCap(16)
+	go v.populateLogBuffer()
+}
+
 func (v *Vhost) NewLogListener() LogListener {
 	logChan := make(LogListener, 100)
 	v.listeners = append(v.listeners, logChan)
@@ -180,6 +182,10 @@ func (v *Vhost) populateLogBuffer() {
 	}
 }
 
+func (v Vhost) String() string {
+	return fmt.Sprintf("%s -> %s:%d", v.Host, v.ServiceHost, v.ServicePort)
+}
+
 // Map given host to 127.0.0.1 in system hosts file (usually /etc/hosts)
 func addToHosts(host string) error {
 	hosts, err := txeh.NewHostsDefault()