diff --git a/download.go b/download.go index 882b66d..63e1064 100644 --- a/download.go +++ b/download.go @@ -99,17 +99,17 @@ func resetDownloadState() { } } -func waitForDownload() string { +func waitForDownload() *dlFile { select { case <-dlNotify: case <-time.After(500 * time.Millisecond): - return "" + return nil } dlTrack.Lock() ev := dlTrack.ev dlTrack.Unlock() if ev == nil { - return "" + return nil } log.Printf("Waiting for download: %s", ev.filename) select { @@ -119,45 +119,49 @@ func waitForDownload() string { dlTrack.Lock() dlTrack.ev = nil dlTrack.Unlock() - return "" + return nil } 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)) + if err != nil { + log.Printf("Failed to read download %s: %v", fpath, err) + return nil + } + os.Remove(fpath) + return &dlFile{name: ev.filename, data: data} +} + +func cacheDownload(f *dlFile) string { + id := shortuuid.New() + dlCache.Lock() + dlCache.files[id] = *f + dlCache.Unlock() + log.Printf("Download cached: /dl/%s (%s, %d bytes)", id, f.name, len(f.data)) return "/dl/" + id } +func writeDownload(w http.ResponseWriter, f *dlFile) { + 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() + log.Printf("Download served inline: %s (%d bytes)", f.name, len(f.data)) +} + 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] + delete(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() + writeDownload(w, &f) } diff --git a/ismap.go b/ismap.go index f6b468a..cd5e300 100644 --- a/ismap.go +++ b/ismap.go @@ -152,8 +152,8 @@ func (rq *wrpReq) action() chromedp.Action { return chromedp.Navigate(rq.url) } -// Navigate to the desired URL, returns download path if a file download was triggered. -func (rq *wrpReq) navigate() string { +// Navigate to the desired URL, returns the downloaded file if one was triggered. +func (rq *wrpReq) navigate() *dlFile { resetDownloadState() ctxErr(chromedp.Run(ctx, rq.action()), rq.w) return waitForDownload() @@ -385,8 +385,12 @@ func mapServer(w http.ResponseWriter, r *http.Request) { rq.printUI(uiParams{}) return } - if dl := rq.navigate(); dl != "" { - http.Redirect(w, r, dl, http.StatusFound) + if dl := rq.navigate(); dl != nil { + if rq.proxy { + writeDownload(w, dl) + } else { + http.Redirect(w, r, cacheDownload(dl), http.StatusFound) + } return } if rq.proxy { diff --git a/wrp.go b/wrp.go index 014526e..f39d121 100644 --- a/wrp.go +++ b/wrp.go @@ -239,8 +239,8 @@ 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) { - if dl := rq.navigate(); dl != "" { - http.Redirect(w, r, dl, http.StatusFound) + if dl := rq.navigate(); dl != nil { + writeDownload(w, dl) return } } @@ -276,8 +276,8 @@ func pageServer(w http.ResponseWriter, r *http.Request) { rq.printUI(uiParams{}) return } - if dl := rq.navigate(); dl != "" { - http.Redirect(w, r, dl, http.StatusFound) + if dl := rq.navigate(); dl != nil { + http.Redirect(w, r, cacheDownload(dl), http.StatusFound) return } if rq.wrpMode == "html" {