Skip to content

Commit c83ee06

Browse files
Xpl0itUXpl0itU
committed
Rewrite download logic
1 parent 4558c9a commit c83ee06

File tree

4 files changed

+69
-82
lines changed

4 files changed

+69
-82
lines changed

‎certificate.go‎

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -28,12 +28,12 @@ func getCert(tmdData []byte, id int, numContents uint16) ([]byte, error) {
2828
}
2929
}
3030

31-
func getDefaultCert(cancelCtx context.Context, progressReporter ProgressReporter, client *http.Client, buffer []byte) ([]byte, error) {
31+
func getDefaultCert(cancelCtx context.Context, progressReporter ProgressReporter, client *http.Client) ([]byte, error) {
3232
if len(cetkData) >= 0x350+0x300 {
3333
return cetkData[0x350 : 0x350+0x300], nil
3434
}
3535
cetkDir := path.Join(os.TempDir(), "cetk")
36-
if err := downloadFile(cancelCtx, progressReporter, client, "http://ccs.cdn.c.shop.nintendowifi.net/ccs/download/000500101000400a/cetk", cetkDir, true, buffer); err != nil {
36+
if err := downloadFile(cancelCtx, progressReporter, client, "http://ccs.cdn.c.shop.nintendowifi.net/ccs/download/000500101000400a/cetk", cetkDir, true); err != nil {
3737
return nil, err
3838
}
3939
cetkData, err := os.ReadFile(cetkDir)
@@ -51,7 +51,7 @@ func getDefaultCert(cancelCtx context.Context, progressReporter ProgressReporter
5151
return nil, fmt.Errorf("failed to download OSv10 cetk, length: %d", len(cetkData))
5252
}
5353

54-
func GenerateCert(tmdData []byte, contentCount uint16, progressReporter ProgressReporter, client *http.Client, cancelCtx context.Context, buffer []byte) (bytes.Buffer, error) {
54+
func GenerateCert(tmdData []byte, contentCount uint16, progressReporter ProgressReporter, client *http.Client, cancelCtx context.Context) (bytes.Buffer, error) {
5555
cert := bytes.Buffer{}
5656

5757
cert0, err := getCert(tmdData, 0, contentCount)
@@ -66,7 +66,7 @@ func GenerateCert(tmdData []byte, contentCount uint16, progressReporter Progress
6666
}
6767
cert.Write(cert1)
6868

69-
defaultCert, err := getDefaultCert(cancelCtx, progressReporter, client, buffer)
69+
defaultCert, err := getDefaultCert(cancelCtx, progressReporter, client)
7070
if err != nil {
7171
return bytes.Buffer{}, err
7272
}

‎cmd/WiiUDownloader/mainwindow.go‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -265,7 +265,7 @@ func (mw *MainWindow) ShowAll() {
265265

266266
wiiudownloader.GenerateTicket(filepath.Join(parentDir, "title.tik"), titleID, titleKey, titleVersion)
267267

268-
cert, err := wiiudownloader.GenerateCert(tmdData, contentCount, mw.progressWindow, http.DefaultClient, context.Background(), make([]byte, 0))
268+
cert, err := wiiudownloader.GenerateCert(tmdData, contentCount, mw.progressWindow, http.DefaultClient, context.Background())
269269
if err != nil {
270270
return
271271
}

‎downloader.go‎

Lines changed: 25 additions & 77 deletions
Original file line numberDiff line numberDiff line change
@@ -14,9 +14,8 @@ import (
1414
)
1515

1616
const (
17-
maxRetries = 5
18-
retryDelay = 5 * time.Second
19-
BUFFER_SIZE = 1048576
17+
maxRetries = 5
18+
retryDelay = 5 * time.Second
2019
)
2120

2221
type ProgressReporter interface {
@@ -30,33 +29,14 @@ type ProgressReporter interface {
3029
AddToTotalDownloaded(toAdd int64)
3130
}
3231

33-
func calculateDownloadSpeed(downloaded int64, startTime, endTime time.Time) int64 {
34-
duration := endTime.Sub(startTime).Seconds()
35-
if duration > 0 {
36-
return int64(float64(downloaded) / duration)
37-
}
38-
return 0
39-
}
40-
41-
func downloadFile(ctx context.Context, progressReporter ProgressReporter, client *http.Client, downloadURL, dstPath string, doRetries bool, buffer []byte) error {
32+
func downloadFile(ctx context.Context, progressReporter ProgressReporter, client *http.Client, downloadURL, dstPath string, doRetries bool) error {
4233
filePath := filepath.Base(dstPath)
4334

4435
startTime := time.Now()
4536
ticker := time.NewTicker(50 * time.Millisecond)
4637
defer ticker.Stop()
47-
isError := false
48-
49-
updateProgress := func(downloaded *int64) {
50-
for range ticker.C {
51-
if progressReporter.Cancelled() {
52-
return
53-
}
54-
progressReporter.UpdateDownloadProgress(*downloaded, calculateDownloadSpeed(*downloaded, startTime, time.Now()), filePath)
55-
}
56-
}
5738

5839
for attempt := 1; attempt <= maxRetries; attempt++ {
59-
isError = false
6040
req, err := http.NewRequestWithContext(ctx, "GET", downloadURL, nil)
6141
if err != nil {
6242
return err
@@ -68,15 +48,19 @@ func downloadFile(ctx context.Context, progressReporter ProgressReporter, client
6848

6949
resp, err := client.Do(req)
7050
if err != nil {
51+
if doRetries && attempt < maxRetries {
52+
time.Sleep(retryDelay)
53+
continue
54+
}
7155
return err
7256
}
7357

7458
if resp.StatusCode != http.StatusOK {
59+
resp.Body.Close()
7560
if doRetries && attempt < maxRetries {
7661
time.Sleep(retryDelay)
7762
continue
7863
}
79-
resp.Body.Close()
8064
return fmt.Errorf("download error after %d attempts, status code: %d", attempt, resp.StatusCode)
8165
}
8266

@@ -86,54 +70,20 @@ func downloadFile(ctx context.Context, progressReporter ProgressReporter, client
8670
return err
8771
}
8872

89-
var downloaded int64
90-
91-
go updateProgress(&downloaded)
92-
93-
Loop:
94-
for {
95-
select {
96-
case <-ctx.Done():
97-
resp.Body.Close()
98-
file.Close()
99-
return ctx.Err()
100-
default:
101-
n, err := resp.Body.Read(buffer)
102-
if err != nil && err != io.EOF {
103-
resp.Body.Close()
104-
file.Close()
105-
if doRetries && attempt < maxRetries {
106-
time.Sleep(retryDelay)
107-
isError = true
108-
break Loop
109-
}
110-
return fmt.Errorf("download error after %d attempts: %+v", attempt, err)
111-
}
112-
113-
if n == 0 {
114-
resp.Body.Close()
115-
file.Close()
116-
break Loop
117-
}
118-
119-
_, err = file.Write(buffer[:n])
120-
if err != nil {
121-
resp.Body.Close()
122-
file.Close()
123-
if doRetries && attempt < maxRetries {
124-
time.Sleep(retryDelay)
125-
isError = true
126-
break Loop
127-
}
128-
return fmt.Errorf("write error after %d attempts: %+v", attempt, err)
129-
}
130-
131-
downloaded += int64(n)
73+
writerProgress := newWriterProgress(file, progressReporter, startTime, filePath)
74+
_, err = io.Copy(writerProgress, resp.Body)
75+
if err != nil {
76+
file.Close()
77+
resp.Body.Close()
78+
if doRetries && attempt < maxRetries {
79+
time.Sleep(retryDelay)
80+
continue
13281
}
82+
return err
13383
}
134-
if !isError {
135-
break
136-
}
84+
file.Close()
85+
resp.Body.Close()
86+
break
13787
}
13888

13989
return nil
@@ -152,10 +102,8 @@ func DownloadTitle(cancelCtx context.Context, titleID, outputDirectory string, d
152102
return err
153103
}
154104

155-
buffer := make([]byte, BUFFER_SIZE)
156-
157105
tmdPath := filepath.Join(outputDir, "title.tmd")
158-
if err := downloadFile(cancelCtx, progressReporter, client, fmt.Sprintf("%s/%s", baseURL, "tmd"), tmdPath, true, buffer); err != nil {
106+
if err := downloadFile(cancelCtx, progressReporter, client, fmt.Sprintf("%s/%s", baseURL, "tmd"), tmdPath, true); err != nil {
159107
if progressReporter.Cancelled() {
160108
return nil
161109
}
@@ -173,7 +121,7 @@ func DownloadTitle(cancelCtx context.Context, titleID, outputDirectory string, d
173121
}
174122

175123
tikPath := filepath.Join(outputDir, "title.tik")
176-
if err := downloadFile(cancelCtx, progressReporter, client, fmt.Sprintf("%s/%s", baseURL, "cetk"), tikPath, false, buffer); err != nil {
124+
if err := downloadFile(cancelCtx, progressReporter, client, fmt.Sprintf("%s/%s", baseURL, "cetk"), tikPath, false); err != nil {
177125
if progressReporter.Cancelled() {
178126
return nil
179127
}
@@ -207,7 +155,7 @@ func DownloadTitle(cancelCtx context.Context, titleID, outputDirectory string, d
207155

208156
progressReporter.SetDownloadSize(int64(titleSize))
209157

210-
cert, err := GenerateCert(tmdData, contentCount, progressReporter, client, cancelCtx, buffer)
158+
cert, err := GenerateCert(tmdData, contentCount, progressReporter, client, cancelCtx)
211159
if err != nil {
212160
if progressReporter.Cancelled() {
213161
return nil
@@ -236,7 +184,7 @@ func DownloadTitle(cancelCtx context.Context, titleID, outputDirectory string, d
236184
return err
237185
}
238186
filePath := filepath.Join(outputDir, fmt.Sprintf("%08X.app", content.ID))
239-
if err := downloadFile(cancelCtx, progressReporter, client, fmt.Sprintf("%s/%08X", baseURL, content.ID), filePath, true, buffer); err != nil {
187+
if err := downloadFile(cancelCtx, progressReporter, client, fmt.Sprintf("%s/%08X", baseURL, content.ID), filePath, true); err != nil {
240188
if progressReporter.Cancelled() {
241189
break
242190
}
@@ -246,7 +194,7 @@ func DownloadTitle(cancelCtx context.Context, titleID, outputDirectory string, d
246194

247195
if tmdData[offset+7]&0x2 == 2 {
248196
filePath = filepath.Join(outputDir, fmt.Sprintf("%08X.h3", content.ID))
249-
if err := downloadFile(cancelCtx, progressReporter, client, fmt.Sprintf("%s/%08X.h3", baseURL, content.ID), filePath, true, buffer); err != nil {
197+
if err := downloadFile(cancelCtx, progressReporter, client, fmt.Sprintf("%s/%08X.h3", baseURL, content.ID), filePath, true); err != nil {
250198
if progressReporter.Cancelled() {
251199
break
252200
}

‎writerProgress.go‎

Lines changed: 39 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,39 @@
1+
package wiiudownloader
2+
3+
import (
4+
"io"
5+
"time"
6+
)
7+
8+
func calculateDownloadSpeed(downloaded int64, startTime, endTime time.Time) int64 {
9+
duration := endTime.Sub(startTime).Seconds()
10+
if duration > 0 {
11+
return int64(float64(downloaded) / duration)
12+
}
13+
return 0
14+
}
15+
16+
type WriterProgress struct {
17+
writer io.Writer
18+
progressReporter ProgressReporter
19+
startTime time.Time
20+
filePath string
21+
totalDownloaded int64
22+
}
23+
24+
func newWriterProgress(writer io.Writer, progressReporter ProgressReporter, startTime time.Time, filePath string) *WriterProgress {
25+
return &WriterProgress{writer: writer, totalDownloaded: 0, progressReporter: progressReporter, startTime: startTime, filePath: filePath}
26+
}
27+
28+
func (r *WriterProgress) Write(p []byte) (n int, err error) {
29+
if r.progressReporter.Cancelled() {
30+
return len(p), nil
31+
}
32+
n, err = r.writer.Write(p)
33+
if err != nil && err != io.EOF {
34+
return n, err
35+
}
36+
r.totalDownloaded += int64(n)
37+
r.progressReporter.UpdateDownloadProgress(r.totalDownloaded, calculateDownloadSpeed(r.totalDownloaded, r.startTime, time.Now()), r.filePath)
38+
return n, err
39+
}

0 commit comments

Comments
 (0)