add download support

download.go

@@ -0,0 +1,163 @@
+package main
+
+import (
+	"fmt"
+	"log"
+	"net/http"
+	"os"
+	"path/filepath"
+	"strconv"
+	"strings"
+	"sync"
+	"time"
+
+	"github.com/chromedp/cdproto/browser"
+	"github.com/chromedp/chromedp"
+	"github.com/lithammer/shortuuid/v4"
+)
+
+var (
+	dlDir    string
+	dlNotify = make(chan struct{}, 1)
+)
+
+type dlEvent struct {
+	guid     string
+	filename string
+	done     chan struct{}
+}
+
+var dlTrack struct {
+	sync.Mutex
+	ev *dlEvent
+}
+
+type dlFile struct {
+	name string
+	data []byte
+}
+
+var dlCache struct {
+	sync.Mutex
+	files map[string]dlFile
+}
+
+func setupDownloads() {
+	if dlDir == "" {
+		var err error
+		dlDir, err = os.MkdirTemp("", "wrp-dl-")
+		if err != nil {
+			log.Fatalf("Failed to create download dir: %v", err)
+		}
+		dlCache.files = make(map[string]dlFile)
+	}
+	chromedp.ListenTarget(ctx, func(ev interface{}) {
+		switch e := ev.(type) {
+		case *browser.EventDownloadWillBegin:
+			log.Printf("Download started: %s (%s)", e.SuggestedFilename, e.GUID)
+			dlTrack.Lock()
+			dlTrack.ev = &dlEvent{
+				guid:     e.GUID,
+				filename: e.SuggestedFilename,
+				done:     make(chan struct{}),
+			}
+			dlTrack.Unlock()
+			select {
+			case dlNotify <- struct{}{}:
+			default:
+			}
+		case *browser.EventDownloadProgress:
+			dlTrack.Lock()
+			ev := dlTrack.ev
+			dlTrack.Unlock()
+			if ev == nil || ev.guid != e.GUID {
+				return
+			}
+			if e.State == browser.DownloadProgressStateCompleted || e.State == browser.DownloadProgressStateCanceled {
+				log.Printf("Download %s: %s", e.State, e.GUID)
+				close(ev.done)
+			}
+		}
+	})
+	err := chromedp.Run(ctx,
+		browser.SetDownloadBehavior(browser.SetDownloadBehaviorBehaviorAllowAndName).
+			WithDownloadPath(dlDir).
+			WithEventsEnabled(true),
+	)
+	if err != nil {
+		log.Printf("Warning: failed to set download behavior: %v", err)
+	}
+}
+
+func resetDownloadState() {
+	dlTrack.Lock()
+	dlTrack.ev = nil
+	dlTrack.Unlock()
+	select {
+	case <-dlNotify:
+	default:
+	}
+}
+
+func waitForDownload() string {
+	select {
+	case <-dlNotify:
+	case <-time.After(500 * time.Millisecond):
+		return ""
+	}
+	dlTrack.Lock()
+	ev := dlTrack.ev
+	dlTrack.Unlock()
+	if ev == nil {
+		return ""
+	}
+	log.Printf("Waiting for download: %s", ev.filename)
+	select {
+	case <-ev.done:
+	case <-time.After(60 * time.Second):
+		log.Printf("Download timed out: %s", ev.guid)
+		dlTrack.Lock()
+		dlTrack.ev = nil
+		dlTrack.Unlock()
+		return ""
+	}
+	fpath := filepath.Join(dlDir, ev.guid)
+	data, err := os.ReadFile(fpath)
+	if err != nil {
+		log.Printf("Failed to read download %s: %v", fpath, err)
+		dlTrack.Lock()
+		dlTrack.ev = nil
+		dlTrack.Unlock()
+		return ""
+	}
+	os.Remove(fpath)
+	id := shortuuid.New()
+	dlCache.Lock()
+	dlCache.files[id] = dlFile{name: ev.filename, data: data}
+	dlCache.Unlock()
+	dlTrack.Lock()
+	dlTrack.ev = nil
+	dlTrack.Unlock()
+	log.Printf("Download cached: /dl/%s (%s, %d bytes)", id, ev.filename, len(data))
+	return "/dl/" + id
+}
+
+func dlServer(w http.ResponseWriter, r *http.Request) {
+	id := strings.TrimPrefix(r.URL.Path, "/dl/")
+	log.Printf("%s Download request for %s", r.RemoteAddr, id)
+	dlCache.Lock()
+	f, ok := dlCache.files[id]
+	dlCache.Unlock()
+	if !ok {
+		http.NotFound(w, r)
+		return
+	}
+	w.Header().Set("Content-Disposition", fmt.Sprintf(`attachment; filename="%s"`, f.name))
+	w.Header().Set("Content-Type", http.DetectContentType(f.data))
+	w.Header().Set("Content-Length", strconv.Itoa(len(f.data)))
+	w.Write(f.data)
+	w.(http.Flusher).Flush()
+	dlCache.Lock()
+	delete(dlCache.files, id)
+	dlCache.Unlock()
+}

ismap.go

@@ -152,9 +152,11 @@ func (rq *wrpReq) action() chromedp.Action {
 	return chromedp.Navigate(rq.url)
 }
 
-// Navigate to the desired URL.
-func (rq *wrpReq) navigate() {
+// Navigate to the desired URL, returns download path if a file download was triggered.
+func (rq *wrpReq) navigate() string {
+	resetDownloadState()
 	ctxErr(chromedp.Run(ctx, rq.action()), rq.w)
+	return waitForDownload()
 }
 
 // Handle context errors
@@ -170,6 +172,7 @@ func ctxErr(err error, w io.Writer) {
 		return
 	}
 	ctx, cncl = chromedp.NewContext(actx)
+	setupDownloads()
 	log.Printf("Created new context, try again")
 	fmt.Fprintln(w, "Created new context, try again")
 }
@@ -382,7 +385,10 @@ func mapServer(w http.ResponseWriter, r *http.Request) {
 		rq.printUI(uiParams{})
 		return
 	}
-	rq.navigate() // TODO: if error from navigate do not capture
+	if dl := rq.navigate(); dl != "" {
+		http.Redirect(w, r, dl, http.StatusFound)
+		return
+	}
 	if rq.proxy {
 		chromedp.Run(ctx, waitForRender())
 		var loc string

wrp.go

@@ -239,7 +239,10 @@ func proxyServer(w http.ResponseWriter, r *http.Request) {
 	chromedp.Run(ctx, chromedp.Location(&currentURL))
 	currentURL = strings.Replace(currentURL, "https://", "http://", 1)
 	if currentURL != strings.Replace(rq.url, "https://", "http://", 1) {
-		rq.navigate()
+		if dl := rq.navigate(); dl != "" {
+			http.Redirect(w, r, dl, http.StatusFound)
+			return
+		}
 	}
 	rq.url = strings.Replace(rq.url, "https://", "http://", 1)
 	if r.Method == "CONNECT" {
@@ -273,7 +276,10 @@ func pageServer(w http.ResponseWriter, r *http.Request) {
 		rq.printUI(uiParams{})
 		return
 	}
-	rq.navigate() // TODO: if error from navigate do not capture
+	if dl := rq.navigate(); dl != "" {
+		http.Redirect(w, r, dl, http.StatusFound)
+		return
+	}
 	if rq.wrpMode == "html" {
 		rq.captureMarkdown()
 		return
@@ -359,12 +365,14 @@ func main() {
 	cncl, acncl = chromedpStart()
 	defer cncl()
 	defer acncl()
+	setupDownloads()
 
 	c := make(chan os.Signal)
 	signal.Notify(c, os.Interrupt, syscall.SIGTERM)
 	go func() {
 		<-c
 		log.Printf("Interrupt - shutting down.")
+		os.RemoveAll(dlDir)
 		cncl()
 		acncl()
 		srv.Shutdown(context.Background())
@@ -380,6 +388,7 @@ func main() {
 	http.HandleFunc(imgZpfx, imgServerTxt)
 	http.HandleFunc("/proxy.pac", pacServer)
 	http.HandleFunc("/shutdown/", haltServer)
+	http.HandleFunc("/dl/", dlServer)
 	http.HandleFunc("/favicon.ico", http.NotFound)
 
 	log.Printf("Default Img Type: %v, Geometry: %+v", *defType, defGeom)