From 0a16ffdb7fd4888fb22e434bfdb701b63f7fdce2 Mon Sep 17 00:00:00 2001
From: Xpl0itU <24777100+Xpl0itU@users.noreply.github.com>
Date: Sun, 10 Sep 2023 15:57:40 +0200
Subject: [PATCH] Improve cancelling the queue

---
 cmd/WiiUDownloader/mainwindow.go | 12 +++++++++++-
 downloader.go                    |  5 ++++-
 2 files changed, 15 insertions(+), 2 deletions(-)

diff --git a/cmd/WiiUDownloader/mainwindow.go b/cmd/WiiUDownloader/mainwindow.go
index e6bf620..a962fbf 100644
--- a/cmd/WiiUDownloader/mainwindow.go
+++ b/cmd/WiiUDownloader/mainwindow.go
@@ -2,6 +2,7 @@ package main
 
 import (
 	"bytes"
+	"context"
 	"encoding/binary"
 	"fmt"
 	"os"
@@ -743,11 +744,20 @@ func (mw *MainWindow) onDownloadQueueClicked(selectedPath string) error {
 	defer close(queueStatusChan)
 	errGroup := errgroup.Group{}
 
+	queueCtx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+
 	for _, title := range mw.titleQueue {
 		errGroup.Go(func() error {
+			select {
+			case <-queueCtx.Done():
+				queueStatusChan <- true
+				return nil
+			default:
+			}
 			tidStr := fmt.Sprintf("%016x", title.TitleID)
 			titlePath := filepath.Join(selectedPath, fmt.Sprintf("%s [%s] [%s]", normalizeFilename(title.Name), wiiudownloader.GetFormattedKind(title.TitleID), tidStr))
-			if err := wiiudownloader.DownloadTitle(tidStr, titlePath, mw.decryptContents, mw.progressWindow, mw.getDeleteEncryptedContents(), mw.logger); err != nil {
+			if err := wiiudownloader.DownloadTitle(cancel, tidStr, titlePath, mw.decryptContents, mw.progressWindow, mw.getDeleteEncryptedContents(), mw.logger); err != nil {
 				return err
 			}
 
diff --git a/downloader.go b/downloader.go
index 234a74e..e239c70 100644
--- a/downloader.go
+++ b/downloader.go
@@ -2,6 +2,7 @@ package wiiudownloader
 
 import (
 	"bytes"
+	"context"
 	"crypto/aes"
 	"crypto/cipher"
 	"encoding/binary"
@@ -69,7 +70,7 @@ func downloadFile(progressReporter ProgressReporter, client *grab.Client, downlo
 	return nil
 }
 
-func DownloadTitle(titleID, outputDirectory string, doDecryption bool, progressReporter ProgressReporter, deleteEncryptedContents bool, logger *Logger) error {
+func DownloadTitle(cancel context.CancelFunc, titleID, outputDirectory string, doDecryption bool, progressReporter ProgressReporter, deleteEncryptedContents bool, logger *Logger) error {
 	titleEntry := getTitleEntryFromTid(titleID)
 
 	progressReporter.SetGameTitle(titleEntry.Name)
@@ -181,12 +182,14 @@ func DownloadTitle(titleID, outputDirectory string, doDecryption bool, progressR
 			}
 			if err := checkContentHashes(outputDirectory, content, cipherHashTree); err != nil {
 				if progressReporter.Cancelled() {
+					cancel()
 					break
 				}
 				return err
 			}
 		}
 		if progressReporter.Cancelled() {
+			cancel()
 			break
 		}
 	}