diff --git a/download.go b/download.go new file mode 100644 index 0000000..882b66d --- /dev/null +++ b/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() +} diff --git a/ismap.go b/ismap.go index e880580..f6b468a 100644 --- a/ismap.go +++ b/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 diff --git a/wrp.go b/wrp.go index 8349693..014526e 100644 --- a/wrp.go +++ b/wrp.go @@ -239,7 +239,10 @@ func proxyServer(w http.ResponseWriter, r *http.Request) { chromedp.Run(ctx, chromedp.Location(¤tURL)) 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)