support multiple hosts - by qwencoder

main.go

@@ -1,4 +1,6 @@
 // Save and serve HTML local storage in server
+// This application provides a wrapper to save and restore localStorage of webapps in a server.
+// It's designed for browsers that clear localStorage when exiting (like Tor Browser) or to share localStorage across multiple browsers.
 package main
 
 import (
@@ -9,6 +11,7 @@ import (
 	"log"
 	"net/http"
 	"os"
+	"strings"
 	"time"
 
 	"go.balki.me/anyhttp"
@@ -17,57 +20,84 @@ import (
 //go:embed wrap.html
 var html []byte
 
-func main() {
+// getDumpFileName generates a hostname-specific dump filename
+func getDumpFileName(host string) string {
+	// Extract hostname from Host header
+	hostname := host
+	if hostname == "" {
+		hostname = "localhost"
+	}
+	// Sanitize hostname to prevent path traversal
+	hostname = strings.ReplaceAll(hostname, "/", "_")
+	hostname = strings.ReplaceAll(hostname, "..", "_")
 
+	// Create dump file name with hostname
+	return hostname + "-dump.json"
+}
+
+func main() {
 	addr := flag.String("address", "localhost:3242", "Listen address. See go.balki.me/anyhttp for usage")
-	dumpFile := flag.String("dumpfile", "dump.json", "Path where local storage is saved in server")
 
 	flag.Parse()
 
+	// Handle root path - show error message
 	http.HandleFunc("/", func(w http.ResponseWriter, _ *http.Request) {
 		if _, err := w.Write([]byte("<h2>Error!</h2><p>Check your webserver config, You should not see this!</p>")); err != nil {
 			log.Panic(err)
 		}
 	})
 
+	// Serve the wrap.html page at /wrap/ path
 	http.HandleFunc("/wrap/", func(w http.ResponseWriter, r *http.Request) {
 		http.ServeContent(w, r, "index.html", time.Now(), bytes.NewReader(html))
 	})
 
+	// Handle saving local storage data to server
 	http.HandleFunc("/wrap/saveLS", func(w http.ResponseWriter, r *http.Request) {
-
+		// Read the request body containing localStorage data
 		data, err := io.ReadAll(r.Body)
 		if err != nil {
 			log.Panic(err)
 		}
 
-		if err := os.WriteFile(*dumpFile, data, 0o600); err != nil {
-			log.Panic(err)
-		}
+		// Generate dump filename
+		dumpFile := getDumpFileName(r.Host)
 
-		if err != nil {
+		// Write the data to the dump file
+		if err := os.WriteFile(dumpFile, data, 0o600); err != nil {
 			log.Panic(err)
 		}
 
 		log.Println("Saved!")
 
+		// Send OK response
 		if _, err := w.Write([]byte("OK")); err != nil {
 			log.Panic(err)
 		}
 	})
 
-	http.HandleFunc("/wrap/loadLS", func(w http.ResponseWriter, _ *http.Request) {
+	// Handle loading local storage data from server
+	http.HandleFunc("/wrap/loadLS", func(w http.ResponseWriter, r *http.Request) {
+		// Set proper content type for JSON response
 		w.Header().Set("Content-Type", "application/json")
-		data, err := os.ReadFile(*dumpFile)
-		if err != nil { // First time
+
+		// Generate dump filename
+		dumpFile := getDumpFileName(r.Host)
+
+		// Read the dump file
+		data, err := os.ReadFile(dumpFile)
+		if err != nil { // First time - return empty JSON
 			data = []byte("{}")
 		}
+
+		// Send the data back to client
 		if _, err := w.Write(data); err != nil {
 			log.Panic(err)
 		}
 		log.Println("Sent!")
 	})
 
+	// Start the HTTP server
 	err := anyhttp.ListenAndServe(*addr, nil)
 	if err != nil {
 		log.Panic(err)