1 // Copyright (c) 2024 Alexander Arkhipov <aa@manpager.org>
2 //
3 // Permission to use, copy, modify, and distribute this software for any
4 // purpose with or without fee is hereby granted, provided that the above
5 // copyright notice and this permission notice appear in all copies.
6 //
7 // THE SOFTWARE IS PROVIDED "AS IS" AND THE AUTHOR DISCLAIMS ALL WARRANTIES
8 // WITH REGARD TO THIS SOFTWARE INCLUDING ALL IMPLIED WARRANTIES OF
9 // MERCHANTABILITY AND FITNESS. IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR
10 // ANY SPECIAL, DIRECT, INDIRECT, OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES
11 // WHATSOEVER RESULTING FROM LOSS OF USE, DATA OR PROFITS, WHETHER IN AN
12 // ACTION OF CONTRACT, NEGLIGENCE OR OTHER TORTIOUS ACTION, ARISING OUT OF
13 // OR IN CONNECTION WITH THE USE OR PERFORMANCE OF THIS SOFTWARE.
14 
15 package main
16 
17 import (
18 	"bufio"
19 	"errors"
20 	"flag"
21 	"fmt"
22 	"io"
23 	"io/fs"
24 	"net/http"
25 	"os"
26 	"strings"
27 )
28 
29 var qflag = flag.Bool("q", false, "be quiet")
30 var pflag = flag.Int("p", 1, "number of parallel downloads")
31 
32 type filename struct {
33 	n        int      // how many times the url has been accessed
34 	name     string   // end-file name
35 	tmpfiles []string // temporary file names
36 }
37 
38 // filemap maps URLs to corresponding filenames
39 var filemap = make(map[string]filename)
40 
41 func getUrl(url, f string, ch chan int) {
42 	defer func() { ch <- 0 }()
43 
44 	rm := func() {
45 		os.Remove(f)
46 	}
47 
48 	if !*qflag {
49 		fmt.Println("GET", url)
50 	}
51 
52 	fp, err := os.Create(f)
53 	if err != nil {
54 		fmt.Fprintln(os.Stderr, err)
55 		rm()
56 		return
57 	}
58 	defer fp.Close()
59 	fmt.Println("created", fp.Name())
60 
61 	resp, err := http.Get(url)
62 	if err != nil {
63 		fmt.Fprintln(os.Stderr, err)
64 		rm()
65 		return
66 	}
67 
68 	buf := make([]byte, 4096)
69 	reader := bufio.NewReader(resp.Body)
70 	writer := bufio.NewWriter(fp)
71 
72 	for readErr := error(nil); readErr == nil; {
73 		n, readErr := io.ReadFull(reader, buf)
74 		if readErr == io.EOF {
75 			break
76 		}
77 		if readErr != nil && readErr != io.ErrUnexpectedEOF {
78 			fmt.Fprintln(os.Stderr, readErr)
79 			rm()
80 			break
81 		}
82 
83 		_, writeErr := writer.Write(buf[:n])
84 		if writeErr != nil {
85 			fmt.Fprintln(os.Stderr, writeErr)
86 			rm()
87 			break
88 		}
89 	}
90 	writer.Flush()
91 }
92 
93 func prepUrl(url, d string) (string, error) {
94 	if !(strings.HasPrefix(url, "http://") || strings.HasPrefix(url, "https://")) {
95 		url = "http://" + url
96 	}
97 
98 	fmentry := filemap[url]
99 	defer func() { filemap[url] = fmentry }()
100 
101 	var fname string
102 
103 	_, fname, _ = strings.Cut(url, "://")
104 	_, fname, _ = strings.Cut(fname, "/")
105 	parts := strings.Split(fname, "/")
106 	fname = parts[len(parts)-1]
107 	if fname == "" {
108 		fname = "index.html"
109 	}
110 
111 	tmpfp, err := os.CreateTemp(d, fname+"*")
112 	if err != nil {
113 		return "", err
114 	}
115 	defer tmpfp.Close()
116 
117 	fmentry.name = fname
118 	fmentry.tmpfiles = append(fmentry.tmpfiles, tmpfp.Name())
119 
120 	return url, nil
121 }
122 
123 func main() {
124 	flag.Parse()
125 
126 	if *pflag < 1 {
127 		fmt.Fprintln(os.Stderr, "can't do less than 1 parallel downloads")
128 		os.Exit(1)
129 	}
130 
131 	var urls []string
132 
133 	tmpdir, err := os.MkdirTemp(".", ".goget*")
134 	if err != nil {
135 		fmt.Fprintln(os.Stderr, err)
136 		os.Exit(1)
137 	}
138 	defer func() {
139 		rename := func(url string) {
140 			fentry := filemap[url]
141 			defer func() {
142 				fentry.tmpfiles = fentry.tmpfiles[1:]
143 				filemap[url] = fentry
144 			}()
145 			os.Rename(fentry.tmpfiles[0], fentry.name)
146 			// Ignoring ErrNotExist since the temporary file might
147 			// have been removed on purpose.
148 			//
149 			// TODO It would be better to have such purposeful
150 			// removals marked explicitly.
151 			if err != nil && !errors.Is(err, fs.ErrNotExist) {
152 				fmt.Fprintln(os.Stderr, err)
153 			}
154 		}
155 
156 		for _, url := range urls {
157 			rename(url)
158 		}
159 
160 		err := os.Remove(tmpdir)
161 		if err != nil {
162 			fmt.Fprintln(os.Stderr, err)
163 			os.Exit(1)
164 		}
165 	}()
166 
167 	for _, arg := range flag.Args() {
168 		url, err := prepUrl(arg, tmpdir)
169 		if err != nil {
170 			fmt.Fprintln(os.Stderr, err)
171 			continue
172 		}
173 		urls = append(urls, url)
174 	}
175 
176 	ch := make(chan int, *pflag)
177 	routines := 0
178 	for _, url := range urls {
179 		if routines >= *pflag {
180 			<-ch
181 			routines--
182 		}
183 		if fmentry, ok := filemap[url]; ok {
184 			go getUrl(url, fmentry.tmpfiles[fmentry.n], ch)
185 			fmentry.n++
186 			filemap[url] = fmentry
187 			routines++
188 		}
189 	}
190 
191 	for routines > 0 {
192 		<-ch
193 		routines--
194 	}
195 }